MatFormer: Nested Transformer for Elastic Inference

We thank the reviewer for their support, and their recognition of the potential inference wins. We address your feedback below:

Saying that MatFormer is “zero additional cost” seems inaccurate given the cost of the additional ‘g - 1’ forward passes during training.

By “zero additional cost”, we mean that the Mix’n’match models are obtained without ever being explicitly optimized for. We will clarify this in the next revision. Moreover, there is no overhead during inference compared to the baseline.

The authors’ don’t explain how the mix’n’match models on the pareto frontier are selected. Given there are many candidates this seems like an important detail.

Currently, we induce nested substructures in the hidden dimensions of FFN blocks. We tested several heuristics to select the best subnetwork, but consistently observed that gradually using larger granularities in deeper layers worked the best. More formally, we use non-decreasing hidden dimensions with the least slope (change in hidden dimensions across consecutive layers) across layers. This choice behaves nearly optimal (performance lies on the pareto-optimal curve). We will add a discussion on search heuristics used for Mix’n’Match in the next revision.

The claim that the FFN is responsible for the largest chunk of latency during inference seems questionable to me. I’d encourage the authors to present a more nuanced perspective on Transformer inference based on previously published data. For example, the study from Kim et. al [1] presents relevant data for CPU serving in Figures 7 and 8.

In Kim et. al [1], Table 4, we note that BERT-Base with 512 sequence length has ~60% of the total FLOPs, whereas attention block takes the remaining. As you recommended, we will refer this paper or other published paper for these numbers in addition to the current ones (Section B and Table 4 in the Appendix)

In the speculative decoding experiments, I’d encourage the authors to briefly describe how consistency between the two models can translate into latency reductions. This will make the impact more clear to readers who aren’t familiar with prior work/do not read the reference you direct them to.

We agree with your suggestion to add more explanation for speculative decoding and will expand upon our explanation in Section 4.1.1 in the final revision.

One ablation I would like to see is how a Transformer performs if the 'g' additional forward passes aren't used during training with your Mix'n'Match procedure. Basically, how bad does a normal Transformer perform with this post-processing if it's not trained appropriately?

Applying Mix'n'Match to the baseline performs poorly because vanilla training does not induce nested structure in the network. We will add this result to the Appendix in the next revision.

We would be happy to discuss any further questions about the work and would appreciate support for acceptance of the paper.

We are happy to discuss if anything in the rebuttal needs more clarification or if the reviewer has further questions regarding the paper.

We thank the reviewer for their review, and recognition of how our work can be beneficial for a wide range of deployment scenarios. We address their feedback and questions below:

Motivation and Efficiency: While the paper does address ... significantly reduced

In Appendix B, we have discussed and empirically demonstrated the training efficiency of MatFormers (our method costs 15% less FLOPs than training independent models). At the same time, unlike the baselines, MatFormer allows extracting 100s of accurate submodels that are not explicitly optimized.

Similarity to Existing Techniques: The paper does ... the novelty of the work.

Conditional computation methods such as MoE [1] are orthogonal to MatFormer, given that in this work, we co-train Transformer feedforward networks within the same representation space as a pre-training method that is suitable for elastic inference, and do not introduce routing mechanisms. In Table 5 in the Appendix, we have compared our training method with training independent MLP modules, and sampling, which are important features of MoEs - we find that MatFormer performs better than these baselines across all granularities. We will add more discussion on this in our next revision.

Scope of Applicability: The focus on FFNs ... Transformers comprehensively.

We have chosen to focus on the FFN of the Matformer, given that it forms the bulk of our computational costs (~60%) at the scale we consider (2.6B) and beyond. In this work, we focus on conducting extensive analysis and ablations on this architecture, and plan extending this technique to other components of the Transformer in future work.

Evaluation and Support for Claims: The paper could benefit from ... asserted lack sufficient empirical validation.

Evaluation and Support for Claims: In our work, we empirically show the effectiveness of the MatFormer architecture across multiple modalities and scales.

• More specifically, in Section 4.1 we evaluate for decoder-only language modeling on 26 different classification and generation NLP tasks and predict scaling trends across 7 different model scales ranging from 70M to 2.6B.
• In Section 4.1.1, we conduct experiments that showcase the elasticity of MatFormers, with evaluations on the models obtained using Mix’n’Match showing the linear scaling trend of the extracted submodel, and experiments showing the effectiveness of MatFormer compared to vanilla transformers for speculative decoding.
• In Section 4.2, we additionally validate our method on image classification and retrieval on ImageNet-1K.

Given our extensive empirical results, we are unsure of which specific claims are unsubstantiated. We are happy to discuss this further and clarify any concerns the reviewer might have.

