Title: MatFormer: Nested Transformer for Elastic Inference
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transformer

Large-scale deployment

Efficiency
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

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 2 , 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. Project website.

1 Introduction

Large Foundation models [49, 45, 17] are deployed in a 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 a smaller (and typically less accurate) model than their latency/cost budget. Alternatively, one can use compression/pruning to fit a bigger model in a given compute budget [19, 36, 53], but that requires additional training.

We introduce MatFormer, a natively elastic Transformer [61] architecture that overcomes this challenge. MatFormer allows for training one universal model which can be used to extract hundreds of smaller submodels without any additional training cost (Figure 1). MatFormer is a general

Training

Add & Norm

𝞼

(

)

Add & Norm

Attention

FFN

Inference

MatFormer Block

MatLM-S

MatLM-M

MatLM-L

MatLM-XL

Mix`n`Match

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 i , such that, T 1 ⊂ T 2 ⊂ ·· · ⊂ T g , where g is the number of nested transformer blocks, and T i ⊂ T i +1 relation indicates that the parameters of T i are contained in those of T i +1 . MatFormer can induce such sub-structure in both the attention and the feedforward network (FFN) blocks of the Transformer (see Figure 1). Consider, say, a FFN block that has d ff neurons in the hidden layer. Then, MatFormer induces matryoshka structure on these neurons, where T i contains the first m i neurons and 1 ≤ m 1 < m 2 · · · < m g = d ff represent the number of neurons for each granularity or sub-model. Intuitively, this implies that the first m 1 neurons are 'most significant' neurons as they belong to all the blocks followed by the next m 2 -m 1 , and so on.

