Matryoshka Representation Learning
Authors: Aditya Kusupati
Gantavya Bhatt
Aniket Rege
Matthew Wallingford
Aditya Sinha
Vivek Ramanujan
William Howard-Snyder
Kaifeng Chen
Sham M. Kakade
Prateek Jain
Ali Farhadi
Venue: NeurIPS 2022 Accept (2022)
Keywords: Representation Learning
Efficient Deployment
Large-scale Retrieval
Large-scale Classification
Deep Learning
Computer Vision

Abstract: 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 $\mathbf{14}\times$ smaller embedding size for ImageNet-1K classification at the same level of accuracy; (b) up to $\mathbf{14}\times$ real-world speed-ups for large-scale retrieval on ImageNet-1K and 4K; and (c) up to $\mathbf{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.

Matryoshka Representation Learning
Aditya Kusupati" Gantavya Bhatt" Aniket Rege",, Matthew Wallingford', Aditya Sinha', Vivek Ramanujan', William Howard-Snyder', Kaifeng Chen', Sham Kakade', Prateek Jain and Ali Farhadi!
University of Washington, "Google Research, *Harvard University (kusupati, acs washington. edu, praj aeggaii@@oogee.co
Abstract
Learned representations are central component in modern ML systems, serv- ing multitude of downstream tasks. When training such representations, it is often the case that computational and statistical constraints for each down- stream task are unknown. In this context, rigid fixed-capacity representations can be either over under-accommodaáing to the task at hand. This leads us to ask: can we design a flexible representation that can adapt to multiple down- stream tasks with varying computational resources? Our main contribution is Matryoshka Representation Learning (MRL) which encodes information different granularities and allows a single embedding to adapt to the computational constraints of downstream tasks. MRI. 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- 1 K classification at the same level of accuracy; (b) up to 14x real-world speed-ups for large-scale retrieval on ImageNet- IK and 4K; and (c) up to 2% accuracy improvements for long-tail tail few-shot classification, all while being as robust as the original representations. Finally, we show that MRL extends seam- lessly to web-scale datasets (ImageNet, JFT) across various modalities vision (ViT, ResNet), vision + language (ALIGN) and language (BERT). MRT code and pretrained models are open-sourced at https: // thub. com wwt.//oo//.......e Bt///RAAVIAAAAA/
Introduction
Learned representations [55] are fundamental building blocks of real-world ML systems [62, 86]. Trained once and frozen, d-dimensional representations encode rich information and can be used to perform multiple downstream tasks [4]. The deployment of deep representations has two steps: (1) an expensive yet constant-cost forward pass to compute the representation [27] and (2) utilization of the representation for downstream applications [48, 84]. Compute costs for the latter part of the pipeline scale with the embedding dimensionality as well as the data size (N) and label space (L). At web-scale [15, 80] this utilization cost overshadows the feature computation cost. The rigidity in these representations forces the use of high-dimensional embedding vectors across multiple tasks despite the varying resource and accuracy constraints that require flexibility.
Human perception the natural world has naturally coarse-to-fine granularity [26, 30]. However, perhaps due to the inductive bias of gradient-based training [79], deep learning models tend to diffuse "information" across the entire representation vector. The desired elasticity usually enabled the existing flat and fixed representations either through training multiple low-dimensional models [27]. jointly optimizing sub-networks of varying capacity [9, 95] or post-hoc compression [36, 57]. Each of these techniques struggle to meet the requirements for adaptive large-scale deployment either
Equal contribution AK led the project with extensive support from GB and AR for experimentation.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

due to training/maintenancee overhead, numerous expensive forward passes through all of the data, storage and memory cost for multiple copies of encoded data, expensive on-the-fly feature selection or significant drop in accuracy. By encoding coarse-to-fine seae--ttffne-griinee representations, which are as accurate as the independently trained counterparts, we learn with minimal overhead a representation that can be deployed adaptively at no additional cost during inference.
We introduce & Matryoshka Representation Learning (MRL) to induce flexibility in the learned representation. MRL learns representations of varying capacities within the same high-dimensional vector through explicit optimization of ((log(d)) lwwer-dimensional vectors in nested fashion, hence the name Matryoshka. MRL can be adapted to any existing representation pipeline and is easily extended to many standard tasks in computer vision and natural language processing. Figure illustrates the core idea of Matryoshka Roprosentation Learning (MRL) and the adaptive deployment settings of the learned Matryoshka Representations.
The first m-dimensions, m E [d]. of the Matryoshka Representation is an information-rich low-dimensional vector, at no additional training cost, that is as accurate as an indepen- dently trained m-dimensional repre- sentation. The information within the Matryoshka Representation in- creases with the dimensionality creat- ing coarse-to to-fine grained represen- tation, all without significant training or additional deployment overhead. MRL equips the representation vector with the desired flexibility and multi- fidelity that can ensure a near-optimal cesuurcccy---oppuue trade-off. With these advantages, MRL enables adap- tive deployment based on accuracy and compute constraints.
Inference
Rd
Training
Shortlistng
4(3)4)4)
Re-ranking
Adaptive Retrieval
D-2(2)
Adaptive Classification
O(Paa)
Matryoshka Representation Learning is adaptable to any representation learning setup and begets a Matryoshka Representation x by optimizing the orig- inal loss C(.) at (log(dd) chosen representation sizes. Matryoshka Representation can be utilized effectively for adap- tive deployment across environments and downstream tasks.
Figure 1:
The Matryoshka Representations improve efficiency for large-scale scale classification and retrieval without any significant loss of accuracy. While there are potentially several applications of coarse-to- fine Matryoshka Representations, in this work focus MI. systems: large-scale classification and retrieval. For classification, we use adaptive cascades with the variable-size representations from a model trained with MRL, significantly reducing the average dimension of embeddings needed to achieve a particular accuracy. For example, on ImageNet-1K: MRL+ adaptive classification results in up to a 14x smaller representation size the same accuracy as baselines (Section 4.2. Similarly, we use MRL in an adaptive retrieval system. Given a query, we shortlist retrieval candidates using the first few dimensions of the query embedding, and then successively use more dimensions to re-rank the retrieved set. A simple implementation of this approach leads to 128x theoretical (in terms FLOPS) and 14x wall-clock time speedups compared to single-shot retrieval system that uses a standard embedding vector; note that MRL's 's retrieval accuracy is comparable that of single-shot retrieval (Section 4.3.1). Finally, as MRT explicitly learns coarse-to-fine representation vectors, intuitively it should share more semantic information among its various dimensions (Figure 5). This is reflected in up to 2% accuracy gains in long-tail continual learning settings while being as robust as the original embeddings. Furthermore, due to its coarse-to- coase-to-fine grained nature, MRL can also be used as method to analyze hardness of classification among instances and information bottlenecks.
key building blocks of real-world
We make the following key contributions:
1. We introduce Matryoshka Representation Learning (MRL) to obtain flexible representa- tions (Matryoshka Representations) for adaptive deployment (Section 3).
2. Up 14x faster yet accurate large-scale classification and retrieval using MRI (Section 4).
3. Seamless adaptation of MRL across modalities (vision ResNet & ViT, vision + language ALIGN, language- BERT) and to web-scale data (ImageNet- K/4K, JFT-300M and ALIGN data).
4. Further analysis of MRL's representations in the context of other downstream tasks (Section 5).

---

2 Related Work
Representation Learning. Large-scale datasets like ImageNet [16, 71] and JFT [80] enabled the learning of general purpose representations for computer vision [4, 93]. These representations are typically learned through supervised and -upeevi-s learning paradigms. Supervised pretraining [27, 49, 77] casts representation learning as multi-class/labe classification problem, while un/self-superviseed learning learns representation via proxy tasks like instance classification [92] and reconstruction [29, 60]. Recent advances [12, 28] in contrastive learning [25] enabled learning from web-scale data [21] that powers large-capacity cross-modal models [18, 44, 67, 96]. Similarly, natural language applications are built [38] on large language models [8] that are pretrained [64, 70]
in un/self-supervised fashion with masked language modelling [19] or autoregressive training [66].
Matryoshka Representation Learning (MRL) is complementary to all these setups and can be adapted with minimal overhead (Section 3). MRL equips representations with multifidelity at no additional cost which enables adaptive deployment based on the data and task (Section 4).
Efficient Classification and Retrieval. Efficiency in classification and retrieval during inference can be studied with respect to the high yet constant deep featurization costs or the search cost which scales with the size of the label space and data. Efficient neural networks address the first issue through variety of algorithms [24, 52] and design choices [37, 51, 82]. However, with strong featurizer, most of the issues with scale are due to the linear dependence on number of labels (L), size of the data (N) and representation size (d), stressing RAM, disk and processor all at the same time.
The sub-linear complexity dependence on number of labels has been well studied in context of compute [3, 41, 65] and memory [20] using Approximate Nearest Neighbor Search (ANNS) [59] or leveraging the underlying hierarchy [17, 53]. In case of the representation size, often dimensionality reduction [72, 83], hashing techniques [14, 50, 73] and feature selection [61] help in alleviating selective aspects of the (d) scaling at a cost of significant drops in accuracy. Lastly, most real-world
search systems [11, 15] are often powered by large-scale embedding based retrieval [10, 62] that scales in cost with the ever increasing web-data. While categorization [84, 94] clusters similar things together, it is imperative to be equipped with retrieval capabilities that can bring forward every instance [7]. Approximate Nearest Neighbor Search (ANNS) [40] makes it feasible with efficient indexing [14] and traversal [5, 6] to present the users with the most similar documents/image: from the database for a requested query. Widely adopted HNSW [59] (o)log(NN) is as accurate as exact retrieval (O(d)) at the cost of a graph-based index overhead for RAM and disk [42].
MRL tackles the linear dependence on embedding size, d. by learning multifidelity Matryoshka Representations. Lower-dimenssonal Matryoshik Representations are as accu- rate as independently trained counterparts without the multiple expensive forward passes. Matryoshka Representations provide an intermediate abstraction between high-dimensional vec- tors and their efficient ANNS indices through the adaptive embeddings nested within the original representation vector (Section 4). All other aforementioned efficiency techniques are complementary and can be readily applied to the learned Matryoshka Representations obtained from MRL.
Several works in efficient neural network literature [9, 88, 95] aim at packing neural networks of varying capacity within the same larger network. However, the weights for each progressively smaller network can be different and often require distinct forward passes to isolate the final representations. This is detrimental for adaptive inference due to the need for re-encoding the entire retrieval database with expensive sub-net forward passes of varying capacities. Finally, ordered representations proposed by Rippel et al. [69] use nested dropout in the context of autoencoders to learn nested representations. MRL differentiates itself in formulation by optimizing only for (log(d)) nesting dimensions instead of Od). Despite this, MRL diffuses information to intermediate dimensions interpolating between the optimized Matryoshka Representation sizes accurately (Figure 5); making web-scale feasible.
3 Matryoshka Representation Learning
For d E N. consider set M [d] of representation sizes. For datapoint in the input do- main X, our goal is to learn a d-dimensional representation vector E IRd, For every m E M, Matsyo Representation Learning (MRT) enables each of the first m dimensions of the em- bedding vector, 21,m E Rd to be independently capable of being transferable and general purpose
3

---

representation of the datapoint 2. We obtain using a deep neural network F(::0p): X' - R parameterized by learnable weights 0r, i.e. 2 := F(x; fr). The multi-gran mualti-gaaullriityy is captured through the set of the chosen dimensions M, that contains less than log( (d) elements, i.e., |M] s [log(d)]. The usual set M consists of consistent halving until the representation size hits a low information bottleneck. We discuss the design choices in Section 4 for each of the representation learning settings.
For the ease of exposition, we present the formulation for fully supervised representation learning via multi-class classification. Matryoshka Representation Learning modifies the typical setting to become a multi-scale representation learning problem on the same task. For example. we train ResNet50 [27] on ImageNet- KK [71] which embeds a 224 x 224 pixel image into a d = 2048 representation vector and then passed through a linear classifier to make a prediction, 9 among the L = 1000 labels. For MRL, we choose M = {8, 16, 1024, 2048} as the nesting dimensions.
Suppose are given a lahelled dataset D { (1), ...................... (IN, ]N where J; E X' is an input point and y/ E [r.] is the label of x; for all i E [N] MRL optimizes the multi-class classification loss for each of the nested dimension m E M using standard empirieal risk minimization using separate linear classifier, parameterized by W(m) R I.x. All the losses are aggregated after scaling with their relative importance (cm 2 0) respectively. That is, we solve
min
Σ AE[N] MEM
Σ
Cm . - (w(m). F(x;0F)2mm) Hk)
(1)
(W(m))
mcM OF N
where L: 12 x [4] +R, + is the multi-class softmax cross-entropy loss function. This is a standard optimization problem that can be solved using sub-gradient descent methods. We set all the impor- lance scales, Cm = 1 for all m E M; see Section 5 for ablations. Lastly, despite only optimizing for (log(d)) nested dimensions, MRL results in accurate representations, that interpolate, for dimensions that fall between the chosen granularity of the representations (Section 4.2).
We call this formulation Matryoshka Representation Learning (MRL). A natural way make this efficient is through weight-tying across all the linear classifiers, i.e., by defining W(m) W I:m for a set of common weights W E 2xd This would reduce the memory cost due to the linear classifiers by almost half, which would be crucial in cases of extremely large output spaces [84, 94]. This variant is called Efficient Matryoshka Representation Learning (MRI-E). Refer to Alg 1 and Alg 2 in Appendix A for the building blocks of Matryoshka Representation Learning (MRL).
Adaptation to Learning Frameworks. MRL can be adapted seamlessly to most representation learning frameworks at web-scale with minimal modifications (Section 4.1). For example, MRL's s adaptation to masked language modelling reduces to MRI-F due to the weight-tying between the input embedding matrix and the linear classifier. For contrastive learning, both in context of vision & vision + language, MRL is applied to both the embeddings that are being contrasted with each other. The presence of normalization on the representation needs to be handled independently for each of the nesting dimension for best results (see Appendix C for more details).
Applications
In this section, we discuss Matryoshka Representation Learning (MRL) for a diverse set of ap- plications along with an extensive evaluation of the learned multifidelity representations. Further, we showcase the downstream applications of the learned Matryoshka Representations for flexible large-scale deployment through (a) Adaptive Classification (AC) and (b) Adaptive Retrieval (AR).
4.1 Representation Learning
We adapt Matryoshka Representation Learning (MRL) to various representation learning setups (a) Supervised learning for vision: ResNet50 [27] on ImageNet-1K IK [71] and ViT-B/16 [22] on JFT-300M [80], (b) Contrastive learning for vision + language: ALIGN model with V/T-B/16 vision encoder and BERT language encoder on ALIGN data [44] and (c) Masked language modelling: BERT [19] on English Wikipedia and BooksCorpus [97]. Please refer to Appendices B and C for details regarding the model architectures, datasets and training specifics.
We do not search for best peramm-prameters for all MRL experiments but use the same hyper- parameters as the independently trained baselines. ResNet50 outputs 2048 digendimelsional repre-

---

80
Top-1 11 Accuracy (%)
1-NN Accuracy (%)
70
70
MRL
MRL
MRL-E
60
MRL-E
60
FF
SVD
SVD
50
Slim. Net
Slim. Net
Rand. LP
Rand. FS
40
40
128
256
512
1024
2048
16
32
128
256
512
1024
2048
Representation Size
Representation Size
Figure 2: ImageNet- 1K linear classification ac- curacy of ResNet50 models. MRL is as accurate the independently trained FF models for every representation size.
Figure 3: ImageNet-IK 1-NN accuracy of ResNet50 models measuring the representation quality for downstream task. MRL outperforms all the baselines across all representation sizes.
sentation while ViT-B/16 and BERT-Base output 768- 768-dimensional embeddings for each data point. We use M {8, 16, 32, 64, 128, 256, 512, 1024, 2048) } and .M- {12, 24, 48, 96, 192, 384, 768} as the explicitly optimized nested dimensions respectively. Lastly, we extensively compare the MRI and MRI-E models to independently trained low-dimensional (fixed feature) representations (FF), dimensionality reduction (SVD), sub-net method (slimmable networks [95]) and randomly selected features of the highest capacity FF model.
In section 4.2, we evaluate the quality and capacity of the learned representations through linear classification/prohhe (LP) and I-nearest neighbour (1-NN) accuracy. Experiments show that MRL models remove the dependence on [M] resource-intensive independently trained models for the coarse-to-fine representations while being as accurate. Lastly, we show that despite optimizing only for JM| dimensions, MRL models diffuse the information, in an interpolative fashion, across all the dimensions providing the finest granularity required for adaptive deployment.
4.2 Classification
Figure 2 compares the linear classification accuracy of ResNet50 models trained and evaluated on ImageNet-IK. ResNet50-MRL, model is least accurate each FF model every rep- resentation size in M while MRL-E within 1% starting from 16-dim. Similarly, Figure showcases the comparison of learned representation quality through I-NN accuracy on ImageNet-IK (trainset with 1.3M samples as the database and validation set with 50K samples as the queries). Matryoshka Representations are up to 2% more accurate than their fixed-feature counterparts for the lower-dimensions while being as accurate elsewhere. I-NN accuracy is an excellent proxy, at no additional training cost, to gauge the utility of learned representations in the downstream tasks.
We also evaluate the quality the representations from training ViT-B/16 on JFT-300M alongside the ViT-B/16 vision encoder of the ALIGN model two web-scale setups. Due the expensive nature of these experiments, we only train the highest capacity fixed feature model and choose random features for evaluation in lower-dimensions. Web-scale is compelling setting for MRL due its relatively inexpensive training overhead while providing multifidelity representations for downstream tasks. Figure evaluated with I-NN on ImageNet- IK, shows that all the MRL models for JFT and ALIGN are highly accurate while providing an excellent cost-vs-accuracy trade-off at nower-dimensions. These experiments show that MRL seamlessly scales large-scale models and web-scale datasets while providing the otherwise prohibitively expensive multi- culli-graauulaityy in the process. We also have similar observations when pretraining BERT; please see Appendix D.2 for more details. Our experiments also show that post-hoc compression (SVD), linear probe on random features, and sub-net style slimmable networks drastically lose accuracy compared to MRL as the representation size decreases. Finally, Figure 5 shows that, while MRL explicitly optimizes ((og(d)) nested representations - removing the (dd) dependence [69]] -, the coarse-to-he grained information interpolated across all d dimensions providing highest flexibility for adaptive deployment.
5

