From MADR to ScaleRank

Summary

To build an alternative to t-SNE and UMAP, I have been experimenting since 2020 with one idea: treat each data point as an agent and decide the layout by moving the agents according to local rules. About 40 notebooks were left in the folder.

In the final 2022 results, MADR beat the other methods by a wide margin on the Q score. In the same notebook, however, a local score restricted to near ranks (k=5) almost reverses the ranking.

The reason lies in how the Q score is defined. The Q score is the sum, over every point, of how far the ranks of all other points as seen from that point differ between the original space and the embedding. Discrepancies at far ranks count as much as those at near ranks, so in practice it is a measure of global structure. Optimise that measure directly and it is no surprise that you win on it.

So, before any new method, I built drbench, a framework for comparing methods under identical conditions. Datasets, methods and metrics are each registered, combinations are run automatically, and results are cached. All distances are Euclidean, and existing methods are called with their default settings.

The central metric: R_NX(K)

For each point, it measures how much the K nearest neighbours in the original space overlap with the K nearest neighbours in 2-D, rescaled so that a random layout scores 0 and a perfect match scores 1. Varying K from 1 to N shows, in a single curve, how well structure is preserved from local (small K) to global (large K).

The R_NX(K) curves clearly separate the scales at which existing methods do well. t-SNE is highest at small K and is overtaken by MDS and PCA once K exceeds about 50.

Across all 8 datasets too, the top method differs between local and global metrics. Switching the metric in the table below shows that t-SNE is strong locally and MDS / PCA are strong globally.

Existing methods carry an implicit, internal choice of which scales to preserve. ScaleRank makes that choice explicit as a user-specified distribution p(K) and directly optimises neighbourhood preservation weighted by that distribution.

Thinking in terms of a single triplet

Fix a point i and a neighbourhood size K. In the original space, there are points j among i's top-K neighbours and points m that are not. In 2-D, the triplet is correct if j is closer to i than m is. Draw K from p(K), draw j and m, and if the triplet is wrong, pull i and j together and push i and m apart. The first version used this triplet as the unit of the loss. The current version draws 20 m's at once (see below).

The knob is γ

If p(K) is proportional to K^−γ, then γ=1 concentrates the weight on small K (local-leaning; the same weighting as AUC_RNX), while γ=0 gives every K the same weight (global-leaning). Moving the γ slider changes which scales receive the effort.

Not an optimisation failure, but a biased surrogate

Even at γ=1, the first version did not reach t-SNE on local metrics (R_NX@10 on digits: 0.44 vs 0.58). So I scored the embeddings of existing methods with ScaleRank's own objective function. If t-SNE scored lower, the optimisation would not be working.

The cause lay in how the objective function was built. Triplet correctness counts a distant non-neighbour being slightly out of order just as heavily, but R_NX(K) only cares about points that intrude into the top K in 2-D. So, for some of the triplets, I chose m from points that are currently close to i in 2-D but are not true neighbours (hard negatives).

In the first version's scatter plots, points were scattered around the clusters. Counting as outliers the points whose 10-NN radius in 2-D exceeds 3 times the median, the first version produced far more of them than t-SNE.

Two kinds of outliers

• The outliers on Pollen were isolated in the original space as well (10-NN radius 3.9 times the median). MDS also produces 16% there. To some extent, these are faithful.
• The outliers on hierarchical clusters (hier_blobs) had ordinary density in the original space (1.06 times). They are artefacts created by the method.

The culprits: hard negatives and anti-hubs

Turning hard negatives off reduced the outliers on hier_blobs from 23% to 2%. The points that became outliers had something in common: they were points that hardly appear in any other point's neighbour list, so-called anti-hubs. In high-dimensional data, some points become neighbours of many points (hubs), while large numbers of points become almost nobody's neighbour (anti-hubs). This is called hubness.

What forces are at work

I decomposed the force each point receives during training into the individual terms of the gradient and measured the component pointing outwards from the centre of the point's own cluster (first version, hier_blobs, epoch 700).

The net force on outliers points inwards, but it is weaker than for the other points (−0.013 vs −0.038). As the cluster cores shrink in the later part of training (distance from the centre 3.5 → 0.6), these points are left behind rather than pushed out. The root cause is the asymmetry of neighbour relations; t-SNE counteracts it with per-point bandwidths (perplexity calibration).