In a departure from related work (Section 2), despite optimizing for only g granularities, we are able to extract exponentially more submodels post-training. Using the trained MatFormer blocks T 1 , . . . , T g at each layer, one can form new models by Mix'n'Match (Section 3.3), i.e., by taking an arbitrary combination of these blocks across layers. For example, in the first layer, one can select T g , the largest block, choose T 2 in the second layer, and so on, forming g l different models (where l 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 g 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'Match lie on the accuracy-vs-parameters trade-off curve generated by the g explicitly trained models, (c) through scaling experiments we observe that the loss vs compute law for different MatFormer models remains similar to vanilla Transformer models across different granularities and (d) the submodels extracted from MatLM have highly consistent behavior that is highly desirable for inference optimizations and deployment across scales.

We further study MatFormer-based ViT models (MatViT) and make similar observations. For example, MatViT-L/16 improves the accuracy of the standard ViT-L/16 model on ImageNet-1K, and the extracted sub-models all match or even perform better than the independently trained baselines. Furthermore, we demonstrate that, due to high consistency, MatViT models can be used as 'elastic encoders' for adaptive image retrieval. That is, the metric-space of an image encoded by the universal (i.e. largest) MatViT model is roughly preserved by the nested submodels. Hence, based on query complexity, system load, and various other considerations, we can use one of the extracted MatViT encoders at inference time for retrieval on a fixed corpus encoded by the universal model - providing over 40% less compute overhead with < 0 . 5% drop in accuracy.

We make these key contributions:

We introduce MatFormer, which incorporates a nested sub-structure within the standard Transformer and optimizes all the g granularities to produce a single, universal elastic model.

We introduce Mix'n'Match, a simple heuristic with no computation overhead that finds optimal submodels within a given parameter budget, outperforming more complex NAS methods. This yields hundreds of accurate and consistent submodels without any training cost (Section 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 a 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, 23, 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 a teacher CNN model, and uses distillation to finetune randomly sampled submodels (not nested) in the universal student CNN model. Moreover, OFA focuses on 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 to 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'n'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 a 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 an 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

recently, some of this work has been extended to Transformer encoders [11, 52] for extracting submodels in both static or dynamic settings. But they either fail at extending further to decoder-only language models ([52]) or perform suboptimal to MatFormer (Section 4) owing to differences in their training methodology. While not in the weight space, matryoshka representation learning [34] & FlexiViT [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 or dynamically) with minimal changes and training overhead (Figure 1). SortedNet [60] is a concurrent work with similar goals which optimizes many sampled submodels (akin to prior work) unlike MatFormer's optimization of a few (typically 4) nested submodels. We also note FLEXTRON [64], a recent work that builds upon MatFormer by extending the nested elasticity in both MLP and Attention Heads simultaneously, tail patches the Matformer style training, and includes a 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 in Section 4.

3 MatFormer

In this section, we define MatFormer's nested substructure (Section 3.1) and discuss its training procedure for a chosen g model granularities (Section 3.2). We then discuss elastic inference using Mix'n'Match models (Section 3.3) from MatFormer along with its deployment considerations.

3.1 MatFormer Structure

MatFormer defines g Transformer blocks T i , such that, T 1 ⊂ T 2 ⊂ · · · ⊂ T g where T i ⊂ T i +1 indicates that the parameters of T i are contained in those of T i +1 . While it is possible to impose such a structure on any part of the Transformer, we select the FFN block to define our method and present a 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 F.2 for experiments applying MatFormer to the attention block of the Transformer). So, in this work, we focus on inducing the MatFormer's nested sub-structure in the FFN block. We then stack individual blocks (for l layers) to form g nested models ( M 1 ...g ) with shared parameters i.e., M i ⊂ M i +1 .

The Transformer FFN block has a single hidden layer with d ff neurons and both input and outputs in R d model , and fixed FFN ratio := d ff /d model (typically ≥ 4 ). MatFormer introduces the matryoshka nested structure with g granularities on the hidden representation of the FFN block. Concretely, a nested sub-block of the Transformer, T i contains the first m i neurons of the FFN and 1 ≤ m 1 < · · · < m g = d ff represent the number of neurons for each granularity or sub-model. So, depending on the chosen granularity the FFN operation of T i i.e., T FFN i on an input x ∈ R d model is:

where the weight matrices of FFN are W 1 , W 2 ∈ R d ff × d model and bias terms are omitted for simplicity. W 1 [0 : k ] denotes the submatrix with the first k rows of W 1 . Finally, σ is a non-linearity often set to GELU[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 are of the sizes { d ff 8 , d ff 4 , d ff 2 , d ff } . We get g nested submodels M 1 ⊂ M 2 . . . , ⊂ M g where M i ← [ T i ] × l , i.e., M i is formed by stacking T i for l layers. The input and output embedding matrices are shared across the models.

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 i attention heads for the i th granularity. We can also introduce this sub-structure in the token embedding ( d model) supplied to each Transformer block.

3.2 Training

For a Transformer model M , the forward pass on an input x is denoted by M ( x ) and let L denote the loss function between the output and the target y : L ( M ( x ) , y ) .

MatFormer relies on a simple training strategy of randomly sampling the g nested submodels across training. To this end, for each step we randomly sample a Matformer granularity i = 1 , 2 ..., g and train for it using the standard stochastic gradient-based optimizers [55]:

where M i is the parameter set of i -th granular submodel, with M i chosen from a probability distribution { p 1 , p 2 ...p g } . For most experiments in this paper, we uniformly sample each submodel in Appendix F.3, we find that tuning this probability distribution can result in stronger submodels.

MatFormer training results in g accurate nested submodels M 1 ...g inside the universal MatFormer model ( M g ), and also enables the extraction of hundreds of smaller submodels along the accuracyvs-compute curve traced by the g explicitly optimized submodels (Section 3.3). These models emerge for free using Mix'n'Match during inference and drastically reduce the amortized training cost per model obtained through MatFormer. This method results in smaller submodels that have highly consistent behavior (Section 3.4) with the universal model.

3.3 Mix'n'Match

At inference time, it is trivial to extract one of the g submodels M 1 ⊂ M 2 . . . , ⊂ M g by stacking the corresponding Transformer block T i across layers. However, by selecting different granularities for each MatFormer layer, it is possible to generate a combinatorially large number of accurate smaller models for free. We call this simple procedure Mix'n'Match and observe that these additional model granularities - which were never explicitly optimized - are highly performant.

For a given compute or 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'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 j th layer is at least that of the i th layer for j > i . To give a concrete example, we find that a submodel that uses granularity g 2 for half the layers then g 3 for the rest will likely be better than a submodel that uses g 1 and g 4 for a 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.1). 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.1, we show that Mix'n'Match works at least as well as using evolutionary search based NAS methods [48] as used by OFA [9].

To summarize, we find that using Mix'n'Match is a simple, cheap and effective hueristic to select a highly performant submodel for a given compute budget (Sections 4.1.1 & 4.2). We provide further details and intuition in Appendix D.1.

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 a 40B parameter from scratch would require 4 . 8 ∗ 10 23 FLOPs, when training a 34B model and 70B model already cost 4 . 08 ∗ 10 23 and 8 . 4 ∗ 10 23 FLOPs respectively. So, one would need to settle for a less accurate model despite the larger latency budget. With Matformer, one could obtain a 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.

Dynamic Workloads For dynamic workloads, where the compute resources or the input hardness change on the fly, we can use the universal MatFormer model to dynamically extract the optimal submodel for each token or 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 can assist in reducing prediction drift between cross platform deployments. We also show that higher model consistency also aids metric-space structure preservation in encoder models (Section 4.2.2). Moreover, given the nested architecture of MatFormer, model colocation can be more memory efficient.

4 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-based models (Sections 4.1.1 & 4.2) for tasks spanning from one-shot generative evals to adaptive image retrieval. Additionally, we also investigate the reliable scaling behavior [29] of MatFormer models (Section 4.1.2).

4.1 MatLM: MatFormer Language Models

Experiment Setting: We build MatFormer-based decoder-only Language Models - MatLMs - and contrast them to their vanilla Transformer counterparts (LMs) [41]. For each MatLM model with fixed d model, we optimize for g = 4 nested granularities represented by FFN ratios of { 0 . 5 , 1 , 2 , 4 } - i.e., only the hidden representation size of the FFN block changes. We denote these submodels as MatLM - {S, M, L, XL} in increasing order of model size and refer to MatLM-XL as the universal MatLM. For baselines, we train vanilla Transformer models with comparable architectures. That is, for each MatLM, we train 4 separate baseline models with FFN ratios of { 0 . 5 , 1 , 2 , 4 } 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 on 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 to independently trained Baselines. Please see Appendix B for further details on training, baselines, evaluation, and the datasets.

2.0

2.5

3.0

3.5

4.0

4.5

N(Non-Embedd. Parameters) 1e8

2.85

2.90

2.95

3.00

Loss

Baseline

DynaBERT

OFA

MatFormer

MatFormer - MnM

(a) Validation loss

2.0

2.5

3.0

3.5

4.0

4.5

N(Non-Embedd. Parameters) 1e8

40

41

42

43

44

45

Accuracy (%)

Baseline

DynaBERT

OFA

MatFormer

MatFormer - MnM

(b) 1-shot Evalulation

2.0

2.5

3.0

3.5

4.0

N(Non-Embedd. Parameters) 1e8

75

80

85

90

95

Consistency (%)

MatFormer

Mix'n'Match

Baseline

(c) Consistency w/ XL model

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. DynaBERT would require more than 15% extra compute to perform as well as MatLM. OFA, similar to MatFormer, maintains a single universal model but

employs random subnetwork sampling during its training. This leads to sampling fewer models close to S and XL granularity, resulting in inferior performance in this regime. The performance gap manifests as a 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 a more detailed discussion of MatFormer performance compared to baselines, and advantages and downsides of each method.

4.1.1 Elastic Inference with MatLM

Accurate MatLM submodels for every constraint for free with Mix'n'Match. Mix'n'Match enables the MatLM to deliver accurate models for any compute constraint between S and XL, 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-vscompute 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'Match submodel approximates the XL's performance with only about a 1% accuracy drop, compared to a 2% drop when using the MatLM-M model. This highlights the efficiency of Mix'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] in 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'Match in Appendix D.1.

MatLM submodels speed up speculative decoding. Speculative decoding leverages an accurate lightweight LM as a draft model to autoregressively generate a few tokens, followed by verifying these drafts with a larger model through parallel decoding on the generated tokens. When the draft is inaccurate, the draft model is rolled back and reset to the larger model's output. This results in considerable inference speed-up for the same accuracy as the large model . We point the reader to the original paper for a more detailed explanation [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 baselines to their corresponding XL model. The significant gap persists even in the KL divergence variant of consistency with the XL model's outputs (see Figure 7 in Appendix). This improved consistency along with the need for only a single universal model positions MatLM favorably to improve techniques that require draft and verifier models such as speculative decoding.

Table 2 shows the inference time speed-ups from speculative decoding using the S and XL submodels of the 850M language model for drafting and verification respectively. Speculative decoding with independently trained baseline LMs results in a speed-up of up to 10% over the standard autoregressive decoding of the 850M-XL model. But MatLM-based speculative decoding is up to 6% faster than traditional speculative decoding. This additional speed-up can be primarily attributed to the more consistent nature

of MatLM-based drafter and verifier models and is further boosted by the ability to share attention cache across models from MatLM which is infeasible for the baselines (see Appendix C.1). Finally, MatLM further reduces the memory overhead for inference by removing the need to have two models during resource-constrained deployment.

4.1.2 MatLM Scales as well as Vanilla Transformer LMs

0

1

2

3

4

N(Non-Embedd. Parameters) 1e8

2.8

3.0

3.2

3.4

3.6

3.8

Loss

MatFormer

Baseline

(a) Validation loss: XL-models

0

1

2

3

4

N(Non-Embedd. Parameters) 1e8

3.00

3.25

3.50

3.75

4.00

Loss

MatFormer

Baseline

(b) Validation loss: all models

0

1

2

3

4

N(Non-Embedd. Parameters) 1e8

36

38

40

42

44

Average Accuracy (%)

MatFormer

Baseline

(c) 1-shot Evals

Now that we have established that a 850M MatLM model and its submodels are at least as accurate as the baseline Transformer LMs, we want to examine the scalability of training MatLM models. So, we study the scaling properties [29, 25] of MatLMs and compare them to vanilla Transformer baseline LMs trained for the same number of tokens. We train models ranging from 78M to 850M parameters on 10B to 80B tokens (per granularity) and plot the validation loss for MatLM - {S, M, L, XL} compared against independently trained baselines in Figure 9.

First, in Figure 3a, we observe that the training of MatLM-XL models across model sizes scale as reliably as the Baseline-XL LMs for loss vs. number of parameters. Figure 3b interestingly shows that all granularities {S, M, L, XL}, of MatLM and Baseline follow the same scaling trend. Therefore, we fit a scaling law according to the number of non-embedding parameters ( N ) and training tokens ( D ) for all possible submodels for both MatLMs and the baselines in

Table 3. 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 are more consistent with their XL 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 E.1. We provide full results split by granularity in Appendix E.

35

45

55

65

75

85

Total Parameters (M)

76

77

78

79

80

Top-1 Accuracy (%)

MatFormer

Mix'n'Match

Baseline

(a) B/16 trained on ImageNet-1K with AugReg

125

150

175

200

225

250

275

300

Total Parameters (M)

84.50

84.75

85.00

85.25

85.50

Top-1 Accuracy (%)

MatFormer

Mix'n'Match

Baseline

(b) L/16 pretrained on IN-21K

→

ImageNet-1K.

35

45

55

65

75

85

Total Parameters (M) in Query Encoder

70

72

74

76

78

1-NN Accuracy (%)

MatFormer

Mix'n'Match

Baseline

125

150

175

200

225

250

275

300

Total Parameters (M) in Query Encoder

82.0

82.5

83.0

83.5

1-NN Accuracy (%)

MatFormer

Mix'n'Match

Baseline

(a) B/16 trained on ImageNet-1K with AugReg

(b) L/16 pretrained on IN-21K

→

ImageNet-1K.

4.2 MatViT: MatFormer Vision Transformers

We extend MatFormer to Vision Transformer (ViT) [21] based computer vision encoder models. MatFormer-based ViT (MatViT) enables elastic inference for fundamental tasks like image classification 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 , 2 , 4 } ). B/16 models are trained on ImageNet-1K [50] with AugReg [57] while L/16 models are pretrained on ImageNet-21K [18] followed by finetuning on ImageNet-1K. All models use the training setup and optimal hyperparameters of standard ViT variants from the Scenic library [16].