---

80
1-NN Accuracy (%)
70
1-NN Accuracy (%)
JFT MRL
60
VIT-ALIGN
40
ALIGN MRL
VIT-JFT
RNSD-INIK
JFT MRL-E
50
VIT-ALIGN-Int
Rand.
VIT-JFT-ntt
ALIGN Rand.
RNSD-NNIK-Int
32
64
128
256
512
1024
2048
12
24
A8
192
384
768
Representation Size
Representation Size
Figure 4: ImageNet-IK I-NN accuracy for V/T-B/16 models trained on JFT-300M as part of ALIGN. MRL scales seamlessly to web-scale with minimal training overhead.
Figure 5: Despite optimizing MRL only for (log(d)) dimensions for ResNet50 and ViT- B/16 models; the accuracy in the intermediate dimensions shows interpolating behaviour.
4.2.1 Adaptive Classification
The flexibility and coarse-to- granularity within Matryoshka Representations allows model cascades [85] for Adaptive Classification (AC) [26]. Unlike standard model cascades [90], MRL does not require multiple expensive neural network forward passes. perform AC with an MRL trained model, we learn thresholds on the maximum softmax probability [31] for each nested classifier a holdout validation set. We then use these thresholds to decide when to transition to the higher dimensional representation (e.g 8 + 16 + 32) of the MRL model. Appendix D.I discusses the implementation and learning of thresholds for cascades used for adaptive classification in detail.
Figure 6 shows the comparison between cascaded MRL representations (MRL-AC) and indepen- dently trained fixed feature (FF) models on ImageNet-I IK with ResNet50. We computed the expected representation size for MRL-AC based on the final dimensionality used in the cascade. We observed that MRL-AC was as accurate, 76.30%, as a 512-dimensional
model but required an expected dimensionality of ~ 37 while being only 0.8% lower than the 2048- 0048-dimensional FF baseline. Note that all MRL-AC models are significantly more accurate than the FF baselines at comparable repre- sentation sizes. MRL-AC uses up to 14x smaller representation size for the same accuracy which affords computational efficiency as the label space grows [84]. Lastly, our results with MRL-AC indicate that instances and classes vary in difficulty which analyze in Section and Appendix J.
4.3 Retrieval
Nearest neighbour search with learned representations powers plethora of retrieval and search appli- cations [15, 86, 11, 62]. In this section, we discuss the image retrieval performance of the pretrained ResNet50 models (Section 4.1) on two large-scale datasets ImageNet- 1K [71] and ImageNet-4K. ImageNet-IK KK has database size of 1.3M and a query set of 50K samples uniformly spanning 1000 classes. We also introduce ImageNet-4K which has a database size of <..22 and query set of 200K samples uniformly spanning 4202 classes (see Appendix B for details), A single forward pass on ResNet50 costs GFLOPs while exact retrieval costs 2.6 GFLOPs per query for ImageNet- I Although retrieval overhead is 40% of the total cost, retrieval cost grows linearly with the size of the database. ImageNet-4K presents a retricval benchmark where the exact search cost becomes the computational bottleneck (8.6 GFLOPs per query). In both these settings, the memory and disk usage are also often bottlenecked by the large databases. However, in most real-world applications exact search, OddN), is replaced with approximate nearest neighbor search (ANNS) method like HNSW [59], (ddogg(N) with minimal accuracy drop at the cost of additional memory overhead.
The goal of image retrieval is to find images that belong to the same class as the query using representations obtained from a pretrained model. this section, we compare retrieval performance using mean Average Precision @ 10 (mAP@ 10) which comprehensively captures the setup of relevant image retrieval at scale. We measure the cost per query using exact search MFL.Ps. All embeddings are unit normalized and retrieved using the L2 distance metric. Lastly, report

---

Top-1 Accuracy (%)
1dx smaller reprosestalton
mAP@10 10 (%)
60
76
MRL
55
MRL-E
FF
75
MRL-AC
50
SVD
74
x
FF
FF 2048
45
Rand. FS
Sii. Net
40
16
32
64
128
256
512
30
128
256
512
1024
2048
(Expected) Representation Size
Representation Size
Figure Adaptive classification on MRL ResNer50 using cascades results in 14x smaller representation size for the same level of accuracy on ImageNet-1K IK (~ 37 vs 512 dims for 76.3%).
Figure 7: mAP@ 10 for Image Retrieval on ImageNet-IK with ResNet50. MRL consistently produces better retrieval performance over the baselines across all the representation sizes.
an extensive set of metrics spanning mAP@@ k and k for k = {10, 25, 50, 100} and real-world wall-clock times for exact search and HNSW. See Appendices and for more details.
Figure compares the mAP@ 10 performance of ResNet50 representations on ImageNet- IK across dimensionalities for MRL, MRI-E, FF, slimmable networks along with post-hoc compression of vectors using SVD and random feature selection. Matryoshka Representations are often the most accurate while being up 3% better than the FF baselines. Similar classification, post-hoc compression and slimmable network baselines suffer from significant drop-off in retrieval mAP 10 with < 256 dimensions. Appendix E discusses the mAP@ 10 of the models on ImageNet-4K.
MRL models are capable of performing accurate retrieval at various granularitics without the additional expense of multiple model forward passes for the web-scale databases. FF models also generate independent databases which become prohibitively expense to store and switch in between. Matryoshka Representations enable adaptive retrieval (AR) which alleviates the need to use full-capacity representations,
2048, for all data and downstream tasks. Lastly, all the vector compression techniques [57, 43] used as part of the ANNS pipelines complimentary to Matryoshka Representations and can further improve the efficiency-vs -i---s--acca vs-accuracy trade-off.
4.3.1 Adaptive Retrieval
We benchmark MRL in the adaptive retrieval setting (AR) [48]. For a given query image, we obtained a shortlist, K = 200, of images from the database using a lower-dimensional representation, e.g. D. 16 followed by reranking with higher capacity representation, e.g. D, 20448. In real-world scenarios where top ranking performance is the key objective, measured with mAP@A k where k covers limited yet crucial real-estate, AR provides significant compute and memory gains over single-shot retrieval with representations of fixed dimensionality. Finally, the most expensive part of AR, as with any retrieval pipeline, the nearest neighbour search for shortlisting. For example, even naive re-ranking of 200 images with 2048 dimensions only costs 400 KFLOPs. While we report exact search cost per query for all AR experiments, the shortlisting component of the pipeline can be sped-up using ANNS (HNSW). Appendix has detailed discussion compute cost for exact search, memory overhead of HNSW indices and wall-clock times for both implementations. We note that using HNSW with 32 neighbours for shortlisting does not decrease accuracy during retrieval.
Figure 8 showcases the compute-vs-accunt Matryoshka Representations compared to single-shot using fixed features with ResNet50 on ImageNet-IK. We observed that all AR settings lied above the Pareto frontier of single-shot retrieval with varying representation sizes. In particular for IK, we show that the AR model with D, 16 & D, 2048 is as accurate as single-shot retrieval with d 2048 while being ~ 128x more efficient in theory and ~ 14x faster in practice (compared using HNSW on the same hardware). We show similar trends with ImageNet-4K, but note that we require D. the increased difficulty of the dataset. This results in 32x and ~ 6x theoretical and in-practice speedups respectively. Lastly, while K 200 works well for our adaptive retrieval experiments, we ablated over the shortlist size & in Appendix K.2 and found that the accuracy gains stopped after
trade-off for adaptive retrieval using
given
7

---

65.3
128x theerstical speedrup
17.5
redbowrddd speedoup
mAP@11 (%)
65.2
2 dhanstica speed up
tortal wuris speed
17.0
mAP@10 (%)
65.1
0
16.5
32
65.0
16.0
Futral
Puseel
64.9
1334
10
10"
10"
10" MFLOPS/Query
10
MFLOPS/Query
(a) ImageNet-IK
(b) ImageNet-4K
Figure& 8: The trade-off between mAp@ 10 vs MFLOPs/Query for Adaptive Retrieval (AR) on ImageNet-IK K (left) and ImageNet-4K (right), Every combination of D, & D, falls above the Pareto line (orange dots) of single-shot retrieval with a fixed representation size while having configurations that are as accurate while being up to 14x faster in real-world deployment. Funnel retrieval is almost as accurate as the baseline while alleviating some of the parameter choices of Adaptive Retrieval.
point, further strengthening the use-case for Matryoshka Representation Learning and adaptive retrieval.
Even with adaptive retrieval, hard to determine the choice of D, & D,- In order to alleviate this issue to extent, we propose Funnel Retrieval, consistent cascade for adaptive retrieval. Funnel thins out the initial shortlist by a repeated re-ranking and shortlisting with a series of increasing capacity representations. Funnel halves the shortlist size and doubles the representation size at every step of re-ranking. For example on ImageNet- 1 a funnel with the shortlist progression of 200 -> 100 + 50 25 - 10 with the cascade of 16 representation sizes within Matryoshka Representation is as accurate as the single-shot 2048-dim retrieval while being 128x more efficient theoretically (see Appendix F for more results). All these results showcase the potential of MRL and AR for large-scale multi-stage search systems [15].
32 64 128 256 -> 2048
Further Analysis and Ablations
Robustness. We evaluate the robustness of the MRL models trained on ImageNet- IK on out-of- domain datasets, ImageNerV2/////SSeeee [68, 32, 33, 89], and compare them to the baselines. Table 17 in Appendix H demonstrates that Matryoshka Representations for classification are at least as robust as the original representation while improving the performance on ImageNet-A by 0.6% a 20% relative improvement. We also study the robustness in the context of retrieval by using Ie2 as the query set for ImageNet-I K database. Table 9 in Appendix shows that MRL models have more robust retrieval compared to the FF baselines by having up to 3% higher mAp 10 performance. This observation also suggests the need for further investigation into robustness using nearest neighbour based classification and retrieval instead of the standard linear probing setup. We also find that the zero-shot robustness of ALIGN-MRL (Table 18 in Appendix H) agrees with the observations made by Wortsman et al. [91]. Lastly, Table 6 in Appendix D.2 shows that MRL also improves the cosine similarity span between positive and random image-text pairs.
Few-shot and Long-tail Learning. We exhaustively evaluated few-shot learning on MRL models using nearest class mean [74]. Table 15 in Appendix G shows that that representations learned through MRL perform comparably to FF representations across varying shots and number of classes.
Matryoshka Representations realize a unique pattern while evaluating FLUID [87], a long-tail sequential learning framework. We observed that MRL provides up to 2% accuracy higher on novel classes in the tail of the distribution, without sacrificing accuracy on other classes (Table 16 in Appendix G). Additionally we find the accuracy between low-dimensional and high-dimensional representations marginal for pretrain classes. We hypothesize that the higher-dimensional represen- tations are required to differentiate the classes when few training examples of each are known. This results provides further evidence that different tasks require varying capacity based on their difficulty.
Disagreement across Dimensions. The information packing in Matryoshka Representation often results in gradual increase of accuracy with increase in capacity. However, we observed that

---

GT: Plastic Bag
Shower Cap
Plastic Bag
2048
GT: Rock Python
Boa Constrictor
Rock Python
e
16
32
2048
GT: Sweatshirt
Sunglasses
Sweatshirt
2048
Figure 9: Grad-CAM [75] progression of predictions in MRL model across 8, 16, 32 and 2048 dimensions. (a) 8-dimensional representation confuses due to presence of other relevant objects (with larger field of view) in the scene and predicts "shower cap"; (b) 8-dim model confuses within the same super-class of "boa" ; (c) 8 and 16-dim models incorrectly focus on the eyes of the doll ("sunglasses") and not the "sweatshirt" which is correctly in focus at higher dimensions; MRL fails gracefully these scenarios and shows potential use cases of disagreement across dimensions.
this trend was not ubiquitous and certain instances and classes were more accurate when evaluated with lower-dimensions (Figure 12 in Appendix J). With perfect routing of instances to appropriate dimension, MRL can gain up to 4.6% classification accuracy. At the same time, the oow-dimensional models are less accurate either due confusion within the same superclass [23] of the ImageNet hierarchy or presence of multiple objects of interest. Figure 9 showcases such examples for dimensional representation. These results along with Appendix put forward the potential for MRL to be a systematic framework for analyzing the utility and efficiency of information bottlenecks.
Superclass Accuracy. As the information bottleneck becomes smaller, the overall accuracy on fine-grained classes decreases rapidly (Figure 3). However, the drop-off is not significant when evaluated at superclass level (Table 24 in Appendix J). Figure 10 presents that this phenomenon
Top-1 Accuracy cy (%)
90
FF
90
88
85
86
masscring
bulling
84
75
Top-1
70
512
1024
2048
Representation Size
32
128
256
512
1024
2048
Figure 10: 31-way ImageNet- IK superclass clas- sification across representation size for MRL & FF models showing the capture of underlying hierarchy through tight information bottlenecks.
Representation Size
Figure 11: Diverse per-superclass accuracy trends across representation sizes for ResNet50- MRL on ImageNet-IK.

---

oceurs with both MRL and FF models; MRL is more accurate across dimensions. This shows that tight information bottlenecks while not highly accurate for fine-grained classification, do capture required semantic information for coarser classification that could be leveraged for adaptive routing for retrieval and classification. Mutifidelity of Matryoshka Representation naturally captures the underlying hierarchy of the class labels with one single model. Lastly, Figure 1 showcases the accuracy trends per superclass with MRL. The utility of additional dimensions in distinguishing class from others within the same superclass is evident for "garment" which has up to 11 1 improvement for 8 + 16 dimensional representation transition. We also observed that superclasses such as "oscine (songbird)" had a clear visual distinction between the object and background and thus predictions using 8 dimensions also led to a good inter-class separability within the superclass.
5.1 Ablations
Table 26 in Appendix K presents that Matryoshka Representations can be enabled within off-the- shelf pretrained models with inexpensive partial finetuning thus paving a way for ubiquitous adoption of MRL. At the same time, Table 27 in Appendix C indicates that with optimal weighting of the nested losses we could improve accuracy of dive--dmmnssoos representations without accuracy loss. Tables 28 and 29 in Appendix C ablate over the choice of initial granularity and spacing of the granularites. Table 28 reaffirms the design choice to shun extremely low dimensions that have poor classification accuracy as initial granularity for MRL. while Table 29 confirms the effectiveness of logarthmic granularity spacing inspired from the behaviour of accuracy saturation across dimensions over uniform. Lastly, Tables 30 and 31 in Appendix K.2 show that the retrieval performance saturates after a certain shortlist dimension and length depending on the complexity of the dataset.
6 Discussion and Conclusions
The results in Section 5.1 reveal interesting weaknesses of MRL that would be logical directions for future work. (I) Optimizing the weightings of the nested losses to obtain a Pareto optimal accuracy-vs-efficienn trade-off - a potential solution could emerge from adaptive loss balancing aspects of anytime neural networks [39]. (2) Using different losses at various fidelities aimed at solving a specific aspect of adaptive deployment c.g. high recall for 8-dimension and robustness for 2048-dimension. (3) Learning search data-structure, like differentiable k-d tree, on top of Matryoshka Representation to enable dataset and representation aware retrieval. (4) Finally, the joint optimization of multi-objective MRL combined with end-to-end learnable search data-structure to have data-driven adaptive large-scale retrieval for web-scale search applications.
In conclusion, we presented Matryoshka Representation Learning (MRL), a flexible represen- tation learning approach that encodes information at multiple granularities in single embedding vector. This enables the MRI to adapt to a downstream task's statistical complexity as well as the available compute resources. We demonstrate that MRL can be used for large-scale adaptive classification as well as adaptive retrieval. On standard benchmarks, MRL matches the accuracy of the fixed-feature baseline despite using 14x smaller representation size on average. Furthermore, the Matryoshka Representation based adaptive shortlisting and re-ranking system ensures comparable mAP@ 10 to the baseline while being 128x cheaper in FLOPs and 14x faster in wall-clock time. Finally, most of the efficiency techniques for model inference and vector search are complementary to MRL further assisting in deployment at the compute- compute-extreme eettute-extreme environments.
Acknowledgments
We are grateful to Srinadh Bhojanapalli, Lovish Madaan, Raghav Somani and Ludwig Schmidt for helpful discussions and feedback. Aditya Kusupati also thanks Tom Duerig and Rahul Sukthankar for their support. Part of the paper's 's large-scale experimentation is supported through a research GCP credit award from Google Cloud and Google Research. Gantavya Bhatt is supported in part by the CONIX Research Center, one of six centers in JUMP, Semiconductor Research Corporation (SRC) program sponsored by DARPA. Sham Kakade acknowledges funding from the NSF award CCF-1703574 and ONR X00014-22-1-2377 awards IIS 1652052, IIS 17303166, DARPA N66001-19-2-4031. DARPA W91 1 and gifts from Allen Institute for Artificial Intelligence.
Ali Farhadi acknowledges funding from the NSF
10

---

References
[1] M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, I. Goodfellow, A. Harp, G. Irving, M. Isard, Y. Jia, R. Joze- fowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, Sutskever, K. Talwar, P. Tucker, V. Vanhoucke, V. Vasudevan, F. Viégas, O. Vinyals, P. Warden, M. Wattenberg, M. Wicke, Y. Yu, and
X. Zheng. TensorFlow: Large-scale machine learning heterogeneous systems, 2015. URL https: hetps://www.. http::///wwwtttsoff..c... tensorflow
org/ Software available from tensorflow.org.
[2] A. Barbu, D. Mayo, J. Alverio, W. Luo, C. Wang, D. Gutfreund, J. Tenenbaum, and B. Katz. Objectnet: large-scale bias bias-controlled dataset for pushing the limits of object recognition models. Advances in neural information processing systems, 32, 2019.
[3] Bengio, Weston, and D. Grangier. Label embedding trees for large multi-class tasks. Advances in Neural Information Processing Systems, 23, 2010.
[4] Y. Bengio. Deep learning of representations for unsupervised and transfer learning. In Proceed- ings of ICML workshop on unsupervised and transfer learning, pages 17-36. JMLR Workshop and Conference Proceedings, 2012.
[5] J. L. Bentley. K-d trees for semidynamic point sets. In Proceedings of the sixth annual symposium on Computational geometry, pages 187-197, 1990.
[6] A. Beygelzimer, S. Kakade, and J. Langford. Cover trees for nearest neighbor. In Proceedings of the 23rd international conference on Machine learning, pages 97-104, 2006.
[7] S. Brin and L. Page. The anatomy of large-scale hypertextual web search engine. Computer networks and ISDN systems, 30(1-7):1071177 1998.
[8] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[9] H. Cai, C. Gan, T. Wang, Z. Zhang, and S. Han. Once-for-all: Train one network and specialize it for efficient deployment. arXiv preprint wrXiv:1908.0979) 2019.
[10] W.-C. Chang, X. Yu, Y.-W. Chang, Y. Yang, and S. Kumar. Pre-training tasks for embedding- based large-scale retrieval. arXiv preprint arXiv:2002.03932, 2020.
[11] W.-C Chang, D. Jiang, H.-F. Yu, C. H. Teo, J. Zhang. K. Zhong, K. Kolluri, Q. Hu, N. Shandilya, V. levgrafov, et al. Extreme multi-label learning for semantic matching in product search. In Proceedings of the 27th ACM SIGKDI Conference Knowledge Discovery & Data Mining, pages 2643-2651, 2021.
[12] T. Chen, S. Korbliihh, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607 PMLR, 2020.
[13) Y. Chen, Z. Liu, H. Xu, T. Darrell, and X. Wang. Meta-baseline: exploring simple meta-learning for few-shot learning. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 9062-9071, 2021.
[14] M. Datar, N. Immorlica. P. Indyk, and V. S. Mirrokni. Locality-sensse sensitive hashing scheme based on p-stable distributions. In Proceedings of the twentieth annual symposium on Computational geometry, pages 253-262, 2004.
[15] J. Dean. Challenges in building large-scale information retrieval systems. In Keynote of the 2nd ACM International Conference on Web Search and Data Mining (WSDM), volume 10, 2009.
[16] J. Deng, W. Dong, R. Socher, Li,J, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255, Icee, 2009.
11

