🤖 machine learning

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

This paper provides a mechanistic analysis of the attractive and repulsive forces in UMAP, explaining how they shape cluster formations and proposing a modification to attraction to improve embedding consistency under random initialization.

Original authors: Mohammad Tariqul Islam, Jason W. Fleischer

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mohammad Tariqul Islam, Jason W. Fleischer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a cosmic architect tasked with taking a massive, chaotic galaxy of stars (your high-dimensional data) and flattening it into a single, beautiful 2D map that humans can actually understand.

To do this, you use a popular tool called UMAP. This paper, written by researchers from MIT and Princeton, acts like a "microscope" for that tool. It looks under the hood to see the invisible "forces" that move the stars around on your map.

Here is the breakdown of their discovery using everyday analogies.

1. The Two Invisible Forces: Magnets and Bumper Cars

The researchers found that UMAP works by using two competing forces to arrange the data:

Attraction (The Magnets): Imagine every star that belongs to the same "family" (a cluster) has a magnet pulling it toward its siblings. This force tries to pull similar things together into tight groups.
Repulsion (The Bumper Cars): Imagine every star that is not in your family is a bumper car. If a stranger gets too close, they hit you and push you away. This force ensures that different groups don't just mush into one giant, unrecognizable blob.

2. The "Glitch" in the Magnets (The Big Discovery)

Usually, we think of magnets as simple: they pull things closer. But the researchers discovered something counterintuitive about UMAP’s magnets.

The Analogy: The "Springy" Magnet.
Imagine a magnet that is actually a very strange, springy coil. If two stars are a medium distance apart, the magnet pulls them together. But if they get too close—closer than a certain "magic distance"—the magnet suddenly flips! Instead of pulling, it starts pushing them apart.

This is why UMAP maps can sometimes look "fuzzy" or "blurry." The stars are caught in a tug-of-war, oscillating back and forth because the very force meant to bring them together is accidentally pushing them away when they get too cozy.

3. Why do we need a "Learning Rate"? (The Cooling Metal)

The paper explains why UMAP uses something called "learning rate annealing" (gradually slowing down the movement).

The Analogy: Forging a Sword.
Imagine you are forging a sword. At first, the metal is glowing red and liquid; you move it around violently to get the general shape right. But if you keep hitting it that hard when it’s almost finished, you’ll shatter it or make it crooked.

In UMAP, the "learning rate" is like the cooling process. You start with big, wild movements to get the clusters in the right place. Then, you slowly "cool down" the movement so the stars can settle into their perfect, sharp positions without the "springy magnet" glitch causing them to bounce around.

4. Fixing the "Near-Sightedness"

The researchers noticed that UMAP is a bit "near-sighted." It’s great at seeing its immediate neighbors, but if two stars that should be together are placed on opposite sides of the map, the magnets are too weak to pull them across the vast distance.

The Analogy: The Local Gossip vs. The Global News.
UMAP is like a person who knows everything about their neighbors but has no idea what’s happening in the next city. To fix this, the researchers experimented with "far-sighted" attraction—giving the magnets a bit more strength over long distances so the "global" structure of the map stays accurate.

Summary: The "Mechanic's Manual"

Before this paper, people used UMAP like a driver uses a car: they knew how to steer, but they didn't really understand how the engine worked.

This paper is the Mechanic’s Manual. It tells us:

If your clusters are fuzzy: Your "magnets" are too springy; you need to slow down the learning rate.
If your clusters are messy: Your "bumper cars" aren't pushing hard enough.
If your map is disorganized: Your "magnets" are too near-sighted.

By understanding these invisible forces, scientists can now create much clearer, more reliable maps of complex data, from human genetics to galaxy formations.

Technical Summary: The Shape of Attraction in UMAP

Title: The Shape of Attraction in UMAP: Exploring the Embedding Forces in Dimensionality Reduction
Authors: Mohammad Tariqul Islam (MIT Media Lab) & Jason W. Fleischer (Princeton University)
Publication: Transactions on Machine Learning Research (04/2026)

1. Problem Statement

