MatFormer: Nested Transformer for Elastic Inference

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 22:30:53

Abstract: Foundation models are applied in a broad spectrum of settings with different inference constraints, from massive multi-accelerator clusters to resource-constrained standalone mobile devices. However, the substantial costs associated with training these models often limit the number of unique model sizes that can be offered. Consequently, practitioners are compelled to select a model that may not be optimally aligned with their specific latency and cost requirements. We present MatFormer, a novel Transformer architecture designed to provide elastic inference across diverse deployment constraints. MatFormer achieves this by incorporating a nested Feed Forward Network (FFN) block structure within a standard Transformer model. During training, we  optimize the parameters of multiple nested FFN blocks with varying sizes, enabling the extraction of hundreds of accurate smaller models without incurring additional computational costs. We empirically validate the efficacy of MatFormer across different model classes (decoders and encoders) and modalities (language and vision), demonstrating its potential for real-world deployment. We show that a 850M decoder-only MatFormer language model (MatLM) allows us to extract multiple smaller models spanning from 582M to 850M parameters, each exhibiting better validation loss and one-shot downstream evaluations than independently trained counterparts. Furthermore, we observe that smaller encoders extracted from a universal MatFormer-based ViT (MatViT) encoder preserve the metric-space structure for adaptive large-scale retrieval. Finally, we showcase that speculative decoding with the accurate and consistent submodels extracted from  MatFormer can lead to significant reduction in inference latency.

Main paper: 
 related work: 
 A standard Transformer (Vaswani et al., 2023) has become the unifying model architecture for foundation models (Bommasani et al., 2021) across modalities like language (Brown et al., 2020) , vision (Dehghani et al., 2023) and audio (Radford et al., 2023) . While extremely powerful, the standard Transformer block is not natively elastic in a way that enables large-scale adaptive and flexible deployment across various resource constraints. To cater to the plethora of deployment requirements, existing solutions include training a family of models of varying sizes (Anil et al., 2023; Touvron et al., 2023b) , post-hoc efficiency techniques like quantization (Dettmers & Zettlemoyer, 2023) , pruning (Lagunas et al., 2021) , distillation (Sanh et al., 2019) and mixture of varying capacity experts (MoE) (Zhang & Ma, 2012) . However, these solutions often are specific to the single constraint at hand, and require additional training or trade-off memory/compute during inference making them far from being a truly elastic solution for adaptive deployment. Lastly, Transformer based LLMs are often sped-up during inference with techniques like speculative decoding (Leviathan et al., 2023; Chen et al., 2023) -that benefits from the smaller draft & the larger verifier models having similar behavior -or early exiting (Schuster et al., 2022) to enable real-time deployment.Obtaining multiple smaller models from a single model has been explored in the past (Yu et al., 2018; Yu & Huang, 2019; Cai et al., 2019; Grimaldi et al., 2022; Cai et al., 2021) with most works focusing on CNN encoders. Specifically, OFA (Cai et al., 2019) creates a universal CNN model which is used to extract and finetune submodels for a handful of deployment constraints while slimmable networks (Yu et al., 2018) optimize for limited preset widths and require explicit training to interpolate for a few more intermediate widths (Yu & Huang, 2019) . NAS techniques that sample random (not nested) subnetworks during training at each step, and then find the subnetwork architecture to retrain from scratch before deployment have been explored (Wang et al., 2020b) . These techniques fall short of being truly elastic and come with significant training overheads. More recently some of them have been extended to Transformer encoders (Chavan et al., 2022; Hou et al., 2020; Salehi et al., 2023) for extracting sub-models in both static or dynamic settings but fail at extending further to decoder-only language models. While not in the weight space, matryoshka representation learning (Kusupati et al., 2022) & FlexiViT (Beyer et al., 2023) showcase elasticity in output & input spaces respectively by smoothly spanning deployment constraints with minimal overhead. MatFormer, in contrast, builds upon these works by nested the weight space instead to enable truly elastic and adaptive Transformer-based (decoder & encoder) models that span all the accuracy-vscompute tradeoff (statically or dynamically) with minimal changes and training overhead (Figure 1 ). Finally, we also point the readers to SortedNet (Valipour et al., 2023) , a concurrent work with similar goals applied to encoders, which optimizes many sampled submodels (akin to prior works) unlike MatFormer's joint optimization of a few (typically 4) nested submodels. 