Prior Work Comparison: The work could ... key prior work, specifically Kusupati et al., 2022, in terms of the nested structure.

As discussed in related work, MRL (Kusupati et al., 2022) focuses on the nested structure of the output representations. In contrast, MatFormer brings the core idea of nesting from MRL into the model weight/parameter space which helps in flexibility not only in embeddings for large-scale retrieval but also for flexibility in model inference. We mention this in the introduction but will make it clearer in the next revision.

The paper seems to use MatLM, MatFormer, and MatViT interchangeably.

We use MatLM and MatViT respectively when talking about MatFormers applied to decoder-only language modeling and to vision transformers. We use the more general term, MatFormer, when our exposition applies to both our vision and language modeling experiments. We will carefully check the next revision to ensure that we have used this terminology as intended.

In Figure 2 and the associated text, what are the differences in model architecture / settings among MatFormer, Mix'n'Match, and baseline?

In Figure 2 the architecture and settings for all three are similar, with the difference between MatFormer and baseline being the training objective. Mix’n’Match models are derived from a trained MatFormer model at inference time by activating different granularities for each layer. We will include these details in the Appendix.

We are happy to discuss if anything in the rebuttal needs more clarification or if the reviewer has further questions regarding the paper.

We thank the reviewer for their response, and for their appreciation of how the flexibility of our design could greatly impact the deployment of large foundation models. We address their concerns below:

The parameters g … which is currently missing.

Thank you for the interesting comment. We note that increasing $g$ increases the number of accurate subnetworks to select from (since we are explicitly optimizing for more models) at the cost of increased training overhead - the inverse holds for decreasing $g$. We refer the reviewer to Table 29 in the MRL paper (Kusupati et al, 2022) for a comparison of logarithmically nested and uniformly distributed granularities.

The Mix'n'Match procedure … the above advantage of MatFormer.

Yes, the Mix'n'Match procedure indeed results in more models than the 9 points we have included for clarity in the results. We will include more data points in the Appendix of the next version.

As mentioned in the first paragraph of the … not a big concern of this work.

Due to resource constraints, we are unfortunately unable to scale our empirical results beyond 2.6B. In lieu of this, we have done extensive scaling law experiments on a range of model sizes (70M-2.6B) to show that our techniques are likely to hold at scale. More specifically, we show that the loss of MatFormer scales as well as that of the baselines in Section 4.1.2, and that the consistency of all MatFormer granularities scales significantly better than that of the baselines in Figure 10.

As authors say in the Introduction: “MatFormer can also form a similar sub-structure on attention heads”. It will be more clear if authors can illustrate the nested structure of attention heads just like Figure 1.

Background: For a given layer let there be $A$ attention heads, with the i-th head projection as $a_i$, each having dimension $d_p$. In a usual Transformer, we concatenate these projections and then project it to output dimension $d$ using $W\in \mathbb{R}^{A\cdot d_p \times d}$. That is: $output = concat(a_1, \cdots, a_{A}) \cdot W$

Nested substructure in attention heads can be introduced using only the first few attention heads. For example, using first $m$ attention heads would result in $output = concat(a_1,\cdots,a_m) \cdot W[:m\cdot d_p]$. Using different nested values of $m$ for different granularities will induce MatFormer nesting in attention heads. We will include this discussion in our next revision.

The core method … this work even more.

Thank you for bringing NetAug [1] to our attention. NetAug aims to train a single small target model. The training involves sampling auxiliary networks from a SuperNetwork in each step, where the auxiliary network contains the target model as a subnetwork, and jointly training the auxiliary and the target network. These sampled auxiliary networks are not nested. Moreover, the auxiliary networks are not trained to be accurate. In comparison, MatFormer training gives accurate subnetworks of various sizes. We will add a discussion about NetAug in our next revision.

Although the experiments are … including other elastic inference techniques

We respectfully disagree with the reviewer that the baselines are weak. The baselines are the corresponding architectures trained from scratch. These will yield the highest performance compared to any other training technique to get a model of the same architecture. We additionally refer the reviewer to Table 5 in the Appendix, where we have compared our method to sampling and training independent MLP modules. Here, we find that our method is the most performant for the same number of tokens.

Inference-time speedup techniques for pretrained networks such as techniques which make inference fast for any pretrained network, like Deja Vu [2], Flash Attention [3], and Speculative Decoding [4], can be applied to our spliced subnetworks as well and are complementary to our contributions.

There are other techniques like distillation that might be stronger baselines, but they create significant training overhead. Moreover, these techniques can also be applied to MatFormer to enhance the quality.

[1] Network augmentation for tiny deep learning, ICLR 2022

