The Shape of Attraction in UMAP: Exploring the Embedding Forces in Dimensionality Reduction

TL;DR

The paper addresses how attraction and repulsion forces govern UMAP embeddings by decomposing low-dimensional updates into shapes $f_a(\zeta)$ and $f_r(\zeta)$ with $\zeta=||y_i-y_j||_2$. It develops analytic results (including contraction/expansion conditions) and derives UMAP-specific shapes, showing that attraction can cause both contraction and expansion and that learning-rate annealing helps drive clusters to concise boundaries, while repulsion modulates inter-cluster distances. It demonstrates that modifying the attraction shape improves consistency under random initialization (via Procrustes analysis on MNIST and other datasets) and analyzes how attraction and repulsion shapes interact to form clusters, comparing UMAP with NEG-$t$-SNE, PaCMAP, TriMap, and LocalMAP. The findings offer a principled lens to interpret and improve DR algorithms, suggesting that shape-mixing and far-distance attraction can enhance robustness and accuracy, with implications for contrastive and representation-learning contexts.

Abstract

Uniform manifold approximation and projection (UMAP) is among the most popular neighbor embedding methods. The method relies on attractive and repulsive forces among high-dimensional data points to obtain a low-dimensional embedding. In this paper, we analyze the forces to reveal their effects on cluster formations and visualization and compare UMAP to its contemporaries. Repulsion emphasizes differences, controlling cluster boundaries and inter-cluster distance. Attraction is more subtle, as attractive tension between points can manifest simultaneously as attraction and repulsion in the lower-dimensional mapping. This explains the need for learning rate annealing and motivates the different treatments between attractive and repulsive terms. Moreover, by modifying attraction, we improve the consistency of cluster formation under random initialization. Overall, our analysis makes UMAP and similar embedding methods more interpretable, more robust, and more accurate.

Figures (31)

• Figure 1: Attraction and repulsion shapes in UMAP. (a) Effect of different values of $f_a$ (top) and $f_r$ (bottom) on a pair. (b) Attraction shape of UMAP. (c) Effective attraction shape $(\lambda f_a)$ for various learning rates $\lambda$. (d) Minimum distance for contraction ($\zeta_{-1}$) as $\lambda$ decreases. (e) Repulsion shape of UMAP. (f) Attraction and (g) repulsion shapes of various embedding methods. (h) Trustworthiness and (i) Silhouette Score of various methods as the constant learning rate is varied for attraction ($\lambda_a$) and repulsion ($\lambda_r$) independently. Default UMAP parameters: $a=1.58$ and $b=0.89$.
• Figure 2: Effect of random initialization on different attraction shapes for the MNIST dataset. (a) Mapping using PCA. (b-d) Four mappings with the lowest Procrustes distance ($p_d$) from the embedding in (a) for (b) UMAP, (c) modified, and (d) composite attraction shapes. (e) Default UMAP and modified attraction shapes. (f-h) Procrustes matrices from 100 runs of (f) UMAP ($0.78\pm0.13$), (g) modified ($0.49\pm0.21)$, and (h) composite attraction ($0.50\pm0.20$) shapes. The diagonal $(i,i)$ entries of the Procrustes matrix are sorted by Procrustes distance ($p_d$) from (a), and the off-diagonal, $(i,j)$, entries correspond to $p_d$ between $i^{th}$ and $j^{th}$ mapping. The matrices and (mean $p_d\pm$std) values show that UMAP's embeddings are not self-similar, while the modified and composite attraction shapes encourage scale-invariant structure.
• Figure 3: Control of inter-cluster distances on the MNIST dataset. (a) Computing $a,b$ by varying the low-dimensional distance $m_d$ restricts exploration. (b-d) UMAP output by setting $m_d$ to $0.1$, $0.01$, and $0.0001$, respectively, shows little improvement in compactness of clusters. (e) Repulsion shapes for different parameters. (f-h) Increasing repulsion by explicitly varying $b$ results in more compact clusters and forms new ones that were absent otherwise. (i) Attraction shapes by varying parameters. (j-l) Increasing repulsion by adding a small positive value ($\varepsilon$) to the repulsion shape increases inter-cluster distance.
• Figure 4: Embedding of the MNIST dataset with (a) default attraction shape but repulsion shape with $b=0.4$ and (b) default repulsion shape but attraction shape with $b=0.4$. The former shows the same clusters of default UMAP with increased compactness. The latter develops new structures within clusters and forms new clusters.
• Figure 5: Values of $g(\zeta,b,a)$ for (a) $a$ fixed at $1.576$ and (b) $b$ fixed at $0.89$.

Theorems & Definitions (12)

• Proposition 4.1
• Proposition 4.2
• Proposition 6.1
• proof
• proof
• proof
• Proposition G.1
• proof
• Proposition G.2
• proof