On UMAP’s True Loss Function
Year: 2022

Abstract:
UMAP has supplanted t-SNE as state-of-the-art for visualizing high-dimensional datasets in many disciplines, but the reason for its success is not well understood. In this work, we investigate UMAP’s sampling based optimization scheme in detail. We derive UMAP’s true loss function in closed form and find that it differs from the published one in a dataset size dependent way. As a consequence, we show that UMAP does not aim to reproduce its theoretically motivated high-dimensional UMAP similarities. Instead, it tries to reproduce similarities that only encode the k nearest neighbor graph, thereby challenging the previous understanding of UMAP’s effectiveness. Alternatively, we consider the implicit balancing of attraction and repulsion due to the negative sampling to be key to UMAP’s success. We corroborate our theoretical findings on toy and single cell RNA sequencing data.

heading: 1 Introduction
text: Today’s most prominent methods for non-parametric, non-linear dimension reduction are t-Distributed Stochastic Neighbor Embedding (t-SNE) [20, 19] and Uniform Manifold Approximation and Projection for Dimension Reduction (UMAP) [13]. The heart of UMAP is claimed to be its sophisticated method for extracting the high-dimensional similarities. However, the reason for UMAP’s excellent visualizations is not immediately obvious from this approach. In particular, UMAP’s eponymous uniformity assumption, see Section 3, is arguably difficult to defend for the wide range of datasets on which UMAP performs well. Therefore, it is not well understood what aspect of UMAP is responsible for its great visualizations.
Both t-SNE and UMAP have to overcome the computational obstacle of considering the quadratic number of interactions between all pairs of points. The breakthrough for t-SNE came with a Barnes-Hut approximation [19]. Instead, UMAP employs a sampling based approach to avoid a quadratic number of repulsive interactions. Other than [4] little attention has been paid to this sampling based optimization scheme. In this work, we fill this gap and analyze UMAP’s optimization method in detail. In particular, we derive the effective, closed form loss function which is truly minimized by UMAP’s optimization scheme. While UMAP’s use of negative sampling was intended to avoid quadratic complexity, we find, surprisingly, that the resulting effective loss function differs significantly from UMAP’s purported loss function. The weight of the loss function’s repulsive term is drastically reduced. As a consequence, UMAP is not actually geared towards reproducing the clever high-dimensional similarities. In fact, we show that most information beyond the shared kNN graph connectivity is essentially ignored as UMAP actually approximates a binarized version of the high-dimensional similarities. These theoretical findings underpin some empirical observations in [4] and demonstrate that the gist of UMAP is not its high-dimensional similarities. This resolves the disconnect between UMAP’s uniformity assumption and its success on datasets of varying density.
From a user’s perspective it is important to gain an intuition for deciding which features of a visualization can be attributed to the data and which ones are more likely artifacts of the visualization method. With our analysis, we can explain UMAP’s tendency to produce crisp, over-contracted substructures, which increases with the dataset size, as a side effect of its optimization.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Without the motivation of reproducing sophisticated high-dimensional similarities in embedding space, it seems unclear why UMAP performs well. We propose an alternative explanation for UMAP’s success: The sampling based optimization scheme balances the attractive and repulsive loss terms despite the sparse high-dimensional attraction. Consequently, UMAP can leverage the connectivity information of the shared kNN graph via gradient descent effectively.