---

[17] J. Deng, A. C. Berg, and L. Fei-Fei. Hierarchical semantic indexing for large scale image retrieval. In CVPR 2011, pages 785-792. IEEE, 2011.
[18] K. Desai and J. Johnson. Virtex: Learning visual representations from textual annotations. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 11162-11173, 2021.
[19] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 000000000000 2018.
[20] G. Dietterich and G. Bakiri. Solving multiclass learning problems via error-corecting output codes, Journal of artificial intelligence research, 2:263-286, 1994.
[21] S. K. Divvala, A. Farhadi, and C. Guestrin. Learning everything about anything: Webly- supervised visual concept learning. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition, pages 3270-3277, 2014.
[22] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenbom, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
[23] L. Engstrom, A. Iyyss, H. Salman, S. Santurkar, and D. Tsipras. Robustness (python library), 2019. URL https: // /gi github. hoa////ddrlaabbbouutteeeee
[24] A. Gholami, S. Kim, Z. Dong, Z. Yao, M. W. Mahoney, and K. Keutzer. A survey of quantization methods for efficient neural network inference. arXiv preprint arXiv:2103.13630 2021.
[25] M. Gutmann and A. Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 297-304. JMLR Workshop and Conference Proceedings, 2010.
[26] M. G. Harris and C. D. Giachritsis. Coarse-grained information dominates fine-grained in- formation in judgments of time-to-contact from retinal flow. Vision research, 4066)160 2000.
[27] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770- 778, 2016.
[28] K. Hc, H. Fan, Y. Wu, S. Xie, and R. Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 9729-9738, 2020.
[29] K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. arXiv preprint arXiv:21 1.06377, 2021.
[30] Hegdé. Time course of visual perception: coarse-to- to-fine processing and beyond. Progress in neurobiology, 84(4):405 2008.
[31] D. Hendrycks and K. Gimpel. A baseline for detecting misclassified and out-of-distribution examples in neural networks. arXiv preprint 1rXiv:1610.02136. 2016.
[32] D. Hendrycks, S. Basart, N. Mu, S. Kadavath, F. Wang, E. Dorundo, R. Desai, T. Zhu, S. Para- juli, M. Guo, et al. The many faces of robustness: A critical analysis of out-of-distribution generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 8340 8349, 2021.
[33] D. Hendrycks, K. Zhao. S. Basart, Steinhardt, and D. Song. Natural adversarial examples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15262-15271 2021.
[34] Hooker, A. Courville, G. Clark, Y. Dauphin, and A. Frome. What do compressed deep neural networks forget? arXiv preprint arXiv:1911.05244 2019.
12

