UMAP: Visualizing High-Dimensional Data's Shape

Why it exists

High-dimensional data is impossible to visualize directly. We need a way to 'flatten' it to 2D or 3D to see its shape, find clusters, and build intuition, without losing the essential relationships between data points. Simple linear methods like PCA can miss complex, non-linear structures that look like tangled curves and surfaces.

The mental model

Imagine your data points are stars in a 3D galaxy. UMAP is like creating a 2D star chart. It tries to place the stars on the paper so that stars that were close neighbors in the galaxy are also close neighbors on the chart. It focuses on preserving these local constellations, even if it has to distort the large-scale distances between constellations to make everything fit on the page.

How it works

UMAP is founded on principles from Riemannian geometry and topology. First, it builds a graph in the original high-dimensional space where edges connect each point to its nearest neighbors. It assumes the data lies on a 'manifold'—a locally smooth surface. Second, it initializes points randomly in a low-dimensional space (e.g., 2D). Finally, it uses an optimization algorithm to move these low-dimensional points around, trying to make their neighborhood graph as structurally similar as possible to the original high-dimensional graph. This process effectively finds a low-dimensional projection that preserves the 'fuzzy topological structure' of the source data.

When to use it

UMAP is a primary tool for exploratory data analysis and visualization. Use it to get a quick, intuitive feel for the structure of high-dimensional datasets, such as identifying potential customer segments, visualizing document similarities, or exploring gene expression patterns. Its high performance makes it a strong choice for large datasets, often faster and better at preserving global structure than its predecessor, t-SNE.

When not to use it

Do not use UMAP if you need to make precise claims about the data's geometry. The size of a cluster, its apparent density, and the absolute distance between two separate clusters in a UMAP plot are often misleading artifacts of the projection. For tasks requiring strict distance preservation, other methods may be more suitable. Also, be mindful that without tuning, it can be memory and CPU intensive on very large datasets.

One canonical example

A classic demonstration is visualizing the MNIST dataset of handwritten digits. Each image is a 784-dimensional vector (28x28 pixels). After running UMAP to project this data down to 2D, the output plot clearly shows ten distinct clusters, one for each digit from 0 to 9. The plot also reveals structural similarities; for instance, the cluster for '4' is often near the cluster for '9', reflecting their visual resemblance, while being far from the cluster for '1'.