anthology-of-the-sota

NOTE-291: Matryoshka Representation Learning

4 min read

Properties2
StatusRead
PaperLIT-547 — Matryoshka Representation Learning

Graph View

Contribution

A training objective that makes an embedding’s prefixes usable on their own, so one vector serves several cost points. Before it, getting a narrower representation meant training another model, compressing after the fact (SVD, which loses a lot), or keeping several databases of embeddings.

Key insight

Put the loss on the prefixes and the model will sort its information by importance. Without it a network spreads information across all coordinates, and truncation destroys it. With it, the first coordinates carry the coarse distinctions and the later ones refine them. The superclass analysis shows exactly this: 8 dimensions lose fine classes and keep superclasses.

Assumptions

  • Nesting set M = {8, 16, …, 2048} for ResNet50 and {12, 24, …, 768} for ViT and BERT, with uniform weights c_m = 1
  • Same hyperparameters as the baselines, untuned for MRL
  • Normalization handled separately for each prefix when the representation is normalized (Appendix C)
  • Retrieval is class-label retrieval on ImageNet-1K/4K, measured with exact search in FLOPs and HNSW in wall-clock time

Key results

  • Linear classification, ResNet50 on ImageNet-1K (Figure 2): MRL is at least as accurate as the separately trained (FF) model at every width. MRL–E is within 1% from 16 dimensions up
  • 1-NN (Figure 3): up to 2% better than FF at low widths, level elsewhere
  • Retrieval mAP@10 (Figure 7): up to 3% better than FF
  • Adaptive classification (Figure 6): 76.30% at an expected ~37 dimensions, where FF needs 512. That is 0.8% under the 2048-d FF model
  • Adaptive retrieval (Figure 8): D_s = 16, D_r = 2048, K = 200 matches single-shot 2048-d at ~128× theoretical and ~14× wall-clock. ImageNet-4K needs D_s = 64, for ~32× and ~6×. Funnel retrieval (halve the shortlist and double the width each step) is as accurate at ~128×
  • Web scale (Table 4): at 12 dimensions, ALIGN-MRL has 43.57 top-1 k-NN against 11.90 for random features, and JFT-ViT-MRL 53.61 against 27.07. At 768: 67.85 against 68.00 and 71.85 against 72.10
  • BERT: MLM validation accuracy within 0.5% of FF
  • Oracle routing per instance to its best width would add up to 4.6% accuracy. Some instances are better at low width

Claims

idclaimstrengthsupport
C1Each nested prefix is as good as a separately trained representation of that widthmoderateResNet50/ImageNet-1K against real FF baselines; web-scale only against random features
C2Shortlist-then-rerank on one MRL vector matches full-width retrieval at much lower costmoderateImageNet-1K/4K image retrieval, one encoder family
C3Intermediate, untrained widths interpolatemoderateFigure 5, Table 5
C4MRL costs nothingweaktrue at inference; at full width web-scale models lose 0.15–0.25 points
C5It transfers to languageweakBERT MLM accuracy only, no text retrieval

Method

min Σ_i Σ_{m∈M} c_m · L(W^(m) · F(x_i)[1:m], y_i), with MRL–E tying W^(m) = W[:, 1:m]. For contrastive training, apply it to both sides of the contrast.

Concepts

  • FF — a fixed-feature baseline: a separately trained model whose representation is m-dimensional
  • Adaptive retrieval (AR) — shortlist with D_s dimensions, re-rank with D_r dimensions of the same vector, so there is one database
  • Funnel retrieval — a cascade of shortlist and re-rank that halves the list and doubles the width at each step

Connections

It descends from nested dropout (Rippel et al.), which ordered all d dimensions. MRL optimizes only O(log d) and interpolates the rest. MatFormer (LIT-546) applies the same nesting to FFN weights.

Recommendations

  • R1 — Train embeddings meant for retrieval with nested prefix losses, store one database, shortlist on a prefix and re-rank on the full vector. Topic: representation-and-encoding. Status: standard in practice. Strength: moderate. Filed as SOTA-327

Bearing on the record

  • SOTA-327 is new
  • SOTA-308 (measure relevance and the end task, not only recall, when changing retrieval) applies directly. MRL’s own metric is class-label mAP@10, which is a recall-style proxy

Limitations

  • Image classification labels as relevance. No text retrieval and no graded relevance
  • Weak web-scale baseline (random features), chosen for cost
  • A small full-width cost at web scale
  • Uniform loss weights and hand-set nesting. The authors list both as weaknesses (§6)

Open questions

  • How it holds up on text retrieval benchmarks with graded relevance. Widespread adoption suggests it does, and the record holds no measurement
  • Whether the full-width cost grows with model scale

Backlinks