matformer: 
 In this section, we define MatFormer's nested substructure (Section 3.1) and discuss its training procedure for a chosen g model granularities (Section 3.2). We then discuss elastic inference using Mix'n'Match models (Section 3.3) from MatFormer along with its deployment considerations.MatFormer defines g Transformer blocks T i , such that,T 1 ⊂ T 2 ⊂ • • • ⊂ T g where T i ⊂ T i+1indicates that the parameters of T i are contained in those of T i+1 . While it is possible to impose such a structure on any part of the Transformer, we select the FFN block to define our method and present our experiments, as the model size and computational cost of a EQUATIONwhere the weight matrices of FFN are W 1 , W 2 ∈ R dff×dmodel and bias terms are omitted for simplicity. W 1 [0 : k] denotes the submatrix with the first k rows of W 1 . Finally, σ is a non-linearity often set to GELU (Hendrycks & Gimpel, 2016) or squared ReLU (So et al., 2021) . In this work, we chose the g = 4 exponentially spaced granularities with FFN ratios of {0.5, 1, 2, 4} i.e., the nested hidden neurons are of the sizes {d f f 8 , d f f 4 , d f f 2 , d f f }.With the nested MatFormer blocks T 1 , T 2 . . . T g , we can combine these to form a MatFormer model, with g nested submodels M 1 ⊂ M 2 . . . , ⊂ M g where M i ← [T i ] ×l , i.e., M i is formed by stacking T i for l layers. The input and output embedding matrices are shared across the models.For a Transformer model M, the forward pass on an input x is denoted by M(x) and let L denote the loss function between the output and the target y: L(M(x), y).MatFormer relies on a simple training strategy of jointly optimizing all the g nested submodels together. To this end, we set the MatFormer loss as a weighted average of loss of g submodels and train for it using the standard stochastic gradient-based optimizers (Shazeer & Stern, 2018) :EQUATIONwhere λ i > 0 is the weight of i-th granular submodel. In this paper, we set {λ i } i=1...g to be uniform i.e., 1/g but explore tuning {λ i } i=1...g in Appendix E.4 to further improve MatFormer.The joint training in MatFormer involves one forward pass per each of the g submodels and benefits from portions of shared computation during backpropagation. MatFormer training results in g accurate nested submodels M 1...g inside the universal MatFormer model (M g ). Note that this simple strategy outperforms various other training techniques (Appendix E.2). Finally, instead of pretraining models with MatFomer structure, we can also induce this structure via finetuning.MatFormer training is ∼ 15% faster (for g = 4) than training all the Transformer based equivalent submodels independently (Appendix B). However, MatFormer also enables the extraction of hundreds of smaller submodels along the accuracy-vs-compute curve traced by the g explicitly optimized submodels (Section 3.3). These models emerge for free using Mix'n'Match during inference and drastically reduce the amortized training cost per model obtained through MatFormer. The joint optimization, even without self-distillation from M g , results in smaller submodels that have highly consistent behavior (Section 3.4) with the universal model. Finally, in Appendix B.1, we argue that the training efficiency of MatFormer can be significantly improved through various optimizations.At inference time, it is trivial to extract one of the g submodels M 1 ⊂ M 2 . . . , ⊂ M g by stacking the corresponding Transformer block T i across layers. However, by selecting different granularity for each MatFormer layer, it is possible to generate a combinatorially large number of accurate smaller models for free. We call this simple procedure Mix'n'Match and observe that these additional model granularities -which were never explicitly optimized -are highly performant.In fact, we can further increase the number of extracted models by generating interpolating blocks between fixed granulaties (Kusupati et al., 2022) . For example, we can generate a T block that uses first 1 2 (m i + m i+1 ) neurons in the FFN layer which still tends to be highly accurate. To summarize, given a computational budget, we can extract a highly accurate model with Mix'n'Match for the constraints rather than using a smaller less accurate model or training a model for this specific constraint (Sections 4.1.1 & 4.2). We note that a compute constraint can be satisfied by various Mix'n'Match models with different accuracies, making identifying the best Mix'n'Match configurations without downstream validation is an exciting direction for future work.During deployment, all we need to store is the single universal MatFormer model for different types of elastic inference depending on the constraints. In the case of static workloads, where compute resources are known beforehand and the inputs remain relatively similar in difficulty, one can choose the most accurate static submodel for the constraints using Mix'n'Match. This eliminates the usage of a less accurate preexisting model or training of a new one for the specific constraints.For dynamic workloads, where the compute resources or the input hardness change on the fly, we can use the universal MatFormer model to dynamically extract the optimal submodel for tokenbased routing in LLMs akin to MoE (Kudugunta et al., 2021; Li et al., 2022) and elastic encoders in dense retrieval (Section 4.2.2). This works largely because all the extracted submodels have high behavioral consistency with universal MatFormer model (Section 4.1) -minimizing the drift across predictions from various submodels. We measure the consistency between two generative models as the percentage of matching tokens generated by them for the same prefix or using the KL divergence of the smaller model outputs with the larger model outputs -this accounts for potential sampling strategies in decoding. This highly consistent nature of MatFormer results in superior inference time speedups for techniques like speculative decoding (Leviathan et al., 2023 ) (Section 4.1.1) and can assist in reducing prediction drift between cross platform deployments. We also show that higher model consistency also aids metric-space structure preservation in encoder models (Section 4.2.2). 
