Title: The Shape of Attraction in UMAP: Exploring the Embedding Forces in Dimensionality Reduction
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: unsupervised learning, neighbor embedding, umap, t-sne, representation learning

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.

TL;DR: This work provides a fresh perspective on non-linear dimensionality reduction by forming attraction-repulsion shapes from attractive and repulsive forces (i.e., gradients), and demystifying many heuristics.