Dimensionality Reduction (DR) methods like UMAP and t-SNE are essential for visualizing high-dimensional data by preserving local and global structures through attraction-repulsion dynamics. However, these algorithms are often treated as "black boxes" governed by heuristic practices (e.g., specific learning rate schedules, initialization, and kernel tuning).

The authors identify a critical gap in understanding: while practitioners know that these forces drive cluster formation, the functional shapes of these forces—and how they interact with hyperparameters like the learning rate—are not well understood. This lack of mechanistic insight makes it difficult to interpret embeddings, avoid spurious structures, and optimize performance.

2. Methodology

The researchers propose a novel analytical framework by decomposing the gradient updates of neighbor embedding algorithms into two independent scalar functions:

Attraction Shape ( $f_a$ ): A function that determines how the distance between similar points changes.
Repulsion Shape ( $f_r$ ): A function that determines how the distance between dissimilar points changes.

Key Analytical Tools:

Decomposition: They reformulate the update equations to isolate the "shape" of the force from its magnitude.
Contraction/Expansion Conditions: They establish mathematical propositions to determine when a force will cause points to move closer (contraction) or further apart (expansion). Specifically, they show that for attraction, contraction only occurs if the effective scaled shape $\lambda f_a$ falls within the range $(-1, 0)$ .
Comparative Analysis: They apply this decomposition to UMAP, Parametric UMAP, NEG-t-SNE, PaCMAP, and TriMap to compare their stability and clustering capabilities.
Interventional Experiments: They modify the attraction/repulsion shapes (e.g., adding a "far-sightedness" term or adjusting the $b$ parameter) to test their impact on cluster formation and consistency under random initialization.

3. Key Contributions

Mechanistic Understanding of UMAP: They reveal that UMAP’s default attraction shape is "unbounded" as distance approaches zero. This leads to a counterintuitive phenomenon where, at certain distances, attractive forces actually cause expansion rather than contraction.
Explanation of Learning Rate Annealing: The paper provides a mathematical justification for UMAP’s learning rate annealing. Annealing is necessary to drive the system into a "contractive regime" where $\lambda f_a > -1$ , thereby eliminating "fuzzy" clusters caused by distance oscillations.
Identification of "Near-Sightedness": They show that UMAP’s attraction force diminishes too quickly at large distances, making the algorithm sensitive to random initialization.
Algorithmic Comparison: They demonstrate that newer methods like PaCMAP and NEG-t-SNE are inherently more stable because their attraction shapes are naturally bounded within the contractive range $[-1, 0]$ , making them less dependent on complex learning rate schedules.
Proposed Improvements: They introduce a "modified attraction shape" that increases attraction at long distances, significantly improving the consistency of embeddings when using random initialization.

4. Results

Cluster Formation: The study proves that attraction governs cluster formation and stability, while repulsion regulates inter-cluster distance and compactness.
Stability & Consistency: Using Procrustes analysis, they showed that default UMAP is highly inconsistent under random initialization. However, by modifying the attraction shape to be "far-sighted," they achieved embeddings that are highly consistent with PCA-initialized baselines.
Empirical Validation: Experiments on MNIST, Fashion-MNIST, and single-cell transcriptomes confirmed that:
- Reducing the learning rate $\lambda$ reduces the "minimum distance for contraction" ( $\zeta^{-1}$ ), leading to sharper boundaries.
- Increasing repulsion (via parameter $b$ ) increases inter-cluster spacing but does not create new internal cluster structures.
- Strengthening attraction (via parameter $b$ ) reveals finer sub-structures within existing clusters.

5. Significance

This paper shifts the study of dimensionality reduction from heuristic optimization to functional analysis. By providing a "mechanistic language" (attraction and repulsion shapes), it allows researchers to:

Diagnose Issues: If an embedding is fuzzy or unstable, practitioners now know to look at the attraction shape rather than just the repulsion.
Design Better Algorithms: It provides a blueprint for designing DR methods that are robust to initialization and do not require heavy hyperparameter tuning.
Bridge Fields: The work connects dimensionality reduction to contrastive learning, suggesting that the principles of attraction-repulsion can be used to enhance general representation learning.