UMAP (Uniform Manifold Approximation and Projection) › Formula

To understand how UMAP achieves its goal of preserving both local and global structure, we need to explore the mathematical machinery that powers it. The algorithm builds on two key ideas from topology: fuzzy simplicial sets (a way to represent how points relate to each other with partial membership) and Riemannian geometry (which lets us think about distances on curved surfaces). Let's walk through the construction step by step, focusing on the intuition behind each piece.

Before diving into formulas, consider the problem UMAP solves. We have high-dimensional data and want to create a low-dimensional representation that preserves important relationships. Linear methods like PCA assume the data lies on a flat hyperplane, but real data often lives on curved manifolds. Imagine a Swiss roll: the data lies on a 2D surface embedded in 3D space. If we simply project it onto a plane, we lose the intrinsic structure. UMAP addresses this by modeling the data as lying on a manifold and preserving the manifold's topology.

The algorithm proceeds in two phases: first, build a rich representation of the high-dimensional structure; second, find a low-dimensional arrangement that matches this structure as closely as possible.

Step 1: Computing Local Distances

We start with the basics. For each point $x_i$ in our dataset, we measure how far it is from other points using Euclidean distance:

$$d_{ij} = \|x_i - x_j\|_2$$

where $d_{ij}$ is the distance between points $x_i$ and $x_j$, and $\|\cdot\|_2$ denotes the Euclidean (L2) norm.

This gives us a distance matrix. UMAP focuses on local neighborhoods because nearby points provide more reliable information about the manifold structure than distant ones.

For each point $x_i$, we identify its $k$ nearest neighbors. These neighbors define a local patch of the manifold around $x_i$. The parameter $k$ (called n_neighbors in practice) controls how many neighbors we consider. A smaller $k$ focuses on very local structure, while a larger $k$ captures more global patterns.

Step 2: Building Fuzzy Membership Strengths

Instead of using a binary neighborhood (where points are either neighbors or not), UMAP assigns a smooth membership strength to each potential neighbor. This membership strength tells us how strongly point $x_j$ belongs to point $x_i$'s local neighborhood.

The membership strength is computed as:

$$\mu_{ij} = \exp\left(-\frac{\max(0, d_{ij} - \rho_i)}{\sigma_i}\right)$$

where:

• $\mu_{ij}$: Membership strength indicating how strongly point $x_j$ belongs to the neighborhood of point $x_i$
• $\rho_i$: Distance from $x_i$ to its nearest neighbor (acts as a baseline)
• $\sigma_i$: Scaling factor specific to point $x_i$ that controls decay rate
• $d_{ij}$: Euclidean distance between points $x_i$ and $x_j$

Any point closer than the nearest neighbor (when $d_{ij} \leq \rho_i$) gets maximum membership strength of 1.0, since $\max(0, d_{ij} - \rho_i) = 0$ and $\exp(0) = 1$. Beyond the nearest neighbor, the membership strength decays exponentially at a rate controlled by $\sigma_i$.

Why use this exponential decay? Because it creates a smooth, continuous transition from "definitely a neighbor" to "definitely not a neighbor." This smoothness is important for the optimization that comes later.

Finding the Right Scale with $\sigma_i$

The parameter $\sigma_i$ is not chosen arbitrarily. For each point $x_i$, we pick $\sigma_i$ so that the effective number of neighbors equals $k$. Mathematically, we want:

$$\sum_{j \neq i} \mu_{ij} = \log_2(k)$$

where:

• $\sum_{j \neq i} \mu_{ij}$: Sum of all membership strengths for point $x_i$ (excluding itself)
• $\log_2(k)$: Target sum based on the number of neighbors $k$

Why $\log_2(k)$ instead of just $k$? This comes from information theory and ensures that the fuzzy simplicial set has the right "size" in terms of entropy. We find $\sigma_i$ using binary search: we try different values until the sum of membership strengths equals $\log_2(k)$.

This adaptive scaling accounts for local density. In dense regions of the data, $\sigma_i$ will be small (neighbors are close, so we don't need to scale distances much). In sparse regions, $\sigma_i$ will be large (neighbors are farther away, so we scale up to maintain connectivity).

Step 3: Creating a Global Graph

Each point $x_i$ now has its own local view of the neighborhood structure, encoded by the membership strengths $\mu_{ij}$. But these local views are asymmetric: point $x_i$ might consider $x_j$ a strong neighbor while $x_j$ considers $x_i$ a weak neighbor (or not at all).

To create a consistent global structure, we symmetrize these relationships using a fuzzy set union:

$$\mu_{ij}^{\text{global}} = \mu_{ij} + \mu_{ji} - \mu_{ij} \cdot \mu_{ji}$$

where:

• $\mu_{ij}^{\text{global}}$: Symmetrized global membership strength between points $i$ and $j$
• $\mu_{ij}$: Local membership strength from point $i$'s perspective
• $\mu_{ji}$: Local membership strength from point $j$'s perspective

This formula is the probabilistic OR operation for fuzzy sets. It is equivalent to:

$$\mu_{ij}^{\text{global}} = 1 - (1 - \mu_{ij})(1 - \mu_{ji})$$

In practice, UMAP often uses a simpler approximation:

$$\mu_{ij}^{\text{global}} \approx \max(\mu_{ij}, \mu_{ji})$$

This approximation says: if either point considers the other a neighbor, preserve that connection in the global graph. The result is a weighted graph where edge weights represent the strength of relationships in the original high-dimensional space.

Step 4: Defining Structure in Low Dimensions

Now we have a target: the fuzzy simplicial set $\mu_{ij}^{\text{global}}$ that describes relationships in high dimensions. We need to find low-dimensional coordinates $y_i$ that recreate these relationships as faithfully as possible.

In the low-dimensional space, we compute similar membership strengths. But instead of using exponential decay, UMAP uses a different functional form:

$$v_{ij} = \frac{1}{1 + a \cdot \|y_i - y_j\|_2^{2b}}$$

where:

• $v_{ij}$: Membership strength between points $i$ and $j$ in the low-dimensional embedding
• $y_i, y_j$: Low-dimensional coordinates of points $i$ and $j$
• $a, b$: Shape parameters that control the distribution (set based on min_dist)

This function is a generalized version of the Cauchy distribution (which t-SNE uses with $b = 1$).

Why a different function in low dimensions? The exponential decay in high dimensions reflects local connectivity on the manifold. In low dimensions, we need a function that allows both tight clusters (when points should be close) and clear separation (when points should be far apart). The form $\frac{1}{1 + a \cdot d^{2b}}$ achieves this balance, where $d = \|y_i - y_j\|_2$ is the Euclidean distance in the low-dimensional space.

The Role of min_dist

The min_dist parameter specifies the minimum distance between points in the low-dimensional embedding. Smaller values allow points to pack tightly (good for revealing fine structure), while larger values spread points out (good for seeing overall organization). The parameters $a$ and $b$ are set to match this desired minimum distance.

Step 5: Optimization via Cross-Entropy

We now have two sets of membership strengths: $\mu_{ij}^{\text{global}}$ (the target from high dimensions) and $v_{ij}$ (the current state in low dimensions). We want to adjust the low-dimensional coordinates $y_i$ to make $v_{ij}$ match $\mu_{ij}^{\text{global}}$ as closely as possible.

UMAP uses cross-entropy as the objective function:

$$C = \sum_{i,j} \mu_{ij}^{\text{global}} \log\left(\frac{\mu_{ij}^{\text{global}}}{v_{ij}}\right) + (1 - \mu_{ij}^{\text{global}}) \log\left(\frac{1 - \mu_{ij}^{\text{global}}}{1 - v_{ij}}\right)$$

where:

• $C$: Cross-entropy loss to be minimized
• $\mu_{ij}^{\text{global}}$: Target membership strength from high-dimensional space
• $v_{ij}$: Current membership strength in low-dimensional space

Cross-entropy measures the difference between two probability distributions. Here we treat $\mu_{ij}^{\text{global}}$ and $v_{ij}$ as probabilities that points $i$ and $j$ are connected.

The first term, $\mu_{ij}^{\text{global}} \log(\mu_{ij}^{\text{global}} / v_{ij})$, penalizes cases where points should be close (high $\mu_{ij}^{\text{global}}$) but are far apart (low $v_{ij}$). The second term, $(1 - \mu_{ij}^{\text{global}}) \log((1 - \mu_{ij}^{\text{global}}) / (1 - v_{ij}))$, penalizes cases where points should be far apart but are close together.

Minimizing this cross-entropy pulls the low-dimensional structure toward the high-dimensional structure. We use stochastic gradient descent (SGD) to optimize the coordinates $y_i$. At each iteration, we sample pairs of points, compute how the cross-entropy would change if we moved them slightly, and update coordinates to reduce the error.

Let's step back and see the full picture:

1. Local structure preservation: By using nearest neighbors and adaptive scaling ($\sigma_i$), UMAP ensures that local neighborhoods are well-represented.

2. Global structure preservation: By connecting local views into a global fuzzy simplicial set, UMAP maintains large-scale relationships.

3. Flexibility: The low-dimensional membership function and min_dist parameter let us control the final embedding's appearance.

4. Mathematical foundation: The use of fuzzy sets and cross-entropy provides a principled objective function that can be efficiently optimized.

The construction above gives UMAP several important properties:

• Topological preservation: By building a fuzzy simplicial set, UMAP preserves topological features such as holes and loops, along with connected components. If two clusters are separate in high dimensions, they remain separate in low dimensions.

• Metric structure preservation: UMAP adapts to the local density of the data. In dense regions, it preserves fine details. In sparse regions, it maintains connectivity without distorting distances too much.

• Parameter interpretability: Each parameter has a clear geometric meaning. The n_neighbors parameter controls the size of local neighborhoods (larger values capture more global structure). The min_dist parameter controls how tightly points can cluster in the embedding (smaller values reveal finer details).

• Computational efficiency: By focusing on local neighborhoods and using efficient optimization, UMAP scales to large datasets. The nearest neighbor search is the main computational bottleneck, but this can be accelerated with approximate methods.