---

[35] S. Hooker, N. Moorosi, G. Clark, S. Bengio, and E. Denton. Characterising bias in compressed models. arXiv preprint arXiv:2010.03058, 2020.
[36] H. Hotelling. Analysis of complex of statistical variables into principal components. Journal of educational psychology, 24(6):417, 1933.
[37] A. G. Howard, M. Zhu, B. Chen, D. Kalenichenko, W. Wang, T. Weyand, M. Andreetto, and H. Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint 1:Xiv:1704... 04861, 2017.
[38] Howard and S. Ruder. Universal language model fine-tuning for text classification. arXiv preprint 1x0iv:1001.06
2018.
[39] H. Hu, D. Dey. M. Hebert, and J. A. Bagnell. Learning anytime predictions in neural networks via adaptive loss balancing. Proceedings of the AAAI Conference on Artificial Intelligence volume 33, pages 3812-3821, 2019.
[40] P. Indyk and R. Motwani. Approximate nearest neighbors: towards removing the curse of dimensionality. In Proceedings of the thirtieth annual ACM symposium on Theory of computing. pages 604 613, 1998.
[41] H. Jain, V. Balasubramanian, B. Chunduri, and M. Varma. Slice: Scalable linear extreme classifiers trained on 100 million labels for related searches. In Proceedings of the Twelfth ACM International Conference on Web Search and Data Mining, pages 528-536 2019.
[42] S. Jayaram Subramanya, F. Devvrit, H. V. Simhadri, R. Krishnawamy, and R. Kadekodi. Diskann: Fast accurate billion-point nearest neighbor search on a single node. Advances in Neural Information Processing Systems, 32, 2019.
[43] H. Jegou, M. Douze, and C. Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128, 2010.
[44] C. Jia, Y. Yang, Y. Xia, Y.-T. Chen, Z. Parekh, H. Pham, Q. Le, Y.-H. Sung, Z. Li, and T. Duerig. Scaling up visual and vision-language representation learning with noisy text supervision. In International Conference on Machine Learning, pages 4904-4916. PMLR, 2021.
[45] Johnson, M. Douze, and H. Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535-544 043355557,45 200,547,2111 2019.
[46] W. B. Johnson. Extensions of lipschitz mappings into hilbert space. Contemp. Math. 26: 189-206, 1984.
[47] N. P. Jouppi, C. Young, N. Patil, D. Patterson, G. Agrawal, R. Bajwa, S. Bates, S. Bhatia, N. Boden, A. Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedingy of the 44th annual international symposium on computer architecture, pages 1-12, 2017.
[48] Kaz Sato. Vertex ai matching engine.
Microsoft Al Blog, 2021.
URL
https://clout google. www.blog//op//////oooooooooooooe http:/oppccs/eeelprrrsssss/ devedoeers pra.pppiorpprttttiieesrsssss/s/e find- anything latgeealllaanaaaaaee fast- fast-oooglss-sogsessse es vector search- technology..
[49] A. Krizhevsky, I. Sutskever, and G. E. Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
[50] B. Kulis, P. Jain, and K. Grauman. Fast similarity search for learned metrics. IEEE Transactions on Pattern Analysis and Machine Intelligence, (112))224332157) 2009.
[51] A. Kusupati, M. Singh, K. Bhatia, A. Kumar, P. Jain, and M. Varma. Fastgrnn: A fast, accurate, stable and tiny kilobyte sized gated recurrent neural network. Advances in Neural Information Processing Systems, 31, 2018.
[52] A. Kusupati, V. Ramanujan, R. Somani, M. Wortsman, P. Jain, S. Kakade, and A. Farhadi. Soft threshold weight reparameterization for learnable sparsity. In International Conference on Machine Learning, pages 5544-5555. PMLR, 2020.
13