experiments: 
 In this section, we empirically evaluate MatFormer across modalities (language and vision), model classes (decoder and encoder), and scales (up to 2.6B parameters). Specifically, we train and analyze MatFormer-based decoder-only Language Models -MatLMs (Section 4.1) -and encoder-only Vision Transformers -MatViT (Section 4.2) models with g = 4 nested granularities across various model sizes. For a fair comparison, we also independently train the Transformer baseline for the submodel of each granularity across model sizes for the same tasks. We primarily focus on the elastic deployment of MatFormer-based models (Sections 4.1.1 & 4.2) for tasks spanning from one-shot generative evals to adaptive image retrieval. Additionally, we also investigate the reliable scaling behavior (Kaplan et al., 2020) of the MatFormer models (Section 4.1.2).We build MatFormer-based decoder-only Language Models -MatLMs -and contrast them to their vanilla Transformer counterparts (LMs) (Liu et al., 2018) . The LMs broadly follow the training pipeline and procedure outlined by Thoppilan et al. (2022) . For each MatLM model with a set d model , we jointly optimize for g = 4 nested granularities represented by FFN ratios of {0.5, 1, 2, 4} -i.e., only the hidden representation size of the FFN block changes. We denote these submodels as MatLM -{S, M, L, XL} in increasing order of model size and refer to MatLM-XL as the universal MatLM. For baselines, we train vanilla Transformer models with comparable architectures. That is, for each MatLM, we train 4 separate baseline models with FFN ratios of {0.5, 1, 2, 4} for a fixed d model denoted as Baseline -{S, M, L, XL}. We evaluate these models on validation loss (= log perplexity) and average accuracy on 26 English tasks similar to (Brown et al., 2020; Du et al., 2022; Anil et al., 2023) . Of these 26 tasks, we group 5 tasks that require generating multiple tokens under "GEN" and the remaining tasks that involve choosing an option from the input text under "RANK". Please see Appendix A for further details on training, evaluation, and the datasets.To showcase elastic inference, we evaluate the 2.6B parameter MatLM models on its ability (a) to provide models spanning the accuracy-vs-compute curve using Mix'n'Match (Section 3.3) and (b) to improve post-hoc inference optimization techniques like Speculative Decoding (Leviathan et al., 2023) to further speed-up accurate auto-regressive generation.Accurate MatLM submodels for every constraint for free with Mix'n'Match. Leveraging Mix'n'Match, a MatLM can provide accurate models for every compute constraint (between S and XL), not just the explicitly optimized granularities {S, M, L, XL}. We evaluate the impact of Mix'n'Match on the 2.6B parameter MatLM in Figure 2 through validation loss and downstream evals and contrast them to four granularities {S, M, L, XL} of the 2.6B baseline LM (all trained independently). In Figures 2a, 2b & 2c, we show that all MatLM -{S, M, L, XL} models all perform as well as their corresponding baselines -with marginal improvements and drops across the scale.In Figure 2a we see that Mix'n'Match helps obtain many models on the optimal loss-vs-compute curve at zero cost. Moreover, downstream eval tasks on these Mix'n'Match models also mimic this trend, as shown in Figures 2c & 2b . In a deployment setting that only has 55% of the required compute resources needed for the MatLM-XL model, it is now possible to have a Mix'n'Match submodel with < 2% accuracy drop on RANK evals. Without elastic deployment due to Mix'n'Match, we would see a > 2.5% accuracy drop due to the use of the MatLM-M model. Note that we highlight only a few of the hundreds of accurate Mix'n'Match models along the curves.We discuss additional details and results on the Mix'n'Match procedure in Appendix C.MatLM submodels speed up speculative decoding. Speculative decoding leverages an accurate lightweight LM as a draft model to autoregressively generate a few tokens, followed by verifying these drafts with a larger model through parallel decoding on the generated tokens. When the draft is inaccurate, the draft model is rolled back and reset to the larger model's output. This results in considerable inference speed-up for the same accuracy as the large model. We point the reader to the original paper for a more detailed explanation (Leviathan et al., 2023) .Slow down of this algorithm stems from cases where the smaller model's predictions disagree with the larger model. A draft model that is significantly more consistent with the larger verifier model would lead to less rollbacks of the draft predictions and therefore lower latency. As seen in Figure 2d the MatLM submodels can be up to 8.5% more consistent than the baselines to their corresponding XL model. The significant gap persists even in the KL divergence variant of consistency with the XL model's outputs (see Figure 8 in Appendix). This improved consistency along with the need for only a single universal model positions MatLM favorably to improve techniques that require draft and verifier models such as speculative decoding.Table 1 shows the inference time speed-ups from speculative decoding using the S and XL submodels of the 2.6B language model for drafting and verification respectively. Speculative decoding with independently trained baseline LMs results in a speed-up of up to 10% over the standard autoregressive decoding of the 2.6B-XL model. But MatLM-based speculative decoding is up to 6% faster than traditional speculative decoding. This additional speed-up can be primarily attributed to the more consistent nature of MatLM-based drafter and verifier models and is further boosted by the ability to share attention cache across models from MatLM which is infeasible for the baselines (see Appendix B.2). Finally, MatLM further reduces the memory overhead for inference by removing the need to have two models during resource-constrained deployment.Now that we have established that a 2.6B MatLM model and its submodels are as accurate as the baseline Transformer LMs, we want to examine the scalability of training MatLM models. So, we study the scaling properties (Kaplan et al., 2020; Hoffmann et al., 2022) Figure 3 : We train various decoder-only MatLM models at a range of sizes from 78M to 2.6B parameters and observe the scaling trends of all granularities (S, M, L, XL) for validation loss and 1-shot downstream evaluation scores. We find that the MatLM-XL models across scales mimic the training trends of Baseline-XL models. Interestingly, we also note that that validation loss and downstream evaluations follow the scaling trends of the XL-models across all granularities.ranging from 78M to 2.6B parameters on 10B to 160B tokens and plot the validation loss for MatLM -{S, M, L, XL} compared against their baselines in Figure 9 .First, in Figure 3a , we observe that the training of MatLM-XL models across model sizes scale as reliably as the Baseline-XL LMs for loss vs. number of parameters. However, Figure 3b interestingly shows that it is not just the XL models but rather all the nested submodels, irrespective of granularity {S, M, L, XL}, of MatLM and Baseline that follow the same scaling trend. Therefore, we fit a scaling law according to the number of non-embedding parameters (N ) and training tokens (D) for all possible submodels for both MatLMs and the baselines in Table 2 . We observe that the fitted parameters are extremely similar, suggesting that MatLMs scale similarly to vanilla Transformer LMs. Table 2: In Figures 3c & 3d we also find that the downstream evals for MatLM are within 0.5% of the baselines, with the smaller submodels even outperforming the baselines at scale. Finally, Figure 9f in the Appendix shows that the MatLM submodels are more consistent with their XL model compared to the baseline counterparts across scales.We note that the scaling law equation does not capture how (1) MatLMs have been optimized for multiple submodels and even have performant submodels that have not been explicitly optimized for (Section 4.1.1), and (2) MatLMs and baselines of the same size have different training FLOPs per step. We leave formulations that capture these subtleties to future work and further discuss this in Appendix D.1. We provide full results split by granularity in Appendix D.In this section, we extend MatFormer to Vision Transformer (ViT) (Dosovitskiy et al., 2020) based computer vision encoder models. MatFormer-based ViT -MatViT -enables elastic inference for fundamental tasks like image classification and retrieval. To this end, we train the MatFormer variant of the standard ViT-B/16 and ViT-L/16 models -MatViT-B/16 and MatViT-L/16 that are trained with g = 4 prechosen nested granularities (FFN ratios of {0.5, 1, 2, 4}). B/16 models are trained on ImageNet-1K (Russakovsky et al., 2015) with AugReg (Steiner et al., 2021) while L/16 models are pretrained on ImageNet-21K (Deng et al., 2009) followed by finetuning on ImageNet-1K. All models are trained with the training setup and optimal hyperparameters of the standard ViT variants from the Scenic library (Dehghani et al., 2022) .For image classification, we evaluate both ViT & MatViT models on ImageNet-1K. Figure 4a shows that the explicitly optimized granularities in MatViT result in as accurate models as the independently trained baselines for the B/16. However for L/16, as shown in Figure 4b , we see that the MatViT models are up to 0.35% more accurate than the baseline for the same inference cost.We then explore using MatFormer at different training stages with a 2 × 2 grid of pretrainingfinetuning pairs ( Figure 5 : MatViT natively enables elastic encoders for adaptive retrieval that can be used for realtime query side computation while retaining strong accuracy on ImageNet-1K, unlike the baselines.bring more accurate and flexible encoders for downstream use. Further, finetuning using MatFormer enhances elastic deployment depending on the constraints at hand through Mix'n'Match.Adaptive Encoders with Mix'n'Match. Furthermore, our Mix'n'match models' accuracy almost lies on the line joining accuracy of explicitly trained granularities. In scenarios where, say, an application can host 50M parameter B/16 model, MatViT can provide 0.8% more accurate model than the current approach which would host the largest baseline model with ≤ 50M parameters.During deployment, the universal MatViT model can be stored in memory and depending on the compute constraints be used to extract an adaptable smaller model to maximize accuracy with the available resources at that moment. Currently, we find the Mix'n'Match models on the accuracycompute curve through a quick inference on the validation set. While relatively scalable, this points to the need for optimal budget allocation across layers in neural networks (Kusupati et al., 2020) .The goal of image retrieval is to find semantically similar images -e.g. images from the same class -using representations obtained from a pretrained encoder (Chen et al., 2022) . Standard approach is to encode the database images as well as query image with same encoder and run nearest neighbor retrieval for the query embedding. While we can embed database images with an expensive encoder, the query encoder generally has to be real-time. Furthermore, the setting of query encoding might be varied, e.g., on-device vs. cloud processing, varying query load and query complexity. Current solutions have to stick to a fixed encoder thus compromising on accuracy or cost for various settings.Given the elastic nature of MatViT, it is a good candidate for query encoder. However, retrieval also requires that submodels preserve distances between fixed database (with large encoder) and query embeddings across all the granularities. If we use smaller baseline ViT models only for query encoding, these distances are not preserved and lead to nearly 0 retrieval accuracy (see Figure 5 ).We evaluate both ViT and MatViT encoders on ImageNet-1K for image retrieval. We compute 1-nearest neighbor (NN) accuracy using the representation vector of the [CLS] token (also see Appendix F.2). Figure 5 shows that submodels extracted from MatViT can approximately preserve distances and provide significantly more flexibility. For example, with a loss of < 0.5% accuracy, MatViT-L/16 can reduce compute cost by 40%. To our knowledge, this is the first result of its kind and opens up a wide variety of adaptive inference strategies for large-scale semantic search. 
