DMAP

§1 Not dimensionality reduction

The companion paper states this plainly: diffusion maps is not dimensionality reduction in the sense of PCA or UMAP. Those methods find a low-dimensional representation that preserves some notion of distance or topology in the ambient space $\mathbb{R}^D$. DMAP does something categorically different — it recovers the intrinsic geometry of the manifold the data lives on.

High-dimensional scientific data — MD trajectories, cryo-EM particle stacks, XFEL diffraction — does not fill $\mathbb{R}^D$ uniformly. It concentrates near a low-dimensional manifold carved by the underlying physics. A protein does not explore all of $\mathbb{R}^{3N}$; it explores a conformational landscape shaped by energy and entropy. The diffusion coordinates are coordinates on that landscape — not a projection of the data, but a discovery of the geometry hidden within it.

§2 The kernel $P = e^{-\beta H}$

The diffusion kernel in DMAP is the quantum Boltzmann operator. Given $N$ data points $\{x_i\}_{i=1}^N \subset \mathbb{R}^D$, we define the Hamiltonian as the simple harmonic oscillator potential:

where $D^2_{ij} = \|x_i - x_j\|^2$ is the squared pairwise distance matrix. The kernel is then the matrix exponential of $-\beta H$:

Entry-wise, this gives the familiar Gaussian kernel $P_{ij} = \exp\!\bigl(-\tfrac{\beta}{2}\|x_i - x_j\|^2\bigr)$, but the derivation from $H = \tfrac{1}{2}D^2$ is not merely notational. It places DMAP firmly within the framework of quantum statistical mechanics: $\beta$ is the inverse temperature, $P$ is the thermal density matrix, and the diffusion coordinates are the low-energy eigenstates of $H$.

Computing $P$ requires all $N^2$ pairwise squared distances — an $O(N^2 D)$ operation. This is done exactly, without approximation. No kNN graph, no pruning, no sparsification. Every pair of points contributes.

§3 Normalization and the operator $P^+$

The raw kernel $P$ conflates geometry with sampling density. Coifman–Lafon normalization removes this bias, producing the row-stochastic diffusion operator $P^+$:

$P^+$ is the transition matrix of a Markov random walk whose stationary distribution reflects the geometry of the manifold, independent of how densely the data was sampled in any region. The normalization parameter $\alpha \in [0,1]$ controls how aggressively density is corrected:

§4 Diffusion coordinates

The diffusion coordinates at time $t$ are the leading eigenvectors of $P^+$, weighted by their eigenvalues raised to power $t$:

where $1 = \lambda_0 > \lambda_1 \geq \lambda_2 \geq \ldots$ are the eigenvalues of $P^+$ and $\psi_j$ the corresponding eigenvectors. The fundamental theorem of diffusion maps: the Euclidean distance in this embedding equals the diffusion distance between points — a geometry-aware, noise-robust intrinsic distance on the manifold.

Choice of $t$

Large $t$ suppresses eigenvectors with $\lambda_j^t \approx 0$, revealing only the slowest, most global modes of the geometry. For MD trajectories $t$ has a physical interpretation as a lag time: fast local transitions stay close, slow conformational changes separate. In the quantum analogy, $t$ plays the role of imaginary time — the path integral representation of the thermal density matrix $e^{-\beta t H}$.

§5 Lanczos eigensolver

Extracting the top $m$ eigenvectors of $P^+$ does not require materializing or factorizing the full $N \times N$ matrix. The Lanczos algorithm only requires matrix-vector products $v \mapsto P^+ v$, each costing $O(N^2 D)$ — one pass over the kernel.

This is the direct analogy to how FlashAttention computes attention: instead of neural network layers, the operator is applied iteratively via Lanczos to extract the geometric eigenmodes. The operator is never stored — only applied.

§6 FlashDiffusion kernel

The Lanczos matvec $v \mapsto P^+ v$ requires computing $P_{ij} = \exp\!\bigl(-\tfrac{\beta}{2}\|x_i - x_j\|^2\bigr)$ for all $i, j$ at each iteration — $O(N^2 D)$ floating point operations. On naive hardware this is memory-bandwidth limited: the $N \times N$ kernel matrix cannot fit in SRAM and must be streamed from HBM, making it intractable beyond $N \sim 10^4$.

FlashDiffusion solves this exactly as FlashAttention solves attention: by tiling. The kernel computation is blocked into tiles that fit in the SRAM of SM100/SM103 streaming multiprocessors on Blackwell GPUs. Each tile computes a block of $P_{ij}$ values, immediately accumulates them into the output vector $P^+ v$, and discards them — the kernel is never written to HBM.

The result: 1M frames $\times$ 64k features computed in approximately 30 minutes on a B300 — a dataset that would require weeks on CPU or be entirely intractable with naive GPU implementations. The geometry is exact. No approximation is made to the kernel or the walk operator.

§7 DMAP vs UMAP vs PCA