Remedies tried

In the first version, the gap in local metrics was largest on the two synthetic datasets built from high-dimensional Gaussians (0.24 vs 0.34 on hier_blobs, 0.31 vs 0.41 on density_blobs). I investigated with a version whose objective function is concentrated around K≈10 (a band profile), putting it on the same footing as t-SNE.

• hier_blobs: With hard negatives, ScaleRank loses to t-SNE on its own objective function. In other words, hard negatives distort the objective itself. Without them, it nearly matches t-SNE and also beats it on the objective.
• density_blobs: It beats t-SNE on the objective function but loses on R_NX@10. A mismatch between the objective and the evaluation metric still remains.
• Allocation across scales: By design, γ=1 also puts a lot of weight on K>10 (Figure 6). On digits, concentrating on K=10 gives an R_NX@10 of 0.599, above t-SNE (0.583).

Rather than tuning hard negatives to each dataset, I asked whether the shape of the loss function could produce the same effect. For one neighbour j, Q non-neighbours m are drawn together and the following loss is used (the same form as InfoNCE in contrastive learning).

• With Q=1, it is the same formula as the first version's triplet loss (agreement confirmed by a test).
• As Q increases, the closer m currently is to i in 2-D, the larger q_im becomes and the stronger the push. Although m is simply drawn uniformly, non-neighbours that have intruded nearby automatically have a stronger effect: a smooth version of hard negatives.
• It is also known that the larger Q is, the closer the result gets to t-SNE (Damrich et al. 2023). Q becomes a second "local or global" knob alongside γ.

Sweeping Q

On the 4 datasets used for tuning so far, starting from the setting without hard negatives, I increased Q from 1 → 5 → 20 → 100.

With Q=20 and 8 groups per point, the quality was almost the same as with 32 groups at 1/4 of the computation. I made this the current default (n_neg=20, per_point=8, hard=0). The table below compares it over 8 datasets × 3 seeds.

Local drops on low-dimensional manifolds

Not every dataset improved. On s_curve (a 2-D surface embedded in 3-D), local fell from 0.830 to 0.797. A simple test of recovering a uniformly random 2-D square shows the same tendency clearly.

I ran the current version with γ = 1 / 0.5 / 0 on 8 datasets × 3 seeds and compared it with the 7 existing methods. With local (R_NX@10) on the vertical axis and global (distance correlation) on the horizontal axis, γ=1 and 0.5 sit, on all 8 datasets, where no existing method beats them on both axes. γ=0 was beaten by an existing method on two datasets (PaCMAP on density_blobs, TriMap on s_curve).

Behaviour as γ is lowered (new experiment)

Averaged over the 8 datasets, the current version's whole path has shifted towards local. At γ=1 it is more local than the first version, but even at γ=0 global does not rise as far as in the first version (distance correlation 0.863 vs 0.906). This is probably because the tendency of many negatives to pull the result towards t-SNE persists even when γ is lowered. Also, in both versions, the lower γ is, the more outliers there are.

In the scatter plots, t-SNE separates the clusters cleanly but does not preserve their positions relative to one another. MDS preserves those positions, but the clusters overlap. ScaleRank sits in between: it separates the clusters while trying to keep their arrangement. The scatter around clusters that stood out in the first version is reduced in the current version.

Even if t-SNE and UMAP at their default settings do not beat ScaleRank, tuning their parameters might bring them up to where ScaleRank is. So I ran 13 t-SNE settings, perplexity {5, 30, 100, 300} × exaggeration {1, 2, 4} plus random initialisation, and 17 UMAP settings, n_neighbors {5, 15, 50, 200} × min_dist {0, 0.1, 0.5, 0.99} plus random initialisation, on 8 datasets × 3 seeds.

• On no dataset did any of the 30 settings beat ScaleRank on both local and global measures. The same holds when the best setting is chosen for each dataset.
• Raising t-SNE's perplexity raises global (mean over 8 datasets 0.52 → 0.71) at the cost of local (0.543 → 0.415). Raising exaggeration lowers local sharply and also increases outliers.
• For UMAP, raising n_neighbors leaves global plateauing at around 0.43, and min_dist has almost no effect.
• The gaps are smallest on s_curve and density_blobs, where tuned t-SNE comes close to ScaleRank.