heading: 2 Related work
text: For most of the past decade t-SNE [20, 19] was considered state-of-the-art for non-linear dimension reduction. In the last years UMAP [13] at least ties with t-SNE. In both methods, points are embedded so as to reproduce high-dimensional similarities. Different from t-SNE, the high-dimensional similarities are sparse for UMAP and do not need to be normalized over the entire dataset. Additionally, t-SNE adapts the local scale of high-dimensional similarities by achieving a predefined perplexity, while UMAP uses its uniformity assumption. The low-dimensional similarity functions also differ. Recently, Böhm et al. [4] placed both UMAP and t-SNE on a spectrum of dimension reduction methods that mainly differ in the amount of repulsion employed. They argue that UMAP uses less repulsion than t-SNE. A parametric version of UMAP was proposed in [17].
UMAP’s success, in particular in the biological community [2, 16], sparked interest in understanding UMAP more deeply. The original paper [13] motivates the choice of the high-dimensional similarities using concepts from algebraic topology and thus justifies UMAP’s transition from local similarities µi→j to global similarities µij . The authors find that while the algorithm focuses on reproducing the local similarity pattern similar to t-SNE, it achieves better global results. In contrast, Kobak and Linderman [10] attribute the better global properties of UMAP visualizations to the more informative initialization and show that t-SNE manages to capture more global structure if initialized in a similar way. Wang et al. [21] also analyze the local and global properties of t-SNE, UMAP and TriMAP [1].
Narayan et al. [14] observe that UMAP’s uniformity assumption leads to visualizations in which denser regions are more spread out while sparser regions get overly contracted. They propose an additional loss term that aims to reproduce the local density around each point and thus spaces sparser regions out. We provide an additional explanation for overly contracted visualizations: UMAP does not reproduce the high-dimensional similarities but exaggerates the attractive forces over the repulsive ones, which can result in overly crisp visualizations, see Figures 1b, 1c and 2a.
Our work aligns with Böhm et al. [4]. The authors conjecture that the sampling based optimization procedure of UMAP prevents the minimization of the supposed loss function, thus not reproducing the high-dimensional similarities in embedding space. They substantiate this hypothesis by qualitatively estimating the relative size of attractive and repulsive forces. In addition, they implement a BarnesHut approximation to the loss function (6) and find that it yields a diverged embedding. We analyze UMAP’s sampling procedure in depth, compute UMAP’s true loss function in closed form and contrast it against the supposed loss in Section 5. Based on this analytic effective loss function, we can further explain Böhm et al. [4]’s empirical finding that the specific high-dimensional similarities provide little gain over the binary weights of a shared kNN graph,1 see Section 6. Finally, our theoretical framework leads us to a new tentative explanation for UMAP’s success in Section 7.

