anthology-of-the-sota

SOTA-327: Train retrieval embeddings with nested losses, so a prefix of the vector can shortlist and the full vector re-rank

2 min read

Properties5
StatusProposed
Consensusemerging
Tagsrepresentation-and-encoding, inference-optimization
SourceLIT-547 — Matryoshka Representation Learning
Introduced byLIT-547 — Matryoshka Representation Learning

Graph View

Source

Kusupati, Bhatt, Rege et al. (2022), LIT-547 — Matryoshka Representation Learning. Read as NOTE-291.

The practice

When training an embedding model, add the loss at about log₂ d prefix lengths, halving from the full width down to a small bottleneck (8 or 12 dimensions in the paper), with a head per prefix or one tied head. Normalize each prefix separately if the representation is normalized. That is the only change.

The objective and the shortlist-then-re-rank deployment both originate in LIT-547, Matryoshka Representation Learning. On ResNet50 ImageNet-1K each prefix there was at least as accurate as a separately trained model of that width in linear classification, and up to 3% better in retrieval mAP@10, while SVD and random features fell off sharply below 256 dimensions.

Then deploy one vector at several costs:

  • Store one database of full-width embeddings. Do not keep one per width
  • Shortlist on a prefix, re-rank on the full vector. A 16-d shortlist of 200 and a 2048-d re-rank matched single-shot 2048-d retrieval at about 14× lower wall-clock on ImageNet-1K. The harder ImageNet-4K needed a 64-d shortlist (6×). So the shortlist width is a per-corpus setting, not a constant
  • Truncate when storage or bandwidth is the constraint. A prefix is at least as good as a separately trained model of that width, and far better than SVD or random dimensions below 256

Conditions

  • Measured on image retrieval with class labels as relevance (ImageNet-1K/4K), on ResNet50, ViT-B/16 (JFT, ALIGN), and BERT for masked LM only. No text retrieval is measured in the source
  • The full-width vector gives up a little. On the web-scale models, 0.15–0.25 top-1 k-NN points at full width. If you only ever use the full width, this is a small cost for nothing
  • Against separately trained narrow models, the evidence is ResNet50 on ImageNet. At web scale the paper compares only against random subsets of the full model’s features, which is a weak baseline
  • Measure the end task, not only the retrieval metric (SOTA-308). The paper’s mAP@10 is label recall

Known implementations

  • OpenAI text-embedding-3-small / -large (dimensions parameter)
  • EmbeddingGemma

Backlinks