UMAP 0.6 races past 0.5: 3 million arXiv papers embedded in under 5 minutes

UMAP, the dimensionality-reduction library behind countless cluster plots in machine learning, just shipped its first major upgrade in years. Version 0.6, announced on November 12 by the library's creator Leland McInnes on r/MachineLearning, rebuilds the initialization and the layout optimization — and on a benchmark of 3 million arXiv papers it finishes in 4 minutes 27 seconds, where the old version needed almost 52. McInnes filmed the two versions racing, and the clip sits at the top of this article.

The race, in 44 seconds#

The clip runs 44 seconds, sped up 24×. On the left, UMAP 0.5 — unseeded, running on all cores — is still hunting for nearest neighbours when the right-hand panel, UMAP 0.6, crosses the line. The 0.6 run is seeded (random_state=42), uses every core, and still finishes first: reproducible and faster at once, a combination 0.5 could never offer. Both panels draw the same 3-million-paper arXiv map, coloured by field — physics, computer science, mathematics, engineering, statistics, biology, finance, economics and the rest — so you can watch the structure form as the clock runs.

Figures from the timing chart in the announcement post. “Reproducible” means a fixed seed: the same input gives the same embedding on every run.

What's new under the hood#

Three changes do the heavy lifting. First, the starting point. UMAP used to begin from a spectral or PCA layout — a blurry first guess the optimizer then had to drag into shape. Version 0.6 can start recursively instead: it lays out a coarse version of the embedding using UMAP itself, then refines it in a tower of steps, arriving at a starting layout far closer to the finished one — and it gets there faster than the old initializations. In the release's own measurements, the optimizer settles to within 3% of the final layout in far fewer epochs.

Second, the optimizers. The release adds Adam and momentum-based layout optimizers alongside the original, plus a much more efficient form of hard-negative mining that makes negative sampling pay off on very large datasets. It also handles high repulsion strengths far better. UMAP was originally tuned for a narrow slice of the attraction–repulsion spectrum, and the new optimizers open up the rest of it — which shows up as better k-nearest-neighbour recall.

Third, raw neighbour search is about four times faster, from better use of PyNNDescent, the approximate-nearest-neighbour engine underneath. The payoff is not only speed: the release resolves finer cluster structure in fewer epochs, with better neighbourhood accuracy than 0.5 despite running faster.

The catch: it's opt-in, for now#

None of this is the default yet. The new initialization and optimizers change the output, so McInnes shipped them behind a flag: set compatibility_layout=False. Leave it unset and 0.6 behaves like 0.5, with a warning that the default will flip in a future release. That caution is the point — embeddings feed downstream models and published figures, and silently changing them would be worse than asking users to opt in.

One trade-off is worth knowing before you switch. A seeded 0.5 run was single-threaded, which is why it took 52 minutes; the new optimizers let seeded runs use every core. And the faster neighbour search is a trade: McInnes notes it can be slightly less accurate on the k-nearest-neighbour graph, a price he judges worth paying for the speed. If you supply your own connectivity map and initialization, don't expect much of a speed gain — the win there is output quality, especially at higher repulsion strengths, and on large datasets, where hard-negative mining pays real dividends.

What the community made of it#

The announcement, posted by McInnes himself in r/MachineLearning, drew 157 upvotes and a 98% upvote rate, and the early field reports are encouraging: one user wrote that they tested 0.6 on their current project the same day and saw a 500% speedup. Another asked how McInnes keeps finding ways to make approximate nearest-neighbour search faster; his answer was unglamorous — small changes that improve data locality and parallelism, plus better-tuned trade-offs. Nobody in the top comments disputed the numbers; the only corrections came from McInnes himself, in the caveats above.

Four more clips from the release#

McInnes posted five videos with the release. The race above is the headline, and the other four each isolate one piece of the new pipeline:

• Recursive initialization for UMAP (0:32) — the coarse-to-fine starting layout, up close.
• Racing different initialization methods for UMAP — PCA, spectral and recursive initialization race to settle within 3% of the final layout over a 1,000-epoch run on the same 3-million-paper map.
• Hard-negative-mining in UMAP 0.6 (1:03) — the negative-sampling change that carries the large-scale gains.
• The repulsion dial for UMAP 0.6 (0:31) — turning the attraction–repulsion dial the new optimizers unlock.

What to watch#

Install it with pip install -U umap-learn (it's on conda-forge too) and try the new pipeline on your own data with compatibility_layout=False — a release candidate has been public since November 3, so the code has already had a shakedown. The interesting question is what happens when the new pipeline becomes the default: embeddings that countless notebooks and pipelines depend on will shift, and the community will find out quickly whether the finer structure and faster runs hold up on data that isn't 3 million arXiv papers. For a library this widely used, a 12× speedup on the benchmark that matters most to its users is the kind of release that quietly changes a lot of workflows.

Sources: the announcement post by u/lmcinnes on r/MachineLearning; “Racing UMAP 0.5 and UMAP 0.6” on the UMAP YouTube channel; the UMAP 0.6 optimizers documentation; the UMAP repository on GitHub; release 0.6.0rc1 notes.