Matryoshka Representation Learning — part 3/5

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

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.