SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbf {P} ^ {\mathrm {s e}} = \exp \left(\left(\boldsymbol {\lambda} ^ {\star} \oplus \boldsymbol {\lambda} ^ {\star} - 2 \mathbf {C}\right) \oslash \left(\boldsymbol {\gamma} ^ {\star} \oplus \boldsymbol {\gamma} ^ {\star}\right)\right). \tag {2}
$$
By defining the symmetric matrix  $\mathbf{P}(\pmb{\gamma},\pmb{\lambda}) = \exp((\pmb{\lambda} \oplus \pmb{\lambda} - 2\mathbf{C}) \odot (\pmb{\gamma} \oplus \pmb{\gamma}))$ , we prove that, when  $\pmb{\gamma} > 0$ ,  $\min_{\mathbf{P} \in S} \mathcal{L}(\mathbf{P},\pmb{\gamma},\pmb{\lambda})$  has a unique solution given by  $\mathbf{P}(\pmb{\gamma},\pmb{\lambda})$  which implies  $q(\pmb{\gamma},\pmb{\lambda}) = \mathcal{L}(\mathbf{P}(\pmb{\gamma},\pmb{\lambda}),\pmb{\gamma},\pmb{\lambda})$ . Thus the proposition shows that when  $\pmb{\gamma}^{\star} > 0$ ,  $\mathbf{P}^{\mathrm{se}} = \mathbf{P}(\pmb{\gamma}^{\star},\pmb{\lambda}^{\star})$  where  $\pmb{\gamma}^{\star},\pmb{\lambda}^{\star}$  solve the following convex problem (as maximization of a concave objective)
$$
\max  _ {\gamma > 0, \boldsymbol {\lambda}} \mathcal {L} (\mathbf {P} (\boldsymbol {\gamma}, \boldsymbol {\lambda}), \boldsymbol {\gamma}, \boldsymbol {\lambda}). \quad \text {(D u a l - S E A)}
$$
Consequently, to find  $\mathbf{P}^{\mathrm{se}}$  we solve the problem (Dual-SEA). Although the form of  $\mathbf{P}^{\mathrm{se}}$  presented in Proposition 5 is only valid when  $\gamma^{\star}$  is positive and we have only proved it for  $n - 1$  indices, we emphasize that if (Dual-SEA) has a finite solution, then it is equal to  $\mathbf{P}^{\mathrm{se}}$ . Indeed in this case the solution satisfies the KKT system associated with (SEA).
Table 1: Properties of  ${\mathbf{P}}^{\mathrm{e}}$  ,  ${\overline{\mathbf{P}}}^{\mathrm{e}}$  ,  ${\mathbf{P}}^{\mathrm{{ds}}}$  and  ${\mathbf{P}}^{\text{se }}$
AFFINITY MATRIX
REFERENCE	Pe[16]	Pē[46]	Pds[29]	Pse(SEA)
P = PT	×	✓	✓	✓
P1 = PT1 = 1	×	×	✓	✓
Hr(P) = (logξ + 1)1	✓	×	×	✓
Numerical optimization. The dual problem (Dual-SEA) is concave and can be solved with guarantees through a dual ascent approach with closed-form gradients (using e.g., SGD, BFGS [28] or ADAM [18]). At each gradient step, one can compute the current estimate  $\mathbf{P}(\boldsymbol{\gamma},\boldsymbol{\lambda})$  while the gradients of the loss w.r.t.  $\boldsymbol{\gamma}$  and  $\boldsymbol{\lambda}$  are given respectively by the constraints  $(\log \xi + 1)\mathbf{1} - \mathrm{H}_{\mathrm{r}}(\mathbf{P}(\boldsymbol{\gamma},\boldsymbol{\lambda}))$  and  $\mathbf{1} - \mathbf{P}(\boldsymbol{\gamma},\boldsymbol{\lambda})\mathbf{1}$  (see e.g., [4, Proposition 6.1.1]). Concerning time complexity, each step can be performed with  $\mathcal{O}(n^2)$  algebraic operations. From a practical perspective, we found that using a change of variable  $\boldsymbol{\gamma} \gets \boldsymbol{\gamma}^2$  and optimize  $\boldsymbol{\gamma} \in \mathbb{R}^n$  leads to enhanced numerical stability.
Remark 6. In the same spirit as Remark 3, one can express  $\mathbf{P}^{\mathrm{se}}$  as a KL projection of  $\mathbf{K}_{\sigma} = \exp (-\mathbf{C} / \sigma)$ . Indeed, we show in Appendix A.5 that if  $0 < \sigma \leq \min_{i}\gamma_{i}^{\star}$ , then  $\mathbf{P}^{\mathrm{se}} = \mathrm{Proj}_{\mathcal{H}_{\xi}\cap \mathcal{S}}^{\mathrm{KL}}(\mathbf{K}_{\sigma})$ .
Comparison between  $\mathbf{P}^{\mathrm{ds}}$  and  $\mathbf{P}^{\mathrm{se}}$ . In Figure 2 we illustrate the ability of our proposed affinity  $\mathbf{P}^{\mathrm{se}}$  to adapt to varying noise levels. In the OT problem that we consider, each sample is given a mass of one that is distributed over its neighbors (including itself since self-loops are allowed). For each sample, we refer to the entropy of the distribution over its neighbors as the spreading of its mass. One can notice that for  $\mathbf{P}^{\mathrm{ds}}$  (DS) (OT problem with global entropy constraint (EOT)), the samples do not spread their mass evenly depending on the density around them. On the contrary, the per-row entropy constraints of  $\mathbf{P}^{\mathrm{se}}$  force equal spreading among samples. This can have benefits, particularly for clustering, as illustrated in the rightmost plot, which shows the eigenvalues of the associated Laplacian matrices (recall that the number of connected components equals the dimension of the null space of its Laplacian [7]). As can be seen,  $\mathbf{P}^{\mathrm{ds}}$  results in many unwanted clusters, unlike  $\mathbf{P}^{\mathrm{se}}$ , which is robust to varying noise levels (its Laplacian matrix has only 3 vanishing eigenvalues). We further illustrate this phenomenon on Figure 3 with varying noise levels.
Figure 3: ARI spectral clustering on the example of three Gaussian clusters with variances:  $\sigma^2$ ,  $2\sigma^2$  and  $3\sigma^2$  (as in Figure 2).
4 Optimal Transport for Dimension Reduction with SNEkhorn
In this section, we build upon symmetric entropic affinities to introduce SNEkhorn, a new DR algorithm that fully benefits from the advantages of doubly stochastic affinities.
SNEkhorn's objective. Our proposed method relies on doubly stochastic affinity matrices to capture the dependencies among the samples in both input and latent spaces. The KL divergence, which is the central criterion in most popular DR methods [44], is used to measure the discrepancy between the two affinities. As detailed in sections 2 and 3,  $\mathbf{P}^{\mathrm{se}}$  computed using the cost  $[\mathbf{C}_{\mathbf{X}}]_{ij} = \| \mathbf{X}_{i}: - \mathbf{X}_{j}\|_{2}^{2}$ , corrects for heterogeneity in the input data density by imposing point-wise entropy constraints. As we do not need such correction for embedding coordinates  $\mathbf{Z}$  since they must be optimized, we opt for the standard affinity (DS) built as an OT transport plan with global entropy constraint (EOT). This OT plan can be efficiently computed using Sinkhorn's algorithm. More precisely, we propose the optimization problem
$$
\min  _ {\mathbf {Z} \in \mathbb {R} ^ {n \times q}} \operatorname {K L} \left(\mathbf {P} ^ {\mathrm {s e}} \mid \mathbf {Q} _ {\mathbf {Z}} ^ {\mathrm {d s}}\right), \tag {SNEkhorn}
$$
where  $\mathbf{Q}_{\mathbf{Z}}^{\mathrm{ds}} = \exp (\mathbf{f}_{\mathbf{Z}}\oplus \mathbf{f}_{\mathbf{Z}} - \mathbf{C}_{\mathbf{Z}})$  stands for the (DS) affinity computed with cost  $[\mathbf{C_{Z}}]_{ij} = \| \mathbf{Z}_{i:}-$ $\mathbf{Z}_{j}\|_{2}^{2}$  and  $\mathbf{f}_{\mathbf{Z}}$  is the optimal dual variable found by Sinkhorn's algorithm. We set the bandwidth to  $\nu = 1$  in  $\mathbf{Q}_{\mathbf{Z}}^{\mathrm{ds}}$  similarly to [46] as the bandwidth in the low dimensional space only affects the scales
Figure 5: From left to right: entries of  $\overline{\mathbf{P}^{\mathrm{e}}}$  (Symmetric-SNE) and associated embeddings generated using  $\overline{\mathbf{P}^{\mathrm{e}}}$ . Then  $\mathbf{P}^{\mathrm{se}}$  (SEA) matrix and associated SNEkhorn embeddings. Perplexity  $\xi = 30$ .
of the embeddings and not their shape. Keeping only the terms that depend on  $\mathbf{Z}$  and relying on the double stochasticity of  $\mathbf{P}^{\mathrm{se}}$ , the objective in (SNEkhorn) can be expressed as  $\langle \mathbf{P}^{\mathrm{se}}, \mathbf{C}_{\mathbf{Z}} \rangle - 2\langle \mathbf{f}_{\mathbf{Z}}, \mathbf{1} \rangle$ .
Heavy-tailed kernel in latent space. Since it is well known that heavy-tailed kernels can be beneficial in DR [21], we propose an extension called t-SNEkhorn that simply amounts to computing a doubly stochastic student-t kernel in the low-dimensional space. With our construction, it corresponds to choosing the cost  $\left[\mathbf{C}_{\mathbf{Z}}\right]_{ij} = \log \left(1 + \left\| \mathbf{Z}_{i:} - \mathbf{Z}_{j:}\right\| _2^2\right)$  instead of  $\| \mathbf{Z}_i:\mathbf{-Z}_j:\| _2^2$
Inference. This new DR objective involves computing a doubly stochastic normalization for each update of  $\mathbf{Z}$ . Interestingly, to compute the optimal dual variable  $\mathbf{f}_{\mathbf{Z}}$  in  $\mathbf{Q}_{\mathbf{Z}}^{\mathrm{ds}}$ , we leverage a well-conditioned Sinkhorn fixed point iteration [19, 13], which converges extremely fast in the symmetric setting:
$$
\forall i, \left[ \mathbf {f} _ {\mathbf {Z}} \right] _ {i} \leftarrow \frac {1}{2} \left(\left[ \mathbf {f} _ {\mathbf {Z}} \right] _ {i} - \log \sum_ {k} \exp \left(\left[ \mathbf {f} _ {\mathbf {Z}} \right] _ {k} - \left[ \mathbf {C} _ {\mathbf {Z}} \right] _ {k i}\right)\right).
$$
On the right side of Figure 4, we plot  $\| \mathbf{Q}_{\mathbf{Z}}^{\mathrm{ds}}\mathbf{1} - \mathbf{1}\|_{\infty}$  as a function of (Sinkhorn) iterations for a toy example presented in Section 5. In most practical cases, we found that about 10 iterations were enough to reach a sufficiently small error.  $\mathbf{Z}$  is updated through gradient descent with gradients obtained by performing backpropagation through the Sinkhorn iterations. These iterations can be further accelerated with a warm start strategy by plugging the last  $\mathbf{f}_{\mathbf{Z}}$  to initialize the current one.
Related work. Using doubly stochastic affinities for SNE has been proposed in [29], with two key differences from our work. First, they do not consider EAs and resort to  $\mathbf{P}^{\mathrm{ds}}$  (DS). This affinity, unlike  $\mathbf{P}^{\mathrm{se}}$ , is not adaptive to the data heterogeneous density (as illustrated
Figure 4: Left: SNEkhorn embedding on the simulated data of Section 5 using  $\widetilde{\mathbf{Q}}_{\mathbf{Z}}$  instead of  $\mathbf{Q}_{\mathbf{Z}}^{\mathrm{ds}}$  with  $\xi = 30$ . Right: number of iterations needed to achieve  $\| \mathbf{Q}_{\mathbf{Z}}^{\mathrm{ds}}\mathbf{1} - \mathbf{1}\|_{\infty} \leq \mathrm{tol}$  with (Sinkhorn).
in Figure 2). Second, they use the affinity  $\bar{\mathbf{Q}}_{\mathbf{Z}}$  in the low-dimensional space and illustrate empirically that matching the latter with a doubly stochastic matrix (e.g.,  $\mathbf{P}^{\mathrm{ds}}$  or  $\mathbf{P}^{\mathrm{se}}$ ) can sometimes impose spherical constraints on the embedding  $\mathbf{Z}$ . This is detrimental for projections onto a  $2D$  flat space (typical use case of DR) where embeddings tend to form circles. This can be verified on the left side of Figure 4. In contrast, in SNEkhorn, the latent affinity is also doubly stochastic so that latent coordinates  $\mathbf{Z}$  are not subject to spherical constraints anymore. The corresponding SNEkhorn embedding is shown in Figure 5 (bottom right).
This section aims to illustrate the performances of the proposed affinity matrix  $\mathbf{P}^{\mathrm{se}}$  (SEA) and DR method SNEkhorn at faithfully representing dependencies and clusters in low dimensions. First, we showcase the relevance of our approach on a simple synthetic dataset with heteroscedastic noise. Then,
we evaluate the spectral clustering performances of symmetric entropic affinities before benchmarking t-SNEkhorn with t-SNE and UMAP [31] on real-world images and genomics datasets.2
Simulated data. We consider the toy dataset with heteroscedastic noise from [23]. It consists of sampling uniformly two vectors  $\mathbf{p}_1$  and  $\mathbf{p}_2$  in the  $10^{4}$ -dimensional probability simplex.  $n = 10^3$  samples are then generated as  $\mathbf{x}_i = \tilde{\mathbf{x}}_i / (\sum_j\tilde{x}_{ij})$  where
$$
\tilde {\mathbf {x}} _ {i} \sim \left\{ \begin{array}{l l} \mathcal {M} (1 0 0 0, \mathbf {p} _ {1}), & 1 \leq i \leq 5 0 0 \\ \mathcal {M} (1 0 0 0, \mathbf {p} _ {2}), & 5 0 1 \leq i \leq 7 5 0 \\ \mathcal {M} (2 0 0 0, \mathbf {p} _ {2}), & 7 5 1 \leq i \leq 1 0 0 0. \end{array} \right.
$$
where  $\mathcal{M}$  stands for the multinomial distribution. The goal of the task is to test the robustness to heteroscedastic noise. Indeed, points generated using  $\mathbf{p}_2$  exhibit different levels of noise due to various numbers of multinomial trials to form an estimation of  $\mathbf{p}_2$ . This typically occurs in real-world scenarios when the same entity is measured using different experimental setups thus creating heterogeneous technical noise levels (e.g., in single-cell sequencing [20]). This phenomenon is known as batch effect [43]. In Figure 5, we show that, unlike  $\overline{\mathbf{P}^{\mathrm{e}}}$  (Symmetric-SNE),  $\mathbf{P}^{\mathrm{se}}$  (SEA) manages to properly filter the noise (top row) to discriminate between samples generated by  $\mathbf{p}_1$  and  $\mathbf{p}_2$ , and represent these two clusters separately in the embedding space (bottom row). In contrast,  $\overline{\mathbf{P}^{\mathrm{e}}}$  and SNE are misled by the batch effect. This shows that  $\overline{\mathbf{P}^{\mathrm{e}}}$  doesn't fully benefit from the adaptivity of EAs due to poor normalization and symmetrization. This phenomenon partly explains the superiority of SNEkhorn and t-SNEkhorn over current approaches on real-world datasets as illustrated below.
Real-world datasets. We then experiment with various labeled classification datasets including images and genomic data. For images, we use COIL 20 [34], OLIVETTI faces [12], UMNIST [15] and CIFAR 10 [22]. For CIFAR, we experiment with features obtained from the last hidden layer of a pre-trained ResNet [38] while for the other three datasets, we take as input the raw pixel data. Regarding genomics data, we consider the Curated Microarray Database (CuMiDa) [11] made of microarray datasets for various types of cancer, as well as the preprocessed SNAREseq (chromatin accessibility) and scGEM (gene expression)
Table 2: ARI  $\left( {\times {100}}\right)$  clustering scores on genomics.
DATA SET	\( \overline{\mathbf{P}}\overline{\mathbf{r}}\overline{\mathbf{s}} \)	\( \mathbf{P}^{\mathrm{ds}} \)	\( \mathbf{P}^{\mathrm{st}} \)	\( \overline{\mathbf{P}}\overline{\mathbf{e}} \)	\( \mathbf{P}^{\mathrm{se}} \)
LIVER (14520)	75.8	75.8	84.9	80.8	85.9
BREAST (70947)	30.0	30.0	26.5	23.5	28.5
LEUKEMIA (28497)	43.7	44.1	49.7	42.5	50.6
COLORECTAL (44076)	95.9	95.9	93.9	95.9	95.9
LIVER (76427)	76.7	76.7	83.3	81.1	81.1
BREAST (45827)	43.6	53.8	74.7	71.5	77.0
COLORECTAL (21510)	57.6	57.6	54.7	94.0	79.3
RENAL (53757)	47.6	47.6	49.5	49.5	49.5
PROSTATE (6919)	12.0	13.0	13.2	16.3	17.4
THROAT (42743)	9.29	9.29	11.4	11.8	44.2
SCGEM	57.3	58.5	74.8	69.9	71.6
SNARESEQ	8.89	9.95	46.3	55.4	96.6
datasets used in [9]. For CuMiDa, we retain the datasets with most samples. For all the datasets, when the data dimension exceeds 50 we apply a pre-processing step of PCA in dimension 50, as usually done in practice [46]. In the following experiments, when not specified the hyperparameters are set to the value leading to the best average score on five different seeds with grid-search. For perplexity parameters, we test all multiples of 10 in the interval  $[10, \min(n, 300)]$  where  $n$  is the number of samples in the dataset. We use the same grid for the  $k$  of the self-tuning affinity  $\mathbf{P}^{\mathrm{st}}$  [53] and for the n_neighbors parameter of UMAP. For scalar bandwidths, we consider powers of 10 such that the corresponding affinities' average perplexity belongs to the perplexity range.
Spectral Clustering. Building on the strong connections between spectral clustering mechanisms and t-SNE [44, 27] we first consider spectral clustering tasks to evaluate the affinity matrix  $\mathbf{P}^{\mathrm{se}}$  (SEA) and compare it against  $\overline{\mathbf{P}^{\mathrm{e}}}$  (Symmetric-SNE). We also consider two versions of the Gaussian affinity with scalar bandwidth  $\mathbf{K} = \exp (-\mathbf{C} / \nu)$ : the symmetrized row-stochastic  $\overline{\mathbf{P}^{\mathrm{rs}}} = \mathrm{Proj}_{\mathcal{S}}^{\ell_2}(\mathbf{P}^{\mathrm{rs}})$  where  $\mathbf{P}^{\mathrm{rs}}$  is  $\mathbf{K}$  normalized by row and  $\mathbf{P}^{\mathrm{ds}}$  (DS). We also consider the adaptive Self-Tuning  $\mathbf{P}^{\mathrm{st}}$  affinity from [53] which relies on an adaptive bandwidth corresponding to the distance from the  $k$ -th nearest neighbor of each point. We use the spectral clustering implementation of scikit-learn [36] with default parameters which uses the unnormalized graph Laplacian. We measure the quality of clustering using the Adjusted Rand Index (ARI). Looking at both Table 2 and Figure 6, one can notice that, in general, symmetric entropic affinities yield better results than usual entropic affinities with significant improvements in some datasets (e.g., throat microarray and SNAREseq). Overall
2Our code is available at https://github.com/PythonOT/SNEkhorn.
Figure 6: ARI spectral clustering score as a function of the perplexity parameter for image datasets.
Table 3: Scores for the UMAP, t-SNE and t-SNEkhorn embeddings.
Silhouette (×100)	Trustworthiness (×100)
UMAP	t-SNE	t-SNEkhorn	UMAP	t-SNE	t-SNEkhorn
COIL	20.4 ± 3.3	30.7 ± 6.9	52.3 ± 1.1	99.6 ± 0.1	99.6 ± 0.1	99.9 ± 0.1
OLIVETTI	6.4 ± 4.2	4.5 ± 3.1	15.7 ± 2.2	96.5 ± 1.3	96.2 ± 0.6	98.0 ± 0.4
UMNIST	-1.4 ± 2.7	-0.2 ± 1.5	25.4 ± 4.9	93.0 ± 0.4	99.6 ± 0.2	99.8 ± 0.1
CIFAR	13.6 ± 2.4	18.3 ± 0.8	31.5 ± 1.3	90.2 ± 0.8	90.1 ± 0.4	92.4 ± 0.3
Liver (14520)	49.7 ± 1.3	50.9 ± 0.7	61.1 ± 0.3	89.2 ± 0.7	90.4 ± 0.4	92.3 ± 0.3
Breast (70947)	28.6 ± 0.8	29.0 ± 0.2	31.2 ± 0.2	90.9 ± 0.5	91.3 ± 0.3	93.2 ± 0.4
Leukemia (28497)	22.3 ± 0.7	20.6 ± 0.7	26.2 ± 2.3	90.4 ± 1.1	92.3 ± 0.8	94.3 ± 0.5
Colorectal (44076)	67.6 ± 2.2	69.5 ± 0.5	74.8 ± 0.4	93.2 ± 0.7	93.7 ± 0.5	94.3 ± 0.6
Liver (76427)	39.4 ± 4.3	38.3 ± 0.9	51.2 ± 2.5	85.9 ± 0.4	89.4 ± 1.0	92.0 ± 1.0
Breast (45827)	35.4 ± 3.3	39.5 ± 1.9	44.4 ± 0.5	93.2 ± 0.4	94.3 ± 0.2	94.7 ± 0.3
Colorectal (21510)	38.0 ± 1.3	42.3 ± 0.6	35.1 ± 2.1	85.6 ± 0.7	88.3 ± 0.9	88.2 ± 0.7
Renal (53757)	44.4 ± 1.5	45.9 ± 0.3	47.8 ± 0.1	93.9 ± 0.2	94.6 ± 0.2	94.0 ± 0.2
Prostate (6919)	5.4 ± 2.7	8.1 ± 0.2	9.1 ± 0.1	77.6 ± 1.8	80.6 ± 0.2	73.1 ± 0.5
Throat (42743)	26.7 ± 2.4	28.0 ± 0.3	32.3 ± 0.1	91.5 ± 1.3	88.6 ± 0.8	86.8 ± 1.0
scGEM	26.9 ± 3.7	33.0 ± 1.1	39.3 ± 0.7	95.0 ± 1.3	96.2 ± 0.6	96.8 ± 0.3
SNAREseq	6.8 ± 6.0	35.8 ± 5.2	67.9 ± 1.2	93.1 ± 2.8	99.1 ± 0.1	99.2 ± 0.1
$\mathbf{P}^{\mathrm{se}}$  outperforms all the other affinities in 8 out of 12 datasets. This shows that the adaptivity of EAs is crucial. Figure 6 also shows that this superiority is verified for the whole range of perplexities. This can be attributed to the fact that symmetric entropic affinities combine the advantages of doubly stochastic normalization in terms of clustering and of EAs in terms of adaptivity. In the next experiment, we show that these advantages translate into better clustering and neighborhood retrieval at the embedding level when running SNEkhorn.
Dimension Reduction. To guarantee a fair comparison, we implemented not only SNEkhorn, but also t-SNE and UMAP in PyTorch [35]. Note that UMAP also relies on adaptive affinities but sets the degree of each node (related to the hyperparameter  $n$  neighbors which plays a similar role to the perplexity) rather than the entropy. All models were optimized using ADAM [18] with default parameters and the same stopping criterion: the algorithm stops whenever the relative variation of the loss becomes smaller than  $10^{-5}$ . For each run, we draw independent
Figure 7: SNAREseq embeddings produced by t-SNE and t-SNEkhorn with  $\xi = 50$
$\mathcal{N}(0,1)$  coordinates and use this same matrix to initialize all the methods that we wish to compare. To evaluate the embeddings' quality, we make use of the silhouette [40] and trustworthiness [48]
scores from scikit-learn [36] with default parameters. While the former relies on class labels, the latter measures the agreement between the neighborhoods in input and output spaces, thus giving two complementary metrics to properly evaluate the embeddings. The results, presented in Table 3, demonstrate the notable superiority of t-SNEkhorn compared to the commonly used t-SNE and UMAP algorithms. A sensitivity analysis on perplexity can also be found in Appendix B. Across the 16 datasets examined, t-SNEkhorn almost consistently outperformed the others, achieving the highest silhouette score on 15 datasets and the highest trustworthiness score on 12 datasets. To visually assess the quality of the embeddings, we provide SNAREseq embeddings in Figure 7. Notably, one can notice that the use of t-SNEkhorn results in improved class separation compared to t-SNE.
We have introduced a new principled and efficient method for constructing symmetric entropic affinities. Unlike the current formulation that enforces symmetry through an orthogonal projection, our approach allows control over the entropy in each point thus achieving entropic affinities' primary goal. Additionally, it produces a DS-normalized affinity and thus benefits from the well-known advantages of this normalization. Our affinity takes as input the same perplexity parameter as EAs and can thus be used with little hassle for practitioners. We demonstrate experimentally that both our affinity and DR algorithm (SNEkhorn), leveraging a doubly stochastic kernel in the latent space, achieve substantial improvements over existing approaches.
Note that in the present work, we do not address the issue of large-scale dependencies that are not faithfully represented in the low-dimensional space [44]. The latter shall be treated in future works. Among other promising research directions, one could focus on building multi-scale versions of symmetric entropic affinities [24] as well as fast approximations for SNEkhorn forces by adapting e.g., Barnes-Hut [45] or interpolation-based methods [26] to the doubly stochastic setting. It could also be interesting to use SEAs in order to study the training dynamics of transformers [54].
The authors are grateful to Mathurin Massias, Jean Feydy and Aurélien Garivier for insightful discussions. This project was supported in part by the ANR projects AllegroAssai ANR-19-CHIA-0009, SingleStatOmics ANR-18-CE45-0023 and OTTOPIA ANR-20-CHIA-0030. This work was also supported by the ACADEMICS grant of the IDEXLYON, project of the Université de Lyon, PIA operated by ANR-16-IDEX-0005.
[1] Mario Beauchemin. On affinity matrix normalization for graph cuts and spectral clustering. Pattern Recognition Letters, 68:90-96, 2015.
[2] Mikhail Belkin and Partha Niyogi. Laplacian eigenmaps for dimensionality reduction and data representation. Neural computation, 15(6):1373-1396, 2003.
[3] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyre. Iterative bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111-A1138, 2015.
[4] Dimitri P Bertsekas. Nonlinear programming. Journal of the Operational Research Society, 48(3):334-334, 1997.
[5] T. Tony Cai and Rong Ma. Theoretical foundations of t-sne for visualizing high-dimensional clustered data. Journal of Machine Learning Research (JMLR), 23(301):1-54, 2022.
[6] Miguel A Carreira-Perpinán. The elastic embedding algorithm for dimensionality reduction. In International Conference on Machine Learning (ICML), volume 10, pages 167–174, 2010.
[7] Fan RK Chung. Spectral graph theory, volume 92. American Mathematical Soc., 1997.
[8] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Neural Information Processing Systems (NeurIPS), 26, 2013.