conclusions: 
 In this work we presented MatFormer, a natively elastic Transformer architecture that allows training a single universal model which can be used to extract hundreds of smaller accurate submodels at zero additional cost at deployment time. We find that the MatFormer Language Model (MatLM) matches the perplexity & 1-shot accuracy of independently trained models. In fact, MatLM demonstrates an interesting loss-vs-compute scaling curve that is nearly independent of trained granularity indicating robust generalization to extremely large models as well. Finally, MatFormer submodels enable diverse inference time speedups like faster autoregressive generation with speculative decoding and elastic query encoders for adaptive dense retrieval across modalities. 
: 
 Large Foundation models (Anil et al., 2023; OpenAI, 2023; Dehghani et al., 2023) are deployed in a variety of settings like real-time response on mobile phones or in batch setting on multi-cluster GPUs for web-scale serving. To handle such varied settings, each model family provides a few independently trained models of different sizes. In order to cover a wide range of applications, typically these models' sizes are nearly linear on log-scale. For example, Llama family provides models with 7B, 13B, 33B and 65B parameters (Touvron et al., 2023a) .Such an approach has two key drawbacks: (a) as the models are independently trained, they incur significant overhead for colocation during inference and are not behaviorally consistent with each other which are detrimental to inference optimization techniques like speculative decoding (Leviathan et al., 2023) and model cascades (Wang et al., 2020c) , and (b) due to training overhead, practitioners typically train only a few models which do not cover the entire set of downstream use-cases. For example, a deployment setup might, say, have the latency budget to support 40B parameter Llama model, but can only host a 33B variant because the next bigger model (65B) has significantly higher latency. So, one would need to settle for a less accurate model despite the larger latency budget. While model compression approaches aim to address this issue, they typically require additional training for each model that needs to be extracted. Furthermore, when applied to LLMs, these techniques are known to significantly drop the accuracy (Jaiswal et al., 2023) .In this paper, we propose MatFormer, a natively elastic Transformer (Vaswani et al., 2023) architecture that allows for training one universal model which can be used to extract hundreds of smaller submodels without any additional training (Figure 1 ). MatFormer is a general architecture that can be applied to both encoders and decoders, is domain agnostic, and is compatible with most design choices and training pipelines of large Transformer-based models -LLMs & ViTs.MatFormer follows the principle of matryoshka representation learning (Kusupati et al., 2022) , to introduce nested substructure inside the standard Transformer block. Formally, MatFormer defines a Transformer blocks T i , such that, T 1 ⊂ T 2 ⊂ • • • ⊂ T g , where g is the number of nested transformer blocks, and T i ⊂ T i+1 relation indicates that the parameters of T i are contained in those of T i+1 . MatFormer can induce such sub-structure in both the attention and the feedforward network (FFN) blocks of the Transformer (see Figure 1 ). Consider a FFN block that has d ff neurons in the hidden layer. Then, MatFormer induces matryoshka structure on these neurons, where T i contains the first m i neurons and1 ≤ m 1 ≤ m 2 • • • ≤ m g = d ffrepresent the number of neurons for each granularity or sub-model. Intuitively, this implies that the first m 1 neurons are "most significant" neurons as they belong to all the blocks followed by the next m 2 -m 1 , and so on. We can form a similar substructure on the attention heads, with the heads being organized from "most" to "least" significant, where the more significant heads are shared by more sub-models. That is, we use only the first m i attention heads for the ith granularity. In fact, we can also introduce this sub-structure in the token embedding (d model ) supplied to each Transformer block.However, in most LLMs and ViTs, the FFN block in the Transformer accounts for more than 60% non-embedding parameters and is responsible for the largest chunk of latency during inference. So, in this work, we focus on inducing the MatFormer's nested sub-structure in the FFN block. We then stack the individual blocks (for l layers) to form g nested models (M 1...g ) with shared parameters i.e., M i ⊂ M i+1 . Finally, we jointly train these g models by combining each model's loss. This leads to a natural question: can one extract more than g models after inducing the MatFormer structure? Yes, in fact, it is possible to extract exponentially many models. Using the trained Mat-Former blocks T 1 , . . . , T g at each layer, one can form new models by Mix'n'Match, i.e., by taking an arbitrary combination of these blocks across layers. For example, in the first layer, one can select T g , the largest block, choose T 2 in the second layer, and so on, forming g l different models. As we explicitly optimized only for g models, instead of the exponentially many models, are the extracted models accurate? Surprisingly, in multiple settings, and for a various model sizes, we observe that the extracted models indeed are accurate, with accuracy scaling with the size of the extracted model.We train Matformer-based decoder-only Language Models (MatLM) up to 2.6B parameters and observe that: (a) MatLMs explicitly trained with g exponentially spaced granularities almost match validation loss and one-shot downstream evals of respective g baseline models trained independently from scratch, (b) our extracted models using Mix'n'Match lie on the accuracy-vs-parameters tradeoff curve generated by the g explicitly trained models, (c) through scaling experiments we observe that the loss vs compute law for different MatFormer models remains similar to vanilla Transformer models across different granularities and (d) the submodels extracted from MatLM have highly consistent behavior that is highly desirable for inference optimizations and deployment across scales.We further studied MatFormer-based ViT models (MatViT) and have similar observations as MatLM. For example, MatViT-L/16 improves the accuracy of the standard ViT-L/16 model on ImageNet-1K, and the extracted sub-models all match or even perform better than the independently trained baselines. Furthermore, we demonstrate that, due to high consistency, MatViT models can be used as "elastic encoders" for adaptive image retrieval. That is, the metric-space of an image encoded by the universal (i.e. the largest) MatViT model is roughly preserved by the nested submodels. Hence, based on query complexity, system load, and various other considerations, we can use one of the extracted MatViT encoders at inference time for retrieval on a fixed corpus encoded by the universal model -providing over 40% lesser compute overhead with < 0.5% drop in accuracy.We make these key contributions:1. We introduce MatFormer, which incorporates a nested sub-structure within the standard Transformer and jointly optimizes all the g granularities to produce a single, universal elastic model. 2. Employing Mix'n'Match of granularities across layers in a universal MatFormer model yields hundreds of accurate and consistent submodels without any additional training cost (Section 3). 3. MatFormer generalizes effectively to both decoder-only language models (MatLM) and vision encoders (MatViT), scaling as reliably and accurately as the standard Transformer, while enabling significantly faster autoregressive generation and large-scale adaptive dense retrieval (Section 4). 
 
 Appendix: MatFormer stands forMatryoshka Transformer due to the model's inherent nested nature. 