To find out how much of ScaleRank is new, I surveyed related work. In short, most of the individual components are already known.

What may be new is the following combination: the user supplies an arbitrary scale distribution p(K); neighbourhood preservation weighted by it is the objective function; and the sampling rules (draw K from p, the positive from the top K and the negatives from outside it) and the triplet weights are derived from that objective. I found no publication that states this combination as such.

However, changing the scale weights of multi-scale t-SNE to follow p(K) is a straightforward extension. Unless experiments can show a difference from it, I consider the novelty weak. The next methods to compare against are multi-scale t-SNE reweighted by p(K), a sweep of the λ in DREAMS, and a sweep of the number of negatives in contrastive learning.

Up to this point, ScaleRank kept a table of the ranks of all points as seen from each point (N × N), which limited it to about 10,000 points. So I dropped the rank table and built a version that optimises the same objective function using only structures of size O(N log N).

• Ranks 1–100 are found with exact k-nearest neighbours.
• Larger ranks are estimated from random subsamples at several levels. The a-th nearest point in a subsample of size n corresponds to a rank of roughly a·N/(n+1) in the full data. This gives each point a "distance → rank" table.
• Negatives are chosen from candidates redrawn uniformly at each epoch: the candidates are ranked with the table, and those whose rank exceeds K are used.

Compared with the original version on 8 datasets × 3 seeds, the differences were within about ±0.01 on every metric (Table 5). Evaluation was also changed to estimate from only 2,000 rows for datasets of more than 10,000 points. Each metric is a mean of per-row values, so this estimate is unbiased.

Table 6 and Figure 21 show the results on all 70,000 points of MNIST and Fashion-MNIST (reduced to 50 dimensions with PCA).

• Globally (distance correlation), ScaleRank is clearly higher than t-SNE and UMAP.
• Locally (R_NX@10), on the other hand, it weakens as N grows. On MNIST at γ=1, it was 0.27 at 10,000 points but 0.10 at 70,000 points: about a quarter of t-SNE (0.39) and about the same as UMAP.

In Figure 21, ScaleRank γ=1 (neg_scale 0.5) keeps the digits apart, but each class breaks into small clumps, with many scattered points around them. γ=0 keeps the overall arrangement but lets the classes overlap.

Why local weakens as N grows

If the number of negatives Q stays fixed while N grows, the Q negatives stand in for an ever larger set of "all other points", and the repulsion becomes relatively weaker. This is a known property in contrastive learning (Damrich et al. 2023), and the behaviour approaches that of UMAP, which draws few negatives.

So I added a setting that multiplies the negative terms by a weight w = ((N−1−K)/Q)^a (neg_scale = a). With a = 0.5, local roughly doubles (MNIST 0.10 → 0.20, Fashion-MNIST 0.12 → 0.20), but it still falls short of t-SNE. On small datasets the density correlation drops (0.66 → 0.48 on digits), so it is not the default.

On 10,000 MNIST points I also tried other settings. Raising γ lifts local to at most 0.36 (γ=2); 100 negatives give 0.35, and neg_scale 0.5 gives 0.36, against 0.28 for the default γ=1 and 0.45 for t-SNE. Four times as many groups per point, or 2,000 epochs instead of 800, change almost nothing (0.29).

So that ScaleRank can be used without data leaving the browser, I rewrote it in Rust and made it into an app that runs in the browser via WebAssembly. For comparison, t-SNE and UMAP are included in the same Rust code and are computed only when the user asks for them. I also implemented approximate nearest neighbours (NN-descent), but with 50 dimensions and 70,000 points it was not faster than brute force (4.3 s on 10 threads), so brute force is used for now.

I checked with drbench that the Rust versions compute the same thing as the reference implementations (Table 7). t-SNE is compared with openTSNE, UMAP with umap-learn (PCA initialisation), and ScaleRank with the Python version.

The same holds at 70,000 points: R_NX@10 is 0.386 for t-SNE (openTSNE 0.391) and 0.108 for UMAP (umap-learn 0.109). With 10 threads, the computation takes 28 s for t-SNE (openTSNE 64 s), 22 s for UMAP (35 s) and 44 s for ScaleRank (Python version 1,347 s).