---

[53] A. Kusupati. M. Wallingford, V. Ramanujan, R. Somani, J. S. Park, K. Pillutla, Jain, S. Kakade, and A. Farhadi. Llc: Accurate, multi-purpose learnt -iwdemmnsional binary codes. Advances in Neural Information Processing Systems, 34, 2021.
[54] G. Leclere, A. Ilyas, L. Engstrom, S. M. Park, H. Salman, and A. Madry. ffev. https: //gihub. com/ wtt.cooooo///eeee htt////fff/fffc////////// 2022. commit 607d117.
[55] Y. LeCun, Y. Bengio, and G. Hinton. Deep learning, nature, 521(7553))436-444 2015.
[56] S. Lee, S. Purushwalkam Shiva Prakash, M. Cogswell, V. Ranjan, D. Crandall, and D. Batra. Stochastic multiple choice learning for training diverse deep ensembles. Advances in Neural Information Processing Systems; 29, 2016.
[57] Y. Linde, A. Buzo, and R. Gray. An algorithm for vector quantizer design. IEEE Transactions on communications, 28( (8844)84 95, 1980.
[58] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05001, 2017.
[59] Y. A. Malkov and D. A. Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence, 42(())824-88) 836, 2018.
[60] J. Masci, U. Meier, D. Ciresan, and J. Schmidhuber. Stacked convolutional auto-encoders for hierarchical feature extraction. In International conference on artificial neural networks, pages 52-59. Springer, 2011.
[61] P. Mitra, C. Murthy, and S. K. Pal. Unsupervised feature selection using feature similarity. IEEE transactions on pattern analysis and machine intelligence, 24(3):30133) 2002.
[62] Nayak. Understanding searches better than ever before. Google Al Blog. 2019. URL https: / / //blog- google/ ppooll//rroducte pttp:ctsooccooeeeoooo httee//earrchreeeeeeeeeeeeeeeeee search wwwechh-annngeeeee wnwee-nannnetttnnnn--eee bert/.
[63] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[64] M. E. Peters, M. Neumann, M. lyyer, M. Gardner, C. Clark, K. Lee, and L. Zettlemoyer. Deep contextualized word representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long Papers), pages 2227-2237, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10.18653/v1/N 1/N/////1/111111 URL https: //aclanthology org/ /t//N18-1221
[65] Y. Prabhu, A. Kusupati, N. Gupta, and M. Varma. Extreme regression for dynamic search advertising. In Proceedings of the 13th International Conference on Web Search and Data Mining, pages 456-464, 2020.
[66] A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever. Improving language understand- ing by generative pre-training. OpenAl Blog, 2018. URL language- wns.gerisee..e...eeee.....
com/bomm/////o. blog/
[67] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pages 8748-8788 8763.876...7 PMLR, 2021.
[68] B. Recht, R. Roelofs, L. Schmidt, and V. Shankar. Do imagenet classifiers generalize to imagenet? In International Conference on Machine Learning, pages 5389-5400. PMLR, 2019.
[69] O. Rippel, M. Gelbart, and R. Adams. Learning ordered representations with nested dropout. In International Conference Machine Learning, pages 1746-1754. PMLR, 2014.
[70] S. Ruder, M. E. Peters, S. Swayamdipta, and T. Wolf. Transfer learning in natural language processing. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: Tutorials, pages 15-18, 2019.
14