4.2.1 Image Classification

For image classification, we evaluate both ViT & MatViT models on ImageNet-1K. Figure 4a shows that the explicitly optimized granularities in MatViT result in as accurate models as the independently trained baselines for the B/16. However for L/16, as shown in Figure 4b, we see that the MatViT models are up to 0 . 35% more accurate than the baseline for the same inference cost.

Wethen explore using MatFormer at different training stages with a 2 × 2 grid of pretraining-finetuning pairs (Table 7 in Appendix G.1) and find that using a MatFormer during pretraining helps bring more accurate and flexible encoders for downstream use. Further, finetuning using MatFormer enhances elastic deployment depending on the constraints at hand through Mix'n'Match.

Adaptive Encoders with Mix'n'Match. Furthermore, our Mix'n'match models' accuracy almost lies on the line joining accuracy of explicitly trained granularities. In scenarios where, say, an application can host 50M parameter B/16 model, MatViT can provide 0 . 8% more accurate model than the current approach which would host the largest baseline model with ≤ 50M parameters.

During deployment, the universal MatViT model can be stored in memory and depending on the compute constraints be used to extract an adaptable smaller model to maximize accuracy with the available resources at that moment. Currently, we find the Mix'n'Match models on the accuracycompute curve through a quick inference on the validation set. While relatively scalable, this points to the need for optimal budget allocation across layers in neural networks [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 is to encode the database images as well as query image with same encoder and run nearest neighbor retrieval for the query embedding. While we can embed database images with an expensive encoder, the query encoder generally has to be real-time. Furthermore, the setting of query encoding might be varied, e.g., on-device vs. cloud processing, varying query load and query complexity. Current solutions have a fixed encoder thus compromising on accuracy or cost for various settings.

Given the elastic nature of MatViT, it is a good candidate for query encoder. However, retrieval also requires that submodels preserve distances between fixed database (with large encoder) and query embeddings across all the granularities. If we use smaller baseline ViT models only for query encoding, these distances are not preserved and lead to nearly 0 retrieval accuracy (see Figure 5).

We evaluate both ViT and MatViT encoders on ImageNet-1K for image retrieval. We compute 1-nearest neighbor (NN) accuracy using the representation vector of the [ CLS ] token (also see Appendix G.2). Figure 5 shows that submodels extracted from MatViT can approximately preserve distances and provide significantly more flexibility. For example, with a loss of < 0 . 5% accuracy, MatViT-L/16 can reduce compute cost by 40% . This corresponds to the 175M Mix'n'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 a wide variety of adaptive inference strategies for large-scale semantic search.

5 Conclusion

In this work we presented MatFormer, a natively elastic Transformer architecture that allows training a single universal model which can be used to extract hundreds of smaller accurate submodels at zero additional cost at deployment time. We find that the MatFormer Language Model (MatLM) matches the perplexity & 1-shot accuracy of independently trained models. In fact, MatLM demonstrates an interesting loss-vs-compute scaling curve that is nearly independent of trained granularity indicating robust generalization to extremely large models as well. Finally, MatFormer submodels enable diverse inference time speedups like faster autoregressive generation with speculative decoding and elastic query encoders for adaptive dense retrieval across modalities. We believe dynamically routing these models to change inference latency [32, 40, 20], and developing the hardware optimizations required is a promising area for future work.

Acknowledgments

We are grateful to Aishwarya P S, Yashas B.L. Samaga, Varun Yerram, Lovish Madaan, Anurag Arnab for assistance in setting up training pipelines, Matthew Wallingford, Praneeth 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 Washington, FAS RC at Harvard University, Kempner Institute and a 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-4031, DARPA W911NF-15-1-0543, and gifts from Allen Institute 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 a 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 IIS2203097, and gift funding from Google, MSR, and OpenAI. Hannaneh Hajishirzi acknowledges funding through a gift from Allen Institute for Artificial Intelligence.

References

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:2110.03036 , 2021.

AI@Meta. Llama 3 model card. 2024. URL https://github.com/meta-llama/llama3/ blob/main/MODEL_CARD.md .

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 arXiv:2305.10403 , 2023.

J. Berant, A. K. Chou, R. Frostig, and P. Liang. Semantic parsing on freebase from questionanswer pairs. In Conference on Empirical Methods in Natural Language Processing , 2013. URL https://api.semanticscholar.org/CorpusID:6401679 .

L. Beyer, P. Izmailov, A. Kolesnikov, M. Caron, S. Kornblith, X. Zhai, M. Minderer, M. Tschannen, I. Alabdulmohsin, and F. Pavetic. Flexivit: One model for all patch sizes. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 14496-14506, 2023.

Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi. Piqa: Reasoning about physical commonsense in natural language, 2019.

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 arXiv:2108.07258 , 2021.

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.

H. Cai, C. Gan, T. Wang, Z. Zhang, and S. Han. Once-for-all: Train one network and specialize it for efficient deployment. arXiv preprint arXiv:1908.09791 , 2019.

H. Cai, C. Gan, J. Lin, and S. Han. Network augmentation for tiny deep learning. arXiv preprint arXiv:2110.08890 , 2021.

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-4941, 2022.

C. Chen, S. Borgeaud, G. Irving, J.-B. Lespiau, L. Sifre, and J. Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318 , 2023.

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.

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.

P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018.

M. Dehghani, A. Gritsenko, A. Arnab, M. Minderer, and Y. Tay. Scenic: A jax library for computer vision research and beyond. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 21393-21398, 2022.

M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. 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.

J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

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.

D. Ding, A. Mallick, C. Wang, R. Sim, S. Mukherjee, V. Ruhle, L. V. Lakshmanan, and A. H. Awadallah. Hybrid llm: Cost-efficient and quality-aware query routing. arXiv preprint arXiv:2404.14618 , 2024.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. 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.11929 , 2020.

N. Du, Y. Huang, A. M. Dai, S. Tong, D. Lepikhin, Y. Xu, M. Krikun, Y. Zhou, A. W. Yu, O. Firat, B. Zoph, L. Fedus, M. Bosma, Z. Zhou, T. Wang, Y. E. Wang, K. Webster, M. Pellat, K. Robinson, K. Meier-Hellstern, T. 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.

M. Grimaldi, L. Mocerino, A. Cipolletta, and A. Calimera. Dynamic convnets on tiny devices via nested sparsity. IEEE Internet of Things Journal , 10(6):5073-5082, 2022.

D. Hendrycks and K. Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415 , 2016.

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 arXiv:2203.15556 , 2022.

S. Hooker, N. Moorosi, G. Clark, S. Bengio, and E. Denton. Characterising bias in compressed models. arXiv preprint arXiv:2010.03058 , 2020.

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.

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: 10.18653/v1/ P17-1147. URL https://aclanthology.org/P17-1147 .

J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. 2020.

S. Kim, C. Hooper, T. Wattanawong, M. Kang, R. Yan, H. Genc, 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.

T. Kudo and J. Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv:1808.06226 , 2018.

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 arXiv:2110.03742 , 2021.

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 on Machine Learning , pages 5544-5555. PMLR, 2020.

A. Kusupati, G. Bhatt, A. Rege, M. Wallingford, A. Sinha, V. Ramanujan, W. Howard-Snyder, K. Chen, S. Kakade, P. Jain, et al. Matryoshka representation learning. Advances in Neural Information Processing Systems , 35:30233-30249, 2022.

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: A benchmark for question answering research. Transactions of the Association for Computational Linguistics , 7:452-466, 2019. doi: 10.1162/tacl_a_00276. URL https://aclanthology.org/Q19-1026 .

F. Lagunas, E. Charlaix, V. Sanh, and A. M. Rush. Block pruning for faster transformers. arXiv preprint arXiv:2109.04838 , 2021.

G. Lai, Q. Xie, H. Liu, Y. Yang, and E. Hovy. Race: Large-scale reading comprehension dataset from examinations, 2017.

H. J. Levesque, E. 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.

Y. Leviathan, M. Kalman, and Y. Matias. Fast inference from transformers via speculative decoding. 2023.

M. Li, S. Gururangan, T. Dettmers, M. Lewis, T. Althoff, N. A. Smith, and L. Zettlemoyer. Branch-train-merge: Embarrassingly parallel training of expert language models. arXiv preprint arXiv:2208.03306 , 2022.

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.10198 , 2018.

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.

N. Mostafazadeh, N. Chambers, X. He, D. Parikh, D. Batra, L. Vanderwende, P. 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.18653/v1/N16-1098. URL https://aclanthology.org/N16-1098 .

Y. Nie, A. Williams, E. Dinan, M. Bansal, J. Weston, and D. Kiela. Adversarial nli: A new benchmark for natural language understanding, 2020.

R. OpenAI. Gpt-4 technical report. arXiv , pages 2303-08774, 2023.

D. Paperno, G. Kruszewski, A. Lazaridou, Q. N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The lambada dataset: Word prediction requiring a broad discourse context, 2016.

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 Learning , pages 28492-28518. PMLR, 2023.

E. Real, A. Aggarwal, Y. Huang, and Q. V. Le. Regularized evolution for image classifier architecture search, 2019.

M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J.-b. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530 , 2024.

O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision , 115:211-252, 2015.

K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019.

M. Salehi, S. Mehta, A. Kusupati, A. Farhadi, and H. Hajishirzi. Sharcs: Efficient transformers through routing with dynamic width sub-networks. Findings of Empirical Methods in Natural Language Processing , 2023.

V. Sanh, L. Debut, J. Chaumond, and T. Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108 , 2019.

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-17472, 2022.

N. Shazeer and M. Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning , pages 4596-4604. PMLR, 2018.

D. R. So, W. Ma´ nke, H. Liu, Z. Dai, N. Shazeer, and Q. V. Le. Primer: Searching for efficient transformers for language modeling. arXiv preprint arXiv:2109.08668 , 2021.

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.10270 , 2021.

R. Thoppilan, D. De Freitas, J. Hall, N. Shazeer, A. Kulshreshtha, H.-T. Cheng, A. Jin, T. Bos, L. Baker, Y. Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239 , 2022.

H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

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.00255 , 2023.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. 2023.

A. Wang, Y. Pruksachatkun, N. Nangia, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems, 2020.

H. Wang, Z. Wu, Z. Liu, H. Cai, L. Zhu, C. Gan, and S. Han. Hat: Hardware-aware transformers for efficient natural language processing. arXiv preprint arXiv:2005.14187 , 2020.

J. Yu and T. Huang. Universally slimmable networks and improved training techniques, 2019. URL https://arxiv.org/abs/1903.05134 .

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.

J. Yu, L. Yang, N. Xu, J. Yang, and T. Huang. Slimmable neural networks. arXiv preprint arXiv:1812.08928 , 2018.

R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can a machine really finish your sentence?, 2019.

B. Zoph and Q. V. Le. Neural architecture search with reinforcement learning, 2017.

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 underrepresented domains/languages [1].

B Implementation Details

B.1 Architecture and Training

For our experiments, we train a 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, S with d ff = 4 ∗ d model , 2 ∗ d model , d model , d model / 2 . All models have 16 layers, 16 attention heads, and a d model : d ff ratio of 1 : 4 . We train a 256k vocabulary using the SentencePiece library [31], use a maximum context length of 1024 tokens, and a 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].

