anthology-of-the-sota

LIT-547: Matryoshka Representation Learning

2 min read

Properties6
StatusActive
Tagsrepresentation-and-encoding, inference-optimization
Extended byLIT-546 — MatFormer: Nested Transformer for Elastic Inference
Cited as “introduced by” bySOTA-327 — Train retrieval embeddings with nested losses, so a prefix of the vector can shortlist and the full vector re-rank
Cited as “paper” byNOTE-291 — Matryoshka Representation Learning
Cited as “source” bySOTA-327 — Train retrieval embeddings with nested losses, so a prefix of the vector can shortlist and the full vector re-rank

Graph View

Kusupati, Bhatt, Rege et al., University of Washington, Google Research, Harvard (2022; NeurIPS 2022) — ARXIV-2205.13147

Key takeaways

  • The objective: sum the usual loss over nested prefixes z[1:m] for m ∈ M = {8, 16, …, d}, each with its own linear head. MRL–E ties those heads into one. Weights are uniform unless tuned. Nothing else in the pipeline changes, and the paper reuses the baseline hyperparameters
  • Each prefix is a working representation. On ResNet50 ImageNet-1K, MRL is “at least as accurate” as a separately trained model of each width in linear classification. In 1-NN it is up to 2% better at low widths, and in retrieval mAP@10 up to 3% better. SVD, random features and slimmable networks fall off sharply below 256 dimensions
  • Widths nobody trained also work. Accuracy interpolates across all d dimensions, not only the O(log d) that were optimized
  • Adaptive retrieval: shortlist 200 with a 16-d prefix and re-rank with 2048. That matches single-shot 2048-d mAP@10 at about 128× fewer theoretical FLOPs and 14× lower wall-clock (HNSW, same hardware). The harder ImageNet-4K needs a 64-d shortlist, for 32× and 6×
  • Adaptive classification: a cascade that escalates width on low confidence reaches 76.30%, a 512-d model’s accuracy, at an expected ~37 dimensions
  • Web scale: ViT-B/16 on JFT-300M and ALIGN. BERT MLM accuracy comes within 0.5% of the fixed-width model

Standing in the anthology

The idea MatFormer (LIT-546) moves into the weights, from the same group. Filed from #163 as “Matryoshka embeddings”, and it sources SOTA-327.

Two hedges the abstract does not carry. At web scale, the lower-width comparison is against random features of the full model, not against separately trained narrow models. That is a much weaker baseline, and the paper says it made the choice for cost. At full width MRL gives something up: Table 4 has ALIGN 67.85 against 68.00 and JFT-ViT 71.85 against 72.10 top-1 k-NN. Small, but “no additional cost” is about inference, not about accuracy at d.

No text retrieval is measured. The BERT result is masked-LM accuracy. The practice’s use in text embedding models is adoption (DP-005), and it is widespread.

Read — NOTE-291.

Backlinks