---

[71] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. KKhola,, M. Bernstein, et Imagenet large scale visual recognition challenge. International journal of computer vision, 115(3):211-253) 2015.
[72] R. Salakhutdinov and G. Hinton. Learning nonlinear embedding by preserving class neigh- bourhood structure. In Artificial Intelligence and Statistics, pages 412 412-419. PMLR, 2007.
[73] R. Salakhutdinov and G. Hinton. Semantic hashing. International Journal of Approximate Reasoning, 50(7);969-778, 2009.
[74] S. Sánchez, Pla, and F. Ferri. On the use of neighbourhood based non- non-parametric classifiers. Pattern Recognition Letters, 18(11-13)1119--11888 1997.
[75] R. R. Selvaraju, M. Cogswell, A. Das, R. Vedantam, D. Parikh, and D. Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization. In Proceedings of the IEEE international conference computer vision, pages 618 1886626,222 2017.
[76] N. Shazeer and M. Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
[77] K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint 0rXiv:1409.1556 2014.
[78] L. N. Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter conference on applications of computer vision (WACV), pages 464-472. IEEE, 2017.
[79] D. Soudry, E. Hoffer, M. S. Nacson, Gunasekar, and N. Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1):2822-2878, 2018.
[80] C. Sun, A. Shrivastava, S. Singh, and A. Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In Proceedings of the IEEE international conference on computer vision, pages 843 852, 2017.
[81] Sutskever, J. Martens, G. Dahl, and G. Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning, pages 1139- 1 1147. PMLR, 2013.
[82] M. Tan and Q Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In International conference machine learning, pages 6105-6114. PMLR, 2019.
[83] L. Van Der Maaten, E. Postma, J. Van den Herik, al. Dimensionality reduction: comparative.
J Mach Learn Res, 10(66-71):13. 2009.
[84] M. Varma. Extreme classification. Communications of the ACM, 62(11):44 2019.
[85] P. Viola and M. Jones. Rapid object detection using boosted cascade of simple features. In Proceedings of the 2001 IEEE computer society conference on computer vision and pattern recognition. CVPR 2001, volume 1, pages I-I. leee, 2001.
[86] C. Waldburger. As search needs evolve, microsoft makes ai tools for better search available to researchers and developers. Microsoft Al Blog, 2019. URL https://blog:: microsoft. com/ htw/ai/bbi httpa////biggvecttoo-oeo
vector- search/.
[87] M. Wallingford, A. Kusupati, K. Alizadeh- A. Walsman, A. Kembhavi, and A. Farhadi. Are we overfitting experimental setups in recognition? arXiv preprint S00iv:2007.025110 2020.
[88] M. Wallingford, H. Li, A. Achille, A. Ravichandran, C. Fowlkes, R. Bhotika, and S. Soatto. Task adaptive parameter sharing for multi-task learning. arXiv preprint a7X.v.2003.16708 2022.
[89] H. Wang, S. Ge, Z. Lipton, and E. P. Xing. Learning robust global representations by penalizing local predictive power. In Advances in Neural Information Processing Systems, pages 10506- 10518 2019.
15

