Matryoshka Representation Learning — part 2/5

Cleaned-text code points [15437, 31390). Bibliography omitted here; source.md includes it.

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.