heading: 3 Background on UMAP
text: UMAP assumes that the data lies on a low-dimensional Riemannian manifold (R, g) in ambient high-dimensional space. Moreover, the data is assumed to be uniformly distributed with respect to the metric tensor g, or put simply, with respect to the distance along the manifold. The metric tensor g in turn is assumed to be locally constant. The key idea of UMAP [13] is to compute pairwise similarities in high-dimensional space which inform the optimization of the low-dimensional embedding. Let x1, ..., xn ∈ RD be high-dimensional, mutually distinct data points for which low-dimensional embeddings e1, ..., en ∈ Rd shall be found, where d D, often d = 2 or 3. First, UMAP extracts high-dimensional similarities between the data points. To do so, the k nearest neighbor graph is computed. Let i1, ..., ik denote the indices of xi’s k nearest neighbors in increasing order of distance to xi. Then, using its uniformity assumption, UMAP fits a local notion of similarity
1The shared k nearest neighbor graph contains an edge ij if i is among j’s k nearest neighbors or vice versa.
for each data point i by selecting a scale σi such that the total similarity of each point to its k nearest neighbors is normalized, i.e. find σi such that
k∑ κ=1 exp ( − ( d(xi, xiκ)− d(xi, xi1) ) /σi ) = log2(k). (1)
This defines the directed high-dimensional similarities
µi→j =
{ exp ( − ( d(xi, xj)− d(xi, xi1) ) /σi ) for j ∈ {i1, . . . , ik}
0 else. (2)
Finally, these are symmetrized to obtain undirected high-dimensional similarities or input similarities between items i and j µij = µi→j + µj→i − µi→jµj→i ∈ [0, 1]. (3) While each node has exactly k non-zero directed similarities µi→j to other nodes which sum to log2(k), this does not hold exactly after symmetrization. Nevertheless, typically the µij are highly sparse, each node has positive similarity to about k other nodes and the degree of each node di = ∑n j=1 µij is approximately constant and close to log2(k), see Supp. Fig. 5 and 6 in Appendix B.
For convenience of notation, we set µii = 0 and define the total similarity as µtot = 12 ∑n i=1 di.
Distance in embedding space is transformed to low-dimensional similarity by a smooth approximation to the high-dimensional similarity function, φ(d; a, b) = (1 + ad2b)−1, for all pairs of points. The shape parameters a, b are essentially hyperparameters of UMAP. We will overload notation and write
νij = φ(||ei − ej ||) = φ(ei, ej) (4)
for the low-dimensional or embedding similarities and usually suppress their dependence on a and b.
Supposedly, UMAP approximately optimizes the following objective function with respect to the embeddings e1, . . . , en:
L({ei}|{µij}) = −2 ∑
1≤i<j≤n
( µij log(νij) +(1− µij) log(1− νij) ) (5)
= −2 ∑
1≤i<j≤n
( µij log(φ(ei, ej))︸ ︷︷ ︸
−Laij
+(1− µij) log(1− φ(ei, ej))︸ ︷︷ ︸ −Lrij
) . (6)
While the high-dimensional similarities µij are symmetric, UMAP’s implementation does consider their direction during optimization. For this reason, our losses in equations (5) and (6) differ by a factor of 2 from the one given in [13]. Viewed through the lens of a force-directed model, the derivative of the first term in each summand of the loss function,−∂Laij/∂ei, captures the attraction of ei to ej due to the high-dimensional similarity µij and the derivative of the second term, −∂Lrij/∂ei, represents the repulsion that ej exerts on ei due to a lack of similarity in high dimension, 1− µij . Alternatively, the loss can be seen as the sum of binary cross entropy losses for each pairwise similarity. Thus, it is minimized if the low-dimensional similarities νij exactly match their high-dimensional counterparts µij , that is, if UMAP manages to reproduce the input similarities in embedding space.
UMAP uses a sampling based stochastic gradient descent to optimize the low-dimensional embedding typically starting from a Laplacian Eigenmap layout [3, 10]. Our main contribution is to show that the sampling based optimization in fact leads to a different objective function, so that UMAP does not reproduce the high-dimensional similarities in low-dimensional space, see Sections 4 to 6.

heading: 4 UMAP does not reproduce high-dimensional similarities
text: UMAP produces scientifically useful visualizations in many domains and is fairly robust to its hyperparameters. Since any visualization of intrinsically high-dimensional data must be somehow unfaithful, it is not straightforward to check the visualization quality other than by its downstream use. Some quantitative metrics such as the correlation between high- and low-dimensional distances [10, 2] exist. We follow a different route to show unexpected properties of UMAP. Consider the toy example of applying UMAP to data that is already low-dimensional, so that no reduction in dimension is
required: D = d = 2. Ideally, the data would be preserved in this situation. One might also expect that UMAP achieves its aim of perfectly reproducing the input similarities. Surprisingly, neither of these expectations is met. In Figure 1, we depict 2D UMAP visualizations of a 2D uniform ring dataset. To ease UMAP’s task, we initialized the embedding with the original data, hoping that UMAP’s optimization would just deem this layout to be optimal. We also used a longer run time to ensure convergence of UMAP’s optimization procedure. Otherwise, we used the default hyperparameters. The results with default initialization and run time are qualitatively similar and shown in Supp. Fig. 20. UMAP manages to capture the ring shape of the data, but changes its appearance significantly. It contracts the width of the ring nearly to a line, see Figure 1b, a phenomenon also observed on real-world datasets, see Figure 2a. Whether this exaggeration of the ring shape is useful depends on the use case. This finding goes beyond Narayan et al. [14]’s observation that over-contraction happens in regions of low density since our toy dataset is sampled uniformly from a circular ring.
Moreover, we embedded the same ring with UMAP as part of a larger dataset of twenty such rings, spaced so that the input similarities for each ring are not influenced by the other rings. According to Section 3, we would expect a qualitatively similar embedding as for the individual ring. However, we see in Figure 1c that UMAP produces an even crisper embedding. In another experiment, we varied the number of samples from a single ring and observed again that the tendency for over-contraction increases with the dataset size. For more information on UMAP’s dependence on the dataset size confer Appendix H. The dataset size dependence is particularly noteworthy, as computing UMAP embeddings of subsets or subsamples of a dataset is common practice.
As described in Section 3, UMAP employs different methods in input and embedding space to transform distances to similarities. In particular, in input space similarities are zero for all but the closest neighbors, while in embedding space they are computed with the heavy-tailed function φ. To test whether this prevents the reproduction of the input data, we also computed the dense similarities on the original data with φ and used these as input similarities for the embedding in Figure 1d. Since we also initialize with the original dataset, a global optimum of the objective function (6) for this choice of input similarity, one would expect no change by UMAP’s optimization scheme. However, we observe that in this setting UMAP produces spurious curves and increases the width of the ring.
Böhm et al. [4] implemented a Barnes-Hut approximation of UMAP’s objective function (6), which produced a diverged embedding. Inspired by this finding, we compute the loss values according to equation (5) for various input and embedding similarities µij and νij in our toy example, see Table 1.
Consider the row with the usual input similarities (µij = µ({x1, . . . , xn})). In a completely diverged embedding, all self similarities are one and all others zero (νij = 1(i == j)). We find that the loss for such an embedding is lower than for the optimized UMAP embedding. This is in
accordance with Böhm et al. [4]’s Barnes-Hut experiment and shows that UMAP does not optimize its supposed objective function (6) as a diverged embedding is approximately feasible in two dimensions. This discrepancy is not just due to the fact that input and embedding similarities are computed differently: The second row of Table 1 contains loss values for the setting in which we use the dense similarities as input similarities, as in Figure 1d. We initialize the embedding at the optimal loss value (νij = φ({x1, . . . , xn}) = µij), but UMAP’s optimization moves away from this layout and towards an embedding with higher loss (νij = φ({e1, . . . , en})) although we always compute similarity in the same manner. Clearly, UMAP’s optimization yields unexpected results.

heading: 5 UMAP’s sampling strategy and effective loss function
text: UMAP uses a sampling based approach to optimize its loss function, to reduce complexity. A simplified version of the sampling procedure can be found in Algorithm 1. Briefly put, an edge ij is sampled according to its high-dimensional similarity and the embeddings ei and ej of the incident nodes are pulled towards each other. For each such sampled edge ij, the algorithm next samples m negative samples s uniformly from all nodes and the embedding of i is repelled from that of each negative sample. Note that the embeddings of the negative samples are not repelled from that of i, see Alg. 1 line 10. So there are three types of gradient applied to an embedding ei during an epoch:
1. ei is pulled towards ej when edge ij is sampled (Alg.1 line 5) 2. ei is pulled towards ej when edge ji is sampled (Alg.1 line 6) 3. ei is pushed away from negative sample embeddings when some edge ij is sampled (Alg.1 line 9).
The full update of embedding ei during epoch t is, according to UMAP’s implementation, given by
gti = n∑ j=1 ( Xtij · ∂Laij ∂ei +Xtji · ∂Laji ∂ei +Xtij · n∑ s=1 Y tij,s · ∂Lris ∂ei ) , (7)
where Xtab is the binary random variable indicating whether edge ab was sampled in epoch t and Y tab,s is the random variable for the number of times s was sampled as negative sample for edge ab in epoch t if ab was sampled in epoch t and zero otherwise. By construction, E(Xtab) = µab and E(Y tab,s|Xtab = 1) = m/n. Taking the expectation over the random events in an epoch, we obtain the expected update of UMAP’s optimization procedure, see Appendix C for more details:
E ( gti ) = 2 n∑ j=1 ( µij · ∂Laij ∂ei + dim 2n · ∂Lrij ∂ei ) . (8)
Comparing the above closed formula for the expectation of the UMAP updates of the low-dimensional embeddings, to the gradient of UMAP’s loss function (6)
∂L ∂ei = 2 n∑ j=1 ( µij · ∂Laij ∂ei + (1− µij) · ∂Lrij ∂ei ) , (9)
we find that the sampling procedure yields the correct weight for the attractive term in expectation, as designed. However, as noticed by Böhm et al. [4], the negative sampling changes the weight for the
Fig. 14. We report quantitative metrics in Appendix G. Figure best viewed in color.
repulsive term significantly. Our closed formula helps to make their qualitative arguments precise: Instead of 1−µij , we have a term dim2n , which depends on the number of negative samplesm per sampled edge. Contrary to the intention of McInnes et al. [13], the repulsive weights are not uniform but vary with the degree of each point di, which, however, is typically close to log2(k), see Appendix B. More practically, since the non-zero high-dimensional similarities are sparse, 1− µij is equal to 1 for most ij. In contrast, the expected repulsive weight is small for large datasets as the numerator dim is of the order of log2(k)m independent of the dataset size n but the denominator scales with n.
Another effect of the negative sampling is that in general the expected update (8) does not correspond to any loss function, see Appendix D. We remedy this by additionally pushing the embedding of a negative sample i away from the embedding ej , whenever i was sampled as negative sample to some edge jk, see Alg.1 line 10. This yields the following gradient at epoch t
g̃ti = n∑ j=1 ( Xtij · ∂Laij ∂ei +Xtji · ∂Laji ∂ei +Xtij · n∑ s=1 Y tij,s · ∂Lris ∂ei + n∑ k=1 XtjkY t jk,i · ∂Lrji ∂ei ) , (10)
corresponding to a loss in epoch t of
L̃t = ∑
1≤i,j≤n
( Xtij · Laij + n∑ s=1 XtijY t ij,s · Lris ) . (11)
Using the symmetry of µij , Laij and Lrij in i and j, we compute the effective loss L̃ = E(L̃t) = 2 ∑
1≤i<j≤n
( µij · Laij + (di + dj)m
2n · Lrij
) . (12)
In fact, the above remedy of also pushing the negative samples does not affect the behavior of UMAP qualitatively, see for instance Supp. Figs. 21 and 15.3 In this light, we can treat L̃ as the effective objective function that is optimized via SGD by UMAP’s optimization procedure. It differs from UMAP’s loss function (6) by having a drastically reduced repulsive weight of (di+dj)m2n instead of 1− µij . We illustrate our analysis on gene expression measurements of 86024 cells of C. elegans [16, 14]. We start out with a 100 dimensional PCA of the data4 and use the cosine metric in high-dimensional space,
2Near the tip of the seam cells, the points seem to lie on a one-dimensional manifold. Closer to the middle of the seam cells the spread appears wider. But there are actually two dense parallel lines of points.
3In fact, the Parametric UMAP [17] does include the update of negative samples. 4 obtained from http://cb.csail.mit.edu/cb/densvis/datasets/. We informed the authors of our
use of the dataset, which they license under CC BY-NC 2.0.
consider k = 30 neighbors and optimize for 750 epochs, similar to [14]. The resulting visualization is depicted in Figure 2a. On this dataset the average value of 1 − µij is 0.9999 but the maximal effective repulsive weight maxij (di+dj)m 2n is 0.0043, showing the dramatic reduction of repulsion due to negative sampling. After each optimization epoch, we log our effective loss L̃ (12), the actual loss L̃t (11) of each epoch computed based on the sampled (negative) pairs as well the purported UMAP loss L (6).5 We always consider the embeddings at the end of each epoch. Note that UMAP’s implementation updates each embedding ei not just once at the end of the epoch but as soon as i is incident to a sampled edge or sampled as a negative sample. This difference does not change the actual loss much, see Appendix F. The loss curves are plotted in Figure 3. We can see that our predicted loss matches its actual counterpart nearly perfectly. While both, L̃ and L̃t, agree with the attractive part of the supposed UMAP loss, its repulsive part and thus the total loss are two orders of magnitude larger. Furthermore, driven by the repulsive part, the total intended UMAP loss increases during much of the optimization process, while the actual and effective losses decrease, exemplifying that UMAP really optimizes our effective loss L̃ (12) instead of its purported loss L (6). Additional loss curves for UMAP on other scRNA-seq datasets and CIFAR-10 [11] can be found in Appendix I.
Recently, a parametric version of UMAP was proposed [17]. Here, a neural network is trained to perform the mapping from high to low dimension. Its effective loss function differs slightly from the one for Non-Parametric UMAP as negative samples only come for the training batch:
Theorem 5.1. The effective loss function of Parametric UMAP with neural network fθ, batch size b and total similarity µtot = 12 ∑n i=1 di is
− 1 (m+ 1)µtot ∑ 1≤i<j≤n µij log ( φ ( fθ(xi), fθ(xj) )) +m b− 1 b didj 2µtot log ( 1− φ ( fθ(xi), fθ(xj) )) .
(13)
Proof. The proof can be found in Appendix A.
While Parametric UMAP’s loss differs slightly from L̃ (12) the same analysis holds unless explicitly mentioned. We mention a relation of Parametric UMAP to modularity clustering in Appendix A.1.
Algorithm 1: UMAP’s optimization input : input similarities µij ,
initial embeddings e1, . . . , en, number of epochs T, learning rate α, number of negative samples m
output :final embeddings e1, . . . , en 1 for t = 0 to T do 2 for ij ∈ {1, . . . , n}2 do 3 r ∼ Uniform(0, 1) 4 if r < µij then 5 ei = ei − α · ∂Laij ∂ei 6 ej = ej − α · ∂Laij ∂ej 7 for l = 1 to m do 8 s ∼ Uniform({1, . . . , n}) 9 ei = ei − α · ∂L r is
∂ei // Next line is omitted in UMAP’s implementation, but included for our analysis
10 /* es = es − α · ∂L r is
∂es */
Figure 3: Loss curves for the optimization leading to Figure 2a. Our effective loss (12) closely matches the actual loss (11), while the purported UMAP loss (6) is two orders of magnitude higher. The total overlays the repulsive purported loss. An average over 7 runs is plotted with shaded areas of one standard deviation, see also Appendix F.1. Figure best viewed in color.
5Our code is publicly available at https://github.com/hci-unihd/UMAPs-true-loss.

heading: 6 True target similarities
text: Since the effective objective function L̃ (12) that UMAP optimizes is different from L (6), we cannot hope that UMAP truly tries to find a low-dimensional embedding whose similarities reproduce the high-dimensional similarities. Nevertheless, using the effective loss L̃, we can compute the true target similarities ν∗ij which UMAP tries to achieve in embedding space. The effective loss L̃ is a sum of non-normalized binary cross entropy loss functions
− ( µij · log(νij) + (di + dj)m
2n · log(1− νij)
) (14)
which are minimal for
ν∗ij = µij
µij + (di+dj)m
2n
{ = 0 if µij = 0 ≈ 1 if µij > 0.
(15)
The approximation holds in the typical case in which (di+dj)m2n ≈ 0, discussed above. In other words, compared to the massively reduced repulsion weight even the smallest high-dimensional similarity appears maximal. Thus, the negative sampling essentially binarizes the input similarities. UMAP’s high-dimensional similarities are non-zero exactly on the shared k nearest neighbor graph edges of the high-dimensional data. Therefore, the binarization explains why Böhm et al. [4] find that using the binary weights of the shared k nearest neighbor graph does not deteriorate UMAP’s performance much.6 The binarization even helps UMAP to overcome disrupted high-dimensional similarities, as long as only the edges of the shared kNN graph have non-zero weight. In Figure 2b, we invert the original positive high-dimensional weights on the C. elegans dataset. That means that the k-th nearest neighbor will have higher weight than the nearest neighbor. The resulting visualization even improves on the original by keeping the layout more compact. This underpins Böhm et al. [4]’s claim that the elaborate theory used to compute the high-dimensional similarities is not the reason for UMAP’s practical success. In fact, we show that UMAP’s optimization scheme even actively ignores most information beyond the shared kNN graph. In Figure 4, we show histograms of the various notions of similarity for the C. elegans dataset. We see in panel 4b how the binarization equalizes the positive target similarities for the original and the inverted high-dimensional similarities.
UMAP’s tendency for over-contraction is already strong for small datasets. Consider UMAP’s default setting of k = 15 and m = 5. Then for datasets with n = 500 points each input similarity µij > 0.2 is mapped to a target similarity ν∗ij > 0.83, see Appendix H.4.
The binary cross entropy terms in the effective loss L̃ (12) are not normalized. This leads to a different weighing of the binary cross entropy terms for each pair ij
L̃ = 2 ∑
1≤i<j≤n
( µij · Laij + (di + dj)m
2n · Lrij
) (16)
= −2 ∑
1≤i<j≤n
(( µij + (di + dj)m
2n
) · ( ν∗ij log(νij) + (1− ν∗ij) log(1− νij) )) . (17)
As (di+dj)m2n is very small for large datasets, the term µij + (di+dj)m
2n is dominated by µij . Hence, the reduced repulsion not only binarizes the high-dimensional similarities, it also puts higher weight on the positive than the zero target similarities. Therefore, we can expect that the positive target similarities are better approximated by the embedding similarities than the zero ones. Indeed, panel 4a shows that the low-dimensional similarities match the positive target similarities very well, as expected from the weighted BCE reading of the effective loss function (17).

heading: 6.1 Explaining artifacts in UMAP visualizations
text: We conclude this section by explaining the observed artifacts of UMAP’s visualization in Figures 1 and 2 in the light of the above analysis. The normal UMAP optimization contracts the ring in Figure 1b even when initialized at the original layout because the reduced repulsion yields nearly binary target similarities. All pairs that are part of the kNN graph not only want to be sufficiently
6Böhm et al. [4] used a scaled version of the kNN graph, but the scaling factor cancels for the target weights.
close that their high-dimensional similarity is reproduced, but so close that their similarity is nearly one. The fact that the effective loss weighs the terms with target similarity near one much more than those with target similarity near zero reinforces this trend. As a result, the ring gets contracted. The same argument applies to the over-contracted parts of the UMAP visualization of the C. elegans dataset in Figure 2. When we increase the dataset size by adding more rings as in Figure 1c, negative sample pairs can come from different rings, reducing the frequency that a pair of points from the same ring is sampled. This reduces the repulsion further. Thus, there is stronger binarization and the embedding looks crisper. If we sample a single ring more densely, over-contraction is also increased. But the visual appearance is more nuanced as over-contraction can appear in the form of densely clustered and line-like substructures within the ring width, see Appendix H.3.
Our framework can also explain the opposite behavior of UMAP when the dense similarities are used as input similarities in Figure 1d. In this setting, the average degree of a node is about 100. With a negative sample ratem = 5 and a dataset size n = 1000 this yields repulsive weights (di+dj)m2n ≈ 0.5. Thus, we increase the repulsion on pairs with high input similarity, but decrease it on pairs with low input similarity. Now, the target similarities are not a binarization of the input similarities, but instead skewed towards 0.5. Hence, we can expect embedding points to increase their distance to nearest neighbors, but distant points to move closer towards each other. This is what we observe in Figure 1d, where the width of the ring has increased and the ring curves to bring distant points closer together.

heading: 7 Discussion

heading: 7.1 Negative sampling in LargeVis
text: The visualization method LargeVis [18] uses negative sampling to optimize its objective function, too. Thus, our analysis of the loss function and the target similarities applies to LargeVis and helps to explain why UMAP and LargeVis visualizations are often similar. For more details confer Appendix E.

heading: 7.2 Balancing attraction and repulsion
text: By deriving UMAP’s true loss function and target similarities, we are able to explain several peculiar properties of UMAP visualizations. According to our analysis, UMAP does not aim to reproduce the high-dimensional UMAP similarities in low dimension but rather the binary weights of the shared kNN graph of the input data. This raises the question just what part of UMAP’s optimization leads to its excellent visualization results. Apparently, the exact formula for the repulsive weights is not crucial as it differs for Non-Parametric UMAP and Parametric UMAP while both produce similarly high quality embeddings. A first tentative step towards an explanation might be the different weighing
of the BCE terms in the effective loss function (17). Focusing more on the similar rather than the dissimilar pairs might help to overcome the imbalance between an essentially linear number of attractive and a quadratic number of repulsive pairs. Inflated attraction was found beneficial for t-SNE as well, in the form of early exaggeration [12].
Put another way, the decreased repulsive weights result in comparable total attractive and repulsive weights, which might facilitate the SGD based optimization. Indeed, the total attractive weight in UMAP’s effective loss function is 2µtot = ∑n i,j=1 µij and the total repulsive weight amounts to
2mµtot = ∑n i,j=1 (di+dj)m 2n for Non-Parametric UMAP. For Parametric UMAP the total attractive and repulsive weights are 1/(m+1) andm/(m+1)·(b−1)/b, respectively. For the default value ofm = 5, the total attractive and repulsive weights are of roughly the same order of magnitude. Moreover, we observe in Figure 3 that the resulting attractive and repulsive losses are also of comparable size. Using UMAP’s purported loss function, however, would yield dominating repulsion.

heading: 7.3 Improved UMAP
text: The reason for the binarization of the input similarities is the reduced repulsion on the pairs of points that are linked in the shared kNN graph. As discussed above, it might not be desirable to increase the repulsion on all pairs of points to 1−µij . But it might be useful to change the optimization procedure of UMAP so as to decrease the repulsion per edge only on the non-kNN graph edges and not on the kNN graph edges. In addition to the positive and negative sampling of UMAP, we suggest to explicitly sample the shared kNN graph edges for repulsion with probability 1− µij − (di + dj)m/(2n). This would correct the repulsion weight for edges in the shared kNN graph. The loss function would read
−2 ∑
ij∈kNN graph
( µij log(νij)+(1−µij) log(1−νij) ) −2 ∑ ij /∈kNN graph ( (di + dj)m 2n log(1− νij) ) . (18)
Now the target similarities are precisely the high-dimensional similarities µij for all pairs i, j. Hence, in an optimal embedding, all µij’s are exactly reproduced and not just a binary version of them. We hope that this will improve the embedding quality and decrease the observation of over-contraction. Crucially, the amount of repulsion per data point is kept constant and is not increasing with n. For the quadratic number of non-kNN graph edges, the individual repulsion pre-factors still decrease with 1/n. The total loss function is a sum of normalized cross-entropy terms for the edges in the kNN graph and of non-normalized and drastically down-weighted cross-entropy terms for the edges not in the kNN graph. Since not all O(n2) negative pairs are considered explicitly and the additional sampling of the kNN graph edges for repulsion only scales with O(n), the complexity of computing the embedding would still scale linearly as O(kmn · numepochs) and not quadratically in n. A more in-depth investigation as to why exactly balanced attraction and repulsion is beneficial for a useful embedding and the implementation of the improved sampling scheme are left for future work.

heading: 8 Conclusion
text: In this work, we investigated UMAP’s optimization procedure in depth. In particular, we computed UMAP’s effective loss function analytically and found that it differs slightly between the nonparametric and parametric versions of UMAP and significantly from UMAP’s alleged loss function. The optimal solution of the effective loss function is typically a binarized version of the highdimensional similarities. This shows why the sophisticated form of the high-dimensional UMAP similarities does not add much benefit over the shared kNN graph. Instead, we conjecture that the resulting balance between attraction and repulsion is the main reason for UMAP’s great visualization capability. Our analysis can explain some artifacts of UMAP visualizations, in particular its tendency to produce over-contracted embeddings which gets stronger for larger datasets.
Acknowledgments and Disclosure of Funding
We would like to thank Dmitry Kobak, Quentin Garrido and the anonymous reviewers for useful discussions and helpful feedback. Moreover, we thank Ashwin Narayan for making the scRNA-seq datasets available.
This work is supported in part by the Deutsche Forschungsgemeinschft (DFG, German Research Foundation) – Projektnummer 240245660 - SFB 1129, the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany’s Excellence Strategy EXC 2181/1 - 390900948 (the Heidelberg STRUCTURES Excellence Cluster) and Informatics for Life funded by the Klaus Tschira Foundation.

Review summaries:
The paper investigated the UMAP's sampling-based optimization. The paper derived the effective objective function and found it to be different from the one proposed in the original paper. The difference is due to using SGD with non-uniformly-sampled positive and negative data points. The paper further derived the actual similarities which UMAP is trying to approximate. Those findings have several important implications: 1) UMAP does not try to reproduce the similarities in the original high dimensional space, 2) instead, it tries to approximate the similarities captured by the KNN graph, 3) points (1) and (2) are more so for large n (# of data points), small m (# of negative samples), and small k (# of neighbors on KNN graph).   -------------------------------- I thank the authors for providing a careful and thorough response to my review. The response has addressed most of my questions and led to new interesting ideas.  

This paper addresses the effects of the negative sampling strategy during UMAP’s optimization procedure. The authors provide a theoretical explanation of how this procedure distorts the objective function portrayed in UMAP’s original paper. The paper derives the actual objective function being optimized when using negative sampling, showing how it corresponds to lower repulsive forces that the ones intended in the cross-entropy loss function. Such behavior causes over contracted embeddings, which do not consider connectivity beyond the Knn neighbors.  

UMAP is a very popular algorithm for dimensionality reduction and data visualisation. This paper analyzes the the sampling based optimisation of UMAP and show that the loss function UMAP implicitly optimises is different from the loss function it tries to optimise (because of the stochastic optimisation algorithm used). The authors derive the qualitative differences between the two and using this show that what UMAP tries to do is approximate binarised similarities of those obtained from the high dimensional setting (similarity 1 if the high dimensional similarity computed is non-zero and 0 otherwise).