---

[90] X. Wang, D. Kondratyuk, K. M. Kitani, Y. Movshovitz-Atiaa, and E. Eban. Multiple networks are more efficient than one: Fast and accurate models via ensembles and cascades. arXiv preprini arXiv:2012.01988, 2020.
[91] M. Wortsman, G. Ilharco, M. Li, J. W. Kim, H. Hajishirzi, A. Farhadi, H. Namkoong, and
L. Schmidt. Robust tinettnning of zero-shot models. arXiv preprint arXiv:2109.01900 2021.
[92] Z. Wu, Y. Xiong. S. Yu, and D. Lin. Unsupervised feature learning via non-parametric instance- level discrimination. arXiv preprint / 01978, 2018.
[93] Yosinski, J. Clune, Y. Bengio, and H. Lipson. How transferable are features in deep neural networks? Advances in neural information processing systems, 27, 2014.
[94] H.-F. Yu, K. Zhong, J. Zhang, W.-C. Chang, and I. S. Dhillon. Pecos: Prediction for enormous and correlated output spaces. Journal of Machine Learning Research, 23(98): 2022.
[95] Yu, L. Yang. N. Xu, Yang, and T. Huang. Slimmable neural networks. arXiv preprint ar arXiv:1812.08928, 2018.
[96] R. Zellers, J. Lu, X. Lu, Y. Yu, Y. Zhao, M. Salehi, A. Kusupati, J. Hessel, A. Farhadi, and Choi. Merlot reserve: Neural script knowledge through vision and language and sound. arXiv preprint arXiv:2201. 02639,, 2022.
[97] Y. Zhu, R. Kiros, Zemel, R. Salakhutdinov, Urtasun, A. Torralba, and Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pages 19 27, 2015.
16

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you descrihe the limitations of your work? [Yes] See Section 6
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Our work does not have any additional negative societal impact on top of the existing impact of representation learning. However, study on the trade-off between representation size and the tendency encode biases an interesting future direction along the lines of existing literature [34, 35]. A part of this is already presented in Section 5.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed reproduce the main perimental results (either in the supplemental material or as a URL)? [Yes] See sup- plemental material and Appendix A. All the code and public models will be open sourced.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section and Appendix C.
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [No] We benchmarked on large-scale datasets like ImageNet- IK, JFT-300M and ALIGN data with models like ResNet and ViT making it extremely expensive to things multiple times.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix C and Appendix I.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license the assets? [No] All the non non-proprietary datasets and code used are public under MIT, BSD CC licenses.
(c) Did you include any new assets either in the supplemental material
URL? [Yes]
We created a new subset of ImageNet-2I1 2K for downstream evaluation of retrieval performance at scale. See Section 4.3 and Appendix B
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
17


