PaperlayerView exact revision on arXiv

Source-linked AI summary

Matryoshka Representation Learning

Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanujan, William Howard-Snyder, Kaifeng Chen, Sham Kakade, Prateek Jain, Ali Farhadi

arXiv:2205.13147v4cs.LGcs.CV

TL;DR

Fixed-capacity embeddings may not match the varying computational and statistical demands of downstream tasks. MRL encodes multiple granularities in one embedding, yielding efficient adaptive deployment with comparable accuracy and extensions across modalities and web-scale datasets.

  • Problem

    Downstream tasks have differing computational and statistical constraints, making rigid fixed-capacity representations potentially over- or under-accommodating.

  • Method

    MRL jointly learns O(log(d)) nested representations within a single embedding and minimally modifies existing representation-learning pipelines.

  • Results

    MRL matches fixed-feature accuracy while using 14× smaller representations on average and achieves comparable retrieval mAP@10 at 128× lower FLOPs and 14× faster wall-clock time.

  • Takeaways & Limitations

    A single MRL embedding can adapt deployment to downstream task complexity and available compute across classification, retrieval, and multiple modalities.

  • Takeaways & Limitations

    The authors identify nested-loss weighting, fidelity-specific losses, and learned search structures as weaknesses and directions for future work.

Abstract

from arXiv · show

Learned representations are a central component in modern ML systems, serving a multitude of downstream tasks. When training such representations, it is often the case that computational and statistical constraints for each downstream task are unknown. In this context rigid, fixed capacity representations can be either over or under-accommodating to the task at hand. This leads us to ask: can we design a flexible representation that can adapt to multiple downstream tasks with varying computational resources? Our main contribution is Matryoshka Representation Learning (MRL) which encodes information at different granularities and allows a single embedding to adapt to the computational constraints of downstream tasks. MRL minimally modifies existing representation learning pipelines and imposes no additional cost during inference and deployment. MRL learns coarse-to-fine representations that are at least as accurate and rich as independently trained low-dimensional representations. The flexibility within the learned Matryoshka Representations offer: (a) up to 14x smaller embedding size for ImageNet-1K classification at the same level of accuracy; (b) up to 14x real-world speed-ups for large-scale retrieval on ImageNet-1K and 4K; and (c) up to 2% accuracy improvements for long-tail few-shot classification, all while being as robust as the original representations. Finally, we show that MRL extends seamlessly to web-scale datasets (ImageNet, JFT) across various modalities -- vision (ViT, ResNet), vision + language (ALIGN) and language (BERT). MRL code and pretrained models are open-sourced at https://github.com/RAIVNLab/MRL.

1 Introduction

MRL addresses the deployment rigidity of fixed-dimensional representations by encoding coarse-to-fine information within one embedding. This supports adaptive accuracy–compute trade-offs for classification and retrieval without additional inference cost.

  • Web-scale downstream utilization costs grow with embedding dimensionality, data size, and label-space size, often overshadowing feature computation.
  • MRL explicitly optimizes O(log(d)) nested lower-dimensional vectors within one high-dimensional representation.The nested representations are designed for adaptive deployment across existing representation-learning pipelines.
  • Each first m-dimensional prefix is intended to match an independently trained m-dimensional representation while retaining increasingly fine-grained information as m grows.This creates a coarse-to-fine representation with no additional training cost for the prefixes.
  • 14× smaller representations achieve the same ImageNet-1K classification accuracy as baselines.MRL uses adaptive cascades to reduce the average embedding dimension required for a target accuracy.
  • 14× wall-clock speedups and 128× theoretical FLOPS speedups are reported for adaptive retrieval relative to single-shot retrieval with standard embeddings.Retrieval shortlists candidates with early dimensions and reranks them using additional dimensions.
  • MRL extends across vision, vision–language, and language modalities and across web-scale datasets.The reported settings include ResNet, ViT, ALIGN, BERT, ImageNet-1K/4K, JFT-300M, and ALIGN data.

2 Related Work

