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