Decision: This submission presents a novel Transformer architecture called MatFormer to provide elastic inference across diverse deployment constraints. Specifically, the authors incorporate a nested Feed Forward Network (FFN) block structure within a standard Transformer model. During training, the authors optimize the parameters of multiple nested FFN blocks with varying sizes, enabling the extraction of hundreds of accurate smaller models without incurring additional computational costs. The authors should consider the reviewers' comments and amend the issues raised by reviewers.

Reviews: review: paper_summary: paper_summary: The authors proposed a novel Transformer architecture called MatFormer to provide elastic inference across diverse deployment constraints. Specifically, the authors incorporate a nested Feed Forward Network (FFN) block structure within a standard Transformer model. During training, the authors optimize the parameters of multiple nested FFN blocks with varying sizes, enabling the extraction of hundreds of accurate smaller models without incurring additional computational costs. Experimental results on different model classes (decoders and encoders) and modalities (language and vision) demonstrate the effectiveness of the proposed method. My detailed comments are as follows.
strength_weakness: strength_weakness: 1. The idea of jointly optimizing a nested sub-structure is similar to Slimmable networks [A] and the supernet in Neural Architecture Search. More explanations are required to clarify the differences between them. What are the new challenges of applying the nested sub-structure to transformer architectures?

2. Some important details of dynamic workloads are missing. It would be better for the authors to show more details about how to use the universal MatFormer model to dynamically extract the optimal submodel for each token or query.

3. Although the paper shows promising results in experimental settings, it lacks extensive evaluation in real-world deployment scenarios. It would be better for the author to deploy the quantized models on GPUs or CPUs and report the memory consumption, inference speed as well as accuracy.

[A] Slimmable neural networks. ICLR 2019.
questions: questions: NA
limitations: limitations: NA
paper_summary: paper_summary: This paper presents MatFormer, a nested Transformer architecture for elastic inference deployment constraints. It follows the principle of matryoshka representation learning and incorporate nested structure in the FFN modules of Transformers. Experiments show that MatFormer can (1) reliably obtain 582-850M model from a single 850M model, (2) preserve the metric-space structure for adaptive large-scale retrieval for extracted encoder, (3) friendly to speculative decoding.
strength_weakness: strength_weakness: The reviewer does not think there are major problems. Please see questions.
questions: questions: (1) Can you add a discussion section with MoE?
(2) Is it possible to obtain a MatFormer model directly from a pretrained checkpoint (e.g. in some form of up-cycling)?
limitations: limitations: Please see question 2.
paper_summary: paper_summary: The paper introduces Matformer, an elastic modeling strategy for inference that provides flexibility for latency and cost requirements. The authors base their method on the recently proposed matryoshka representation learning (MRL) paradigm to introduce nested substructures in the transformer blocks. Specifically, they introduce this in the FFN layers of the transformer and allocate a hyperparameter $g$ to design the number of granular blocks. In the paper, they use $g=4$, where the granularities are mapped as $\{d_{ff}, d_{ff}/2, d_{ff}/4, d_{ff}/8 \}$, where $d_{ff}$ represents the inner dimension of the FFN blocks. The models are trained following the original loss functions used to train baseline models (i.e., models without MRL blocks). The paper showcases results for decoder-only LLMs, encoder-only ViTs, and adaptive retrieval following a simple NN retrieval metric. The decoder-only LLMs also show that the models exhibit scaling laws similar to the baseline models (with slightly shifted constants) and show how a single model can be used efficiently for speculative decoding scenarios, with an optional attention-cache sharing mechanism.
strength_weakness: strength_weakness: 1. Scaling the proposed Matformer models will be difficult. For a given model scale, to keep the models "equivalent" to the original baselines (for downstream accuracy/log perplexity loss), the MatLM models, for example, are trained with 4x more tokens. Training for such long durations is computationally intractable at model scales that are deployed today.
    - The scaling trends (from the trained models) also do not look too promising. For example, at the XL scale for MatLM or the B-16 scale for ViTs, the MatLM models seem to perform very similarly to the baseline models. Also, if you look at the loss/accuracy for the MatLM-XL models, you get very similar losses/accuracy to the base model - which "sees" 4x fewer tokens, so overall gains are minimal.
3. While the authors show through ablations that their "least slope" method of Mix-n-Match works best, in practice, it seems difficult to select granularities unless some quick experiments are not done. For e.g., at the model depths presented in the paper, it is to figure out to use M+L blocks to outperform an L baseline potentially, but as depth increases, will relying on only two granularities work - or will there be a need to define a fine-grained slope (still following the least-slope method).
4. While the speculative decoding scenario shown is interesting, generally for effective speedups (1.5x ~ 2.5x) from speculative decoding rely on much smaller draft models relative to target models (for e.g., the original paper referred to in the paper - there is ~10x gap in model sizes). Even if the MatLM models are scaled up, with the current granularities, the smallest model size will be 50%-60% smaller than the largest model size).
    - This may not pave the way to more effective speedups unless more granular blocks $g$ are introduced. However, that is counter-intuitive to the training FLOPs spent as you'd effectively scale up more tokens $\propto$ number of $g$ blocks.
    - Is the intuition here to rely on better consistency and attention-cache sharing capabilities to scale for more effective speedups?
questions: questions: 1. Can the authors elaborate more on the attention-cache sharing mechanism? Is the sharing speed because the draft models see a better cache from the largest model's inference pass (because the largest model has to generate 1 final token) during the subsequent generation phase? And does that eventually help improve the consistency of generated tokens? If so, this favors setups where more tokens are speculated and verified. Can the authors comment on how many tokens they speculate per step?

2. In some places, it is a little difficult to understand what the results correlate to in graphs. For eg, in Section 4.2.2, the authors state "For example, with a loss of < 0.5% accuracy, MatViT-L/16 can reduce compute cost by 40%" - but it is unclear where in the graph this result is evident. There are many other instances like this in the paper. Can the authors either point to such results in the graphs? Or be more verbose in the Appendix about how these results are inferred for the readers to understand.

3. For some of the results, where the authors claim better results - the "betterness" claims are very weak. For e.g., the authors claim 0.35% better accuracy for the L/16 ViT models, but this also seems very close to the seed range for some models of this size. Similarly, for the MatLM results in Section 4.1 - they claim 0.01 log perplexity is a better result than the DynaBERT approach. This is, again, a very minor difference. Can the authors expand on this by showing either seed numbers (which might be difficult in rebuttal time - I understand this) or highlighting other papers that talk about this to understand the significance of the results?

4. For many of the downstream results in the Appendix, several results in Tables 9 and 10 are very close to random chance accuracy - making it a little difficult to gauge the gains from using the MatLM training recipe. Can the authors comment on this? Are some tasks more representative of the model's performance at these scales, which can help readers understand the significance of the results?
limitations: limitations: Yes, the limitations of most results and methods are presented throughout the paper. Nothing extra is needed.
paper_summary: paper_summary: The paper proposes Matformer, a technique to achieve elastic inference where one model is trained, encapsulating sub-models that can be extracted on demand. The main idea to achieve this is to apply a Matryoshka "self-stacking" of hidden states in the FNN blocks of transformers, which are randomly sampled at train time to for all sub-models to be trained concurrently. Experimental results on Million-parameter transformers show that this idea works very well in producing low complexity models at inference time.
strength_weakness: strength_weakness: I have identified a few weaknesses in the paper, most of which are included in the "questions" section below. Please address those:
- I have doubts on the application of the Matryoshka approach to the attention block (see questions)
- In terms of training time, it is unclear that the proposed approach would be better than training small models, when taking into account FLOPs per iteration and speed of convergence (larger models need more tokens).
- Experiments are limited to models with less than a billion parameters, while the description uses Llama 34B, 40B, and 70B as motivating examples.
-Minor: There are no theoretical results.
questions: questions: I have several questions to the authors:

Issues in the motivation of the work:
- The authors identify compression/pruning as one option to fit a bigger model to a compute budget, but argue that such scheme requires additional training. This is not very fair, since Matformer itself trains from scratch, one could also do quantization-aware, pruning-aware, compression-aware training from scratch. While this approach would still necessitate one trianing session for one model size, I disagree with the author's assertion that compression "requires additional training".
- At line 40, the authors claim that Matformer is applied to both attention and FFN blocks. However, Figure 1 shows that this is only applied to the hidden state of the FFN, and also later sections in the work (Section 3 where the scheme is presented, and Section 4 where results are reported) only use Matformwer on the FFN blocks. It seems like this in an overclaim in the introduction which should be removed. Unless Matformer can be applied to the Attention. If so, how do we define Matformer for Attention? In the FFN, it is very clear how the hidden state is arbitrarily truncated to smaller chunks in a Matryoshka fasion. What is the corresponding operation in Attention, is it applied to the attention heads? How would that impact LLM serving optimization such as KV cache optimization, FlashAttantion and etc...? And how would it differ from grouped query attention? A very vague hint to the potential of doing that is given at lines 147. Please discuss further, or just drop the issue of attention and explicitly state that the method is applied to FNN only. Right now, it's a bit confusing what is or can be done for the attention block.


On the related work section: Overall this is a good survey of similar works. I would suggest comparing to Flextron [1]. This work essentially does the same thing as Matformer as far as I can tell (authors, please let me know what the differences are). It is contemporary since it came out around the same time as the Neurips deadline, so this does not in any way impact my recommendation score for this paper. However, I think it is still good to compare to Flextron, who does elastic inference on Billion-parameter models.
[1] Flextron: Many-in-One Flexible Large Language Model, Cai et al., ICML2024


Questions on the training scheme proposed:
- Does the random sampling strategy affect convergence time? Since at each iteration, we are sampling one of exponentially many sub-networks, does it take more time to train (in terms of iterations or "epochs"). 
- How does the computational cost of training (measured in FLOPs) compares to training small models from scratch. I'd assume that training a small model from scratch would converge faster since parameter volume is known to be correlated with training token availability. And also, a smaller model would have a smaller computational cost per iteration. So, all in all, is it really easier to use Matformer, rather than training small models. I would appreciate a quantitative answer to this question, if possible.
P.S., the example of Llama3 34B, 40B, 80B is specifically talking to what I am referring to (even though it doesn't take into account the algorithmic convergence). We need to train for a number of FLOPs that is needed by Llama3-80B in order to produce a 40B model.


Comments on the experiments:
- In the earlier sections, a lot is being made about Llama models, yet all reported results are on Million-parameter models. I understand that Meta did not provide a training recipe for Llamas (nor a training dataset) - but there are open source implementations of Billion-parameter models, such as Megatron-LM. It would be good to see how Matformer scales beyond 850M parameters. E.g., Flextron does produce elastic inferencing up to 8B parameters. Currently, while all results are good, I am worried that using such small models may be deriving conclusions from toy examples.
limitations: limitations: Yes, there are no worries on that, and the authors did attach the Neurips checklist to the back of their paper.

Comments: title: Official Comment by Reviewer 51tp
Question is being side-stepped
Response to the follow up
Responses to rebuttal
comment: Thanks for your detailed responses. After reading all reviews and responses, I will maintain my score to accept.
Thank you for the support of the work. We are excited about the future applications of MatFormer as well!
Thank you! This is great. I am particularly interested in the idea of up-cycling for a potential future work - I will keep the original score of 9. Please consider accept this paper!
What I am curious about is a convergence curve showing after how many tokens consumed does the small model's loss flatten, and compare that to the many-in-one. Fig 2 (b) is showing something completely unrelated (zero-shot accuracies). Please respect the reviewer's time and don't make them re-check data that was not asked for.

Anyway, the authors should take this as something to study in the future and improve their work. I maintain my score of 5. Good job.
We thank the reviewer for the prompt reply. We answer their follow up question here:

> For instance a smaller models would consume fewer tokens until its accuracy saturates ... do you need to consume more tokens until the accuracy of the many-in-one models saturate?

We do not require more tokens during elastic training, regardless of whether the baseline models to compare against are small or large, rather _**less**_ tokens compared to baselines to reach the same quality and saturation as them. Let’s consider two baseline models, A and B, with model A being smaller or equal in size to model B. Our MatFormer model, MatFormer-B, is of the same size as the larger model B. Within MatFormer-B, consider a subnetwork, MatFormer-A, which mirrors the size and architecture of Baseline-A.
During training, we optimize MatFormer-A for the same number of tokens and steps as Baseline-A and MatFormer-B for the same number of tokens and steps as Baseline-B. Consequently, each granularity in MatFormer is trained for the same number of steps and tokens as its corresponding Baseline model. And the total compute (FLOPs) and total memory required is also same or less than training the two baseline models (Baseline-B, Baseline-A). 

Due to parameter sharing within MatFormer, optimizing either of the granularity MatFormer-A or MatFormer-B also partially trains the other granularity for free in addition, without requiring additional FLOPs. This happens due to parameter sharing, and results in MatFormer-A and MatFormer-B achieving performance _**at least**_  as that of Baseline-A and Baseline-B (Fig. 2 (a))

In other words, to achieve saturation and quality levels equivalent to Baseline-A and Baseline-B, MatFormer requires using fewer tokens and less compute compared to the Baselines. While using the same compute and tokens results in better performance.

Note this is not the main advantage of MatFormer. The primary advantage of MatFormer lies in its ability to provide thousands of intermediate models between MatFormer-A and MatFormer-B, all while using the same total FLOPs, tokens, and compute as the Baseline models, which only provide two fixed-size models.

----
We welcome further questions about the work, and would really appreciate an appropriate increase in score if reviewers’ concerns are adequately addressed
Thank you for the detailed response to the reviews and additional experimental results. After reading all reviews and response, and the overall global response, I will update my score to accept (score: 7).

Please incorporate necessary changes for the unclear sections in the revised draft.
Thanks for the rebuttal. I would like to continue the conversation:
1 & 3 & 6. The authors are side-stepping the fact that smaller models converge faster. For instance a smaller models would consume fewer tokens until its accuracy saturates. As such, training small models is much cheaper than training large models. On the other hand, how is convergence speed impacted with elastic training? Since you're training many-in-one, do you need to consume more tokens until the accuracy of the many-in-one models saturate?

4. OK and thanks for acknowledging that you only have preliminary results for the attention part.

5. Sounds good. Note that Flextron has experiments for Billion-scale models, as opposed to Matformer's results using Million-scale models. Also Flextron uses an auto-router. But again, because this work is contemporary, I do not take into consideration this in my recommendation.

7. Sounds good, thank you for agreeing that the results need to be scaled up. Good luck trying to add those to the final draft.

In conclusion, I appreciate the responses, and I am glad the authors mostly agree with my initial assessment. I maintain that this is a borderline paper leaning on the accept side. But I certainly do not think that the contribution warrants a very strong accept.