Prior efficiency methods address representation and search costs but often trade accuracy for reduced dimensionality or require multiple models and forward passes. MRL instead embeds nested multifidelity representations in one vector and complements existing efficiency techniques.

  • MRL is complementary to supervised, self-supervised, and contrastive representation-learning setups and adds multifidelity with minimal overhead.
  • Large-scale classification and retrieval costs depend on label-space size, data size, and representation dimensionality, stressing RAM, disk, and processors.
  • Approximate search and hierarchical methods reduce dependence on label count, while dimensionality reduction, hashing, and feature selection address O(d) scaling with possible accuracy loss.
  • MRL learns lower-dimensional representations that match independently trained counterparts without requiring multiple expensive forward passes.
  • MRL provides an intermediate abstraction between high-dimensional vectors and efficient approximate-nearest-neighbor indices through nested adaptive embeddings.Other efficiency techniques remain complementary and can be applied to Matryoshka Representations.
  • Varying-capacity neural-network approaches can require distinct forward passes to isolate representations, increasing the cost of re-encoding retrieval databases.

Matryoshka Representation Learning

MRL turns one d-dimensional embedding into nested prefixes that are independently usable at selected capacities. It trains these prefixes jointly through multi-scale losses and can be adapted to standard representation-learning frameworks.

  • For each selected m, the prefix z_1:m is trained to serve as a transferable representation of the input.
  • MRL selects fewer than log(d) nesting dimensions and obtains the full representation from a neural network F(x; θ_F).
  • In supervised classification, a d-dimensional image representation is paired with separate linear classifiers for the chosen nesting dimensions.The ImageNet-1K example uses d = 2048 and M = {8, 16, . . . , 1024, 2048}.
  • MRL aggregates separate classification losses across all nested dimensions, weighting each by its importance coefficient c_m.
  • Despite optimizing only selected nested dimensions, MRL produces accurate representations at intermediate dimensions between the chosen granularities.
  • Weight tying across classifiers yields Efficient MRL, reducing classifier memory cost by almost half.
  • MRL adapts to web-scale frameworks with minimal modifications, including masked language modeling and contrastive vision or vision–language learning.

4 Applications

MRL provides coarse-to-fine representations that support flexible classification and retrieval across representation sizes, datasets, and modalities. Its adaptive deployments reduce representation or query cost while retaining comparable accuracy.

  • Representation Learning: MRL removes the need for independently trained models at each representation size while preserving comparable coarse-to-fine representation quality.It explicitly optimizes nested dimensions and can be adapted to standard vision and language representation pipelines.
  • Representation Learning: MRL matches independently trained models for linear classification at every representation size and is up to 2% more accurate in lower-dimensional 1-NN evaluation.The 1-NN evaluation measures downstream representation quality on ImageNet-1K.
  • Representation Learning: MRL scales to JFT-300M and ALIGN web-scale setups, providing accurate lower-dimensional representations with favorable cost-versus-accuracy trade-offs.The evaluation uses ViT-B/16 models and 1-NN accuracy on ImageNet-1K.
  • Adaptive Classification: Adaptive classification achieves 76.30% accuracy with an expected representation size of approximately 37 dimensions, matching a 512-dimensional fixed-feature model.This corresponds to roughly 14× smaller representation size, while remaining only 0.8% below the 2048-dimensional fixed-feature baseline.
  • Adaptive Retrieval: MRL improves ImageNet-1K retrieval by up to 3% over fixed-feature baselines across dimensionalities and supports adaptive retrieval on ImageNet-1K and ImageNet-4K.Adaptive retrieval with Ds = 16 and Dr = 2048 is as accurate as fixed d = 2048 retrieval, while being approximately 14× faster in practice on ImageNet-1K.
  • Adaptive Retrieval: Adaptive retrieval configurations lie above the fixed-size Pareto frontier, offering comparable accuracy with lower MFLOPs per query.Funnel retrieval is reported as almost as accurate as the baseline while reducing some adaptive-retrieval parameter choices.

5 Further Analysis and Ablations

MRL retains robustness across out-of-domain classification and retrieval settings, improves long-tail performance, and reveals how representation capacity affects errors and superclass recognition. These analyses show that different inputs and tasks can benefit from different dimensionalities.

  • Robustness: Up to 3% higher mAP@10 demonstrates more robust MRL retrieval on ImageNetV2 queries against an ImageNet-1K database.Classification robustness is at least as strong as the original representation, with a 0.6% ImageNet-A improvement, equivalent to a 20% relative improvement.
  • Few-shot and Long-tail Learning: Up to 2% higher accuracy on novel tail classes occurs without sacrificing accuracy on other classes in FLUID long-tail sequential learning.Accuracy differences between low- and high-dimensional representations are marginal for pretrain classes.
  • Disagreement across Dimensions: 4.6% additional classification accuracy is possible with perfect routing of instances to their most suitable representation dimension.Some instances and classes are more accurate at lower dimensions, while low-dimensional errors often involve same-superclass confusion or multiple objects of interest.
  • Superclass Accuracy: MRL is more accurate than FF models across dimensions for 31-way ImageNet-1K superclass classification under tight information bottlenecks.Fine-grained accuracy decreases more rapidly than superclass accuracy as the bottleneck narrows, indicating that smaller representations retain coarser semantic information.
  • Superclass Accuracy: 11% improvement for garment classification occurs when increasing representation size from 8 to 16 dimensions.Superclass trends vary: oscine benefits from visual separability between object and background even at 8 dimensions.
  • Ablations: Inexpensive partial finetuning enables MRL in off-the-shelf pretrained models, while loss weighting can improve lower-dimensional accuracy without reducing overall accuracy.Additional ablations examine the initial granularity and spacing of nested representation sizes.

6 Discussion and Conclusions

The paper presents MRL as a flexible single-embedding approach for adaptive classification and retrieval. It matches fixed-feature accuracy with smaller representations and enables substantially cheaper, faster retrieval, while identifying several directions for improving the accuracy-efficiency trade-off.

  • Discussion and Conclusions: Future work targets Pareto-optimal loss weighting, fidelity-specific losses, learned search structures, and joint adaptive retrieval optimization.These directions address weaknesses identified in the reported results.
  • Discussion and Conclusions: MRL encodes information at multiple granularities in one embedding, adapting representation capacity to downstream statistical complexity and available compute.The approach supports adaptive classification and retrieval.
  • Discussion and Conclusions: 14× smaller representation size on average matches fixed-feature baseline accuracy on standard benchmarks.The adaptive shortlisting and re-ranking system also achieves comparable mAP@10 while reducing retrieval cost.
  • Discussion and Conclusions: 128× cheaper FLOPs and 14× faster wall-clock time are achieved for adaptive retrieval while maintaining comparable mAP@10 to the baseline.The conclusion positions MRL alongside complementary inference and vector-search efficiency techniques.

Checklist

The paper reports checklist compliance for claims, limitations, reproducibility materials, and data practices, while noting that repeated benchmarking was not performed because of computational expense.

  • Research transparency: The authors state that the paper’s claims accurately reflect its contributions and scope, and that limitations are discussed in Section 6.The checklist marks both items as satisfied.
  • Societal impact: The work reports no additional negative societal impact beyond existing representation-learning impacts and identifies representation size versus encoded bias as future work.The checklist marks personally identifiable information, offensive content, and human-subject procedures as not applicable.
  • Reproducibility: The code and public models are reported as open sourced, with training details, compute resources, and reproducibility materials provided in the appendices.The checklist also confirms that code, data, and instructions needed to reproduce the main experiments are provided.
  • Experimental caveats: The authors did not report error bars because repeated runs on ImageNet-1K, JFT-300M, and ALIGN data were considered extremely expensive.The checklist states that the experiments used large-scale datasets and models such as ResNet and ViT.

B Datasets

Experiments cover standard, robust, and web-scale vision and vision-language datasets, using ResNet50, ViT-B/16, ALIGN, and BERT configurations with fixed-feature and nested-representation baselines.

  • Datasets: ImageNet-1K contains 1,281,167 labeled training images and 50,000 labeled validation images across 1,000 classes.The images use standard transformations detailed by FFCV.
  • Datasets: ImageNet-4K contains 4,202 non-overlapping ImageNet-21K classes, totaling approximately 4.2M training and 200K query or validation examples.Each class has 1,000 training examples and 50 query or validation examples.
  • Web-scale datasets: The web-scale datasets include JFT-300M with 300M images across 18,291 categories and ALIGN with 1.8B image-text pairs.These setups evaluate vision and vision-language representations.
  • Robustness datasets: Robustness evaluation uses ImageNetV2, ImageNet-A, ImageNet-R, and ImageNet-Sketch alongside ObjectNet.These datasets vary in sampling time, adversarial filtering, artistic renditions, sketches, and object classes.
  • Experimental models: The experiments compare MRL and MRL-E against fixed-feature models with representation sizes from 8 to 2048 dimensions.The models include ResNet50 configurations, with additional ViT-B/16, ALIGN, and BERT training setups.
  • Experimental protocol: MRL models use the same representation normalization as fixed-feature implementations when applicable, supporting fair comparisons.The paper reports minimal modifications to the FFCV training pipeline.

