# Matryoshka Representation Learning

###### 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 {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ) which encodes information at different granularities and allows a single embedding to adapt to the computational constraints of downstream tasks. MRL {\rm MRL} minimally modifies existing representation learning pipelines and imposes no additional cost during inference and deployment. MRL {\rm 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 {\rm 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 {\rm MRL} extends seamlessly to web-scale datasets (ImageNet, JFT) across various modalities – vision (ViT, ResNet), vision + language (ALIGN) and language (BERT). MRL {\rm MRL} code and pretrained models are open-sourced at https://github.com/RAIVNLab/MRL .

## 1 Introduction

Learned representations [ 57 ] are fundamental building blocks of real-world ML systems [ 66 , 91 ] . Trained once and frozen, d 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 [ 29 ] and (2) utilization of the representation for downstream applications [ 50 , 89 ] . Compute costs for the latter part of the pipeline scale with the embedding dimensionality as well as the data size ( N N ) and label space ( L L ). At web-scale [ 15 , 85 ] 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 of the natural world has a naturally coarse-to-fine granularity [ 28 , 32 ] . However, perhaps due to the inductive bias of gradient-based training [ 84 ] , deep learning models tend to diffuse “information” across the entire representation vector. The desired elasticity is usually enabled in the existing flat and fixed representations either through training multiple low-dimensional models [ 29 ] , jointly optimizing sub-networks of varying capacity [ 9 , 100 ] or post-hoc compression [ 38 , 60 ] . Each of these techniques struggle to meet the requirements for adaptive large-scale deployment either due to training/maintenance 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 a significant drop in accuracy. By encoding coarse-to-fine-grained 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 {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ) to induce flexibility in the learned representation. MRL {\rm MRL} learns representations of varying capacities within the same high-dimensional vector through explicit optimization of O ⁡ ( log ⁡ ( d ) ) O(\log(d)) lower-dimensional vectors in a nested fashion, hence the name Matryoshka {\rm Matryoshka} . MRL {\rm 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 1 illustrates the core idea of Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ) and the adaptive deployment settings of the learned Matryoshka ​ Representations {\rm Matryoshka~Representations} .

The first m m -dimensions, m ∈ [ d ] m\in[d] , of the Matryoshka ​ Representation {\rm Matryoshka~Representation} is an information-rich low-dimensional vector, at no additional training cost, that is as accurate as an independently trained m m -dimensional representation. The information within the Matryoshka ​ Representation {\rm Matryoshka~Representation} increases with the dimensionality creating a coarse-to-fine grained representation, all without significant training or additional deployment overhead. MRL {\rm MRL} equips the representation vector with the desired flexibility and multifidelity that can ensure a near-optimal accuracy-vs-compute trade-off. With these advantages, MRL {\rm MRL} enables adaptive deployment based on accuracy and compute constraints.

The Matryoshka ​ Representations {\rm Matryoshka~Representations} improve efficiency for large-scale classification and retrieval without any significant loss of accuracy. While there are potentially several applications of coarse-to-fine Matryoshka ​ Representations {\rm Matryoshka~Representations} , in this work we focus on two key building blocks of real-world ML systems: large-scale classification and retrieval. For classification, we use adaptive cascades with the variable-size representations from a model trained with MRL {\rm MRL} , significantly reducing the average dimension of embeddings needed to achieve a particular accuracy. For example, on ImageNet-1K, MRL {\rm MRL} + adaptive classification results in up to a 14 × 14\times smaller representation size at the same accuracy as baselines (Section 4.2.1 ). Similarly, we use MRL {\rm 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 128 × 128\times theoretical (in terms of FLOPS) and 14 × 14\times wall-clock time speedups compared to a single-shot retrieval system that uses a standard embedding vector; note that MRL {\rm MRL} ’s retrieval accuracy is comparable to that of single-shot retrieval (Section 4.3.1 ). Finally, as MRL {\rm MRL} 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 % 2\% accuracy gains in long-tail continual learning settings while being as robust as the original embeddings. Furthermore, due to its coarse-to-fine grained nature, MRL {\rm MRL} can also be used as method to analyze hardness of classification among instances and information bottlenecks.

We make the following key contributions: 1. We introduce Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ) to obtain flexible representations ( Matryoshka ​ Representations {\rm Matryoshka~Representations} ) for adaptive deployment (Section 3 ).

2. Up to 14 × 14\times faster yet accurate large-scale classification and retrieval using MRL {\rm MRL} (Section 4 ).

3. Seamless adaptation of MRL {\rm MRL} across modalities (vision - ResNet & ViT, vision + language - ALIGN, language - BERT) and to web-scale data (ImageNet-1K/4K, JFT-300M and ALIGN data).

4. Further analysis of MRL {\rm MRL} ’s representations in the context of other downstream tasks (Section 5 ).

## 2 Related Work

##### Representation Learning.

Large-scale datasets like ImageNet [ 16 , 76 ] and JFT [ 85 ] enabled the learning of general purpose representations for computer vision [ 4 , 98 ] . These representations are typically learned through supervised and un/self-supervised learning paradigms. Supervised pretraining [ 29 , 51 , 82 ] casts representation learning as a multi-class/label classification problem, while un/self-supervised learning learns representation via proxy tasks like instance classification [ 97 ] and reconstruction [ 31 , 63 ] . Recent advances [ 12 , 30 ] in contrastive learning [ 27 ] enabled learning from web-scale data [ 21 ] that powers large-capacity cross-modal models [ 18 , 46 , 71 , 101 ] . Similarly, natural language applications are built [ 40 ] on large language models [ 8 ] that are pretrained [ 68 , 75 ] in a un/self-supervised fashion with masked language modelling [ 19 ] or autoregressive training [ 70 ] .

Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ) is complementary to all these setups and can be adapted with minimal overhead (Section 3 ). MRL {\rm 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 a variety of algorithms [ 25 , 54 ] and design choices [ 39 , 53 , 87 ] . However, with a strong featurizer, most of the issues with scale are due to the linear dependence on number of labels ( L L ), size of the data ( N N ) and representation size ( d 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 , 43 , 69 ] and memory [ 20 ] using Approximate Nearest Neighbor Search (ANNS) [ 62 ] or leveraging the underlying hierarchy [ 17 , 55 ] . In case of the representation size, often dimensionality reduction [ 77 , 88 ] , hashing techniques [ 14 , 52 , 78 ] and feature selection [ 64 ] help in alleviating selective aspects of the O ⁡ ( d ) O(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 , 66 ] that scales in cost with the ever increasing web-data. While categorization [ 89 , 99 ] 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) [ 42 ] makes it feasible with efficient indexing [ 14 ] and traversal [ 5 , 6 ] to present the users with the most similar documents/images from the database for a requested query. Widely adopted HNSW [ 62 ] ( O ⁡ ( d ​ log ⁡ ( N ) ) O(d\log(N)) ) is as accurate as exact retrieval ( O ⁡ ( d ​ N ) O(dN) ) at the cost of a graph-based index overhead for RAM and disk [ 44 ] .

MRL {\rm MRL} tackles the linear dependence on embedding size, d d , by learning multifidelity Matryoshka ​ Representations {\rm Matryoshka~Representations} . Lower-dimensional Matryoshka ​ Representations {\rm Matryoshka~Representations} are as accurate as independently trained counterparts without the multiple expensive forward passes. Matryoshka ​ Representations {\rm Matryoshka~Representations} provide an intermediate abstraction between high-dimensional vectors 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 {\rm Matryoshka~Representations} obtained from MRL {\rm MRL} .

