normrouter

Training-free per-sample routing over pools of frozen pre-trained encoders.

Given a pool of (backbone, layer) candidates, normrouter scores each test sample with training-free embedding-geometry statistics and routes it to one candidate or to a weighted combination of all of them. The router is not trained and uses neither test labels nor per-candidate calibration.

Two scores are provided: Levy, the radial typicality of the embedding norm, and LSNG, a local spectral score that is comparable across embedding dimensions.

Installation

pip install normrouter

Route with a pool of encoders

from normrouter import NormRouter

router = NormRouter.from_backbones(["dinov2_vitl14", "mae_vit_large"])
router.fit(train_images, train_labels)
predictions = router.predict(test_images)

fit embeds the training samples once per encoder, then fits one scorer and one linear probe per (encoder, layer) candidate. predict scores every candidate for each test sample and combines the probes' predictions according to the routing mode: "hard" picks the argmax candidate, "soft" (the default) averages the full pool with softmax weights, and "topk-uniform" averages the top_k best-scored candidates uniformly. Scores are quantile-calibrated against each candidate's training scores by default, which keeps them comparable across encoders with different embedding scales; pass calibrate="none" for the raw statistics. A custom pool, including finetuned checkpoints, goes in through NormRouter(pool, ...).

Score your own embeddings

The functional API works on plain numpy arrays. train_embeddings are the embeddings a candidate produced for the training set, test_embeddings the ones to score.

from normrouter import compute_levy, compute_lsng, softmax_route

levy = compute_levy(test_embeddings, train_embeddings)
lsng = compute_lsng(test_embeddings, train_embeddings, seed=0)

Scores from several candidates turn into routing weights:

import numpy as np

scores = np.stack([lsng_a, lsng_b, lsng_c])   # (n_candidates, n_samples)
weights = softmax_route(scores, temperature=1.0)

Hard routing is the argmax over candidates; soft ensembling is the weights-weighted average of the candidates' predicted class probabilities.

Compose a candidate pool

from normrouter import EncoderEntry, Pool

base = EncoderEntry(key="dinov2_vitl14", base_key="dinov2_vitl14", num_layers=24)
expert = EncoderEntry(
    key="dinov2_vitl14_galaxy10",
    base_key="dinov2_vitl14",
    num_layers=24,
    checkpoint_path="checkpoints/dinov2_vitl14_galaxy10/backbone.pt",
)

pool = Pool.of(base, layers="last:4") + Pool.of(expert, layers="all")
pool = pool.subset(encoders=["dinov2*"])
pool.candidate_ids()

Pools refuse duplicate candidates and conflicting encoder definitions, so a checkpoint can never silently shadow another one under the same key.

Research code

The ssl_model_audit package and the scripts/ directory contain the full experiment pipeline behind the package: embedding extraction with on-disk caching, linear and kNN probes, the scorer registry, routing strategies, and the figure and table generators. The main entry point is:

python scripts/run_experiment.py --dataset galaxy10 --backbones dinov2_vitl14 mae_vit_large

Supported backbones include DINOv2 (S/B/L/G), MAE, CLIP, BEiT, I-JEPA, MoCo v3, SimSiam, SimCLR, SwAV, BYOL, Barlow Twins, VICReg, and a supervised ResNet-50. The transformer backbones declare their layer counts and work directly with NormRouter.from_backbones; the ResNet-50 ones (SimCLR, SwAV, BYOL, Barlow Twins, VICReg, supervised) need a Pool with an explicit num_layers. Finetuned experts are registered from local checkpoints. Text encoders (Pythia) are supported through the same pipeline.

Development

pip install -e ".[dev]"
pytest -m "not slow and not network"

Tests marked slow or network download model weights and are excluded from the default CI run.

License

Attribution-NonCommercial 4.0 International (CC BY-NC 4.0), see LICENSE.md.