[2] Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time, ICML 2023

[3] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022

[4] Fast Inference from Transformers via Speculative Decoding, ICML 2023

We hope that the rebuttal clarifies the questions raised by the reviewer. We would be happy to discuss any further questions about the work, and would appreciate an appropriate increase in the score if the reviewer’s concerns are adequately addressed.

We are happy to discuss if anything in the rebuttal needs more clarification or if the reviewer has further questions regarding the paper.

We thank the reviewer for their comments and suggestions. We are glad that the reviewer found the paper well-written, the idea to be novel and effective for the challenges at hand, and the experiments to be extensive. We address the concerns raised by the reviewer below:

Weak baselines … the performance of MatFormer

Thank you for highlighting Deja Vu [1] which is a complementary approach to MatFormer. MatFormer’s main focus is designing transformer blocks that have nested structures so that smaller models can be extracted without any further training. In contrast, the goal of Deja Vu is to route tokens through different parts of the model based on sparsity. Deja Vu can be applied to any pretrained network, including spliced subnetworks of MatFormer, to speed up their inference.

Lack of technical … of their technical contributions.

Thank you for bringing HAT [2] to our attention. The contributions of work like Once-for-all [3] and HAT differ in numerous ways compared to MatFormer:

• HAT vs MatFormer: HAT samples random (not nested) subnetworks during training at each step to optimize. During inference, it conducts NAS to find the subnetwork architecture and then trains it from scratch before deployment. In contrast, MatFormer trains one model and reads off accurate nested models without any further training.
• Once-for-all (OVA) vs MatFormer: OVA first trains a teacher model, and fine-tunes subnetworks using knowledge distillation. In contrast, MatFormer doesn’t require any finetuning and maintains just a single model.

In addition, both OVA and HAT focus on smaller-scale end devices. OVA trains models of size <=10M params, while HAT trains <=60M param models. In comparison, we train models at a much larger scale (up to 2.6B parameters). We will add the comparison with HAT in the next revision.

Lack of search algorithm: MatFormer … best combination among pre-trained blocks

We experimented with several heuristics to select the best subnetwork, but consistently observed that gradually using larger granularities in deeper layers worked the best. More formally, we use non-decreasing hidden dimensions with the least slope (change in hidden dimensions across consecutive layers) across layers. Given that this choice behaves nearly optimally (performance lies on the pareto-optimal curve), we did not focus on search techniques. In future work, we plan to extend the nested substructure to other components of the Transformer like attention heads, model dimensions, and n(layers), and will explore NAS methods then. We will add a discussion on search heuristics used for Mix’n’Match in the next revision.

Simple search … designed by NAS methods

We would like to clarify that we do not choose among the number of FFN blocks, but choose the hidden dimension in each FFN block. Hence, simple search techniques like MixnMatch work well, unlike existing works that require careful NAS. In the future when we extend our work to induce elasticity in other components of the Transformer, we will explore more complex techniques like NAS.

The reviewer … used for comparison are weak

We respectfully disagree with the reviewer that the baselines are weak. The baselines are the corresponding architectures trained from scratch. These will yield the highest performance compared to any other training technique to get a model of the same architecture. We additionally refer the reviewer to Table 5 in the Appendix, where we have compared our method to sampling and training independent MLP modules. Here, we find that our method is the most performant for the same number of tokens. Inference-time speedup techniques for pretrained networks such as techniques which make inference fast for any pretrained network, like Deja Vu [1], Flash Attention [4], and Speculative Decoding [5], can be applied to our spliced subnetworks as well and are complementary to our contributions.

There are other techniques like distillation that might be stronger baselines, but they create significant training overhead. Moreover, these techniques can also be applied to MatFormer to enhance the quality.

We are happy to discuss if anything in the rebuttal needs more clarification or if the reviewer has further questions regarding the paper.

Thank you to the authors for their response. I have carefully reviewed both the authors' reply and the comments from other reviewers.

It seems that the reviewers share common concerns, namely: 1) the lack of empirical evidence to support that Mix'n'Match selected pareto-optimal or necessity to improve Mix'n'Match algorithm. 2) weak baselines 3) suboptimal search space - dimension of FFN blocks. However, I did not find any indication in the authors' response that they suggested improved algorithms or the definition of a new search space or enough empiricial evidence. These aspects appear to be the main limitations that need to be addressed.

Therefore, my concerns are still remained and I am inclined to maintain the original score.

We have updated the draft in response to your helpful feedback. In particular, we have:

• Clarified your questions on training cost in Appendix B
• Added details on Mix'n'Match in Appendix C.
• Discussed HAT in Related Work (Section 2)