The browser version is currently single-threaded and takes about 17 s for ScaleRank, 11 s for UMAP and 7 s for t-SNE on 6,000 points. At 20,000 points, ScaleRank takes about a minute; I plan to shorten this with multithreading.

Caveats

• On local metrics alone, t-SNE remains first on every dataset (mean 0.543 vs 0.511 for the current version).
• On low-dimensional manifolds (s_curve, the 2-D square), the current version scores lower locally than the first version.
• Lowering γ to lean towards global structure increases outliers and does not reach the global scores of the first version.
• Outliers are still far more frequent (8.8%) than with t-SNE (0.3%). Some are genuinely isolated points, but not all.
• AUC_RNX measures something close to the γ=1 objective function, so the lead on it must be discounted. R_NX@10 and distance correlation are not optimised directly.
• Hyperparameters were tuned on digits, Pollen, hier_blobs and density_blobs. The 4 datasets not used for tuning showed the same tendencies.
• At 70,000 points, local drops to about a quarter of t-SNE's. neg_scale roughly doubles it, but it still falls short.
• Most of the individual components exist in prior work. ScaleRank has not yet been compared with multi-scale t-SNE reweighted by p(K) or with DREAMS.
• The browser version is single-threaded, and ScaleRank takes about a minute on 20,000 points.

Candidate next steps

• Estimate ranks in 2-D directly. With Q negatives drawn, the rank of j can be estimated from the fraction of non-neighbours that are closer than j in 2-D. Penalising only when this exceeds K would optimise R_NX(K) almost as it stands.
• Vary Q with K. Use fewer negatives when leaning global (large K) and more when leaning local. This may help with global scores falling short when γ is lowered.
• Fix the asymmetry on the positive side. As t-SNE does with per-point bandwidths, even out how strongly each point is pulled as a neighbour, so that anti-hubs also receive enough attraction.
• Keep local quality at large N. Set the default neg_scale according to N. Try a stage corresponding to t-SNE's early exaggeration.
• Compare directly with prior work. Add multi-scale t-SNE reweighted by p(K), a sweep of the λ in DREAMS and a sweep of the number of negatives in contrastive learning to drbench.
• Make it faster in the browser. Aim for 100,000 points with multithreading and SIMD. Add an FFT-based version of t-SNE.

The results can be reproduced with the following commands on the proto/scalerank branch of the repository tsjshg/madr. The original pickle files are assumed to be in the parent folder.

uv sync
uv run python scripts/convert_legacy.py      # 2020 pickles → data/raw/*.npz
uv run drbench run                           # existing methods (8 datasets × 3 seeds)

# current version (softmax loss, default settings)
uv run drbench run -m scalerank:gamma=1.0 scalerank:gamma=0.5 scalerank:gamma=0.0
# first version (triplets + hard negatives)
uv run drbench run -m "scalerank:gamma=1.0,hard=0.3,n_neg=1,per_point=32"
uv run drbench report                        # figures and tables in results/report/

# investigation scripts
uv run python scripts/sweep_scalerank.py -d digits --grid n_neg=1,5,20 hard=0.0
uv run python scripts/score_under_objective.py digits --gammas 1
uv run python scripts/diagnose_outliers.py hier_blobs -m tsne "scalerank:gamma=1.0"
uv run python scripts/force_decomposition.py hier_blobs --epochs 700
uv run python scripts/plot_scale_profile.py

# comparison with tuned t-SNE / UMAP (after drbench run with 30 settings × 8 datasets × 3 seeds)
uv run python scripts/compare_tuned_baselines.py

# N²-free version and Rust version (after installing scalerank-web's Python bindings)
uv run drbench run -m scalerank_fast:gamma=1.0 scalerank_fast:gamma=0.5 scalerank_fast:gamma=0.0
uv run drbench run -m scalerank_rs:gamma=1.0 tsne_rs umap_rs umap:init=pca
uv run python scripts/compare_fast.py
uv run python scripts/compare_rs.py
uv run drbench run -d mnist70k fmnist70k -s 0 -m tsne umap scalerank_rs:gamma=1.0 scalerank_rs:gamma=0.0 "scalerank_rs:gamma=1.0,neg_scale=0.5"

The code for the Rust version and the app is in tsjshg/scalerank-web.