D Classification Results

MRL improves classification and retrieval at small representation sizes while supporting adaptive deployment across modalities and datasets. Its adaptive cascades preserve accuracy while reducing the representation or computation required.

  • At d ≤128, MRL significantly outperforms all classification baselines, showing stronger low-dimensional representations.
  • MRL improves k-NN classification at low dimensions compared with feature selection, SVD, random projection, and slimmable-network baselines.
  • Adaptive Classification: ∼37 expected dimensions achieve 76.3% ImageNet-1K accuracy, roughly 14× smaller than the FF–512 baseline.
  • Cross-modal Evaluation: MRL improves classification accuracy at nearly all dimensions for ALIGN and JFT-ViT, and remains within 0.5% of FF BERT representations for MLM validation accuracy.
  • Retrieval: MRL outperforms FF retrieval models at nearly all representation sizes, especially when Ds ≤32, and remains robust on ImageNetV2.
  • Adaptive Retrieval: Funnel retrieval reaches within 0.1% of Ds = 2048 top-1 accuracy on ImageNet-1K using 128× fewer MFLOPs, and within 0.15% on ImageNet-4K using 64× fewer MFLOPs.

G Few-shot and Sample Efficiency

MRL representations retain few-shot performance while adapting capacity to task difficulty. They improve novel-tail accuracy in long-tailed classification without sacrificing accuracy on other classes.

  • Few-shot Classification: MRL matches FF performance across ImageNetV2 shot numbers and representation sizes, while lower-shot tasks reach optimal accuracy at smaller dimensions.
  • Long-tailed Classification: MRL provides up to 2% higher accuracy on novel-tail FLUID classes without sacrificing accuracy on other classes.
  • Long-tailed Classification: The 64-dimensional MRL is approximately 1% below its 2048-dimensional counterpart on pretrain-head classes, but the novel-tail gap is larger.
  • Sample Efficiency: These results support that different tasks require varying representation capacity according to their difficulty.
  • Few-shot Classification: For 1-shot ImageNetV2 classification, 32-dimensional representations perform comparably to 2048-dimensional representations.

J Analysis of Model Disagreement

MRL representations exhibit structured disagreement across dimensions: some classes improve with more capacity, while others benefit from smaller representations. This disagreement motivates adaptive routing and exposes an oracle ceiling for classification.

  • Class Trends: 244 ImageNet-1K classes improve monotonically with dimension, while independently trained FF models show no clear trend for 950 classes.
  • Class Trends: Some instances and classes benefit from lower-dimensional representations, so increasing dimension does not uniformly improve classification.
  • Oracle Accuracy: 18.46% of ImageNet-1K validation samples cannot be correctly predicted at any representation size, leaving an 81.54% oracle accuracy.
  • Adaptive Routing: Adaptive cascading achieved 76.30% top-1 accuracy with expected dimensionality d ∼37, while more optimal routing remains future work.
  • Grad-CAM Examples: Grad-CAM analysis links smaller-model errors to cluttered scenes and confusions among visually similar classes within the same superclass.

K.1 MRL Training Paradigm

MRL trains nested representations at selected granularities and can be modified or extended through alternative training choices. Logarithmic nesting is effective at low dimensions, while very small dimensions are difficult to optimize.

  • Finetuning: MRL can induce nesting in pretrained FF-2048 models by adding an MRL layer and unfreezing backbone layers.
  • Relative Importance: Boosting the loss importance of lower nesting dimensions provides an explicit way to target low-dimensional performance.
  • Arbitrary Granularities: Logarithmic nesting dimensions are used because accuracy gains increase more logarithmically than linearly with representation size.
  • Arbitrary Granularities: Uniformly spaced nesting suffers at low dimensions compared with logarithmic spacing, despite interpolating between trained granularities.
  • Lower Dimensionality: Training below 8 dimensions yields very low accuracy and optimization difficulty, although higher granularities remain largely unaffected.
  • Adaptive Retrieval Ablation: Adaptive retrieval performance depends on shortlist length: larger shortlists saturate ImageNet-1K performance but continue improving ImageNet-4K results through k = 2048.
Loading 2205.13147v4…