Reviews:
[{"review_id": "jyVSDZCGFKj", "reviewer": "Reviewer_GFsE", "summary": "The authors propose a flexible representation learning approach MRL for an adaptive deployment on downstream tasks, where a single embedding can encode information at multiple granularities. An efficient variant MRL-E is also introduced with a set of common weights for nested dimensions. The coarse -to-fine representations are then applied to adaptive classification and retrieval tasks with various settings, where the method matches the performances of fixed-feature baselines. They provide as well further analysis in terms of robustness, few-shot learning and different dimensions.", "questions": "I am wondering if there exist any difficulties in training by introducing the nested loss.", "limitations": "I think the limitations mentioned in Section 6 are quite complete.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1)    The paper is well-written and quite easy for readers to follow the principal ideas, though switching between appendix and main paper from time to time troubles a bit.\n2)    The author provides a surprising number of experiments, covering different aspects of modalities, datasets and applications, which is quite impressive and complete.\n3)    Coming to the classification task, in Tables 1,2,4, for the high dimensions, it’s interesting to see that MRL sometimes cannot achieve the same results as FF. I am wondering if some explanations could be provided. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "QfDyZT6Va4c", "reviewer": "Reviewer_pjGy", "summary": "In this work, the authors propose Matryoshka Representations, a flexible coarse-to-fine representation that can be applied to multiple downstream tasks based on the computational budget. The authors demonstrate that utilizing the proposed Matryoshka Representations leads to a substantial reduction in embedding size for large-scale retrieval tasks. Finally, the authors have shown that the proposed representations can be incorporated into web-scale dataset experiments.     ", "questions": "- Have the authors experimented with arbitrary granularities? For instance, m random numbers can be generated up to the highest embedding size and these random numbers can be used for nesting dimensions. It will be interesting to see how the embeddings interpolate if the features are learned using such an arbitrary granularity.  \n\n- Have the authors tried lower than 8 dimensions?\n\n- Any insight into why tuning the relative importance weight of dimension 16 did not lead to improvement of dimension 32 whereas only increasing the weight of dimension 8 led to a noticeable improvement of dimension 16.", "limitations": "The limitations of the proposed work have been adequately addressed. However, no discussion about the potential negative soceital impact has been provided in the main text and supplementary materials. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The proposed solution is simple and intuitive.\n- The learned low-dimensional representations are as good as the models trained independently with that particular dimension without any additional computational overhead.\n- Significantly improves the efficiency of large-scale classification and retrieval tasks.\n- Allows for adaptive classification with variable embedding size.\n- Achieves promising improvement on long-tail continual learning experiments.\n- Matryoshka Representations can be easily extended to different modalities and even to web-scale datasets.\n\nWeaknesses:\n- No major weakness. The idea is simple and the results are compelling.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Of6FCavBFPp", "reviewer": "Reviewer_DT9D", "summary": "This paper proposes the so-called matryoshka representation learning, which can learns coarse-to-fine representations by modifying the regular learning pipeline. The matryoshka representation can encode information at different granularities, which achieves the same level of performance while being of much smaller sizes. The authors show the effectiveness of the matryoshka representations for image classification and retrieval tasks. Extensive ablation studies demonstrates the robustness and the ability of learning with limited data (few shots) and long-tailed data settings.", "questions": "1. It is unclear to me why the proposed algorithm can deal with out-of-distribution data and long-tailed data.", "limitations": "I did not see any discussion about the limitations of this work.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n1. it is an interesting idea to learn representations that can encode information of different granularities. \n\n2. The learned representations are deployed to different downstream tasks, e.g., classification and retrieval.\n\n3. The experiments and ablation studies are thorough and comprehensive.\n\nWeakness\n\n1. The idea is simple and straightforward. The major weakness of this work is the technical novelty, which makes the contribution a bit weak.\n\n2. The proposed algorithm does not show performance advantages when the representation size is large although it performs much better when size is very small. It seems the proposed algorithm also relies on a large representations to produce SOTA performance. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "2zGRRWU9vn", "reviewer": "Reviewer_HvJ7", "summary": "The paper introduces \"Matryoshka Representations\" which are a way to encode dimensions over multiple scales such that for a chosen embedding dimension d, there are log(d) embeddings of increasingly smaller size that can encode the same information as the original d-dimensional vector, with minimal loss in performance of the downstream task. The paper thoroughly validates MRs on an impressively large variety of tasks spanning over challenging and large scale datasets, and different modalities (vision+language, vision, language).", "questions": "Please see strengths/ weaknesses.", "limitations": "The authors do not explicitly discuss societal impacts of their work, but its unclear if the work has direct impacts. It would be useful for the community if the authors commented/showed if the lower dimensional representations are more or less susceptible of encoding biases in their representations, as its important to discuss such details. For practical purposes, is there a tradeoff between the dimensionality of the representations and the tendency to encode biases.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- the main strength of the paper is that it proposes a simple and practical solution for the problem, which can be used in many downstream tasks\n- The paper is written very clearly.\n- 14 x faster retrieval is a very significant improvement\n- The experiments are exhaustive and satisfying, for an approach that offers a lot of generality. \n\nWeaknesses:\n- Outside of increasing the speed for retrieval tasks, I am not sure what other practical applications of MRs are there. \n- I would like to see the authors discuss the research done on anytime neural networks[1]. They are not similar in methodology, but try to solve a very similar problem.\n- The improvement over FF is impressive, albeit unclear to me, why there is still improvements are higher dimensions (~64+). I would think/hope that at around 64 dim, the amount of information starts to become noisy. I would also \n- I would be very interested in these two baselines on any task: \n    - VIB [2] reformulated for this task, to see if its possible that there are any direct or indirect connections if the low-dim representations learned using VIB can be useful. \n    - Taking the full d-dim embedding and projecting it to a lower dim space using a random projection and seeing if that works. According to Gram-Schmidt, the low-dim representations should be usefully informative, and it was tried in [3] to some success. It can be a useful baseline for the paper as well. \n\n\n[1] https://arxiv.org/abs/1708.06832\n[2] https://arxiv.org/abs/1612.00410\n[3] https://arxiv.org/abs/1910.13540", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept