MatFormer: Nested Transformer for Elastic Inference
Authors: Fnu Devvrit
Sneha Kudugunta
Aditya Kusupati
Tim Dettmers
Kaifeng Chen
Inderjit S Dhillon
Yulia Tsvetkov
Hannaneh Hajishirzi
Sham M. Kakade
Ali Farhadi
Prateek Jain
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformer
Large-scale deployment
Efficiency
Primary area: deep_learning_architectures
TL;DR: We present 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.

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.

MatFormer: Nested Transformer for Elastic Inference
Devvrit Ao
Sneha Kudugunta
Aditya Kusupati
Tim Dettmers"
Kaifeng Chen
Inderjit D
Yulia Tsvetkoy
Hannaneh Hajishirzi
Sham Kakade
Ali Farhadi Prateek Jain
Google DeepMind
University of Texas at Austin
Unniversity of Washington Harvard University
Abstract
Foundation models applied in broad spectrum of settings with different infer- ence constraints, from massive multi-accelerator clusters to resource-constraine standalone mobile devices. However, the substantial costs associated with training these models often limit the number of unique model sizes that can be offered. Con- sequently, practitioners aligned with their specific latency and cost requirements. We present MatFormer", novel Transformer architecture designed provide elastic inference across di- verse deployment constraints. MatFormer achieves this by incorporating a nested Feed Forward Network (FFN) block structure within 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 with- out incurring additional computational costs. We empirically validate the efficacy of MatFormer across different model classes (decoders and encoders) and modali- ties (language and vision), demonstrating its potential for real-world deployment. We show that a 850M decoder-only MatFormer language model (MatLM) allows 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, extracted from a universal HatFormer-bassee VIT (MT) encoder preserve the metric-space structure for adaptive large-scale retrieval. Finally, showcase that speculative decoding with the accurate and consistent submodels extracted from MatFormer can lead to significant reduction in inference latency. Project website.
compelled select model that may not be optimally
observe that smaller encoders
Introduction
Large Foundation models [49, 45, 17] are deployed in variety of settings with different compute and accuracy demands like real-time response on mobile phones or on multi-cluster GPUs for web-scale batch serving. However, typical model families provide only a few independently trained models of different sizes. For example, the Llama-2 family provides models with 7B, 13B, 34B, and 70B parameters [59]. So practitioners are forced to choose smaller (and typically less accurate) model than their latency/cost budget. Alternatively, one can use compression/pruning to fit a bigger model given compute budget [19, 36, 53], but that requires additional training.
We introduce MatFormer, a natively elastic Transformer [61] architecture that overcomes this chal- lenge. MatFormer allows for training universal model which can be used extract hundreds of smaller submodels without any additional training cost (Figure 1). MatFormer is a general
'Equal technical contribution. *Aditya Kusupati and Prateek led project.
Correspondence: devvritdcs.ute utexas. edu, (snehakudugunta kusupati. prajainn[googll... com "MttFormer stands for Matryoshka Transformer, reflecting the model's inherent nested structure.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Add Norm
Training
Wg&3 Relind
? [ [Pffwssrrdc
ManLM-S
MaiLMMM
MatLMH+
MatLM-XL
Inference
Add & Norm
Allention
MatFormer Block
Mixin Match
Figure 1: MatFormer introduces nested structure into the Transformer's FFN block & trains all the submodels, enabling free extraction of hundreds of accurate submodels for elastic inference.
architecture that can be applied to encoders and decoders, is domain agnostic, and is compatible with standard foundation model training pipelines.
MatFormer follows the principle of matryoshka representation learning [34], to introduce nested substructure inside the standard Transformer block. Formally, MatFormer defines Transformer blocks T, such that, Ti C T2 C ... C Tg. where is the number of nested transformer blocks, and T, C. T relation indicates that the parameters of are contained in those of T MatFormer can induce such sub-structure in both the attention and the feedforward network (FFN) blocks of the Transformer (see Figure 1). Consider, say, a FFN block that has du neurons in the hidden layer. Then, MatFormer induces matryoshka structure on these neurons, where T, contains the first neurons and 1 S 1 < m2 ... < m, = du represent the number of neurons for each granularity or sub-model. Intuitively, this implies that the first my neurons are "most significant" neurons they belong to all the blocks followed by the next m2 - m., and so on.
In a departure from related work (Section 2), despite optimizing for only granularities, we are able to extract exponentially more submodels post-training. Using the trained MatFormer blocks ,, at each layer, one can form new models by Mix n'Match (Section 3.3), i.c., by taking an arbitrary combination of these blocks across layers. For example, the first layer, one can select Tg. the largest block, choose T2 in the second layer, and on, forming g different models (where I is the number of layers). Surprisingly, in multiple settings, and for 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 850M parameters and observe that: (a) MatLMs explicitly trained with y exponentially spaced granularities outperform validation loss and one-shot downstream evals of respective g baseline models trained independently from scratch, (b) our extracted models using Mix' 'n' 'aMtch lie on the accuracy-ve wwsrracyy-s-paaaae parameters trade-off curve generated by the 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 highly desirable for inference optimizations and deployment across scales.
We further study MaaFordmr- ViT models (MatViT) and make similar observations. For example, MatViT-L/16 improves the accuracy of the standard ViT-L/16 model 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 "elastic encoders" for adaptive image retrieval. That is, the metric-space of an image encoded by the universal (i.e. 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 inference time for retrieval on a fixed corpus encoded by the universal model providing over 40% less 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 Trans- former and optimizes all the g granularities to produce a single, universal elastic model.
We introduce Mix'n n' Match, a simple heuristic with no computation overhead that finds optimal submodels within given parameter budget, outperforming more complex NAS methods. This yields hundreds of accurate and consistent submodels without any 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).
2 Related Work
Transformers [61] have become the unifying model architecture for foundation models [7] across modalities like language [8], vision [17] and audio [47]. While 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 [49, 2]. post-hoc efficiency techniques like quantization [19], pruning [36], and distillation [53]. However, these solutions often are specific to the single constraint at hand, and require additional training for each new downstream usecase. This makes them far from being truly elastic solution for adaptive deployment. Lastly, Transformer based LLMs are often sped-up during inference with techniques like speculative decoding [39, 12] that benefits from the smaller draft & the larger verifier models having similar behavior or early exiting [54] to enable real- time deployment.
Obtaining multiple smaller models from a single model has been explored in the past [66, 65, 9,29,, 10] with most work focusing on CNN encoders. In this work, we focus on Transformers in decoder-only language models and pretrained vision models. Specifically, OFA [9] trains teacher CNN model, and uses distillation to finetune randomly sampled submodels (not nested) in the universal student CNN model. Moreover, OFA focuses small scale models to be deployed on end devices. In contrast, MatFormer doesn't require distillation, thereby using substantially less memory, and uses nested models. Using nested models allows us host multiple models together without significantly increasing the model's memory footprint, which is advantageous for scenarios where we want to route queries through different sub-networks. Slimmable networks [66] jointly optimizes and provide limited preset widths. Universal Slimmable network [65] extends this to sample from a continuous search space of submodels and optimizes them jointly. In contrast, MatFormer samples and only optimizes one of the preset granularities. HAT [63] trains a universal network only to learn relative performance for different architectures. For deployment, the authors use NAS to find the optimal architecture and train it from scratch before serving. In contrast, MatFormer requires no additional training and accurate subnetworks can be obtained using Mix") "i''n'T Match (Section 3.3) yielding results as good as NAS without the additional complexity. DynaBERT [27] jointly trains a fixed set of submodels, doesn't introduce any search strategy and discusses only using explicitly trained granularities as submodels. As result of joint optimization of all granularities, DynaBERT yields fewer gradient updates and hence suboptimal performance compared to MatFormer while using same compute and memory (Section 4).
Moreover, we emphasize that while most work in this area optimizes for exponential number of models, we optimize a small number of subnetworks (g = 4) to obtain an exponential number of models at inference time which leads to significantly better accuracy for large datasets. More
Table 1: Comparison of MatFormer with comparable techniques across training and inference. We emphasize that in contrast with earlier work, MatFormer requires optimizing fewer models to obtain an exponential number of models at inference time without the need for post-training or NAS. Moreover, MatFormer subnetworks are nested, allowing for adaptive retrieval & colocation of models during inference. Here, l is the number of layers in the model and exp(l) refers to exponential in
N(Models Optimizd)
NCModels Oleainal) Noaee?
Model Selection
Post-tmining Ncoled?
Architacture
Decader Modd??
MatFormer
(X))
exptl)
Mix'n'Match
Trandfrmer
expli)
expil)
NAS
Neeworks (66)
Q(1)
Q(1)
expli
exill)
NAS
CNN
ene-dos X
Somed Newwork [60]
expl])
erpil)
X
Beth
DynaBuRI [27]
0(1)
O(1)
X
Transformer
3

---

recently, some of this work has been extended to Transformer encoders [11, 52] for extracting sub- models in both static or dynamic settings. But they either fail extending further to decoder-only language models ([52J) or perform suboptimal to MatFormer (Section 4) owing to differences in their training methodology. While not the weight space, matryoshka representation learning [34] & ViT [5] showcase elasticity in output & input spaces respectively by smoothly spanning deployment constraints with minimal overhead. MatFormer, in contrast, builds upon these works by creating nested structure within the weight space instead to enable truly elastic and adaptive Transformer-based (decoder & encoder) models that span all the accuracy-Vs-compute tradeoff (statically dynamically) with minimal changes and training overhead (Figure 1). SortedNet [60] a concurrent work with similar goals which optimizes many sampled submodels (akin to prior work) unlike MatFormer's optimization of few (typically 4) nested submodels. We also note FLEXTRON [64], recent work that builds upon MatFormer by extending the nested elasticity in both MI.P and Attention Heads simultaneously, tail patches the Matformer style training, and includes router to automatically route each token within the different granularities in each layer.
In Table 1, we summarize the differences between MatFormer and the related work discussed. Among these works, we identified two central ideas commonly employed: jointly training multiple submodels and sampling random submodels. We consider DynaBERT [27] and Once-for-All (OFA) [9] as the respective most relevant prior works and compare them to MatFormer Section 4.
MatFormer
In this section, we define MatFormer's nested substructure (Section 3.1) and discuss its training procedure for chosen model granularities (Section 3.2). We then discuss elastic inference using Mix'n M'Mcch models (Section 3.3) from MatFormer along with its deployment considerations.
3.1 MatFormer Structure
MatFormer defines g Transformer blocks T,, such that, Ti C T2 C.... C Tg where T₁ C T++1 indicates that the parameters of Ti are contained in those of T.+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 majority of our experiments, as the model size and computational cost of a Transformer is dominated (around 60% for LLMs and ViTs) by the FFN block (see Appendix C. and Appendix F22 for experiments applying MatFormer to the attention block of the Transformer). So, in this work, we focus on inducing the MatFormer's 's nested sub-structure in the FFN block. We then stack individual blocks (for / layers) to form y nested models (Mi...)) with shared parameters i.e., M, C M++1.
The Transformer FFN block has single hidden layer with dy neurons and both input and outputs in IR dhasee and fixed FFN ratio := dg/ dmodel (typically 2 4). MatFormer introduces the matryoshka nested structure with g granularities on the hidden representation of the FFN block. Coneretely, nested sub-block of the Transformer, T; contains the first m, neurons of the FFN and 1 S m2 < my = dp represent the number of neurons for each granularity or sub-model. So, depending on the chosen granularity the FFN operation of T, i.e., TPPN on an input x E Ruos is:
T.FN (x) = o W0 (:00mm])) m.]")- W2[ m.].
where the weight matrices of FFN are W,. W₂ Rede danei and bias terms are omitted for simplicity.
[ k] denotes the submatrix with the first is rows of W, Finally, is a non-linearity often set to
GELU [24] or squared ReLU [56]. 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 of the sizes { 441 4 di d. We get nested submodels Mj C M2.... C Mg where M; < [T,]xl,
xl, i.e., M, formed by stacking T; for I layers. The input and output embedding matrices are shared across the models.
We note that we can form a similar sub-structure 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 the first m. attention heads for the ith granularity. We can also introduce this sub-structure in the token embedding (dmoda) supplied to each Transformer block.

---

3.2 Training
For Transformer model M, the forward pass on an input r is denoted by M(x) and let L denote the loss function between the output and the target y: L(M(x), v).
MatFormer relies on simple training strategy of randomly sampling the g nested submodels across training. To this end, for each step we randomly sample a Matformer granularity = 1,2... 9 and train for it using the standard stochastic gradient-based optimizers [55]:
SSAMPLING (z. y) ((M.(x).)).
(2)
where M, is the parameter set of i-th granular submodel, with M. chosen from a probability distribution 1P1.P2.-P93 J. For most experiments in this paper, we uniformly sample each submodel in Appendix F.3, we find that tningg this probability distribution can result in stronger submodels.
MatFormer training results in accurate nested submodels Mi...g inside the universal MatFormer model (Mg). and also enables the extraction of hundreds of smaller submodels along the accuracy- vs-compute curve traced by the explicitly optimized submodels (Section 3.3). These models emerge for free using Mix'n n Match during inference and drastically reduce the amortized training cost per model obtained through MatFormer. This method results in smaller submodels that have highly consistent behavior (Section 3.4) with the universal model.
3.3 Mix'n'Mateh
At inference time, it is trivial to extract one of the submodels Mi C M2....,.M, by stacking the corresponding Transformer block T; across layers. However, by selecting different granularities for each MatFormer layer, it is possible to generate combinatorially large number of accurate smaller models for free. We call this simple procedure Mix's 'n Match and observe that these additional model granularities - which were never explicitly optimized are highly performant.
For given compute parameter budget, there are multiple possible submodels. A common strategy to select an optimal submodel is Neural Architecture Search (NAS) [48, 68). This, however, is computationally expensive (Appendix D.2). With, Mix'n 'n'' Match, we propose gradually increasing the sub-block size with the "least slope". More concretely, we recommend selecting sub-blocks with minimal granularity changes across layers, ensuring that the size of the jth layer is at least that of the 7th layer for j > i. To give a concrete example, we find that a submodel that uses granularity d2 for half the layers then 93 for the rest will likely be better than submodel that uses 91 and 9ı for similar model size. Our heuristic is underpinned by the training methodology. where each sampled subnetwork maintains consistent layer granularity across the model. Consequently, the model adapts best to configurations where layer granularities are either uniform or display minimal variation. This intution is also backed by NAS, which predicts balanced configurations over skewed (Appendix D. Among these "balanced" configurations, we empirically found the increasing with minimum slope configuration to perform the best. In Section 4. 1.1 and Appendix D. we show that Mix' n ''MMcch works at least as well as using evolutionary search based NAS methods [48] as used by OFA [9].
To summarize, find that using Mix'n n'k''nnn MaMcc is simple, cheap and effective hueristic to select a highly performant submodel for given compute budget (Sections 4.1.1 & 4.2), We provide further details and intuition in Appendix D.I.
3.4 Deployment
The design of MatFormer is beneficial for both static and dynamic workloads:
Static Workloads To expand upon the example of Llama-2 models in Section 1, a deployment setup might, say, have the latency budget to support 40B parameter Llama model, but can only host a 34B variant because the next bigger model (70B) has significantly higher latency. Training 40B parameter from scratch would require 4.8 * 1023 FLOPs, when training a 34B model and 70B model already cost 4.08 * 10 and 8.4 20 FLOPs respectively. So, one would need to settle for a less accurate model despite the larger latency budget. With Matformer, one could obtain highly accurate 40B model for 0 additional training FLOPS. More precisely, for 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.
5

---

Dynamic Workloads For dynamic workloads, where the compute resources or the imput hardness change on the fly, we can use the universal MatFormer model dynamically extract the optimal submodel for each token query. This works especially well for MatFormer 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 high consistency results in superior inference time speedups for techniques like speculative decoding [39] (Section 4.1.1) and 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). Moreover, given the nested architecture of MatFormer, model colocation can be more memory efficient.
Experiments
In this section, we empirically evaluate MatFormer across modalities (language in Section 4.1 and vision in Section 4.2) and model classes (decoder and encoder). We demonstrate the elastic deployment of MatFormer-basee 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 [29] of MatFormer models (Section 4.1.2).
4.1 MatLM: MatFormer Language Models
Experiment Setting: We build MatFormer-based decoder-only Language Models MatL.Ms and contrast them their vanilla Transformer counterparts (L.ss) [41]. For each Matl LM model with fixed dmwdel. we optimize for g 4 nested granularities represented by FFN ratios of (0.5, 1, 2, 4] -i.e., only the hidden representation size 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 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, 4} denoted as Baseline (S, M, L. XL]. In addition, we adapt OFA [9] and DynaBERT [27] to our language modeling setup, and compare those to MatFormer at the same model size. We evaluate these models validation loss and average accuracy on 25 English tasks [8, 22, 3]. We note that no additional memory and compute is used during training these methods compared independently trained Baselines. Please see Appendix B for further details on training, baselines, evaluation, and the datasets.
45
95
Daseline
3.00
Dyna
44
Consistency [%]
90
OTA MatFormer
2.95
Accuracy
Baseline DynsBERT DFA
MirM
2.90
MatFormer Mix'n Match
MatFormer MatFormer MnM
80
Baseline
2.0
NINon-Embedd. Parameters)
2.5
3.0
3.5
4.0
4.5
1eB
40
2.0
N(Non-Em Embbdd. Parameters)
2.5
3.0
3.5
4.5
75
N(Non-Embedd. Parameters)
2.5
(a) Validation loss
(b) 1-shot Evalulation
(c) Consistency w/ XL model
Figure 2: Validation loss & one-shot downstream evaluation scores for the 850M MatLM & baseline models. Mix' nn' Match helps generate accurate and more consistent models from MatLM that lie on the performance-vs-compite curve spanned by the explicitly optimized submodels.
Results compared to baselines: To showcase efficacy of MatFormer over baselines, we evaluate 850M MatLM model with the corresponding baseline counterparts in Figure 2.
Overall, in Figures 2a and 2b we observe all granularity submodels of MatLM outperform their baseline counterparts. Specifically, we find that DynaBERT exhibits a significant 0.01 log perplexity gap compared to MatFormer on the 850M model. The underlying reason is DynaBERT employs joint optimization of all granularities, which yields to fewer gradient updates and hence suboptimal performance compared to MatFormer. DERT would require more than 15% extra compute to perform well MatLM. OFA, similar to MatFormer, maintains single universal model but

---

employs random subnetwork sampling during its training. This leads to sampling fewer models close and XL granularity. resulting in inferior performance in this regime. The performance gap manifests as bell-shaped loss curve (Figure 2a), highlighting OFA shortcomings in handling the trade-off between maintaining universal (XL.) model quality and model elasticity. Additionally, OFA's training necessitates complicated NAS strategy for optimal submodel selection. However, using NAS at scale is costly and erroneous, which we further discuss in Appendix D.2. We refer the reader to Appendix B.4 for more detailed discussion of MatFormer performance compared to baselines, and advantages and downsides each method.
4.1.1 Elastic Inference with MatLM
Accurate MatLM submodels for every constraint for free with Mix'n MMMcch. Mix 'n' Match enables the MaiLM to deliver accurate models for any compute constraint between and XI., beyond the fixed granularities (S. M. L,, XL.. We assess the efficacy of Mix' n Match on the 850M parameter MatLM, comparing validation loss and downstream performance against independently trained baseline models (S, M, L, XL). Figure 2a demonstrates that Mix 'n' Match achieves optimal loss-vs- compute trade-offs at no additional cost. Additionally, downstream evaluations in Figure 2b reinforce this trend. In deployment scenarios with only 55% of the compute resources for a MatLM-XL model, a Mix'n' n' Match submodel approximates the XL's performance with only about a 1 1% accuracy drop, compared to 2% drop when using the MatLM-M model. This highlights the efficiency of Mix'n 'n' Match in creating numerous optimal models, as exemplified by selected instances along the performance curves.
We experimented with several heuristics to select the best subnetwork, but consistently observed that gradually using larger granularities in deeper layers worked the best (Section 3.3). We find that this heuristic better than the evolutionary search based techniques [48]. used in OFA [9] Figures 2a
& 2b. We also find that applying NAS to MatFormer provides no benefit over Mix n Match in Figure 6. We discuss additional details on Mix'n nn' n'MMcch in Appendix D.1.
MatLM submodels speed up speculative decoding. Speculative decoding leverages an accurate lightweight LM as draft model to autoregressively generate a few tokens, followed by verifying these drafts with larger model through parallel decoding on the generated tokens. When the draft inaccurate, the draft model is rolled back and reset 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 more detailed explanation [39].
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 2c, MatLM submodels can be up to 11.5% more consistent than the basclines their corresponding XL model. The significant gap persists even in the KL divergence variant of consistency with the XL model's outputs (see Figure in Appendix). This improved consistency along with the need for only a single universal model positions MaILM favorably to improve techniques that require draft and verifier models such as speculative decoding.
Table 2 shows the inference time speed-ups from speculative decoding using the S and XL sub- models of the 850M language model for drafting and verification respectively. Speculative decod- ing with independently trained baseline LMs results in speed-up of up to 10% over the stan- dard autoregressive decoding of the 850M-XI model. But MatLM- speculative decoding is up to 6% faster than traditional speculative decoding. This additional speed-up can be pri- marily attributed to the more consistent nature of MatLM M-based drafter and verifier models and further boosted by the ability to share attention cache across models from MatLM which is infeasible for the baselines (see Appendix C. Finally, Matl M further reduces the memory overhead for inference by removing the need to have two models during essuucce-cossrraine deployment.
Table 2: Inference time speed-ups over standard 850M model through speculative decoding using a 393M (S) draft and 850M (XI.) verifier model.
Speculative Decoding
LAMBADA
TriviaQA
Baseline MaILM + shared attention cache
1.10x
1.08x
1.14x
1.11x
1.16x
1.1x
7

---

4.1.2
MatLM Scales as well as Vanilla Transformer LMs
MatFormer Baseline
4.00
MatFormer Baseline
3.75
Average Accuracy
8
g
3:50
3:25
38
3.0
3.00
-6
MatFormer
Baseline
288
N(Non-Embedd,
NINon Embedd. Porameters) 148
NNNon-Embedd. Parameters)
(a) Validation loss: XL-models
(b) Validation loss: all models
(c) 1-shot Evals
Figure 3: We train various decoder-only MatLM models range of sizes from 78M 850M 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-XX XL models. Interestingly. we also note that that validation loss and downstream evaluations follow the scaling trends of the XL-models across granularities.
Now that we have established that 850M MatLM model and its submodels are at least accurate as the baseline Transformer LMs, we want to examine the scalability of training MatLM models. So. study the scaling properties [29, 25] of MatLMs and compare them to vanilla Transformer baseline LMs trained for the same number of tokens. parameters 1OB to 8OB tokens (per granularity) and plot the validation loss for MatLM (S, M. L, XL) compared against independently trained baselines in Figure
train models ranging from 78M to 850M
First, in Figure 3a, we observe that the training of MatLM-XL models across model sizes scale as re- liably as the Baseline-XL XL. LMs for loss vs. number of parameters. Figure interestingly shows that all granularities (S. M, I., XI.), of MatLM and Baseline follow the same scaling trend. Therefore, we fit a scal- ing law according to the number of non-embedding parameters (/V and training tokens (D) for all possi- ble submodels for both MatLMs and the baselines in
Table 3: Fitted parameters for the scaling equa-
tion: Loss(N, NN, D) = a (ND)*
b
Baseline
14.08
-0.10
0.89
Matformer
21.60
-0.13
1.33
Table We observe that the fitted parameters are extremely similar, suggesting that MatLMs scale similarly to vanilla Transformer LMs. In Figure 3c we also find that the downstream evals for MatLM 0.3% better than the baselines, with the smaller submodels further outperforming the baselines at scale by upto 1.4%. Finally, Figure 9f in the Appendix shows that the MatLM submodels more consistent with their XI. model compared to the baseline counterparts across scales.
We note that the scaling laws do not capture how MatLMs have been optimized for multiple submodels and even have performant submodels that have not been explicitly optimized for (Section 4.1.1) We leave formulations that capture these subtleties to future work and further discuss this in Appendix 1. We provide full results split by granularity in Appendix E.
80
Top-1 Accuracy (96
85.50
Accuracy
79
85.25
78
85.00
77
MatFormer
84.75
MatFormer
Mix'n Match
Mix Matchh
76
Baseline
125
150 175 200 225 250 275 300 Total Parameters (M)
Total Parameters (M)
(a) B/16 trained on ImageNet-1K with AugReg
(b) T./16 pretrained on IN-21K + ImageNet-IK.
Figure 4: MatViT variants match or outperform standard ViT models on ImageNet- K classification and provide free extracted models that span the accuracy-compute curve through Mix 'n' Match.

---

1-NN Accuracy (%)
83.55
Accuracy
MatFormer Mixn'MMach
MatFormer
82.5
Mix'nIMatch
Baseline
125 150 200
Total Parameters (M) in Query Encoder
Total Parameters (M) in Query Encoder
(a) B/16 trained on ImageNet-IK with AugReg
(b) L/16 pretrained IN-21K ImageNet-IK.
Figure 5: MatViT natively enables elastic encoders for adaptive retrieval that can be used for real-time query side computation while retaining strong accuracy ImageNet- K, unlike the baselines.
4.2 MatViT: MatFormer Vision Transformers
We extend MatFormer Vision Transformer (ViT) [21] based computer vision encoder models. MadFormer-based ViT (MatViT) enables elastic inference for fundamental tasks like image classifica- tion and retrieval. 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 nested granularities (FFN ratios of (0.5, 1, 43). B/16 models are trained on ImageNet- 1K [50] with AugReg [57] while L/16 models are pretrained on ImageNet-21K [18] followed by finetuming on ImageNet- 1K. All models use the training setup and optimal hyperparameters of standard ViT variants from the Scenie library [16].
4.2.1 Image Classification
For image classification, we evaluate both ViT & MatViT models on ImageNet-IK. IK. 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 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 pretraining finetuning pairs (Table 7 in Appendix G.I) and find that using MatFormer during pretraining helps bring more accurate and flexible encoders for downstream use. Further, finetuning using MatFormer enhances elastic deployment depending on the constraints hand through Mix MaMtc..
Adaptive Encoders with Mix' ' Match.. Furthermore, our Mix n 'mtcch models' accuracy almost lies the line joining accuracy of explicitly trained granularities. In scenarios where, say, 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 that moment. Currently, we find the Mix' n' Match models on the accuracy- compute curve through quick inference the validation set. While relatively scalable, this points to the need for optimal budget allocation across layers in neural networks [33].
4.2.2 Adaptive Image Retrieval
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 [13]. Standard approach 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 a fixed encoder thus compromising accuracy cost for various settings.
Given the elastic nature MatViT, also requires that submodels preserve distances between fixed database (with large encoder) and query embeddings across all the granularities. 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).
good candidate for query encoder. However, retrieval

---

We evaluate both VIT and MatViT encoders on ImageNet- IK for image retrieval. We compute 1-nearest neighbor (NN) accuracy using the representation vector of the [CLS] token (also see Appendix G.2). Figure 5 shows that submodels extracted from MatViT can approximately preserve distances and provide significantly more flexibility. For example, with loss of < 0.5% accuracy. MatViT-LL/16 can reduce compute cost by 40%. This corresponds to the 175M Mix ' Match parameter model in Fig 5b, which is 40% smaller than the 300M XL model, and has < 0.5% accuracy drop. To our knowledge, this is the first result of its kind and opens up wide variety of adaptive inference strategies for large-scale semantic search.
Conclusion
In this work we presented MatFormer, natively elastic Transformer architecture that allows training 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 clastic query encoders for adaptive dense retrieval across modalities. We believe dynamically routing these models to change inference latency [32, 40, 20], and developing the hardware optimizations required a promising area for future work.
Acknowledgments
We are grateful to Aishwarya P S. Yashas B.I.. Samaga, Varun Yerram, Lovish Madaan, Anurag Arnab for assistance in setting up training pipelines, Matthew Wallingford, Pranceth Netrapalli, Orhan Firat, Rohan Anil, Tom Duerig, Luke Zettlemoyer, Manish Gupta, Rahul Sukthankar and Jeff Dean for helpful discussions, support and feedback.
We also acknowledge the computing resources and support from HYAK at the University of Wash- ington, FAS RC at Harvard University, Kempner Institute and GCP credit award for the early-stage exploration of this project. Ali Farhadi acknowledges funding from the NSF awards IIS 1652052, IIS 1703166, DARPA N66001-19-2-400 for Artificial Intelligence and Google. Sham Kakade acknowledges funding from the Office of Naval Research under award N00014-22-1-2377. This work has been made possible in part by gift. from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. Yulia Tsvetkov acknowledges support from the National Science Foundation under CAREER Grant No. IIS2142739, NSF grants No. IIS2125201 and ISS2203097, and gift funding from Google, MSR, and OpenAL. Hannanch Hajishirzi acknowledges funding through a gift from Allen Institute for Artificial Intelligence.
DARPA W91 1911NFF15-1-0544 and gifts from Allen Institute
References
[1] O. Ahia, J. Kreutzer, and S. Hooker. The low-resource double bind: An empirical study of pruning for low-resource machine translation. arXiv preprint arXiv:21 2021.
[2] AI@ Meta. Ilama 3 model card. 2024. blob/main//ODEL_.CO
md.
[3] R. Anil, A. M. Dai, O. Firat, M. Johnson, D. Lepikhin, A. Passos, S. Shakeri, E. Taropa, P. Bailey, Z. Chen, et al. Palm 2 technical report. arXiv preprint 2rX1v:2305.104403 2023.
[4] J. Berant, A. K. Chou, R. Frostig, and P. Liang. Semantic parsing on freebase from question- answer pairs. In Conference on Empirical Methods in Natural Language Processing, 2013. URL https://app. semanticschon
pegg/orpuuID. 6401679.
[5] L. Beyer, P. Izmailov, A. Kolesnikov, M. Caron, S. Korblithh, X. Zhai, M. Minderer, M. Tschan- nen, I. Alabdulmohsin, and F. Pavetic. Flexivit: One model for all patch sizes. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 14496-14506.
2023.
10

---

[6] Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi. Piqa: Reasoning about physical common- sense in natural language, 2019.
[7] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint S.8i::210.0022:: 2021.
[8] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[9] H. Cai, C. Gan, T. Wang, 2. Zhang, and S. Han. Once-for-all: Train one network and specialize it for efficient deployment. arXiv preprint arXiv: 1rX1vv:908.0979/ 2019.
[10] H. Cai, C. Gan, Lin, and S. Han. Network augmentation for tiny deep learning. arXiv preprint arXiv:2110.08890, 2021.
[11] A. Chavan, Z. Shen, Z. Liu, Z. Liu, K.-T. Cheng, and E. P. Xing. Vision transformer slimming: Multi-dimension searching in continuous optimization space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4931 1931-994, 2022.
[12] C. Chen, S. Borgeaud, G. Irving, J.-B. Lespiau, L. Sifre, and J. Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint (rXiiv2302.01316 2023,
[13] W. Chen, Y. Liu, W. Wang, E, M. Bakker, T. Georgiou, P. Fieguth, L. Liu, and M. S. Lew, Deep learning for instance retrieval: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
[14] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung,
C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. Rao, P. Barnes, Y. Tay, N. Shazeer, V. Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. Bradbury,
J. Austin, M. Isard, G. Gur-Ari, P. Yin, et al. Palm: Scaling language modeling with pathways, 2022.
[15] P. Clark, I. Cowhey, O, Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge, 2018
[16] M. Dehghani, A. Gritsenko, A. Arnab, M. Minderer, and Y. Tay. Scenie: A jax library for computer vision research and beyond. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 21393-21398, 2022.
[17] M. Dehghani, Djolonga, B. Mustafa, P. Padlewski, Heek, Gilmer, A. P. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pages 7480-7512. PMLR, 2023.
[18] J. Deng, W. Dong. R. Socher, L.-J. Li, K. Li, and Fei-Fei. Imagenet: large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 248-255 Ieee, 2009.
[19] T. Dettmers and L. Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws. In International Conference on Machine Learning, pages 7750-7774, PMLR, 2023.
[20] D. Ding, A. Mallick, C. Wang, R. Sim, S. Mukherjee, V. Ruhle, L. V. Lakshmanan, and A. H. Awadallah. Hybrid Ilm: Cost-efficient and quality-aware query routing, arXiv preprint arXiv:2404. 14618, 2024.
[21] A. Dosovitskiy, L. Beyer, Kolesnikov, D. Weissenborn, Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
[22] N. Du, Y. Huang, A. M. Dai, Tong, D. Lepikhin, Y. Xu, M. Krikun, Y. Zhou, A. Yu, O. Firat, B. Zoph, L. Fedus, M. Bosma, Z. Zhou, T. Wang. Y. E. Wang, K. Webster, M. Pellat, Duke, L. Dixon, K. Zhang, Q. V. Le, Y. Wu, Z. Chen, and C. Cui. Glam: Efficient scaling of language models with mixture-of-experts, 2022.
K. Robinson, K. Meier-Hellstem,,
11

---

[23] M. Grimaldi, L. Mocerino, A. Cipolletta, and A. Calimera. Dynamic convnets on tiny devices via nested sparsity. IEEE Internet of Things Journal, 10(6):5033-5082, 2022.
[24] D. Hendrycks and K. Gimpel. Gaussian error linear units (gelus). arXiv preprint
1rXiv:1606.08415, 2016.
[25] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. arXiv preprint urXiv:2203. 15556, 2022.
[26] S. Hooker, N. Moorosi, G. Clark, S. Bengio, and E. Denton. Characterising bias in compressed models, arXiv preprint arXiv:2010.03058, 2020.
[27] L. Hou, Z. Huang. L. Shang, X. Jiang, X. Chen, and Q. Liu, Dynabert: Dynamic bert with adaptive width and depth. Advances in Neural Information Processing Systems, 33:9782-9793, 2020.
[28] M. Joshi, E. Choi, D. Weld, and L. Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1601-1611, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 18653vv// P17-1147. URL https: / ac alaahhologggy org/P17-1147.
[29] Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. 2020.
[30] S. Kim, C. Hooper, T. Wattanawong, M. Kang, R. Yan, H. Gene. G. Dinh, Q. Huang, K. Keutzer,
M. W. Mahoney, et al. Full stack optimization of transformer inference: a survey. arXiv preprint arXiv:2302. 14017, 2023.
[31] T. Kudo and J. Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv: wwXivv/8008 (08.8888..22.0. 2018.
[32] S. Kudugunta, Y. Huang, A. Bapna, M. Krikun, D. Lepikhin, M.-T. Luong, and O. Firat. Beyond distillation: Task-level mixture-of-experts for efficient inference. arXiv preprint 10X.v:210..03742 2021.
[33] A. Kusupati, V. Ramanujan, R. Somani, M. Wortsman, P. Jain, S. Kakade, and A. Farhadi. Soft threshold weight reparameterization for learnable sparsity. In International Conference Machine Learning, pages 5544-5555. PMLR, 2020.
[34] A. Kusupati, G. Bhatt, A. Rege, M. Wallingford, A. Sinha, V. Ramanujan, W. Hloard--nydde K. Chen, S. Kakade, Jain, et Matryoshka representation learning. Advances in Neural Information Processing Systems, 35:30233-30249, 2022.
[35] T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, J. Devlin, K. Lee, K. Toutanova, L. Jones, M. Kelcey, M.-W. Chang, A. M. Dai, J. Uszkoreit, Q. Le, and S. Petrov. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452-466, 2019, doi: (1162//acl_a0002777 URL https: taaps://acllee http///aaaaatttooy.ooooo org/ 1tt//Q99-022
[36] Lagunas, E. Charlaix, V. Sanh, and A. M. Rush. Block pruning for faster transformers. arXiv preprint arXiv:2109.0458 04838, 2021.
[37] G. Lai, Q. Xie, H. Liu, Y. Yang, and E. Hovy. Race: Large-scale reading comprehension dataset from examinations, 2017.
[38] H. J. Levesque, F. Davis, and L. Morgenstern. The winograd schema challenge. In Proceedings of the Thirteenth International Conference on Principles of Knowledge Representation and Reasoning, KR 12, page 552-561. AAAI Press, 2012. ISBN 9781577355601.
[39] Y. Leviathan, M. Kalman, and Y. Matias. Fast inference from transformers via speculative decoding. 2023.
12

---

[40] M. Li, S. Gururangan, T. Dettmers, M. Lewis, T. Althoff, N. A. Smith, and L. Zettlemoyer. Franch-train merge: Embarrassingly parallel training of expert language models. arXiv preprint SBti:22008003:000 2022.
[41] P. J. Liu, M. Saleh, E. Pot, B. Goodrich, R. Sepassi, L. Kaiser, and N. Shazeer. Generating wikipedia by summarizing long sequences. arXiv preprint arXiv:1801.101901 2018.
[42] T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering, 2018.
[43] N. Mostafazadeh, N. Chambers, X. He, D. Parikh, D. Batra, L. Vanderwende, Kohli, and J. Allen. A corpus and cloze evaluation for deeper understanding of commonsense stories, In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 839 849, San Diego, California, June 2016. Association for Computational Linguistics. doi: 10. 10.188533/1/NB 1098.
URL https: http:////aaaaaahhooooooeooo orr/N16- 1098.
[44] Y. Nie, A. Williams, E. Dinan, M. Bansal, J. Weston, and D. Kiela, Adversarial nli: A new benchmark for natural language understanding, 2020.
[45] R. OpenAL. Gpt-4 technical report. arXiv,, pages 2303-08774, 2023.
[46] D. Paperno, G. Kruszewski, A. Lazaridou, Q. N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and Fernández. The lambada dataset: Word prediction requiring a broad discourse context, 2016.
[47] A. Radford, J. W. Kim, T. Xu, G. Brockman, C. McLeavey, and I. Sutskever. Robust speech recognition via large-scale weak supervision. In International Conference on Machine Leurning, pages 28492-285 20992-281188 18. PMLR, 2023.
[48] E. Real, A. Aggarwal, Y. Huang, and Q. V. Le. Regularized evolution for image classifier architecture search, 2019.
[49] M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J.-b. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schritwiieerr, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.0553 2024.
[50] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, et Imagenet large scale visual recognition challenge. International journal of computer vision, 1 115:211-255 2015.
[51] K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019.
[52] M. Salehi, S. Mehta, A. Kusupati, A. Farhadi, and H. Hajishirzi. Shares: Efficient transformers through routing with dynamic width sub-networks. Findings of Empirical Methods in Natural Language Processing, 2023.
[53] V. Sanh, L. Debut, J. Chaumond, and T. Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01 108, 2019.
[54] T. Schuster, A. Fisch, J. Gupta, M. Dehghani, D. Bahri, V. Tran, Y. Tay, and D. Metzler. Confident adaptive language modeling, Advances in Neural Information Processing Systems, 35:17456-17444 17472, 2022.
[55] N. Shazeer and M. Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference Machine Learning, pages 4596-4604, PMLR, 2018.
[56] D. R. So, W. Manke, H. Liu. Z. Dai, N. Shazeer, and Q. V. Le. Primer: Searching for efficient transformers for language modeling. arXiv preprint arXiv:2109.08668; 2021.
[57] A. Steiner, A. Kolesnikov, X. Zhai, R. Wightman, J. Uszkoreit, and L.. Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint arXiv:2106.102711 10270, 2021.
13

---

[58] R. Thoppilan, D. De Freitas, J. Hall, N. Shazeer, A. Kulshreshtha, H.-T. Cheng, A. Jin, T. Bos, L. Baker, Du. et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239, 2022.
[59] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.099288 20..099.8.8 2023.
[60] M. Valipour, M. Rezagholizadeh, H. Rajabzadeh, M. Tahaei, B. Chen, and A. Ghodsi. Sortednet, a place for every network and every network in its place: Towards a generalized solution for training many-in-one neural networks. arXiv preprint arXiv:2309.00233
2023.
[61] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and
I. Polosukhin. Attention is all you need. 2023.
[62] A. Wang. Y. Pruksachatkun, N. Nangia, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. Superglue: stickier benchmark for general-purpose language understanding systems, 2020.
[63] H. Wang. Z. Wu, Z. Liu, H. Cai, L. Zhu, C. Gan, and Han. Hat: Hardware-awarc transformers for efficient natural language processing. arXiv preprint arXiv:2005. 2020.
[64] Yu and T. Huang. Universally slimmable networks and improved training techniques, 2019.
URL https: homps://aaae
orggabs////e 19t../99...000 05134.
[65] J. Yu and T. S. Huang. Universally slimmable networks and improved training techniques, In Proceedings of the IEEE/CVF international conference on computer vision, pages 1803-1811, 2019.
[66] J. Yu, L. Yang, N. Xu, J. Yang, and T. Huang. Slimmable neural networks. arXiv preprint 181.v:1812.088988, 2018.
[67] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can machine really finish your sentence?, 2019.
[68] B. Zoph and Q. V. Le. Neural architecture search with reinforcement learning, 2017.
14

---

A Broader Impact Statement
The elasticity of MatFormer enables its use in diverse deployment scenarios. This lowers the barrier for practitioners to use foundation models tailored to their deployment scenarios. Training foundation models remains expensive, with the largest models we discuss trained on 256 TPU-v4 cores for 3 days. Moreover, while we report validation loss and downstream evaluation scores on a variety of tasks, we acknowledge the possibility that MatFormer can have adverse effects on bias [26] or undersepresented domains/languages [1].
B Implementation Details
B.1 Architecture and Training
For our experiments, we train range of MatLMs varying from size 78M to 850M for 10B-80B tokens we scale model size equally with the number of training tokens [25]. For each MatLM granularity, we also train a corresponding baseline vanilla Transformer model. That is, for each model size we train Baseline- XL, L, M, with djf 4 * dmoddl dmodet dmndel dmodel /2. All models have 16 layers, 16 attention heads, and dmnddl dhf ratio of 1 4. We train a 256k vocabulary using the SentencePiece library [31], use a maximum context length of 1024 tokens, and batch size of 1M tokens. We pretrained the 850M models on 256 v3 TPU chips. We provide further details on these models in Table 4. For further details on training setup, we point the reader to [58].
Table 4: Model details for the models scales used to conduct the experiments described in Section 4.1, with breakdown of total parameter counts, non-embedding parameter counts and FFN parameter counts for each model granularity.
Parameter Count (full / spliced)
Non- Non-Embedding Params (full spliced)
FFN Params (full)
danciel
N(tokens)
78M (74M/ / 72M 71M)
12.6M (8. 8..4//6.3MM 5.3M)
8.4M
256
10B
180M (164M/ 157M 152M)
50M 233.7M225... IM)
33.6M
512
20B
310M (272M/ / 253M 244M)
113M (75M/56M/47M)
75.6M
768
30B
463M (397M/ / 363M / 346M)
201M (134M/100M/84M)
134M
1024
40B
850M (696M / 620M 582M)
453M (302M/227M/189M)
302M
1536
80B
B.2 Downstream Evaluation
We evaluate all the LM models trained on set 25 English tasks similar to [8, 22, 14, 3J, including:
1. Open-Domain Closed-Book Question Answering tasks: TriviaQA [28], Natural Questions [35], and WebQuestions [4].
Cloze and completion tasks: LAMBADA [46], HellaSwag [67], and StoryCloze [43].
3. Winograd-style tasks: Winograd [38] and WinoGrande [51].
4. Reading comprehension: RACE [37].
Common sense reasoning: PIQA [6], ARC [15], and OpenBookQA [42].
SuperGLUE [62]
7. Natural language inference: Adversarial II.I [44].
For all the granularities corresponding to each model, we present evaluation numbers along with development set log perplexity loss on all the 25 tasks in Tables 9 to 13.
B.3 Baseline Implementation Details
Independently trained Baseline: Baselines are trained from scratch independently. where each granularity in given model size uses X tokens mentioned in Table 4. For example, for 850M model = 80B tokens. Therefore, Baseline- M, L, XL. process 4X tokens in total. Baselines trained from scratch result in the same number of models that are explicitly trained for. That is, there is no way to directly get a submodel of size between the trained granularities without using additional methods like distillation, pruning, etc.
15

---

OFA: OFA [9] trains a big model from scratch, then freezes it and starts another run of similar sized "universal model" where samples random subnetworks and optimizes them with distillation and ground truth loss. The central idea of this work is sampling random submodels and optimizing them. Distillation is orthogonal component which can be added to MatFormer as well. In order to use comparable compute and memory to the Baseline (Section 4.1), we modify the method to only optimize w.r.t ground truth loss as also done in the case of MatFormer. Overall, OFA uses 4X tokens, same all the baseline granularities trained form scratch. In each step, OFA samples a random model, where each of its layer has (S, M. L, XL) granularity sampled randomly, and optimizes it.
To obtain submodels according to a given cost constraint, OFA proposes using NAS [48] to search for optimal architecture. Specifically, OFA samples a random set of (architecture, loss) pairs from the universal model. Then, it trains a predictor on this pair, taking architecture as input and predictinn the loss. Finally, it uses an evolutionary search based routine [48] to search for the optimal architecture within specified constraint, using the predictor in its search routine. More details of NAS and comparison with MatFormer Mix n Match is presented in Section D.
DynaBERT [27] has g fixed granularities of submodels- similar to MatFormer. Similar to OFA, DynaBERT employs distillation loss as well. For a comparable analysis and to maintain compute/memory comparable to baseline, we only optimize DynaBERT wrt ground truth loss, similar to MatFormer.
DynaBERT:
The crucial difference between DynaBERT and MatFormer lies in its training methodology. In each step. DynaBERT takes 4 batches of data, optimizes each batch with each of the (S, M, L., XL.) granularity and averages these losses. Similar to Matformer and OFA, it processes total of 4X tokens. Though, since it optimizes on 4 batches of data in single step, it runs for one fourth the steps of MatFormer or OFA, and same number of steps as the baseline.
An important distinction between contribution of DynaBERT work and MatFormer is the search strategy. In MatFormer, we introduced the Mix n Match heuristic, which calculates optimal submodel with a cost constraint without any overhead. In contrast, DynaBERT doesn't 't introduce any search strategy and discusses only using explicitly trained granularities as submodels.
B.4
Comparison with Baselines
Here we do more detailed discussion of comparison between MatFormer and the baselines used in Section 4.
1. MatFormer vs trained from scratch Baseline: Let Baseline- M, L,, XL) combined process 4X tokens, where each granularity is trained on X tokens independently. Since MatFormer maintains just one model, we train it for 4X tokens. in each step sampling either of MatLM-(S. M, L, XLJ. As a result, the total compute and peak memory is still the same as Baseline since we effectively train MatLM-(S, M. L, XL) for X tokens each - matching Baseline training configuration. But since we maintain a single model, the shared parameters get trained for up to 4X tokens for free. For example, each time we train XL/L/M granularity, since S granularity is contained in them, it also receives gradients updates. Therefore, the MatLM-S effectively processes 4X tokens, MatLM-M processes 3X tokens, MatLM-L processes 2X tokens, and MatLM-XL processes X tokens. We can also see this translating to Validation loss (Figure 2a) the smaller granularities outperform Baseline by large margin. Whereas the larger granularity match in performance.
2. MatFormer vs DynaBERT: DynaBERT, similar to MatFormer, trains g nested subnetworks, but differs critically in its training methodology. It performs joint optimization - simultancously processing batches of data- 1 batch through each (S, M. L, XL granularity. The final loss being optimized is the average of each granularity's 's loss. Conversely, MatFormer samples and optimizes individual granularities in separate steps, resulting four times more gradient updates for the same amount of data and compute. This difference in training methodology translates to a substantial performance gap: DynaBERT exhibits a significant 0.01 log perplexity gap compared to MatFormer on the 850M model. Even with 15% increase in compute, DynaBERT still falls short of MatFormer's performance. Moreover, unlike MatFormer, which offers flexibility in submodel size selection through its Mix nn Match approach, DynaBERT lacks a dedicated sampling strategy.
16

---

Baseline
3.00
MatFormer
MatFormer MnM
Loss
2.95
MatFormer NAS
2.90
2.85
2.0
2.5
3.0
3.5
4.0
4.5
Parameters)
1e8
Figure 6: We search for optimal architectrure configuration using evolutionary search [48] and with Mix's n' Match. We find that Mix' n Match yields architectures lying on pareto-optimal curve, and are least good as those found using the evolutionary search.
MatFormer vs OFA: OFA maintains one universal model similar to MatFormer, but randomly samples subnetworks such that each layer is one of {S. M. L, XL). Similar to MatFormer and DynaBERT, it takes advantage of processing 4X data by a single model. But due to sampling, not enough small or large models (close to S and XL granularity model) are sampled, leading to inferior performance these regions. That is, sampling models ranging from S model size to XL model size, most of the sampled models close M-L model sizes. This can also be seen in the Loss curve, where we see a bell like curve having much higher loss close to and XL model sizes, and matching MatLM's performance in between. Another disadvantage of OFA is since random model are sampled, once can't employ simple Mix n' Match technique like in FarForrer. Rather a more complicated NAS strategy is needed find optimal submodels. To train NAS when searching for such large models is costly and erroneous, as we discuss in Appendix D.2.
C Training and Inference Costs
We currently make minimal changes and optimizations to the training scripts of vanilla Transformer architecture. other words, we use the same implementation for both Baselime and MatFormer, except using different sized splices of FFN block for each forward pass. The wall-clock time for MatLM training is the same cost as training serving, observe the 850M model FFN latency to attention latency ratio 56 44. We note that this FFN:MHA latency ratio depends highly on scale and sequence length. More specifically, for given sequence length FFN latency dominates the overall latency at scale, while the attention heads' cost increases with sequence length. We refer the reader to Kim et al. [30] for a more extensive illustration of this. We emphasize that though trained one MatFormer and compare its training time with Baselines combined, we can extract many more models than the model granularities we explicitly trained.
the granulatities baseline counterparts. During
C.1 Speculative Decoding Attention Sharing
An additional benefit of MatLM is that the attention cache is shared between the draft and verifier model. When the XL model verifies model's draft, it overwrites the attention cache with its richer latent representation compared to the one generated by the drafter model. Note that 1) this does not involve extra computation since MatLM has single universal model including both draft and verifier model; 2) attention sharing isn't possible in the Baseline since they are not explicitly trained together. Hence, latent representation of one model is quite meaningless the other model. Thus, attention sharing gives further improvement over vanilla speculative decoding shown in Table 2.
17

---

Table 850M MatFormer model, while running NAS we observe prefers balanced granularities across layers rather than skewed. On few parameter constraints we list the MnM heuristic configuration, and configuration predicted with NAS.
FFN Params budget
Mix n' Match configuration & loss
NAS predicted configuration & loss
226M
2.931
A.M.M.M...MMMMMMMMMMMMM.M..MMMMM M.M.MMMMMM.M.... MMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMM 2.922
2.94
245M
M.M.UA.M.LLMMMMMMMMMMMMMMMMMMMMM MMM.M.MAALLLLLLLLLLLLLLLLLL W.L.LXXXXXXXXXX.......
2,926
278M
[MI.J.J.ILL M.J. M.M.M.M....... 2.91
377M
I.L.XLXLLXXXL
D Search Techniques
D.1
Mix'n n'tcch
For given compute or parameter budget multiple submodels may meet the constraint. A common strategy could involve employing Neural Architecture Search (NAS) to identify the optimal architec- ture within this subset; however, this approach can be prohibitively expensive. Instead, we propose simpler Mixn' Match heuristic, which identifies optimal submodels that lie on the Pareto-optimal curve.
The Mix' n Match heuristic recommends selecting layer granularities with minimal changes across layers, ensuring that the granularity of the jth layer at least that of the jth layer for > i. This means that gradual increase in granularity with the "least slope" empirically yields the best performance among all tested mix-and-match configurations. Our heuristic is underpinned by the training methodology, where each sampled granularity configuration maintains consistent layer granularity across the model. Consequently, the model adapts best to configurations where layer granularities are either uniform display minimal variation. For instance, a configuration of [M. M, M, M, L, L, L] across layers is more effective than [S, S, S, S, S, XL, XL] despite having similar number of parameters, as it maintains& more balanced distribution of granularity opposed to skewed one.
This intuition is also supported by results from evolutionary search-based Neural Architecture Search (NAS) method [48], similar to that employed in [9]. Our NAS experiments, few of the runs present in Table 5, indicated preference for balanced configurations under various constraints, aligning with the findings from our Mix' 'n' Match heuristic.
We explored various balanced configurations, including:
1. Serrasiing-Decressiin Granularity increases until the midpoint and then decreases, e.g., [M, M, L. L, M, MJ.
2. Decreasing- Docreasing-Incrrasing:: The reverse of the previous heuristic, where granularity decreases then increases.
3. Increasing: non-decreasing sequence of layer granularities, such as [M, M, M, M, I,, L, L].
4. Decreasing: The opposite of the Increasing heuristic.
Among these configurations, the Increasing heuristic consistently outperformed the others. Therefore, we recommend adopting strategy of selecting increasing granularities across layers with the minimum slope for effective model performance.
On the 850M MatLM, we compare architectures predicted by Mix 'n' match and the NAS algorithm, in Figure 6. We find that our heuristic works at least as well as NAS, resulting in models that lie on pareto-optimal curve. Note that Mix'n n' match requires additional training of module overhead in calculating optimal submodel within a cost constraint.
D.2 NAS
Neural Architecture Search [68, 48] is a natural method to search an optimal network among possible architectures within constraint budget. In fact, OFA [9] uses an evolutionary search based technique [48] along with architecture loss predictor (Appendix B.3) in its search routine. While on smaller scale models, it might be easy to deploy NAS method similar to what OFA does, on larger scale it becomes prohibitively costly and erroncous.
18

---

To get loss corresponding to an architecture, one needs to run on a sufficiently large held out set to get an average value. For example, in our runs we used held out set of around 200M tokens. On large scale, it becomes costly to repeatedly run large models on this held out set to collect loss values. For example, to collect 16k (architecture, loss) pairs, we used 256 TPUv-4 chips for around days.
Once we have the dataset, we split it into 60%-40% train-eval split. After training the architecture loss predictor, we tested out its performance on the held out set and noticed ~ 5e 5 mean squared error.
As we move large models regime, the performance gap between different sized models diminish. For example, the gap between 78M parameter Baseline-S S and Baseline- XL, which are of size 7 7IM and 78M. is 4.01 vs 3.83. Which a gap of 0.18 loss for just 7M parameter difference. But case of 850M model, Baseline-S and Baseline-XX XL has size 582M vs 850M, the corresponding loss is 3.017 vs 2.84. Thus, a difference of 0.177 for 268M parameter difference. This trend scales to larger models, and the gap reduces between various model sizes. So, training a architecture accuracy predictor becomes increasingly challenging and erroneous, since has to learn differentiating architectures with very small differences in their loss values.
In our experiment for OFA, we take the best of the model configs predicted by NAS and the model config among the sampled points for a corresponding budget.
E
Scaling Laws for Language Decoders
We provide results split by granularities for validation loss, average score on evaluation tasks, and consistency in Figures 9, 10, and 11 respectively. We observe that the gap in validation loss and evaluation accuracy between Ma and Baselines appears to be constant. For consistency, the gap appears to reduce with scale, but one would need to scale the models by many orders of magnitude beyond what's possible today for baselines have comparable consistency with MatLMs.
E.1 Scaling laws of MatFormers vs Transformers.
Scaling laws are essential tools estimate quality as the cost of training or inference is increased. Scaling laws can help us consider various nuances of training and deployment such as overall training cost in FLOPS, training data and parameter efficiency, and inference mean FLOPS utilization vs latency for deployments.
The scaling relationship of MatFormers versus Transformers is both simple and complex. Simple, because MatFormers scaling curves for pretraining are similar to that of Transformers thus Mat- Formers only require similar amount of compute and the same hyperparameters that work for Transformers are effective for MatFormers. The complex scaling relationship comes from the fact that MatFormers allow the training of multiple models with single training run which is qualitative different from Transformers and difficult to factor into scaling equations. In terms of efficiency, if we compare the training FLOPs equivalent of all the extractable models from MatFormers, then MatFormer training alone has a clear advantage in any case where all parameters used to train standard Transformer models on the same dataset exceed 2.58P, where the MatFormer and the largest Transformer model. This is so because MatFormer uses 2.58 times more FLOPs per token for a training run than a Transformers: 4x more FLOPs for attention layers parameters and {1 + 1/2 + 1/4 + 1/8- 1.875) x more FLOPs for MLP layers.
the number of parameters of
F Further Analysis on Language Decoders
F.
KL Divergence Between S, M, L and XL Models
Figure showcases the smoother consistency calculation between two generative models measured with KL--iiergence of the smaller model's outputs with the larger model outputs. Similar to the exact match style hard consistency metric used in the main paper, there is a significant gap between the consistency of MaiLM's submodels with the MatLM-XL. model and between that of the corresponding baseline models. This points to how sampling strategies based on the output probabilities do not change the behavioral consistency between two models and that still follows the trend of generating the token with the highest probability. This smoother notion of consistency argues for the metric-space
19

---

Table For 850M model, we experiment with modifying {ps. PMr PL. PXL) sample submodels from non-umiform distribution during training and report the results across all granularities. We find that all strategies that upweight the loss for the largest granularity perform well, with modest degradation on the M and S granularties.
Model
Probabilities
M
XI
Bascline
N/A
3.017
2.971
2.910
2.840
0.44/0.3 15/0.10
2.963
2.925
2.899
2.877
0.31/0.27/0.22/00.2
2.977
2.929
2.890
2.857
MatFormer
0.25/0. 2500 25.50.25/00.2/0.25
2.970
2.934
2.886
2.846
0.20/0.22/0. 200.022///2//.344
3.000
2.939
2.885
2.836
0.17/0.20/0.22/00 2210..4
3.010
2.943
2.884
2.829
preservation given that the output classifier/embeddiing matrix shared across all the submodels of MatLM.
KL div w/ XL Model
0.5
MatFormer Mixch'Match Baseline
0.4
0.3
0.2
0.1
2.0
N(Non-Embedd. Parameters) 1e8
The smoother variant of consistency measures the KL divergence between the smaller models and the corresponding XL model. This metric, unlike the exact match accuracy variant, also accounts for different sampling strategies on the output distribution during deployment. In this figure, we plot KL divergence of S, M, L granularities with respect to XL for the 850M parameter model.
Figure 7:
F.2 Using MatFormers for Attention Sub-Block
We experiment with applying MRL to the attention sub-block of the Transformer block. More specifically, we jointly optimize 4 subnetworks with varying number of attention heads aattn 3 nattn. natin 2,, neattn/4 of size d/ Nattn. each. Figure 8, we plot the validation loss for these models (MatLM- AAttn,, their corresponding baselines and Mix Matched subnetworks. We find that similar to our experiments on Mix 'n' Matched MatLMs, Mix n 'atch helps obtain many MatLM-Attn models the optimal loss-vs-compute curve.
F33 Tuning Sampling Probability
In Table 6, we experiment with tuning the sampling probabilities for individual granularities in order investigate the trade-off between granularity size and performance. More specifically, we tune {pi.P2 P2, Ps, Pa). and find that all strategies that upweight the loss for the largest granularity perform well, with modest degradation on the M and granularties.
G Further Analysis on Vision Encoders
G.1 Decoupling Effect MatFormer on Pretraining and Finetuning
Table 7 investigates the effect of MatFormer pretaining and finetuning phases of VIT-L/16 model. ViT-I/16 is typically pretrained on ImageNet-21K and then finetuned on ImageNet- for the final evaluation. Table 7 shows that having MatFormer during pretraining generates a better model for
20

---

3.0
2,6
MatFormer
2.4
Baseline
2.2
7,50
7.75
8.00
0.25
0.50
NNNon-Embeda. Parameters) IeB
(a) Validation loss
Figure Validation loss for the 850M MatLM-Attn & baseline models.
downstream finetuning compared to regular ViT pertaining. At the same time, finetuning a vanilla pretrained VIT with MatFormer results in flexibility being induced into the model. Despite being up to 2% less accurate than its counterparts at some granularities, fine-tuned MatViT learned to reallocate the information to provide strong nested models. Considering that this is insignificant compared to pretaining costs, possible to take the largest pretrained ViT model and finetune with MatFormer to obtain deployable MatViT variant.
Table x grid of pairs to evaluate (top-1 I accuracy (%) the effects of MatFormer and standard training the pretraining (PT) ImageNet-21 K and finetuning (FT) ImageNet-1K K using L/16 architecture. Using a MatFormer during pretraining helps bring more accurate, and elastic encoders for downstream uses.
PT
# Params (M)
VIT
MatViT
306
85.26
85.57
206
85.12
84.27
VIT
156
85.02
82.79
131
84.42
82.1
306
85.58
85.61
206
85.40
MatVIT
156
85.02
131
84.41
G.2 Traditional Image Retrieval Evaluation
Table 8 showcases traditional image retrieval evaluation on ImageNet-1K where the query and the document encoders are the same for nearest neighbor retrieval. The 1-nearest neighbor (NN) based evaluation closely follows one-vs-all classification results shown Figure 4. Both MatViT variants B/16 and L/16 have submodels that have good or better retrieval performance compared to their independently trained counterparts. Concretely, MatViT-based retrieval can be up to 0.5% more accurate than the baselines while 200M parameter MatViT submodel can be more accurate than the 300M parameter ViT baseline.
21

---

Table 8: Image retrieval I-NN accuracy (%) when the query and document encoders are the same model. Similar to the image classification results, MatViT variants either match or outperform the corresponding standard ViT counterparts. Note that all the smaller models of given model in Vat are extracted for free while the baselines have to be explicitly trained for the constraints.
Encoder
# Params (M)
ViT
MatViT
85
77.46
77.38
57
76.58
76.41
B/16
43
74.90
74.49
36
71.44
71.72
300
83.17
83.67
200
82.92
83.23
L/16
150
82.81
82.89
125
82.22
82.14
3.8
MatFormer Baseline
3.6
2
3.4
3.4
3.2
2.8
N(Non-Embedd, Parameters)
N(Non-Embedd, Parameters) 1e8
(a) XL-model Loss
(h) I-model Loss
4.0
MatFormer
MatFormer
3.8
3.8
Baseline
3.6
3.4
3.4
3.2
3.2
3.0
3.0
0.0
0.5
1.0
1.5
2.0
0.D
0.5
1.0
1.5
Embend. Parameters) 1c8
N(Non-Embedd, Parameters) 1eB
(c) M-model Loss
(d) S-model Loss
90
4.00
MatFormer
Baseline
85
Consistency
3.50
80
75
3.00
MatFormer Basell
70
2
N(Non-Embedd, Parameters) 1c8
NINon-Embedd, Parameters) 1eB
(c) Loss for all granularities XL, L,, M,
(f) Consistency with the XL-models
Figure 9: We train various decoder-only MatLM models at a range of sizes from 78M to 850M parameters and observe the scaling trends each model granularity on validation loss. We observe that the gap between MatLM and the baseline appears to be constant at each granularity. The consistency between the submodels of granularities and the XI. models shows the effect of MatFormer joint training natively ensuring similar behavior across submodels.
22

---

Average Accuracy (96
Matformer
Average (%6
MatFormer
42
Baseline
42
Baseline
40
40
36
36
0
2 NINon-Embedd, Parameters) 1e8
N(Non-Embedd. Parameters) 1e8
(a) XL-model Average Score on Evals
(b) L-model Average Score on Evals
42
Average Accuracy (%)
MatFormer
Average Accuracy (%)
Baseline
40
40
38
36
MatFormer
0.0
0.5
1.0
1.5
2.0
N(Non-Embedd. Parameters) 1e8
N(Non-Embedd, Parameters) 1e8
(c) M-model Average Score on Evals
(d) S-model Average Score on Evals
Average Accuracy (96)
44
42
40
38
MatFormer Baseline
36
0
1
4
NINon-Embedd. Parameters) 1e8
(c) Average Score on Evals for all granularitics
XL, L, M, S
Figure 10: We train various decoder-only MatLM models at a range of sizes from 78M 850M parameters and observe the scaling trends for each model granularity for the average score on 1-shot evaluation. We observe that the gap between MatLM and the baseline stays the same with scale, outperforming the baselines for on all granularities for the largest models.
23

---

8500
Consistency
8500
Consistency
8000
8000
7500
MatFormer
MatFormer Baseline
Basellie
N(Non-Embedd. Parameters) leB
N(Non-Embedd. Parameters)
(a) Consistency of L-model with XL-model
(b) Consistency of L-model with XL-model
90
Consistency
Consistency
7750
80
7500
75
7250
Mat Former Baseline
MatFormer Baseline
7000
70
0.0
1.5 N/Non-Embedd, Parameters) 1e8
0.5
1.0
NINon-Embedd. Parameters) 1e8
(c) Consistency L-model with XL-model
(d) Consistency of all model granularities with XL- model L, M. S
Figure 11:
We train various decoder-only MatLM models at a range of sizes from 78M to 850M parameters and observe the scaling trends for each submodel S, M, L for the consistency with the XL model. We observe that the gap between MatL. and the baseline reduces with scale, but one would need to scale the baseline by many orders of magnitude have consistency comparable to that of MatLMs.
24

---

Table 9: Downstream Eval numbers and development log perplexity loss on 78M model size
granularities.
Downstream Task
Baseline S
MaiLMM-S
Baseline
MaiL.M-M
Baseline-I.
MaI..M-I
Baseline XI.
Mail-M-XI.
TriviaQA (EM)
0.14
0.13
0.19
0.14
0.14
0.11
0.19
0.10
NaturalQuestions (EM)
0.06
0.03
0.03
0.00
WebQuestions (EM)
0.10
0.10
0.15
0.10
0.20
0.15
0.30
0.15
LAMBADA
0.06
0.00
0.02
0.02
0.02
0.04
HellaSwag
25.42
26.59
26,00
26.28
25.95
25.95
25.95
26.32
StoryCloze
52.81
53.82
53.13
54.14
$4,46
55.32
$4,46
55.16
WSC
52.98
52.98
53.68
54.74
55.79
55.09
$2.28
57.19
WinoGrande
48.46
47.91
51.34
50.36
50.99
51.38
48,36
50.20
Winugrad
52.38
52.75
RACE-H
25.53
27.22
24.71
27.13
26.07
26.70
25.96
26.73
29.18
32.80
28.83
33.15
28.83
32.94
29.74
32.38
PIQA
55.77
$6.47
54.62
56.31
54,52
56.42
56.86
56.09
ARC-C
21.50
22.01
21.08
21.50
21.59
21.76
22.35
21.76
ARC-E
34.55
34,55
34.30
34.68
34.89
34,97
34.55
34.55
OpenBookQA
25.40
28.00
27.60
29.00
28.20
30.60
29,80
30.60
BoolQ
48.72
52.14
51.87
52.08
51.28
52.14
52.11
52.20
COPA
62.00
$8.00
62.00
59.00
63.00
58.00
60.00
59.00
RTE
53.79
53.43
52.35
53.07
51.26
52.35
51.99
52.35
WIC
49.53
41.49
49.06
47.18
47.34
47,34
47.65
47.18
MuliRC
RECORD
39.52
41.60
40.03
41.75
40.55
42.77
40,80
43.72
41.07
41.07
44.64
41.07
44.64
41.07
42.86
42.86
ANLI-RI
30,90
32.60
32.30
32.30
32.50
32.40
32.50
32.70
ANLI-R2
31.10
30.60
31.10
30.60
30.70
30.60
30.60
30.70
ANI.IR3
31.75
31.17
30.58
30.58
30.33
30.67
30.00
31,00
Average
34,59
35.17
35.02
34.86
35.41
35.20
35.12
35,53
Loss
4.011
3.874
3.966
3.82
3.905
3.776
3.83
3.74
25

---

Table 10: Downstream Eval numbers development set log perplexity loss on 180M model size granularities.
Downstream Task
Baseline-S
MaILM-S
Baseline-M
MaIL.M-M
Baseline-L.
MarLM-L
Baseline-XI.
Marl M-XI.
TriviQQA (EM)
1.04
0.79
0,98
1.05
1.66
1.11
1.66
1.21
NaturalQuestion (EM)
0.08
0.17
0.14
0.30
0.33
0.39
WebQuestions (EM)
0.59
0.44
0.44
0.64
1.28
0.79
1.33
0.94
LAMBADA
0.16
1.13
0.43
1,28
1.51
1.20
0.49
1.28
HellaSwag
27.77
27.99
27.45
28.29
27.58
28.53
28.86
28.95
StoryCloze
56.33
$7.03
57.03
57
$7.30
58.31
$8.63
58.90
WSC
55.44
56.14
56.49
58.25
58.25
57.54
57.89
WinoGrande
52.01
$2.09
50,28
$4.30
51.22
$2.25
51.54
51.54
Winograd
54.21
RACE-H
27.93
28.82
27.50
28.47
28.70
29.33
28.73
28.99
RACE-M
33.29
34,05
34.19
33.98
34,54
34,89
33.29
35.58
PIQA
57.13
59.30
56,91
59.85
57.94
59.52
59.52
60.50
ARC-C
22.53
22.10
23.63
22.87
24.06
23.55
24.66
22.95
ARC-E
40.24
40.19
40.19
41.46
41.71
41.71
41.62
42.72
OpenBookQA
30.60
34.20
30.80
33.60
31.00
33.60
34.00
35.00
BoolQ
54.13
$2.48
52.45
53.00
55.63
52.57
53.90
52.94
COPA
62.00
65.00
61.00
63.00
61.00
66.00
64.00
64.00
52.71
50.18
$2.35
50.90
50.54
47.65
52.71
48.38
WiC
47.34
47.34
47.34
49.84
47.96
47.96
47.65
47.96
RECORD
48.58
50.56
48,99
52.11
50.56
53.21
52.82
54.09
42.86
42.86
42.86
42.86
39.29
44.64
42.86
42.86
ANLI-RI
31.80
32.20
31.80
31.50
32.40
32.50
32.20
32,60
ANII-R2
30.50
29.40
31.10
29.70
32.00
29.30
30.50
30.30
ANLI-R3
30.08
31.33
30.50
30.67
33.50
30.75
30.67
30.42
Average
36.43
37.02
36,56
37.37
37.24
37.57
37.78
37.75
Loss
3,548
3.484
3.513
3.43
3,456
3.387
3.354
3.348
Table 11: Downstream Eval numbers and development set log perplexity loss
310M model size granularities.
Downstream Task
Baseline-S
MaILM-S
Baseline-M M
MaiLM-I M
Baseline- L
MaiLM-L
Baseline XL
MatLM XL
TriviaQA (EM)
2.09
3.77
2.20
4,92
5.18
5.40
NaturalQuestions (EM)
0.11
0.69
0.28
0.72
0.58
0.72
0.91
1.00
WehQuestions (FM)
2.12
2.41
1.08
2,41
1.67
2.71
2.11
3.35
LAMBADA
0.29
1.28
0.66
1.00
2.66
2.76
3.10
HellaSwag
29.89
30.49
30.05
31.28
31.18
31.87
32,52
32.65
StoryClove
60.24
60.88
60.24
61.79
62.48
WSC
61.05
$9.30
59.30
60.00
61.75
63.16
$8.95
62.46
WinoGrande
51,46
52.25
49,57
50.75
$2.41
52.72
50.91
52.33
Winograd
55.68
62.27
57.88
64.10
63.00
67.77
61.17
67.40
RACE-H
29.45
29.50
28.90
29.42
29.22
29.87
29.67
29.42
RACE-M
35.31
37.19
36.14
37.26
36.42
37.53
37.60
38.23
PIQA
58.98
60,83
59.58
61.92
59.79
63.00
62.19
62.62
ARC-C
23.38
25.17
23,21
24.06
23.81
24.49
25.00
23.98
ARC-E
42,30
43.86
44.11
45,50
44.53
47.05
46.30
48.36
OpenBookQA
BoolQ
53.43
$4.56
55,32
$3.79
$2.87
$1.87
$4.22
51.31
COPA
RTE
52,71
49.10
53.43
49.46
$1.62
48.74
$4.
51.26
47.18
48.75
47.65
47.65
47.65
47.34
47.34
47.34
MulüRC
51.67
$1.57
$2.70
51.16
53.84
$2.58
$3,28
52.76
RECORD
54.34
55.93
55.18
56.80
56.75
58.21
58.39
58.97
42.86
44.64
42.86
51.79
42,86
50.00
50.00
44.64
ANLI-RI
32.00
33.20
32.00
33.40
32.50
33.30
32.20
32.40
ANLI-R2
32.60
30.60
30.10
30.60
30.70
29.80
ANLI-RO R3
32.08
31.58
30.75
31.83
32.17
32.00
31,50
31.17
Average
37.75
38.67
37.95
39.12
38.77
40.00
39.41
39.81
3.299
3.235
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All claims in abstract are backed up by experimental results in Section 4.
2. Limitations
26

---

Table 12: Downstream Eval numbers and development set log perplexity loss on 463M model size granularities.
Downstream
Baseline-S
MatLM-S
Baseline-M
MaILM-M
Buseline-L
MatLM-L
Baseline-XL
MatLM-XL
TriviaQA (EM)
4,63
5.67
4.87
6.78
6.11
7.51
8.09
7.76
NaturalQuestion (EM)
0.61
1.19
0.80
1.33
0.94
1.41
1.66
1.33
WebQuestions
2.26
2.71
3.20
LAMBADA
2,10
1.82
2.60
1.98
3.94
2.60
3.49
3.71
HellaSwag
36.21
StoryCloze
61.25
61.46
61.36
62.21
63.66
63.23
64.24
64.35
62.81
62.11
66.32
62.46
61.05
WinoGrande
52.33
51.46
49,09
51,78
52.64
50.20
53.12
51.85
Winograd
60.07
61.17
60.07
63.37
67.40
64.47
68.50
64.84
RACE-
29.85
29.65
29.47
30.10
30.56
30.33
30.70
30.93
RACE-M
37.53
38.23
37.33
39.28
40.39
39.62
40,95
40.67
PIQA
61.26
61.33
61.48
62.40
60,99
62.79
63.17
63.44
ARC-C
23.04
24.23
24.06
24.32
24.49
25.43
23.72
24.91
ARC-E
47.35
OpenBookQA
37.20
37.40
37.00
38.60
36.40
38.80
41.00
39.20
BoolQ
52.39
51.77
56.12
51.96
50.28
53.52
54.98
55.17
COPA
67.00
65.00
73,00
64.00
71.00
63.00
67.00
70.00
RTE
52.35
55.23
53.43
50.54
$2.35
49.46
$2.35
52.35
WIC
47.34
47.34
47.34
47.34
47.34
47.34
47.34
47.34
MultRC
54.93
53.01
50.89
53.11
$2.62
$2.10
$2.33
52.89
RECORD
57.58
59.93
59.31
61.06
62.16
63.42
63.51
CB
42,86
41.07
44.64
46.43
44.64
37.50
42,86
35.71
ANLI-RI
32.60
32.90
31.70
32.30
31.40
32.40
32.50
32,30
ANLI-R2
30.70
33.70
28.40
30.50
30.40
30.90
31.20
31.90
Average
39.05
39.64
39.43
39.91
40.40
39.83
41.01
40.83
Loss
3.205
3.162
3.16
3.107
3.096
3.06
3.023
3.02
Table 13: Downstream Eval numbers and development set log perplexity loss on 850M model size granularities.
Downstream Task
Baseline S
MaIlMMS
Baseline
Baseline-I. Mat.MM-M
MaI.M
Baseline XI.
XI.
TriviaQA (EM)
6.62
10.50
9.78
11.61
11.72
13.15
13.31
15.15
NaturalQuestion (EM)
0.89
1.97
1,58
1.97
2.38
2.33
2.66
2.52
WehQuestions (EM)
3.35
3.69
4.18
2.95
4.43
3.94
4,08
4.82
LAMBADA
8.25
6.89
10.83
7.51
10.44
9.55
14.03
9.96
HellaSwag
36.64
37.42
37.70
39.21
39.64
41.48
43.40
43.43
StoryCloze
65.26
65.10
66.17
67.08
67.13
68.47
71.25
70.12
WSC
65.96
65.96
64.21
67.02
62.12
69.82
70.53
67.72
WinoGrande
51.54
$2.57
52.57
$4.46
$2.96
$4.54
$4.14
$4.85
Winugrad
69.23
64.84
71.43
67.77
70.33
71.79
72.16
73.26
RACE-H
RACE-M
40.95
41.85
41,16
42.76
42.55
44,43
44.64
44,57
PIQA
63.98
64.04
64.91
65.02
65.23
66.81
67.25
67.19
ARC-C
24.15
26.02
24.91
27.65
26.54
28.84
27.13
30.38
ARC.E
51.01
52.69
52.95
55.13
$4.92
56.27
57.11
59.01
OpenBookQA
40.40
41.60
41.20
42,20
40,80
42.20
43.00
43,40
BoolQ
50.31
$2.48
47.80
$4.80
$0.15
$4.59
$5.60
54.46
COPA
73,00
74.00
73.00
73.00
73.00
75.00
73.00
78.00
RTE
51.99
$4.87
52.35
54.15
$1.99
$4.15
$3.07
53.43
WIC
47.34
47.18
47.34
MuliRC
51.79
55.01
53,09
$4.81
$3.36
$5.61
$6.08
56.19
CB
37.50
41.07
42,86
44.64
42.86
48.21
46.43
50.00
ANLI-RI
31.80
32.70
32.10
31.80
32.20
31.20
32.60
31,30
ANLIIR2
31.50
30.20
30.90
30.00
30.60
30.40
30.40
30.90
ANI.I-S
30.25
30.92
30.17
31.33
30.00
32.08
30.58
31.92
42.58
43.34
43.35
44.23
44.01
45.42
45.83
46.11
ANLI-R3
30.83
33.17
30.08
32.50
30.83
31.58
30.92
31.25
Average
41.14
42.03
42.01
42.92
42.72
44.13
44.53
44.87
3.017
2.987
2.971
2.934
2.91
27

---

Table 14: Downstream Eval numbers development set log perplexity loss on 850M model size granularities for DynaBERT and OFA
Downstream
DynaBERT-S
DynaBERT-M
OFA-M
DynaBERT-L
OFA-L
DynaBERT-XI.
OFA-XI
TriviaQA (FM)
8.52
9.96
9.57
10.60
13.14
11.35
14.55
13.34
NaturalQuestions (EM)
1.39
1.52
1.83
2.24
2.55
2.22
2.96
2.60
WebQuestions (EM)
3.74
LAMBADA
8,85
4.72
11.84
5,74
14,88
5.98
14.65
8.31
HellaSwag
37.86
37.22
39.33
41.15
41.26
43.39
43.04
StoryCloze
65.42
64.40
66.54
66.49
68.36
68.57
69.05
69.75
WSC
64.91
68.07
68.77
65.96
71.58
68.77
72.63
68.42
WinoGrande
55.72
52.09
55.72
56.04
57.30
55.33
59.35
55.33
Winograd
69.23
68.86
69.60
68.50
72.53
70.33
75.46
72.53
RACE-H
32.70
32.62
33.19
RACE
41.43
42.41
43.04
43.04
45.54
43.73
46.52
43.73
PIQA
ARC-C
24.66
24.74
25.94
26.62
26.62
27.22
28.50
29.27
ARC-E
52.44
52.78
54.46
54.46
57.41
56.78
58.54
58.96
OpenBookQA
39.60
39.60
39.40
42.00
41.40
42,40
42.00
42.40
BoolQ
54.56
52.11
52.81
52.17
50.55
54.17 13
$1.28
54,92
COPA
70.00
68,00
72.00
72.00
75.00
71.00
73.00
74.00
RTE
53.07
52.71
54.15
52.71
49.10
52.71
46.21
52.71
WIC
47.34
47.02
47.34
MuliRC
$1.73
54.41
$3.01
54.93
55.42
55.18
$5.18
55.42
RECORD
64.60
68.91
CB
41.07
42.86
41.07
42.86
41.07
46.43
41.07
44.64
ANLI-RI
32.60
32.50
31.70
32.50
32.60
32.30
32.10
32.10
ANII R2
30.60
30.60
30.60
30.50
31.00
30.80
31.70
30.70
ANLI-R3
30.42
30.67
31.08
30.67
32.00
30.58
31.58
30.75
Average
41.71
42.72
42.64
43.9
43.44
44.47
44.17
Loss
2,993
3.01
2,942
2.935
2.895
2.89
2.854
2.863
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discuss this throughout the draft.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA].
Justification: No theoretical results.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Implementation details provided in Appendix B. Moreover, code to repro- duce Section 4.2 will be been open-sourced at link.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: Implementation details are provided in Appendix B. While code reproduce Section 4.2 has been open-sourced, the language model has been trained on proprietary data.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
28

---

Justification: Implementation details are provided in Appendix B, with citation of the models upon which these models built.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: This has not been done due to the computational expense of training language models from scratch.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: This information is provided in Appendix B.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : http:://eeuuueeeeeee
Answer: [Yes]
Justification: Broader impacts have been discussed in Appendix
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Broader impacts have been discussed in Appendix A.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: No datasets or models are released with this paper.
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
The vision models and data used in this have been public for a long time. The language models inherit the appropriate licenses of the Lamda [58] paper while we train on proprietary data.
Justification:
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
Code to reproduce experiments will be released for camera ready.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
29

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjeets
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "btZ6DGNYgy", "reviewer": "Reviewer_51tp", "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.", "questions": "NA", "limitations": "NA", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of incorporating a nested sub-structure within the standard Transformer and optimizing all the g granularities to produce a single, universal elastic model is interesting.\n\n2. The paper introduces Mix’n’Match, a simple heuristic with no computation overhead that finds optimal submodels within a given parameter budget, outperforming more complex NAS methods without any training cost.\n\n3. The results show the proposed method generalizes well to both decoder-only language models (MatLM) and vision encoders (MatViT), which has great potential in practice.\n\n4. The paper is easy to read and provides enough experimental details to reproduce.", "weaknesses": "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?\n\n2. 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.\n\n3. 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.\n\n[A] Slimmable neural networks. ICLR 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments and support towards the paper. We clarify the main concerns raised by the reviewer:\n\n-----------\n\n1. 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?\n\n**Answer**: Slimmable networks optimizes all models simultaneously. This idea is also used in DynaBERT, where it applies this training technique on transformer based architecture rather than on CNNs like Slimmable Networks. Hence, as a baseline we compare with DynaBERT (which is also more recent) and demonstrate the advantages of the MatFormer training methodology, which optimizes only one model at a time. This results in more gradient updates for the same FLOPs and memory. Furthermore, neither Slimmable Networks nor DynaBERT incorporate a strategy for selecting submodels beyond the few explicitly trained submodels. In contrast, MatFormer employs Mix’n’Match as an efficient and near-optimal model selection strategy, offering 1000+ models at no additional cost. We will clarify these differences further in the revised paper draft. \n\n-----------\n\n2. 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.\n\n**Answer**: In this paper, we primarily focus on pretraining a single model that can result in multiple performant submodels without any additional training. It is possible to use these resulting submodels with many types of algorithms that are geared towards using multiple models to improve latency such as speculative decoding (Leviathan et al., 2023; Kim et al,. 2023), model routing (Ong et al., 2024; FrugalGPT) and cascade algorithms geared towards inference efficiency  (Narasimhan et al 2024, Kolawole et al 2024).\n\nThere are recent works which build on top of our work and apply routing for dynamic workloads. We skip mentioning them to not break double blind policy, but will certainly include a discussion in our revised version. We believe that developing new algorithm variants to use MatFormer for dynamic workloads is a promising area for future work, and will add a discussion of this to the final draft.\n\n- Y. Leviathan, M. Kalman, and Y. Matias. Fast inference from transformers via speculative 479 decoding. 2023.\n- Kim, Sehoon, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W. Mahoney, Amir Gholami, and Kurt Keutzer. \"Speculative decoding with big little decoder.\" Advances in Neural Information Processing Systems 36 (2024).\n- Ong, I., Almahairi, A., Wu, V., Chiang, W. L., Wu, T., Gonzalez, J. E., ... & Stoica, I. (2024). RouteLLM: Learning to Route LLMs with Preference Data. arXiv preprint arXiv:2406.18665.\n- Chen, L., Zaharia, M., & Zou, J. (2023). Frugalgpt: How to use large language models while reducing cost and improving performance. arXiv preprint arXiv:2305.05176.\n- Narasimhan, H., Jitkrittum, W., Rawat, A. S., Kim, S., Gupta, N., Menon, A. K., & Kumar, S. (2024). Faster Cascades via Speculative Decoding. arXiv preprint arXiv:2405.19261.\n- Kolawole, S., Dennis, D., Talwalkar, A., & Smith, V. (2024). Revisiting Cascaded Ensembles for Efficient Inference. arXiv preprint arXiv:2407.02348.\n\n-----------\n\n3. 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.\n\n**Answer**: We thank the reviewer for their feedback, and will include these details in the final draft.\n\nWe welcome further questions about the work, and if key issues are addressed, we would greatly appreciate an appropriate increase in score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0nIisbAQd9", "reviewer": "Reviewer_zcbm", "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.", "questions": "(1) Can you add a discussion section with MoE?\n(2) Is it possible to obtain a MatFormer model directly from a pretrained checkpoint (e.g. in some form of up-cycling)?", "limitations": "Please see question 2.", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "(1) Important problem and a very interesting solution.\n\n(2) The writing is straightforward and easy to understand.\n\n(3) Results are comprehensive and multi-dimensional. In additional to the main claim, its ability to be reliably used in speculative decoding shows the generalization ability of the methodology", "weaknesses": "The reviewer does not think there are major problems. Please see questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their appreciation of our contributions, and answer the questions asked below:\n\n--------\n\n1. Can you add a discussion section with MoE?\n\n**Answer**: We agree that a discussion section on MoE would be appropriate, and will include a discussion section on this in the final draft. Both Matformer and MoE models are conditional computation architectures that can activate certain portions of the model depending on the input. For dynamic workloads, where the compute resources or the input hardness changes for each model query, we can use the universal MatFormer model to dynamically extract the optimal submodel for token-based routing in LLMs, similar to MoE models that focus on inference efficiency (Kudugunta et al., 2021; Li et al., 2022).\n\n- Kudugunta, S., Huang, Y., Bapna, A., Krikun, M., Lepikhin, D., Luong, M. T., & Firat, O. (2021). Beyond distillation: Task-level mixture-of-experts for efficient inference.\n- Li, M., Gururangan, S., Dettmers, T., Lewis, M., Althoff, T., Smith, N. A., & Zettlemoyer, L. (2022). Branch-train-merge: Embarrassingly parallel training of expert language models. \n\n---------\n\n2. Is it possible to obtain a MatFormer model directly from a pretrained checkpoint (e.g. in some form of up-cycling)?\n\n**Answer**: It is indeed possible to obtain a MatFormer model directly from a pretrained checkpoint. In Table 7 (Appendix G1), we validate this with MatViT. While training MatViT from scratch results in more accurate submodels, we are still able to obtain deployable submodels by finetuning ViT with the MatViT objective using 2% of the training budget.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MPcV8HvAfl", "reviewer": "Reviewer_vEVM", "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.", "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?\n\n2. 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.\n\n3. 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?\n\n4. 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": "Yes, the limitations of most results and methods are presented throughout the paper. Nothing extra is needed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper presents a simple yet effective method to introduce elasticity into existing transformer architectures without relying on compute-intensive methods such as NAS (though the method is composable with NAS).  \n2. To simplify the model design, The model relies on the efficient MRL formulation (ie, using nested sub-structures) over $g$ different layers per FFN block. \n3. For inference, instead of relying on NAS, the authors propose a simple Mix-n-Match strategy to combine all the different nested sub-substructures. They propose a simple \"least slope\" increasing method for the granularity of the model for best performance (simplifying the design space for the users) \n3. The proposed method applies to different transformer networks (both decoder-only and encoder-only). The authors train decoder-only LLMs and encoder-only ViTs to showcase the method's flexibility.\n4. The authors back their paper with comprehensive evaluations, ablation studies, results on scaling trends, and comparisons against relevant baselines (traditional training and other elastic inference methods).\n5. For the LLM inference use case, the authors showcase how additional capabilities, such as speculative decoding, are inherent to the model's design.", "weaknesses": "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.\n    - 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.\n3. 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).\n4. 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).\n    - 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.\n    - Is the intuition here to rely on better consistency and attention-cache sharing capabilities to scale for more effective speedups?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments and support. We clarify the main concerns below:\n\n1. \"Scaling ... will be difficult...Training .. intractable at model scales that are deployed today.\"\n\nWe clarify that MatFormer does not use more data or compute compared to the baselines trained separately and can be scaled similarly to the baselines. Rather than training different sized baseline models from scratch and separately, our approach trains them together and thereby enjoys the benefits of training the common parameters jointly.\n\nConsider 4 baselines trained using FLOPs-S/M/L/XL for Baseline-S/M/L/XL, resulting in total FLOPs_TOTAL=FLOPs_S+...+FLOPs_XL. These separate baseline models, collectively use memory equivalent to Memory-S+Memory-M+Memory-L+Memory-XL.\nMatFormer, in comparison, has memory=max(Memory-S,...,Memory-XL)=Memory-XL. The FLOPs used by MatFormer does not exceed FLOPs-TOTAL. \n\n2. \"The scaling trends... do not look too promising...\"\n\nThe aim of MatFormer is to provide 1000s of models along a #param-vs-performance curve that is at least as good as the baseline #param-vs-performance curve, while using less memory and the same FLOPs as training a fixed number of baselines. Note that MatFormer gives 1000s of models on this optimal curve for free, whereas the baseline provides only a fixed number of explicitly trained models (4 in our paper). MatFormer achieves this (Figure 2), where its #param-vs-performance curve matches around the XL scale but outperforms the baseline significantly at smaller scales. \n\nAdditionally, it is possible to adjust the sampling probability while maintaining the same total FLOPs, thereby achieving higher performance across models from S to XL (as shown in Table 6). We will add this discussion to the paper to clarify this.\n\n3. \"While the authors show through ablations their \"least slope\" method of Mix-n-Match...\"\n\nIn initial experiments, we found that this trend was consistent at different model scales, and the least slope method can be calculated mathematically for a given latency budget. The intuition behind this is that since \"uniform\" submodels are used during training, i.e. all layers are S/M/L/XL, making the least change to these explicitly trained granularities ought to produce the best results. We further discuss this in Appendix D1, and believe same heuristic should work when generalizing MatFormer further.\n\n4. \"While the speculative decoding...effective speedups unless more granular blocks are introduced.\"\n\nAns: Thank you for the insightful comment. We believe that the smallest model can be much smaller than 50% of the XL model, especially as we scale to much larger sizes. Consider the Llama-3.1 70B model, with model_dim=8k. We can have the smallest model as 1/16th of the total parameters, resulting in the S model having approximately ((70-2.5)/16+2.5)≈7B parameters. This is ~10x smaller than the 70B XL model.\n\nYes, having more granularities would allow for a much smaller S model. To avoid higher training costs, the smallest granularity can be quite small while maintaining the same n(granularities) - we could have 1/16, 1/4, 1/2 as granularity ratios for the S, M, L models relative to the XL model. Thank you for highlighting this, we will incorporate this discussion into the paper. \n\n5. Is the intuition here to rely on better consistency ... more effective speedups?\n\nYes, better consistency will contribute to more effective speedups. We presented significantly higher consistency as a key advantage of MatFormer and illustrated two of the many possible use cases for it. For Decoder LMs, speculative decoding benefits from this consistency (Section 4.1.1). For Encoders, it enables adaptive retrieval (Section 4.2.2).\n\n6. \"Can the authors elaborate more on the attention-cache sharing...\"\n\nYou are correct. Since S,XL models share attention blocks, when the XL model verifies the S model’s drafts, it performs a forward pass, resulting in the generation of the XL model’s KV Cache. This cache can overwrite the S model’s cache. This is not possible with baselines because their embedding spaces are different. We verified this empirically and found that baselines failed completely. \n\nWe observed higher speedups when sharing the attention-cache (Table 2). For our experiments we speculated 3 draft tokens before verification through the larger model.\n\n7. \"...little difficult to understand what the results correlate to in graphs...\"\n\nThank you for pointing this out. In the MatViT instance, we were referring to Fig. 5(b). Here, the Mix’n’Match model having 175M parameters has less than 0.5% accuracy drop compared to the XL model which has 300M parameters, resulting in ~40% speedup. We’ll clarify this and go over the paper to ensure there are no ambiguities.\n\n8. \"... seed numbers... the significance of the results?\"\n\nFor ImageNet-ViT models, the variance between runs is marginal – which we observed across multiple seeds. 0.35% gain in the range of 85% accuracy is statistically significant. We also want to mention that this gain is a bonus and is not the main takeaway for MatFormer. For LMs, 0.01 log perplexity is quite significant. On two different seeds we gave ~15% more FLOPs to DynaBERT than MatFormer, but the 0.01 log perplexity gap was not bridged. We’ll clarify this further in the final draft.\n\n9.\"... readers understand the significance of the results?\"\n\nAt small scales, the models may not be sufficient to provide meaningful comparisons for individual evaluation tasks, which can result in noise. To mitigate this, we use 25 different evaluation tasks and average the performance over them to judge the final performance. This provides a more reliable assessment of the model's capabilities, which strongly correlate with loss/perplexity. In future work, we aim to scale to even larger models.\n\n---\n\nWe welcome further questions about the work, and would really appreciate an appropriate increase in score if reviewers’ concerns are adequately addressed.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jpEf0GiqjM", "reviewer": "Reviewer_6C9L", "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.", "questions": "I have several questions to the authors:\n\nIssues in the motivation of the work:\n- 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\".\n- 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.\n\n\nOn 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.\n[1] Flextron: Many-in-One Flexible Large Language Model, Cai et al., ICML2024\n\n\nQuestions on the training scheme proposed:\n- 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\"). \n- 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.\nP.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.\n\n\nComments on the experiments:\n- 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": "Yes, there are no worries on that, and the authors did attach the Neurips checklist to the back of their paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The idea proposed is very original: train one model and obtain many sub-models at inference time.\n\nExperimentally, the approach is shown to work well with LLMs and ViTs, and shows promising results in terms of scaling with vanilla transformers. In addition to compression, it is shown that an actual speed-up in execution can be achieved in tandem with speculative decoding.\n\nFinally, the paper is very well written; thank you to the authors for making everything so clear that I was able to understand everything in my review.", "weaknesses": "I have identified a few weaknesses in the paper, most of which are included in the \"questions\" section below. Please address those:\n- I have doubts on the application of the Matryoshka approach to the attention block (see questions)\n- 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).\n- Experiments are limited to models with less than a billion parameters, while the description uses Llama 34B, 40B, and 70B as motivating examples.\n-Minor: There are no theoretical results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments and support. We clarify the main concerns below:\n\n1.\"...unclear that the proposed approach would be better than training small...\"\n\nWe clarify that MatFormer does not use more data or compute compared to the baselines trained separately and can be scaled like the baselines. Rather than training different sized baseline models from scratch and separately, our approach trains them together and enjoys the benefits of training common parameters jointly.\n\nConsider 4 baselines trained using FLOPs-S/M/L/XL for Baseline-S/M/L/XL, resulting in FLOPs_TOTAL=FLOPs_S+...+FLOPs_XL. These separate baselines, collectively use memory Memory-S+Memory-M+Memory-L+Memory-XL. MatFormer, in comparison, has memory=max(Memory-S,...,Memory-XL)=Memory-XL. The FLOPs used by MatFormer does not exceed FLOPs-TOTAL.\n\nMoreover, Baselines offer only the fixed few models they were explicitly trained for, whereas MatFormer provides thousands of accurate models that were never explicitly trained for.\n\n3. \"...I disagree...that compression \"requires additional training\".\"\n\nWe'd like to clarify what we meant by “additional training”. MatFormer, once trained (using the same FLOPs as the 4 baseline models), provides 1000s of models that can be extracted without any post-hoc training. In contrast, with compression, after training a large model, additional post-hoc processes are required to obtain smaller models for each new latency budget. To generate 1000s of models, each of these processes would need to be repeated individually. We will add this nuance in the final draft.\n\n4. \"At line 40, the authors claim that Matformer is applied to both attention and FFN...\"\n\nWe have preliminary results on applying MatFormer to the attention block in Figure 8 (Appendix F2). Similar to the FFN case, we can get many submodels for free using Mix’n’Match. As you correctly referenced, we apply MatFormer structure to the number of attention heads n. Specifically, we use the “first” n/8 heads in MatFormer-S, first n/4 of heads (superset of n/8 heads) in MatFormer-M, first n/2 heads in MatFormer-L and all n heads in MatFormer-XL. This results in the attention KV cache being reduced proportionally—⅛ / ¼ / ½ for MatFormer-S/M/L. Optimization is the same as optimizing a smaller model having n/8, n/4, or n/2 heads respectively. Flash Attention works on individual heads independently, so it will act on whatever heads are being considered for a given model. For grouped query attention, MatFormer could be applied group-wise, prioritizing the first few groups and so on.\n\n5. \"...would suggest comparing to Flextron [1]...\"\n\nWe thank the reviewer for their feedback, and will include these details in the final draft.\n\n6. \"Does the random sampling .. convergence time .. sampling one of exponentially many sub-networks...\"\n\nWe would like to clarify that we sample only one of g models per step, not one of the exponentially many possible sub-networks. Convergence is not affected by this. We have also experimented with alternative sampling methods, such as rotating sequentially through Model-S, M, L, and XL. These methods yielded similar performance to random sampling. The total training time for MatFormer is always less than or equal to the time required to train the individual baseline models separately.\n\n6. \"How does the computational cost of training (measured in FLOPs) compares to training...\"\n\nConsider the Llama3 model family as an example —34B, 40B, 80B models. Assume each model is trained for X, Y, and Z steps with the same batch size.\n\nIn contrast, MatFormer trains a single 80B model, thus taking significantly less memory than maintaining three separate 34B, 40B, and 80B models. During training, MatFormer samples and optimizes the 34B submodel for X steps, the 40B submodel for Y steps, and the entire 80B model for Z steps. This approach consumes the same total FLOPs as training the individual 34B, 40B, and 80B models separately. The advantages of MatFormer include:\n\nEnhanced Model Quality: MatFormer submodels, particularly at smaller scales, benefit from extended training of shared parameters, resulting in higher quality compared to baseline models of the same size (Appendix B4, 1st point).\n\nMix’n’Match: MatFormer provides 1000s of models through Mix’n’Match without requiring additional training. In comparison, baseline models only provide the explicitly trained 34B, 40B, and 80B models. Post-hoc techniques such as pruning or distillation adds substantial computational overhead, as discussed previously.\n\nThus, MatFormer leverages the same amount of FLOPs as training individual models but achieves greater efficiency and flexibility by consolidating the training process, avoiding the need for extensive post-hoc adjustments.\n\nWe hope this explanation clarifies the computational cost and efficiency benefits of using MatFormer relative to training smaller models from scratch.\n\n7. \"... limited to models with less than a billion parameters...\" \"...worried that using such small models...\"\n \nWe share your interest in scaling up MatFormer. Although we show this for 850M parameter models due to resource constraints, using MatFormer becomes critical as we scale to larger sizes, where training models with intermediate sizes becomes increasingly challenging. We use these larger models as motivating examples for elastic inference algorithms that enable users to obtain any sized model during inference without any post-training. We hope this clarifies our intention behind mentioning Llama models.\n\nWe will try to include billions parameter sized models in the final draft. Additionally, our scaling laws show that MatFormer scales as well as baselines. This suggests that larger models will behave similarly\n\n---\nWe would be very happy to discuss any further questions about the work, and would really appreciate an appropriate increase in score if reviewers’ concerns are adequately addressed.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (poster)