B.2 Downstream Evaluation

We evaluate all the LM models trained on set of 25 English tasks similar to [8, 22, 14, 3], including:

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].

Winograd-style tasks: Winograd [38] and WinoGrande [51].

Reading comprehension: RACE [37].

Common sense reasoning: PIQA [6], ARC [15], and OpenBookQA [42].

SuperGLUE [62]

Natural language inference: Adversarial NLI [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 a given model size uses X tokens mentioned in Table 4. For example, for 850M model X = 80B tokens. Therefore, Baseline-{S, 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 in between the trained granularities without using additional methods like distillation, pruning, etc.

OFA: OFA [9] trains a big model from scratch, then freezes it and starts another run of a similar sized "universal model" where it 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 a 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 as 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 predictiing the loss. Finally, it uses an evolutionary search based routine [48] to search for the optimal architecture within a 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 : 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.

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 a total of 4X tokens. Though, since it optimizes on 4 batches of data in a 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 an optimal submodel with a cost constraint without any overhead. In contrast, DynaBERT doesn't introduce any search strategy and discusses only using explicitly trained granularities as submodels.

B.4 Comparison with Baselines

Here we do a more detailed discussion of comparison between MatFormer and the baselines used in Section 4.

1. MatFormer vs trained from scratch Baseline : Let Baseline-{S, 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, XL}. 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 a 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 - simultaneously processing 4 batches of data 1 batch through each {S, M, L, XL} granularity. The final loss being optimized is the average of each granularity's loss. Conversely, MatFormer samples and optimizes individual granularities in separate steps, resulting in 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 a 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'n'Match approach, DynaBERT lacks a dedicated sampling strategy.

2.0

2.5

3.0

3.5

4.0

4.5

N(Non-Embedd. Parameters) 1e8

2.85

2.90

2.95

3.00

Loss

Baseline

MatFormer

MatFormer - MnM

MatFormer - NAS

3. 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 in these regions. That is, sampling models ranging from S model size to XL model size, most of the sampled models are close to 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 S 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 matFormer. Rather a more complicated NAS strategy is needed to find optimal submodels. To train a 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. In 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 all the 4 granulatities baseline counterparts. During serving, we 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 a 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 we trained one MatFormer and compare its training time with Baselines combined, we can extract many more models than the 4 model granularities we explicitly trained.

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 S 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 a 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 to the other model. Thus, attention sharing gives further improvement over vanilla speculative decoding as shown in Table 2.

D Search Techniques

D.1 Mix'n'Match

For a 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 architecture within this subset; however, this approach can be prohibitively expensive. Instead, we propose a simpler Mixn'Match heuristic, which identifies optimal submodels that lie on the Pareto-optimal accuracy-vs-parameters curve.

The Mix'n'Match heuristic recommends selecting layer granularities with minimal changes across layers, ensuring that the granularity of the j th layer is at least that of the i th layer for j > i . This means that a 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 or 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 a more balanced distribution of granularity as opposed to a skewed one.

This intuition is also supported by results from an evolutionary search-based Neural Architecture Search (NAS) method [48], similar to that employed in [9]. Our NAS experiments, few of the runs we present in Table 5, indicated a preference for balanced configurations under various constraints, aligning with the findings from our Mix'n'Match heuristic.

We explored various balanced configurations, including:

Increasing-Decreasing: Granularity increases until the midpoint and then decreases, e.g., [M, M, L, L, L, M, M].

Decreasing-Increasing: The reverse of the previous heuristic, where granularity decreases then increases.

Increasing: A non-decreasing sequence of layer granularities, such as [M, M, M, M, L, L, L].

Decreasing: The opposite of the Increasing heuristic.

Among these configurations, the Increasing heuristic consistently outperformed the others. Therefore, we recommend adopting a 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'match requires no additional training of a module or 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 a constraint budget. In fact, OFA [9] uses an evolutionary search based technique [48] along with a 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 erroneous.

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 a 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 2 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 ∼ 5 e -5 mean squared error.

As we move to large models regime, the performance gap between different sized models diminish. For example, the gap between 78M parameter Baseline-S and Baseline-XL, which are of size 71M and 78M, is 4.01 vs 3.83. Which is a gap of 0.18 loss for just 7M parameter difference. But in case of 850M model, Baseline-S and Baseline-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 it 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 MatLMs 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 to have comparable consistency with MatLMs.

E.1 Scaling laws of MatFormers vs Transformers.

Scaling laws are essential tools to 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 MatFormers only require a 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 a single training run which is a 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 . 58 P , where P is the number of parameters of 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: 4 × more FLOPs for attention layers parameters and { 1 + 1 / 2 + 1 / 4 + 1 / 8 = 1 . 875 }× more FLOPs for MLP layers.

F Further Analysis on Language Decoders

F.1 KL Divergence Between S, M, L and XL Models

Figure 7 showcases the smoother consistency calculation between two generative models measured with KL-divergence 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 MatLM'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 it still follows the trend of generating the token with the highest probability. This smoother notion of consistency argues for the metric-space

preservation given that the output classifier/embedding matrix is shared across all the submodels of MatLM.

2.0

2.5

3.0

3.5

4.0

N(Non-Embedd. Parameters) 1e8

0.0

0.1

0.2

0.3

0.4

0.5

KL div w/ XL Model

MatFormer

Mix'n'Match

Baseline

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 a varying number of attention heads n attn , 3 ∗ n attn / 4 , n attn / 2 , n attn / 4 of size d/n attn each. In Figure 8, we plot the validation loss for these models (MatLM-Attn), their corresponding baselines and Mix'n'Matched subnetworks. We find that similar to our experiments on Mix'n'Matched MatLMs, Mix'n'Match helps obtain many MatLM-Attn models on the optimal loss-vs-compute curve.

F.3 Tuning Sampling Probability

In Table 6, we experiment with tuning the sampling probabilities for individual granularities in order to investigate the trade-off between granularity size and performance. More specifically, we tune { p 1 , p 2 , p 3 , p 4 } , and find that all strategies that upweight the loss for the largest granularity perform well, with modest degradation on the M and S granularties.

G Further Analysis on Vision Encoders

G.1 Decoupling Effect of MatFormer on Pretraining and Finetuning

Table 7 investigates the effect of MatFormer on pretaining and finetuning phases of ViT-L/16 model. ViT-L/16 is typically pretrained on ImageNet-21K and then finetuned on ImageNet-1K for the final evaluation. Table 7 shows that having a MatFormer during pretraining generates a better model for

7.50

7.75

8.00

8.25

8.50

N(Non-Embedd. Parameters) 1e8

2.2

2.4

2.6

2.8

3.0

Loss

MatFormer

Mix'n'Match

Baseline

(a) Validation loss

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, a 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 a deployable MatViT variant.

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 in Figure 4. Both MatViT variants B/16 and L/16 have submodels that have as 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 a 200M parameter MatViT submodel can be more accurate than the 300M parameter ViT baseline.

0

1

2

3

4

N(Non-Embedd. Parameters) 1e8

2.8

3.0

3.2

3.4

3.6

3.8

Loss

MatFormer

Baseline

(a) XL-model Loss

0

1

2

3

N(Non-Embedd. Parameters) 1e8

3.0

3.2

3.4

3.6

3.8

Loss

MatFormer

Baseline

(b) L-model Loss

0.0

0.5

1.0

1.5

2.0

N(Non-Embedd. Parameters) 1e8

3.0

3.2

3.4

3.6

3.8

4.0

Loss

MatFormer

Baseline

(c) M-model Loss

0.0

0.5

1.0

1.5

N(Non-Embedd. Parameters) 1e8

3.0

3.2

3.4

3.6

3.8

4.0

Loss

MatFormer

Baseline

(d) S-model Loss

0

1

2

3

4

N(Non-Embedd. Parameters) 1e8

3.00

3.25

3.50

3.75

4.00

Loss

MatFormer

Baseline

(e) Loss for all granularities - XL, L, M, S.

0

1

2

3

N(Non-Embedd. Parameters) 1e8

70

75

80

85

90

Consistency

MatFormer

Baseline

(f) Consistency with the XL-models

0

1

2

3

4

N(Non-Embedd. Parameters) 1e8

36

38

40

42

44

Average Accuracy (%)

MatFormer

Baseline

(a) XL-model Average Score on Evals

0

1

2

3

N(Non-Embedd. Parameters) 1e8

36

38

40

42

44

Average Accuracy (%)

MatFormer

Baseline

(b) L-model Average Score on Evals

0.0

0.5

1.0

1.5

2.0

N(Non-Embedd. Parameters) 1e8

36

38

40

42

Average Accuracy (%)

MatFormer

Baseline

(c) M-model Average Score on Evals

0.0

0.5

1.0

1.5

N(Non-Embedd. Parameters) 1e8

36

38

40

42

Average Accuracy (%)

MatFormer

Baseline

(d) S-model Average Score on Evals

0

1

2

3

4

N(Non-Embedd. Parameters) 1e8

36

38

40

42

44

Average Accuracy (%)

MatFormer

Baseline

(e) Average Score on Evals for all granularities -

XL, L, M, S

0

1

2

3

N(Non-Embedd. Parameters) 1e8

7500

8000

8500

9000

Consistency

MatFormer

Baseline

(a) Consistency of L-model with XL-model

0.0

0.5

1.0

1.5

2.0

N(Non-Embedd. Parameters) 1e8

7500

8000

8500

Consistency

MatFormer

Baseline

(b) Consistency of L-model with XL-model

0.0

0.5

1.0

1.5

N(Non-Embedd. Parameters) 1e8

7000

7250

7500

7750

8000

Consistency

MatFormer

Baseline

0

1

2

3

N(Non-Embedd. Parameters) 1e8

70

75

80

85

90

Consistency

MatFormer

Baseline

(c) Consistency of L-model with XL-model

(d) Consistency of all model granularities with XL- model - L, M, S

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

Question: Does the paper discuss the limitations of 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 a 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Implementation details are provided in Appendix B. Moreover, code to reproduce Section 4.2 will be been open-sourced at link.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Implementation details are provided in Appendix B. While code to 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 (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Implementation details are provided in Appendix B, with citation of the models upon which these models are 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Broader impacts have been discussed in Appendix A.

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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: No datasets or models are released with this paper.

12. 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]

Justification: 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.

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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an 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 to 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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 Subjects

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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

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 to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nYes, there are no worries on that, and the authors did attach the Neurips checklist to the back of their paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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?\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes, the limitations of most results and methods are presented throughout the paper. Nothing extra is needed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: 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.

STRENGTHS:
(1) Important problem and a very interesting solution.

(2) The writing is straightforward and easy to understand.

(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.

QUESTIONS:
(1) Can you add a discussion section with MoE?
(2) Is it possible to obtain a MatFormer model directly from a pretrained checkpoint (e.g. in some form of up-cycling)?

LIMITATIONS:
Please see question 2.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\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\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.\n\nQUESTIONS:\nNA\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-19"}