On UMAP's True Loss Function

1 Introduction . 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 . 2 Related work . 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 . 3 Background on UMAP . 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 .

Summary: 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.