Several works in efficient neural network literature [ 9 , 93 , 100 ] 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. Several works [ 23 , 26 , 65 , 59 ] investigate the notions of intrinsic dimensionality and redundancy of representations and objective spaces pointing to minimum description length [ 74 ] . Finally, ordered representations proposed by Rippel et al. [73] use nested dropout in the context of autoencoders to learn nested representations. MRL {\rm MRL} differentiates itself in formulation by optimizing only for O ⁡ ( log ⁡ ( d ) ) O(\log(d)) nesting dimensions instead of O ⁡ ( d ) O(d) . Despite this, MRL {\rm MRL} diffuses information to intermediate dimensions interpolating between the optimized Matryoshka ​ Representation {\rm Matryoshka~Representation} sizes accurately (Figure 5 ); making web-scale feasible.

## 3 Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning}

For d ∈ ℕ d\in\mathbb{N} , consider a set ℳ ⊂ [ d ] \mathcal{M}\subset[d] of representation sizes. For a datapoint x x in the input domain 𝒳 \mathcal{X} , our goal is to learn a d d -dimensional representation vector z ∈ ℝ d z\in\mathbb{R}^{d} . For every m ∈ ℳ m\in\mathcal{M} , Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ) enables each of the first m m dimensions of the embedding vector, z 1 : m ∈ ℝ m z_{1:m}\in\mathbb{R}^{m} to be independently capable of being a transferable and general purpose representation of the datapoint x x . We obtain z z using a deep neural network F ⁡ ( ⋅ , θ F ) : 𝒳 → ℝ d F(\,\cdot\,;\theta_{F})\colon\mathcal{X}\rightarrow\mathbb{R}^{d} parameterized by learnable weights θ F \theta_{F} , i.e., z ≔ F ⁡ ( x , θ F ) z\coloneqq F(x;\theta_{F}) . The multi-granularity is captured through the set of the chosen dimensions ℳ \mathcal{M} , that contains less than log ⁡ ( d ) \log(d) elements, i.e., | ℳ | ≤ ⌊ log ⁡ ( d ) ⌋ \lvert\mathcal{M}\rvert\leq\left\lfloor\log(d)\right\rfloor . The usual set ℳ \mathcal{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 {\rm Matryoshka~Representation~Learning} modifies the typical setting to become a multi-scale representation learning problem on the same task. For example, we train ResNet50 [ 29 ] on ImageNet-1K [ 76 ] which embeds a 224 × 224 224\times 224 pixel image into a d = 2048 d=2048 representation vector and then passed through a linear classifier to make a prediction, y ^ \hat{y} among the L = 1000 L=1000 labels. For MRL {\rm MRL} , we choose ℳ = { 8 , 16 , … , 1024 , 2048 } \mathcal{M}=\{8,16,\ldots,1024,2048\} as the nesting dimensions.

Suppose we are given a labelled dataset 𝒟 = { ( x 1 , y 1 ) , … , ( x N , y N ) } \mathcal{D}=\{(x_{1},y_{1}),\ldots,(x_{N},y_{N})\} where x i ∈ 𝒳 x_{i}\in\mathcal{X} is an input point and y i ∈ [ L ] y_{i}\in[L] is the label of x i x_{i} for all i ∈ [ N ] i\in[N] . MRL {\rm MRL} optimizes the multi-class classification loss for each of the nested dimension m ∈ ℳ m\in\mathcal{M} using standard empirical risk minimization using a separate linear classifier, parameterized by 𝐖 ( m ) ∈ ℝ L × m \mathbf{W}^{(m)}\in\mathbb{R}^{L\times m} . All the losses are aggregated after scaling with their relative importance ( c m ≥ 0 ) m ∈ ℳ \left(c_{m}\geq 0\right)_{m\in\mathcal{M}} respectively. That is, we solve min { 𝐖 ( m ) } m ∈ ℳ , θ F 1 N ∑ i ∈ [ N ] ∑ m ∈ ℳ c m ⋅ ℒ ( 𝐖 ( m ) ⋅ F ( x i ; θ F ) 1 : m ; y i ) , \min_{\left\{\mathbf{W}^{(m)}\right\}_{m\in\mathcal{M}},\ \theta_{F}}\frac{1}{N}\sum_{i\in[N]}\sum_{m\in\mathcal{M}}c_{m}\cdot{\cal L}\left(\mathbf{W}^{(m)}\cdot F(x_{i};\theta_{F})_{1:m}\ ;\ y_{i}\right)\ , (1) where ℒ : ℝ L × [ L ] → ℝ + {\cal L}\colon\mathbb{R}^{L}\times[L]\to\mathbb{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 importance scales, c m = 1 c_{m}=1 for all m ∈ ℳ m\in\mathcal{M} ; see Section 5 for ablations. Lastly, despite only optimizing for O ⁡ ( log ⁡ ( d ) ) O(\log(d)) nested dimensions, MRL {\rm 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 as Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ). A natural way to make this efficient is through weight-tying across all the linear classifiers, i.e., by defining 𝐖 ( m ) = 𝐖 1 : m \mathbf{W}^{(m)}=\mathbf{W}_{1:m} for a set of common weights 𝐖 ∈ ℝ L × d \mathbf{W}\in\mathbb{R}^{L\times d} . 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 [ 89 , 99 ] . This variant is called Efficient Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL ​ – ​ E {\rm MRL\text{--}E} ). Refer to Alg 1 and Alg 2 in Appendix A for the building blocks of Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ).

##### Adaptation to Learning Frameworks.

MRL {\rm MRL} can be adapted seamlessly to most representation learning frameworks at web-scale with minimal modifications (Section 4.1 ). For example, MRL {\rm MRL} ’s adaptation to masked language modelling reduces to MRL ​ – ​ E {\rm MRL\text{--}E} 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 {\rm 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).

## 4 Applications

In this section, we discuss Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ) for a diverse set of applications along with an extensive evaluation of the learned multifidelity representations. Further, we showcase the downstream applications of the learned Matryoshka ​ Representations {\rm 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 {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ) to various representation learning setups (a) Supervised learning for vision: ResNet50 [ 29 ] on ImageNet-1K [ 76 ] and ViT-B/16 [ 22 ] on JFT-300M [ 85 ] , (b) Contrastive learning for vision + language: ALIGN model with ViT-B/16 vision encoder and BERT language encoder on ALIGN data [ 46 ] and (c) Masked language modelling: BERT [ 19 ] on English Wikipedia and BooksCorpus [ 102 ] . Please refer to Appendices B and C for details regarding the model architectures, datasets and training specifics.

We do not search for best hyper-parameters for all MRL {\rm MRL} experiments but use the same hyper-parameters as the independently trained baselines. ResNet50 outputs a 2048 2048 -dimensional representation while ViT-B/16 and BERT-Base output 768 768 -dimensional embeddings for each data point. We use ℳ = { 8 , 16 , 32 , 64 , 128 , 256 , 512 , 1024 , 2048 } \mathcal{M}=\{8,16,32,64,128,256,512,1024,2048\} and ℳ = { 12 , 24 , 48 , 96 , 192 , 384 , 768 } \mathcal{M}=\{12,24,48,96,192,384,768\} as the explicitly optimized nested dimensions respectively. Lastly, we extensively compare the MRL {\rm MRL} and MRL ​ – ​ E {\rm MRL\text{--}E} models to independently trained low-dimensional (fixed feature) representations (FF), dimensionality reduction (SVD), sub-net method (slimmable networks [ 100 ] ) 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/probe (LP) and 1-nearest neighbour (1-NN) accuracy. Experiments show that MRL {\rm MRL} models remove the dependence on | ℳ | |\mathcal{M}| resource-intensive independently trained models for the coarse-to-fine representations while being as accurate. Lastly, we show that despite optimizing only for | ℳ | |\mathcal{M}| dimensions, MRL {\rm MRL} models diffuse the information, in an interpolative fashion, across all the d d dimensions providing the finest granularity required for adaptive deployment.

### 4.2 Classification

Figure 3 compares the linear classification accuracy of ResNet50 models trained and evaluated on ImageNet-1K. ResNet50– MRL {\rm MRL} model is at least as accurate as each FF model at every representation size in ℳ \mathcal{M} while MRL ​ – ​ E {\rm MRL\text{--}E} is within 1 % 1\% starting from 16 16 -dim. Similarly, Figure 3 showcases the comparison of learned representation quality through 1-NN accuracy on ImageNet-1K (trainset with 1.3M samples as the database and validation set with 50K samples as the queries). Matryoshka ​ Representations {\rm Matryoshka~Representations} are up to 2 % 2\% more accurate than their fixed-feature counterparts for the lower-dimensions while being as accurate elsewhere. 1-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 of 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 to 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 a compelling setting for MRL {\rm MRL} due to its relatively inexpensive training overhead while providing multifidelity representations for downstream tasks. Figure 5 , evaluated with 1-NN on ImageNet-1K, shows that all the MRL {\rm MRL} models for JFT and ALIGN are highly accurate while providing an excellent cost-vs-accuracy trade-off at lower-dimensions. These experiments show that MRL {\rm MRL} seamlessly scales to large-scale models and web-scale datasets while providing the otherwise prohibitively expensive multi-granularity 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 {\rm MRL} as the representation size decreases. Finally, Figure 5 shows that, while MRL {\rm MRL} explicitly optimizes O ⁡ ( log ⁡ ( d ) ) O(\log(d)) nested representations – removing the O ⁡ ( d ) O(d) dependence [ 73 ] –, the coarse-to-fine grained information is interpolated across all d d dimensions providing highest flexibility for adaptive deployment.

#### 4.2.1 Adaptive Classification

The flexibility and coarse-to-fine granularity within Matryoshka ​ Representations {\rm Matryoshka~Representations} allows model cascades [ 90 ] for Adaptive Classification (AC) [ 28 ] . Unlike standard model cascades [ 95 ] , MRL {\rm MRL} does not require multiple expensive neural network forward passes. To perform AC with an MRL {\rm MRL} trained model, we learn thresholds on the maximum softmax probability [ 33 ] for each nested classifier on a holdout validation set. We then use these thresholds to decide when to transition to the higher dimensional representation (e.g 8 → 16 → 32 8\to 16\to 32 ) of the MRL {\rm MRL} model. Appendix D.1 discusses the implementation and learning of thresholds for cascades used for adaptive classification in detail.

Figure 7 shows the comparison between cascaded MRL {\rm MRL} representations ( MRL {\rm MRL} –AC) and independently trained fixed feature (FF) models on ImageNet-1K with ResNet50. We computed the expected representation size for MRL {\rm MRL} –AC based on the final dimensionality used in the cascade. We observed that MRL {\rm MRL} –AC was as accurate, 76.30 % 76.30\% , as a 512-dimensional FF model but required an expected dimensionality of ∼ 37 \sim 37 while being only 0.8 % 0.8\% lower than the 2048-dimensional FF baseline. Note that all MRL {\rm MRL} –AC models are significantly more accurate than the FF baselines at comparable representation sizes. MRL {\rm MRL} –AC uses up to ∼ 14 × \sim 14\times smaller representation size for the same accuracy which affords computational efficiency as the label space grows [ 89 ] . Lastly, our results with MRL {\rm MRL} –AC indicate that instances and classes vary in difficulty which we analyze in Section 5 and Appendix J .

### 4.3 Retrieval

Nearest neighbour search with learned representations powers a plethora of retrieval and search applications [ 15 , 91 , 11 , 66 ] . In this section, we discuss the image retrieval performance of the pretrained ResNet50 models (Section 4.1 ) on two large-scale datasets ImageNet-1K [ 76 ] and ImageNet-4K. ImageNet-1K has a database size of ∼ \sim 1.3M and a query set of 50K samples uniformly spanning 1000 classes. We also introduce ImageNet-4K which has a database size of ∼ \sim 4.2M and query set of ∼ \sim 200K samples uniformly spanning 4202 classes (see Appendix B for details). A single forward pass on ResNet50 costs 4 GFLOPs while exact retrieval costs 2.6 GFLOPs per query for ImageNet-1K. Although retrieval overhead is 40 % 40\% of the total cost, retrieval cost grows linearly with the size of the database. ImageNet-4K presents a retrieval benchmark where the exact search cost becomes the computational bottleneck ( 8.6 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, O ⁡ ( d ​ N ) O(dN) , is replaced with an approximate nearest neighbor search (ANNS) method like HNSW [ 62 ] , O ⁡ ( d ​ log ⁡ ( N ) ) O(d\log(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. In this section, we compare retrieval performance using mean Average Precision @ 10 (mAP@ 10 10 ) which comprehensively captures the setup of relevant image retrieval at scale. We measure the cost per query using exact search in MFLOPs. All embeddings are unit normalized and retrieved using the L2 distance metric. Lastly, we report an extensive set of metrics spanning mAP@ k k and P@ k k for k = { 10 , 25 , 50 , 100 } k=\{10,25,50,100\} and real-world wall-clock times for exact search and HNSW. See Appendices E and F for more details.

Figure 7 compares the mAP@ 10 10 performance of ResNet50 representations on ImageNet-1K across dimensionalities for MRL {\rm MRL} , MRL ​ – ​ E {\rm MRL\text{--}E} , FF, slimmable networks along with post-hoc compression of vectors using SVD and random feature selection. Matryoshka ​ Representations {\rm Matryoshka~Representations} are often the most accurate while being up to 3 % 3\% better than the FF baselines. Similar to classification, post-hoc compression and slimmable network baselines suffer from significant drop-off in retrieval mAP@ 10 10 with ≤ 256 \leq 256 dimensions. Appendix E discusses the mAP@ 10 10 of the same models on ImageNet-4K.

MRL {\rm MRL} models are capable of performing accurate retrieval at various granularities 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 {\rm Matryoshka~Representations} enable adaptive retrieval (AR) which alleviates the need to use full-capacity representations, d = 2048 d=2048 , for all data and downstream tasks. Lastly, all the vector compression techniques [ 60 , 45 ] used as part of the ANNS pipelines are complimentary to Matryoshka ​ Representations {\rm Matryoshka~Representations} and can further improve the efficiency-vs-accuracy trade-off.

#### 4.3.1 Adaptive Retrieval

We benchmark MRL {\rm MRL} in the adaptive retrieval setting (AR) [ 50 ] . For a given query image, we obtained a shortlist, K = 200 K=200 , of images from the database using a lower-dimensional representation, e.g. D s = 16 D_{s}=16 followed by reranking with a higher capacity representation, e.g. D r = 2048 D_{r}=2048 . In real-world scenarios where top ranking performance is the key objective, measured with mAP@ k k where k covers a 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, is 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 I has a detailed discussion on 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-accuracy trade-off for adaptive retrieval using Matryoshka ​ Representations {\rm Matryoshka~Representations} compared to single-shot using fixed features with ResNet50 on ImageNet-1K. We observed that all AR settings lied above the Pareto frontier of single-shot retrieval with varying representation sizes. In particular for ImageNet-1K, we show that the AR model with D s = 16 D_{s}=16 & D r = 2048 D_{r}=2048 is as accurate as single-shot retrieval with d = 2048 d=2048 while being ∼ 𝟏𝟐𝟖 × \mathbf{\sim 128\times} more efficient in theory and ∼ 𝟏𝟒 × \mathbf{\sim 14\times} faster in practice (compared using HNSW on the same hardware). We show similar trends with ImageNet-4K, but note that we require D s = 64 D_{s}=64 given the increased difficulty of the dataset. This results in ∼ 32 × \sim 32\times and ∼ 6 × \sim 6\times theoretical and in-practice speedups respectively. Lastly, while K = 200 K=200 works well for our adaptive retrieval experiments, we ablated over the shortlist size k k in Appendix K.2 and found that the accuracy gains stopped after a point, further strengthening the use-case for Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} and adaptive retrieval.

Even with adaptive retrieval, it is hard to determine the choice of D s D_{s} & D r D_{r} . In order to alleviate this issue to an extent, we propose Funnel Retrieval , a 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-1K, a funnel with the shortlist progression of 200 → 100 → 50 → 25 → 10 200\to 100\to 50\to 25\to 10 with the cascade of 16 → 32 → 64 → 128 → 256 → 2048 16\to 32\to 64\to 128\to 256\to 2048 representation sizes within Matryoshka ​ Representation {\rm Matryoshka~Representation} is as accurate as the single-shot 2048-dim retrieval while being ∼ 128 × \sim 128\times more efficient theoretically (see Appendix F for more results). All these results showcase the potential of MRL {\rm MRL} and AR for large-scale multi-stage search systems [ 15 ] .

## 5 Further Analysis and Ablations

##### Robustness.

We evaluate the robustness of the MRL {\rm MRL} models trained on ImageNet-1K on out-of-domain datasets, ImageNetV2/R/A/Sketch [ 72 , 34 , 35 , 94 ] , and compare them to the FF baselines. Table 17 in Appendix H demonstrates that Matryoshka ​ Representations {\rm Matryoshka~Representations} for classification are at least as robust as the original representation while improving the performance on ImageNet-A by 0.6 % 0.6\% – a 20 % 20\% relative improvement. We also study the robustness in the context of retrieval by using ImageNetV2 as the query set for ImageNet-1K database. Table 9 in Appendix E shows that MRL {\rm MRL} models have more robust retrieval compared to the FF baselines by having up to 3 % 3\% higher mAP@ 10 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 {\rm MRL} (Table 18 in Appendix H ) agrees with the observations made by Wortsman et al. [96] . Lastly, Table 6 in Appendix D.2 shows that MRL {\rm 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 {\rm MRL} models using nearest class mean [ 79 ] . Table 15 in Appendix G shows that that representations learned through MRL {\rm MRL} perform comparably to FF representations across varying shots and number of classes.

Matryoshka ​ Representations {\rm Matryoshka~Representations} realize a unique pattern while evaluating on FLUID [ 92 ] , a long-tail sequential learning framework. We observed that MRL {\rm MRL} provides up to 2 % 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 is marginal for pretrain classes. We hypothesize that the higher-dimensional representations 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.

(a)

(b)

(c)

##### Disagreement across Dimensions.

The information packing in Matryoshka ​ Representations {\rm Matryoshka~Representations} often results in gradual increase of accuracy with increase in capacity. However, we observed that 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 {\rm MRL} can gain up to 4.6 % 4.6\% classification accuracy. At the same time, the low-dimensional models are less accurate either due to confusion within the same superclass [ 24 ] of the ImageNet hierarchy or presence of multiple objects of interest. Figure 9 showcases 2 such examples for 8 8 -dimensional representation. These results along with Appendix J put forward the potential for MRL {\rm 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 as significant when evaluated at a superclass level (Table 24 in Appendix J ). Figure 11 presents that this phenomenon occurs with both MRL {\rm MRL} and FF models; MRL {\rm 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 {\rm Matryoshka~Representation} naturally captures the underlying hierarchy of the class labels with one single model. Lastly, Figure 11 showcases the accuracy trends per superclass with MRL {\rm MRL} . The utility of additional dimensions in distinguishing a class from others within the same superclass is evident for “garment” which has up to 11% improvement for 8 → \to 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 {\rm Matryoshka~Representations} can be enabled within off-the-shelf pretrained models with inexpensive partial finetuning thus paving a way for ubiquitous adoption of MRL {\rm MRL} . At the same time, Table 27 in Appendix C indicates that with optimal weighting of the nested losses we could improve accuracy of lower-dimensions representations without accuracy loss. Tables 29 and 29 in Appendix C ablate over the choice of initial granularity and spacing of the granularites. Table 29 reaffirms the design choice to shun extremely low dimensions that have poor classification accuracy as initial granularity for MRL {\rm 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 {\rm MRL} that would be logical directions for future work. (1) Optimizing the weightings of the nested losses to obtain a Pareto optimal accuracy-vs-efficiency trade-off – a potential solution could emerge from adaptive loss balancing aspects of anytime neural networks [ 41 ] . (2) Using different losses at various fidelities aimed at solving a specific aspect of adaptive deployment – e.g. high recall for 8 8 -dimension and robustness for 2048 2048 -dimension. (3) Learning a search data-structure, like differentiable k-d tree, on top of Matryoshka ​ Representation {\rm Matryoshka~Representation} to enable dataset and representation aware retrieval. (4) Finally, the joint optimization of multi-objective MRL {\rm 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 {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} ), a flexible representation learning approach that encodes information at multiple granularities in a single embedding vector. This enables the MRL {\rm MRL} to adapt to a downstream task’s statistical complexity as well as the available compute resources. We demonstrate that MRL {\rm MRL} can be used for large-scale adaptive classification as well as adaptive retrieval. On standard benchmarks, MRL {\rm MRL} matches the accuracy of the fixed-feature baseline despite using 14 × 14\times smaller representation size on average. Furthermore, the Matryoshka ​ Representation {\rm Matryoshka~Representation} based adaptive shortlisting and re-ranking system ensures comparable mAP@ 10 10 to the baseline while being 128 × 128\times cheaper in FLOPs and 14 × 14\times faster in wall-clock time. Finally, most of the efficiency techniques for model inference and vector search are complementary to MRL {\rm MRL} further assisting in deployment at the compute-extreme environments.

## Acknowledgments

We are grateful to Srinadh Bhojanapalli, Lovish Madaan, Raghav Somani, Ludwig Schmidt, and Venkata Sailesh Sanampudi for helpful discussions and feedback. Aditya Kusupati also thanks Tom Duerig and Rahul Sukthankar for their support. Part of the paper’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, a Semiconductor Research Corporation (SRC) program sponsored by DARPA. Sham Kakade acknowledges funding from the NSF award CCF-1703574 and ONR N00014-22-1-2377. Ali Farhadi acknowledges funding from the NSF awards IIS 1652052, IIS 17303166, DARPA N66001-19-2-4031, DARPA W911NF-15-1-0543 and gifts from Allen Institute for Artificial Intelligence.

## Appendix A Code for Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} ( MRL {\rm MRL} )

We use Alg 1 and 2 provided below to train supervised ResNet50– MRL {\rm MRL} models on ImageNet-1K. We provide this code as a template to extend MRL {\rm MRL} to any domain.

## Appendix B Datasets

ImageNet-1K [ 76 ] contains 1,281,167 labeled train images, and 50,000 labelled validation images across 1,000 classes. The images were transformed with standard procedures detailed by FFCV [ 56 ] .

ImageNet-4K dataset was constructed by selecting 4,202 classes, non-overlapping with ImageNet-1K, from ImageNet-21K [ 16 ] with 1,050 or more examples. The train set contains 1,000 examples and the query/validation set contains 50 examples per class totalling to ∼ \sim 4.2M and ∼ \sim 200K respectively. We will release the list of images curated together to construct ImageNet-4K.

JFT-300M [ 85 ] is a large-scale multi-label dataset with 300M images labelled across 18,291 categories.

ALIGN [ 46 ] utilizes a large scale noisy image-text dataset containing 1.8B image-text pairs.

##### ImageNet Robustness Datasets

We experimented on the following datasets to examine the robustness of MRL {\rm MRL} models:

ImageNetV2 [ 72 ] is a collection of 10K images sampled a decade after the original construction of ImageNet [ 16 ] . ImageNetV2 contains 10 examples each from the 1,000 classes of ImageNet-1K.

ImageNet-A [ 35 ] contains 7.5K real-world adversarially filtered images from 200 ImageNet-1K classes.

ImageNet-R [ 34 ] contains 30K artistic image renditions for 200 of the original ImageNet-1K classes.

ImageNet-Sketch [ 94 ] contains 50K sketches, evenly distributed over all 1,000 ImageNet-1K classes.

ObjectNet [ 2 ] contains 50K images across 313 object classes, each containing ∼ \sim 160 images each.

## Appendix C Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} Model Training

We trained all ResNet50– MRL {\rm MRL} models using the efficient dataloaders of FFCV [ 56 ] . We utilized the rn50_40_epochs.yaml configuration file of FFCV to train all MRL {\rm MRL} models defined below:

• MRL {\rm MRL} : ResNet50 model with the fc layer replaced by MRL_Linear_Layer ( efficient = False )

• MRL ​ – ​ E {\rm MRL\text{--}E} : ResNet50 model with the fc layer replaced by MRL_Linear_Layer ( efficient = True )

• FF–k: ResNet50 model with the fc layer replaced by torch . nn . Linear ( k , num_classes ) , where k ∈ [ 8 , 16 , 32 , 64 , 128 , 256 , 512 , 1024 , 2048 ] \in[8,16,32,64,128,256,512,1024,2048] . We will henceforth refer to these models as simply FF, with the k value denoting representation size.

We trained all ResNet50 models with a learning rate of 0.475 0.475 with a cyclic learning rate schedule [ 83 ] . This was after appropriate scaling (0.25 × \times ) of the learning rate specified in the configuration file to accommodate for 2xA100 NVIDIA GPUs available for training, compared to the 8xA100 GPUs utilized in the FFCV benchmarks. We trained with a batch size of 256 per GPU, momentum [ 86 ] of 0.9, and an SGD optimizer with a weight decay of 1e-4.

Our code (Appendix A ) makes minimal modifications to the training pipeline provided by FFCV to learn Matryoshka ​ Representations {\rm Matryoshka~Representations} .

We trained ViT-B/16 models for JFT-300M on a 8x8 cloud TPU pod [ 49 ] using Tensorflow [ 1 ] with a batchsize of 128 and trained for 300K steps. Similarly, ALIGN models were trained using Tensorflow on 8x8 cloud TPU pod for 1M steps with a batchsize of 64 per TPU. Both these models were trained with adafactor optimizer [ 81 ] with a linear learning rate decay starting at 1e-3.

Lastly, we trained a BERT-Base model on English Wikipedia and BookCorpus. We trained our models in Tensorflow using a 4x4 cloud TPU pod with a total batchsize of 1024. We used AdamW [ 61 ] optimizer with a linear learning rate decay starting at 1e-4 and trained for 450K steps.

In each configuration/case, if the final representation was normalized in the FF implementation, MRL {\rm MRL} models adopted the same for each nested dimension for a fair comparison.

## Appendix D Classification Results

We show the top-1 classification accuracy of ResNet50– MRL {\rm MRL} models on ImageNet-1K in Table 1 and Figure 3 . We compare the performance of MRL {\rm MRL} models ( MRL {\rm MRL} , MRL ​ – ​ E {\rm MRL\text{--}E} ) to several baselines:

• FF: We utilize the FF-k models described in Appendix C for k ∈ { 8 , … ​ 2048 } k\in\{8,...2048\} .

• SVD : We performed a low rank approximation of the 1000-way classification layer of FF-2048, with rank = 1000.

• Rand. LP : We compared against a linear classifier fit on randomly selected features [ 30 ] .

• Slim. Net : We take pretrained slimmable neural networks [ 100 ] which are trained with a flexible width backbone (25%, 50%, 75% and full width). For each representation size, we consider the first k k dimensions for classification. Note that training of slimmable neural networks becomes unstable when trained below 25% width due to the hardness in optimization and low complexity of the model.

At lower dimensions ( d ≤ 128 d\leq 128 ), MRL {\rm MRL} outperforms all baselines significantly, which indicates that pretrained models lack the multifidelity of Matryoshka ​ Representations {\rm Matryoshka~Representations} and are incapable of fitting an accurate linear classifier at low representation sizes.

We compared the performance of MRL {\rm MRL} models at various representation sizes via 1-nearest neighbors (1-NN) image classification accuracy on ImageNet-1K in Table 2 and Figure 3 . We provide detailed information regarding the k-NN search pipeline in Appendix E . We compared against a baseline of attempting to enforce nesting to a FF-2048 model by 1) Random Feature Selection (Rand. FS): considering the first m m dimensions of FF-2048 for NN lookup, and 2) FF+SVD: performing SVD on the FF-2048 representations at the specified representation size, 3) FF+JL: performing random projection according to the Johnson-Lindenstrauss lemma [ 48 ] on the FF-2048 representations at the specified representation size. We also compared against the 1-NN accuracy of slimmable neural nets [ 100 ] as an additional baseline. We observed these baseline models to perform very poorly at lower dimensions, as they were not explicitly trained to learn Matryoshka ​ Representations {\rm Matryoshka~Representations} .

### D.1 Adaptive Classification ( MRL {\rm MRL} –AC)

In an attempt to use the smallest representation that works well for classification for every image in the ImageNet-1K validation set, we learned a policy to increase the representation size from m i m_{i} to m i + 1 m_{i+1} using a 10K sized subset of the ImageNet-1K validation set. This policy is based on whether the prediction confidence p i p_{i} using representation size m i m_{i} exceeds a learned threshold t i ∗ t_{i}^{\ast} . If p i ≥ t i ∗ p_{i}\geq t_{i}^{\ast} , we used predictions from representation size m i m_{i} otherwise, we increased to representation size m i + 1 m_{i+1} . To learn the optimal threshold t i ∗ t_{i}^{\ast} , we performed a grid search between 0 and 1 (100 samples). For each threshold t k t_{k} , we computed the classification accuracy over our 10K image subset. We set t i ∗ t_{i}^{\ast} equal to the smallest threshold t k t_{k} that gave the best accuracy. We use this procedure to obtain thresholds for successive models, i.e., { t j ∗ ∣ j ∈ { 8 , 16 , 32 , 64 , … , 2048 } } \{t_{j}^{\ast}\mid j\in\{8,16,32,64,\ldots,2048\}\} . To improve reliability of threshold based greedy policy, we use test time augmentation which has been used successfully in the past [ 82 ] .

For inference, we used the remaining held-out 40K samples from the ImageNet-1K validation set. We began with smallest sized representation ( m = 8 m=8 ) and compared the computed prediction confidence p 8 p_{8} to learned optimal threshold t 8 ∗ t_{8}^{\ast} . If p 8 ≤ t 8 ∗ p_{8}\leq t_{8}^{\ast} , then we increased m = 16 m=16 , and repeated this procedure until m = d = 2048 m=d=2048 . To compute the expected dimensions, we performed early stopping at m = { 16 , 32 , 64 , … ​ 2048 } m=\{16,32,64,\ldots 2048\} and computed the expectation using the distribution of representation sizes. As shown in Table 3 and Figure 7 , we observed that in expectation, we only needed a ∼ 37 \sim 37 sized representation to achieve 76.3 % 76.3\% classification accuracy on ImageNet-1K, which was roughly 14 × 14\times smaller than the FF–512 baseline. Even if we computed the expectation as a weighted average over the cumulative sum of representation sizes { 8 , 24 , 56 , … } \{8,24,56,\ldots\} , due to the nature of multiple linear heads for MRL {\rm MRL} , we ended up with an expected size of 62 62 that still provided a roughly 8.2 × 8.2\times efficient representation than the FF–512 baseline. However, MRL ​ – ​ E {\rm MRL\text{--}E} alleviates this extra compute with a minimal drop in accuracy.

### D.2 JFT, ALIGN and BERT

We examine the k-NN classification accuracy of learned Matryoshka ​ Representations {\rm Matryoshka~Representations} via ALIGN– MRL {\rm MRL} and JFT-ViT– MRL {\rm MRL} in Table 4 . For ALIGN [ 46 ] , we observed that learning Matryoshka ​ Representations {\rm Matryoshka~Representations} via ALIGN– MRL {\rm MRL} improved classification accuracy at nearly all dimensions when compared to ALIGN. We observed a similar trend when training ViT-B/16 [ 22 ] for JFT-300M [ 85 ] classification, where learning Matryoshka ​ Representations {\rm Matryoshka~Representations} via MRL {\rm MRL} and MRL ​ – ​ E {\rm MRL\text{--}E} on top of JFT-ViT improved classification accuracy for nearly all dimensions, and significantly for lower ones. This demonstrates that training to learn Matryoshka ​ Representations {\rm Matryoshka~Representations} is feasible and extendable even for extremely large scale datasets. We also demonstrate that Matryoshka ​ Representations {\rm Matryoshka~Representations} are learned at interpolated dimensions for both ALIGN and JFT-ViT, as shown in Table 5 , despite not being trained explicitly at these dimensions. Lastly, Table 6 shows that MRL {\rm MRL} training leads to a increase in the cosine similarity span between positive and random image-text pairs.

We also evaluated the capability of Matryoshka ​ Representations {\rm Matryoshka~Representations} to extend to other natural language processing via masked language modeling (MLM) with BERT [ 19 ] , whose results are tabulated in Table 7 . Without any hyper-parameter tuning, we observed Matryoshka ​ Representations {\rm Matryoshka~Representations} to be within 0.5 % 0.5\% of FF representations for BERT MLM validation accuracy. This is a promising initial result that could help with large-scale adaptive document retrieval using BERT– MRL {\rm MRL} .

## Appendix E Image Retrieval

We evaluated the strength of Matryoshka ​ Representations {\rm Matryoshka~Representations} via image retrieval on ImageNet-1K (the training distribution), as well as on out-of-domain datasets ImageNetV2 and ImageNet-4K for all MRL {\rm MRL} ResNet50 models. We generated the database and query sets, containing N N and Q Q samples respectively, with a standard PyTorch [ 67 ] forward pass on each dataset. We specify the representation size at which we retrieve a shortlist of k-nearest neighbors (k-NN) by D s D_{s} . The database is a thus a [ N N , D s D_{s} ] array, the query set is a [ Q Q , D s D_{s} ] array, and the neighbors set is a [ Q Q , k] array. For metrics, we utilized corrected mean average precision (mAP@k) [ 55 ] and precision (P@k): P ​ @ ​ k = c ​ o ​ r ​ r ​ e ​ c ​ t ​ _ ​ p ​ r ​ e ​ d k P@k=\dfrac{correct\_pred}{k} where c ​ o ​ r ​ r ​ e ​ c ​ t ​ _ ​ p ​ r ​ e ​ d correct\_pred is the average number of retrieved NN with the correct label over the entire query set using a shortlist of length k k .

We performed retrieval with FAISS [ 47 ] , a library for efficient similarity search. To obtain a shortlist of k-NN, we built an index to search the database. We performed an exhaustive NN search with the L2 distance metric with faiss . IndexFlatL2 , as well as an approximate NN search (ANNS) via HNSW [ 47 ] with faiss . IndexHNSWFlat . We used HNSW with M = 32 M=32 unless otherwise mentioned, and henceforth referred to as HNSW32. The exact search index was moved to the GPU for fast k-NN search computation, whereas the HNSW index was kept on the CPU as it currently lacks GPU support. We show the wall clock times for building the index as well as the index size in Table 20 . We observed exact search to have a smaller index size which was faster to build when compared to HNSW, which trades off a larger index footprint for fast NN search (discussed in more detail in Appendix K ). The database and query vectors are normalized with faiss . normalize_L2 before building the index and performing search.

Retrieval performance on ImageNet-1K, i.e. the training distribution, is shown in Table 8 . MRL {\rm MRL} outperforms FF models for nearly all representation size for both top-1 and mAP@10, and especially at low representation size ( D s D_{s} ≤ 32 \leq 32 ). MRL ​ – ​ E {\rm MRL\text{--}E} loses out to FF significantly only at D s D_{s} = 8 =8 . This indicates that training ResNet50 models via the MRL {\rm MRL} training paradigm improves retrieval at low representation size over models explicitly trained at those representation size (FF- 8 ​ … ​ 2048 8...2048 ).

We carried out all retrieval experiments at D s D_{s} ∈ { 8 , 16 , 32 , 64 , 128 , 256 , 512 , 1024 , 2048 } \in\{8,16,32,64,128,256,512,1024,2048\} , as these were the representation sizes which were a part of the nesting_list at which losses were added during training, as seen in Algorithm 1 , Appendix A . To examine whether MRL {\rm MRL} is able to learn Matryoshka ​ Representations {\rm Matryoshka~Representations} at dimensions in between the representation size for which it was trained, we also tabulate the performance of MRL {\rm MRL} at interpolated D s D_{s} ∈ { 12 , 24 , 48 , 96 , 192 , 384 , 768 , 1536 } \in\{12,24,48,96,192,384,768,1536\} as MRL {\rm MRL} –Interpolated and MRL ​ – ​ E {\rm MRL\text{--}E} –Interpolated (see Table 8 ). We observed that performance scaled nearly monotonically between the original representation size and the interpolated representation size as we increase D s D_{s} , which demonstrates that MRL {\rm MRL} is able to learn Matryoshka ​ Representations {\rm Matryoshka~Representations} at nearly all representation size m ∈ [ 8 , 2048 ] m\in[8,2048] despite optimizing only for | ℳ | |\mathcal{M}| nested representation sizes.

We examined the robustness of MRL {\rm MRL} for retrieval on out-of-domain datasets ImageNetV2 and ImageNet-4K, as shown in Table 9 and Table 10 respectively. On ImageNetV2, we observed that MRL {\rm MRL} outperformed FF at all D s D_{s} on top-1 Accuracy and mAP@10, and MRL ​ – ​ E {\rm MRL\text{--}E} outperformed FF at all D s D_{s} except D s D_{s} = 8 =8 . This demonstrates the robustness of the learned Matryoshka ​ Representations {\rm Matryoshka~Representations} for out-of-domain image retrieval.

## Appendix F Adaptive Retrieval

The time complexity of retrieving a shortlist of k-NN often scales as O ⁡ ( d ) O(d) , where d = d= D s D_{s} , for a fixed k and N N . We thus will have a theoretical 256 × 256\times higher cost for D s D_{s} = 2048 =2048 over D s D_{s} = 8 =8 . We discuss search complexity in more detail in Appendix I . In an attempt to replicate performance at higher D s D_{s} while using less FLOPs, we perform adaptive retrieval via retrieving a k-NN shortlist with representation size D s D_{s} , and then re-ranking the shortlist with representations of size D r D_{r} . Adaptive retrieval for a shortlist length k = 200 k=200 is shown in Table 11 for ImageNet-1K, and in Table 12 for ImageNet-4K. On ImageNet-1K, we are able to achieve comparable performance to retrieval with D s D_{s} = 2048 =2048 (from Table 8 ) with D s D_{s} = 16 =16 at 128 × 128\times less MFLOPs/Query (used interchangeably with MFLOPs). Similarly, on ImageNet-4K, we are able to achieve comparable performance to retrieval with D s D_{s} = 2048 =2048 (from Table 10 ) with D s D_{s} = 64 =64 on ImageNet-1K and ImageNet-4K, at 32 × 32\times less MFLOPs. This demonstrates the value of intelligent routing techniques which utilize appropriately sized Matryoshka ​ Representations {\rm Matryoshka~Representations} for retrieval.

##### Funnel Retrieval.

We also designed a simple cascade policy which we call funnel retrieval to successively improve and refine the k-NN shortlist at increasing D s D_{s} . This was an attempt to remove the dependence on manual choice of D s D_{s} & D r D_{r} . We retrieved a shortlist at D s D_{s} and then re-ranked the shortlist five times while simultaneously increasing D r D_{r} (rerank cascade) and decreasing the shortlist length (shortlist cascade), which resembles a funnel structure. We tabulate the performance of funnel retrieval in various configurations in Table 13 on ImageNet-1K, and in Table 14 on ImageNet-4K. With funnel retrieval on ImageNet-1K, we were able to achieve top-1 accuracy within 0.1 % 0.1\% of retrieval with D s D_{s} = 2048 =2048 (as in Table 8 ) with a funnel with D s D_{s} = 16 =16 , with 128 × 128\times less MFLOPs. Similarly, we are able to achieve equivalent top-1 accuracy within 0.15 % 0.15\% of retrieval at D s D_{s} = 2048 =2048 (as in Table 10 ) with funnel retrieval at D s D_{s} = 32 =32 on ImageNet-4K, with 64 × 64\times less MFLOPs. This demonstrates that with funnel retrieval, we can emulate the performance of retrieval with D s D_{s} = 2048 =2048 with a fraction of the MFLOPs.

## Appendix G Few-shot and Sample Efficiency

We compared MRL {\rm MRL} , MRL ​ – ​ E {\rm MRL\text{--}E} , and FF on various benchmarks to observe the effect of representation size on sample efficiency. We used Nearest Class Means [ 79 ] for classification which has been shown to be effective in the few-shot regime [ 13 ] .

##### ImageNetV2.

Representations are evaluated on ImageNetV2 with the n-shot k-way setup. ImageNetV2 is a dataset traditionally used to evaluate the robustness of models to natural distribution shifts. For our experiments we evaluate accuracy of the model given n n examples from the ImageNetV2 distribution. We benchmark representations in the traditional small-scale (10-way) and large-scale (1000-way) setting. We evaluate for n ∈ 1 , 3 , 5 , 7 , 9 n\in{1,3,5,7,9} with 9 being the maximum value for n n because there are 10 images per class.

We observed that MRL {\rm MRL} had equal performance to FF across all representation sizes and shot numbers. We also found that for both MRL {\rm MRL} and FF, as the shot number decreased, the required representation size to reach optimal accuracy decreased (Table 15 ). For example, we observed that 1-shot performance at 32 32 representation size had equal accuracy to 2048 2048 representation size.

##### FLUID.

For the long-tailed setting we evaluated MRL {\rm MRL} on the FLUID benchmark [ 92 ] which contains a mixture of pretrain and new classes. Table 16 shows the evaluation of the learned representation on FLUID. We observed that MRL {\rm MRL} provided up to 2% higher accuracy on novel classes in the tail of the distribution, without sacrificing accuracy on other classes. Additionally we found the accuracy between low-dimensional and high-dimensional representations was marginal for pretrain classes. For example, the 64-dimensional MRL {\rm MRL} performed ∼ 1 % \sim 1\% lower in accuracy compared to the 2048-dimensional counterpart on pretrain-head classes (84.46% vs 85.60%). However for novel-tail classes the gap was far larger (6.22% vs 12.88%). We hypothesize that the higher-dimensional representations are required to differentiate the classes when few training examples of each are known. These results provide further evidence that different tasks require varying capacity based on their difficulty.

## Appendix H Robustness Experiments

We evaluated the robustness of MRL {\rm MRL} models on out-of-domain datasets (ImageNetV2/R/A/Sketch) and compared them to the FF baseline. Each of these datasets is described in Appendix B . The results in Table 17 demonstrate that learning Matryoshka ​ Representations {\rm Matryoshka~Representations} does not hurt out-of-domain generalization relative to FF models, and Matryoshka ​ Representations {\rm Matryoshka~Representations} in fact improve the performance on ImageNet-A. For a ALIGN– MRL {\rm MRL} model, we examine the the robustness via zero-shot retrieval on out-of-domain datasets, including ObjectNet, in Table 18 .

## Appendix I In Practice Costs

All approximate NN search experiments via HNSW32 were run on an Intel Xeon 2.20GHz CPU with 24 cores. All exact search experiments were run with CUDA 11.0 on 2xA100-SXM4 NVIDIA GPUs with 40G RAM each.

##### MRL {\rm MRL} models.

As MRL {\rm MRL} makes minimal modifications to the ResNet50 model in the final fc layer via multiple heads for representations at various scales, it has only an 8MB storage overhead when compared to a standard ResNet50 model. MRL ​ – ​ E {\rm MRL\text{--}E} has no storage overhead as it has a shared head for logits at the final fc layer.

##### Retrieval

Exact search has a search time complexity of O ⁡ ( d ​ k ​ N ) O(dkN) , and HNSW has a search time complexity of O ⁡ ( d ​ k ​ log ⁡ ( N ) ) O(dk\log(N)) , where N N is the database size, d d is the representation size, and k k is the shortlist length. To examine real-world performance, we tabulated wall clock search time for every query in the ImageNet-1K and ImageNet-4K validation sets over all representation sizes d d in Table 19 for both Exact Search and HNSW32, and ablated wall clock query time over shortlist length k k on the ImageNet-1K validation set in Table 21 . The wall clock time to build the index and the index size is also shown in Table 20 .

## Appendix J Analysis of Model Disagreement

##### Class Trends

Does increasing representation size necessarily help improve classification performance across all classes in ImageNet-1K? We studied this question by examining trends in performance with increasing representation size from d = 8 , … ​ 2048 d={8,...2048} . For MRL {\rm MRL} models, we observed that 244 244 classes showed a monotonic improvement in performance with increasing d d , 177 177 classes first improved but then observed a slight dip (one or two misclassifications per class), 49 49 classes showed a decline first and then an improvement, and the remaining classes did not show a clear trend. When we repeated this experiment with independently trained FF models, we noticed that 950 950 classes did not show a clear trend. This motivated us to leverage the disagreement as well as gradual improvement of accuracy at different representation sizes by training Matryoshka ​ Representations {\rm Matryoshka~Representations} . Figure 12 showcases the progression of relative per-class accuracy distribution compared to the Matryoshka ​ Representation ​ Learning {\rm Matryoshka~Representation~Learning} -2048 dimensional model. This also showed that some instances and classes could benefit from lower-dimensional representations.

##### Discussion of Oracle Accuracy

Based on our observed model disagreements for different representation sizes d d , we defined an optimal oracle accuracy [ 58 ] for MRL {\rm MRL} . We labeled an image as correctly predicted if classification using any representation size was correct. The percentage of total samples of ImageNet-1K that were firstly correctly predicted using each representation size d d is shown in Table 22 . This defined an upper bound on the performance of MRL {\rm MRL} models, as 18.46 % 18.46\% of the ImageNet-1K validation set were incorrectly predicted ∀ d ∈ { 8 , 16 , … , 2048 } \forall d\in\{8,16,\ldots,2048\} . We show the oracle performance on MRL {\rm MRL} models for ImageNet-1K/V2/A/R/Sketch datasets in Table 23 .

In an attempt to derive an optimal routing policy to emulate oracle accuracy, we designed the adaptive classification via cascading method as discussed in Appendix D.1 . This led to an interesting observation on the expected dimensionality for 76.30 % 76.30\% top-1 classification accuracy being just d ∼ 37 d\sim 37 . We leave the design and learning of a more optimal policy for future work.

##### Grad-CAM Examples

We analyzed the nature of model disagreement across representation sizes with MRL {\rm MRL} models with the help of Grad-CAM visualization [ 80 ] . We observed there were certain classes in ImageNet-1K such as "tools", "vegetables" and "meat cutting knife" which were occasionally located around multiple objects and a cluttered environment. In such scenarios, we observed that smaller representation size models would often get confused due to other objects and fail to extract the object of interest which generated the correct label. We also observed a different nature of disagreement arising when the models got confused within the same superclass. For example, ImageNet-1K has multiple "snake" classes, and models often confuse a snake image for an incorrect species of snake.

##### Superclass Performance

We created a 30 superclass subset of the validation set based on wordnet hierarchy (Table 24 ) to quantify the performance of MRL {\rm MRL} model on ImageNet-1K superclasses. Table 25 quantifies the performance with different representation size.

## Appendix K Ablation Studies

### K.1 MRL {\rm MRL} Training Paradigm

##### Matryoshka ​ Representations {\rm Matryoshka~Representations} via Finetuning.

To observe if nesting can be induced in models that were not explicitly trained with nesting from scratch, we loaded a pretrained FF-2048 ResNet50 model and initialized a new MRL {\rm MRL} layer, as defined in Algorithm 2 , Appendix C . We then unfroze different layers of the backbone to observe how much non-linearity in the form of unfrozen conv layers needed to be present to enforce nesting into a pretrained FF model. A description of these layers can be found in the ResNet50 architecture [ 29 ] . All models were finetuned with the FFCV pipeline, with same training configuration as in the end-to-end training aside from changing lr = 0.1 =0.1 and epochs = 10 =10 . We observed that finetuning the linear layer alone was insufficient to learn Matryoshka ​ Representations {\rm Matryoshka~Representations} at lower dimensionalities. Adding more and more non-linear conv+ReLU layers steadily improved classification accuracy of d = 8 d=8 from 5 % 5\% to 60 % 60\% after finetuning, which was only 6 % 6\% less than training MRL {\rm MRL} end-to-end for 40 epochs. This difference was successively less pronounced as we increased dimensionality past d = 64 d=64 , to within 1.5 % 1.5\% for all larger dimensionalities. The full results of this ablation can be seen in Table 26 .

##### Relative Importance.

We performed an ablation of MRL {\rm MRL} over the relative importance, c m c_{m} , of different nesting dimensions m ∈ ℳ m\in\cal{M} , as defined in Sec. 3 . In an attempt to improve performance at lower dimensionalities, we boosted the relative importance c m c_{m} of training loss at lower dimensions as in Eq. 1 with two models, MRL {\rm MRL} -8boost and MRL {\rm MRL} -8+16boost. The MRL {\rm MRL} -8boost model had c m ∈ ℳ = [ 2 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 ] c_{m\in\cal M}=[2,1,1,1,1,1,1,1,1] and the MRL {\rm MRL} -8+16boost model had c m ∈ ℳ = [ 2 , 1.5 , 1 , 1 , 1 , 1 , 1 , 1 , 1 ] c_{m\in\cal M}=[2,1.5,1,1,1,1,1,1,1] . The relative importance list c m ∈ ℳ c_{m\in\cal M} had a 1-to-1 correspondence with nesting dimension set ℳ \mathcal{M} . In Table 27 , we observed that MRL {\rm MRL} -8boost improves top-1 accuracy by 3 % 3\% at d = 8 d=8 , and also improves top-1 accuracy of all representation scales from 16 to 256 over MRL {\rm MRL} , while only hurting the performance at 512 to 2048 representation scales by a maximum of 0.1%. This suggests that the relative importance c m c_{m} can be tuned/set for optimal accuracy for all m ∈ ℳ m\in\mathcal{M} , but we leave this extension for future work.

##### Matryoshka ​ Representations {\rm Matryoshka~Representations} at Arbitrary Granularities.

To train MRL {\rm MRL} , we used nested dimensions at logarithmic granularities ℳ = { 8 , 16 , … , 1024 , 2048 } \mathcal{M}=\{8,16,\ldots,1024,2048\} as detailed in Section 3 . We made this choice for two empirically-driven reasons: a) The accuracy improvement with increasing representation size was more logarithmic than linear (as shown by FF models in Figure 3 ). This indicated that optimizing for granularities increasing in a non-logarithmic fashion would be sub-optimal both for maximum performance and expected efficiency; b) If we have m m arbitrary granularities, the expected cost of the linear classifier to train MRL {\rm MRL} scales as O ⁡ ( L ∗ ( m 2 ) ) O(L*(m^{2})) while logarithmic granularities result in O ⁡ ( L ∗ 2 ​ l ​ o ​ g ​ ( d ) ) O(L*2log(d)) space and compute costs.

To demonstrate this effect, we learned Matryoshka ​ Representations {\rm Matryoshka~Representations} with uniform ( MRL {\rm MRL} -Uniform ) nesting dimensions m ∈ ℳ = { 8,212,416,620,824 , 1028 , 1232 , 1436 , 1640 , 1844 , 2048 } m\in\mathcal{M}=\{8,212,416,620,824,1028,1232,1436,1640,1844,2048\} . We evaluated this model at the standard ( MRL {\rm MRL} -log ) dimensions m ∈ ℳ = { 8 , 16 , 32 , 64 , 128 , 256 , 512 , 1024 , 2048 } m\in\mathcal{M}=\{8,16,32,64,128,256,512,1024,2048\} for ease of comparison to reported numbers using 1-NN accuracy (%). As shown in Table 29 , we observed that while performance interpolated, MRL {\rm MRL} -Uniform suffered at low dimensions as the logarithmic spacing of MRL {\rm MRL} -log resulted in tighter packing of information in these initial dimensions. The higher nesting dimensions of MRL {\rm MRL} -Uniform did not help in significant accuracy improvement due to accuracy saturation, which is often logarithmic in representation size as shown by FF models. Note that the slight improvement at dimensions higher than 512 for MRL {\rm MRL} -Uniform is due to multiple granularities around them compared to just three for MRL {\rm MRL} -log , which are not useful in practice for efficiency.

##### Lower Dimensionality.

We experimented with training MRL {\rm MRL} with smaller nesting dimension than m = 8 m=8 , as shown in Table 29 , with two models: MRL-4 and MRL-6. We found that using lower than 8-dimensions to train MRL {\rm MRL} , i.e. m 0 ∈ { 4 , 6 } m_{0}\in\{4,6\} for MRL-4 and MRL-6 respectively, did not affect the top-1 accuracy of other granularities significantly. However, granularities smaller than 8-dimensions had very low accuracy and were often unusable for deployment along with additional training difficulty. We also observed a small dip in accuracy at higher dimensions which we attribute to the joint loss that now also included the harder optimization of the smallest dimension. Lastly, we hypothesize the dimensionality of 8 is an empirically validated design choice due to the considerable accuracy it provided along with the ease of training.

### K.2 Retrieval

##### Adaptive Retrieval.

To examine the effect of increasing shortlist lengths on search time, we performed a reranking ablation over shortlist lengths for D s D_{s} = 16 and D r D_{r} = 2048 over ImageNet-1K in Table 30 , and over ImageNet-4K in Table 31 . We observed that using a larger shortlist k k saturated ImageNet-1K performance at k k =200. But using larger shortlists until k = 2048 k=2048 , the maximum value supported by the FAISS framework, steadily improved performance on ImageNet-4K. This is likely due to the increased database size, but could also indicate a correlation with ImageNet-4K being slightly out-of-distribution making the task at hand harder.

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .