Matryoshka Representation Learning — part 4/5
Cleaned-text code points [46490, 62453). Bibliography omitted here; source.md includes it.
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