Manifold Learning: t-SNE & UMAP Mathematics

Linear dimensionality reduction techniques like PCA maximize global variance along orthogonal projection axes. When high-dimensional data resides on non-linear topological manifolds (such as Swiss rolls, spirals, or multi-cluster embeddings), linear methods flatten distinct geodesic folds onto each other.

Non-Linear Manifold Learning preserves local neighborhood topology, allowing complex geometries to unfold into lower dimensions.

1. High-Dimensional Similarities & The Crowding Problem#

In t-Distributed Stochastic Neighbor Embedding (t-SNE) (van der Maaten and Hinton, 2008), pairwise Euclidean distances are converted into conditional probabilities representing neighborhood affinities.

In high-dimensional space $\mathbb{R}^D$, the probability that sample $x_i$ picks $x_j$ as its neighbor under a Gaussian centered at $x_i$ with variance $\sigma_i^2$ is:

To ensure symmetry and robust handling of outliers, affinities are symmetrized:

The Crowding Problem in Low Dimensions#

In high-dimensional space $\mathbb{R}^D$, the volume of a hypersphere of radius $r$ scales as $r^D$. In low-dimensional projection space $\mathbb{R}^2$, volume scales only as $r^2$.

In $D = 100$, an enormous amount of volume is available at moderate Euclidean distances. When mapped to 2D, the available area is drastically smaller. If we use a Gaussian distribution in low dimensions, points that are moderately distant in $\mathbb{R}^D$ must be pushed far apart in $\mathbb{R}^2$ to match small probabilities, crowding the center and causing distinct clusters to overlap.

2. Derivation of the Student-t Heavy-Tailed Kernel#

To solve the crowding problem, t-SNE models low-dimensional affinities using a Student-t distribution with 1 degree of freedom (Cauchy distribution):

Because the Cauchy distribution features heavy polynomial tails ($1/r^2$) rather than exponential Gaussian tails ($e^{-r^2}$), low-dimensional points must separate by significantly larger distances to reproduce small probabilities, naturally segregating clusters into visually distinct islands.

3. Analytical Gradient Derivation of t-SNE#

t-SNE minimizes the Kullback-Leibler (KL) Divergence between high-dimensional distribution $P$ and low-dimensional distribution $Q$:

Let $d_{ij} = \|y_i - y_j\|_2$ and define the unnormalized Cauchy kernel $w_{ij} = (1 + \|y_i - y_j\|_2^2)^{-1}$ with partition function $Z = \sum_{k \ne l} w_{kl}$, so that $q_{ij} = \frac{w_{ij}}{Z}$.

4. UMAP: Topological Foundations#

Uniform Manifold Approximation and Projection (UMAP) (McInnes et al., 2018) provides faster computation and superior preservation of global structure compared to t-SNE.

Rather than minimizing KL-divergence, UMAP models data as a Riemannian manifold using fuzzy simplicial sets and minimizes fuzzy set cross-entropy:

The second term $(1 - p_{ij}) \ln\left( \frac{1 - p_{ij}}{1 - q_{ij}} \right)$ explicitly penalizes placing dissimilar points close together in low dimensions, maintaining large-scale global cluster relationships that t-SNE frequently shuffles.

5. From-Scratch Implementation: Zero Dependencies#

The following implementations compute a minimal t-SNE embedding with gradient descent from scratch in Python and Rust.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import math
import random
from typing import List

class TSNEFromScratch:
    def __init__(self, n_components: int = 2, learning_rate: float = 10.0, n_iter: int = 300):
        self.n_components = n_components
        self.learning_rate = learning_rate
        self.n_iter = n_iter

    @staticmethod
    def _dist_sq(a: List[float], b: List[float]) -> float:
        return sum((x - y) ** 2 for x, y in zip(a, b))

    def fit_transform(self, X: List[List[float]]) -> List[List[float]]:
        num_samples = len(X)
        random.seed(42)

        # 1. Compute high-dimensional affinities P with fixed bandwidth sigma=1.0
        P = [[0.0] * num_samples for _ in range(num_samples)]
        for i in range(num_samples):
            for j in range(num_samples):
                if i != j:
                    P[i][j] = math.exp(-self._dist_sq(X[i], X[j]) / 2.0)

        # Symmetrize and normalize P
        total_p = sum(sum(row) for row in P)
        for i in range(num_samples):
            for j in range(num_samples):
                P[i][j] = max(P[i][j] / total_p, 1e-12)

        # 2. Initialize low-D coordinates Y with small Gaussian noise
        Y = [[random.gauss(0.0, 1e-4) for _ in range(self.n_components)] for _ in range(num_samples)]

        # 3. Gradient Descent Loop
        for _ in range(self.n_iter):
            # Compute low-D Cauchy kernel affinities Q
            Q = [[0.0] * num_samples for _ in range(num_samples)]
            inv_dist = [[0.0] * num_samples for _ in range(num_samples)]
            total_q = 0.0

            for i in range(num_samples):
                for j in range(i + 1, num_samples):
                    val = 1.0 / (1.0 + self._dist_sq(Y[i], Y[j]))
                    inv_dist[i][j] = val
                    inv_dist[j][i] = val
                    total_q += 2.0 * val

            for i in range(num_samples):
                for j in range(num_samples):
                    if i != j:
                        Q[i][j] = max(inv_dist[i][j] / total_q, 1e-12)

            # Compute gradients: dL/dy_i = 4 * sum (p_ij - q_ij) * (y_i - y_j) * inv_dist
            dY = [[0.0] * self.n_components for _ in range(num_samples)]
            for i in range(num_samples):
                for j in range(num_samples):
                    if i != j:
                        mult = 4.0 * (P[i][j] - Q[i][j]) * inv_dist[i][j]
                        for d in range(self.n_components):
                            dY[i][d] += mult * (Y[i][d] - Y[j][d])

            # Update coordinates
            for i in range(num_samples):
                for d in range(self.n_components):
                    Y[i][d] -= self.learning_rate * dY[i][d]

        return Y

if __name__ == "__main__":
    X = [
        [1.0, 1.2], [1.1, 1.0], [0.9, 1.1],
        [8.0, 8.2], [8.1, 8.0], [7.9, 8.1]
    ]

    tsne = TSNEFromScratch(n_components=2, learning_rate=50.0, n_iter=150)
    embedding = tsne.fit_transform(X)
    print("t-SNE 2D Coordinates:")
    for pt in embedding:
        print(f"  [{pt[0]:.4f}, {pt[1]:.4f}]")

6. Active Knowledge Checks#