LLaMaFlex: Many-in-one LLMs via Generalized Pruning and Weight Sharing

Authors: Ruisi Cai, Saurav Muralidharan, Hongxu Yin, Zhangyang Wang, Jan Kautz, Pavlo Molchanov
Primary area: foundation or frontier models, including LLMs
Keywords: large language models, elastic networks, training efficiency, inference efficiency

Abstract:
Large Language Model (LLM) providers typically train a family of models, each of a different size targeting a specific deployment scenario. Models in the family are all trained from scratch, making the process extremely resource intensive.
Recent work has successfully reduced the cost of training model families through a combination of structured pruning and knowledge distillation; here, only the largest model in the family is trained from scratch, and smaller models are obtained via pruning. We observe that while effective, this strategy must still perform pruning and distillation with hundreds of billions of training tokens for every new model, keeping overall training costs high.
In this work, we introduce a novel nested weight-shared architecture named LLaMaFlex that can be pruned across both width and depth dimensions in a zero-shot manner to instantly yield a large number of highly accurate compressed models.
LLaMaFlex starts from a pretrained model, and only requires a single continued training phase consisting of ~60B tokens, which trains the elastic network and an end-to-end Gumbel Softmax-based router; this router is able to interpolate smoothly across model sizes, enabling the "train once, deploy many'' paradigm.
We train LLaMaFlex on Llama 3.1 8B and use it to zero-shot generate a family of compressed models that achieves accuracy on par with or better than state-of-the-art pruned, elastic/flexible, and trained-from-scratch models.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730624865881

thread_type:
official_review

<thread_object>
creator: reviewer_geny

<content>
summary:
The paper presents LLaMaFlex, an elastic architecture for large language models that enables zero-shot resizing across both width and depth dimensions, allowing for the instant generation of multiple compressed models without additional fine-tuning. LLaMaFlex utilizes a nested weight-shared network architecture and a Gumbel Softmax-based router for smooth interpolation between model sizes, achieving a "train once, deploy many" paradigm. It introduces policy-aware modulation to enhance the expressivity of nested architectures and produces a family of compressed models from Llama 3.1 8B that outperform state-of-the-art compressed, elastic/flexible, and trained-from-scratch models in accuracy. The uniform architectures generated by LLaMaFlex are also easier to deploy using existing LLM frameworks, offering a significant advancement in training efficiency and deployment flexibility for large language models.

strengths:
1. The authors introduce a novel elastic architecture for large language models that enables zero-shot resizing across both width and depth, addressing a key challenge for efficient model deployment.

2. The end-to-end trained router could dynamically determine the optimal network architecture under a given parameter budget.

3. It achieves a promising performance of the compressed model in various of size, which outperform the same size model trained from scratch.

weaknesses:
1. The authors validate their proposed method in only one model, Llama 3.1 8B. The serviceability to other models and architecture should be considered.

2. I understand that the router is used to determine the optimal architecture, that is the hyperparams of the network architecture, under a given parameter budget. Once the hyperparams of the network architecture are determined, the algorithm will select the top-k importance submodule, which is ranked by a pre-defined metric, e.g. the accumulated magnitude of activations as stated in Section 2.3. There is no ablation on the pre-defined metric, which might limit the rigor and integrity of this work.

3. How to guarantee that the sampled submodel (i.e., the combination of N, H, D, N_A, lambda) satisfies the parameter budget (i.e., the second equation in Eq. 4), given that N, H, D, N_A, and lambda are modeled independently in the routers (i.e., Eq. 5)? What is the explicit joint distribution of Q in Eq. 6?

4. The authors did not explicitly discuss the limitations of the proposed method.

questions:
1. Instead of interpolating the router input h_k in Eq. 7, is it possible to interpolate the router output in a similar way given h_{n+1} and h_{n}?

2. In Line 196, the author indicates that \mathcal{D}, \mathcal{N_A}, \mathcal{N}, \mathcal{H} are predefined, how to define them?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730694338671

thread_type:
official_review

<thread_object>
creator: reviewer_xxxs

<content>
summary:
This paper presents LLaMaFlex, a novel approach to create elastic Large Language Models (LLMs) that can be dynamically resized without additional fine-tuning. The key innovation is a nested weight-shared architecture combined with a Gumbel Softmax-based router that allows zero-shot pruning across both width (attention heads, hidden dimensions, MLP intermediate size) and depth (layers) dimensions. Unlike previous approaches that require separate training or distillation for each model size, LLaMaFlex needs only a single continued training phase of ~60B tokens to enable "train once, deploy many" capabilities. The authors also introduce a policy-aware modulation technique inspired by diffusion models to enhance the expressiveness of nested architectures. When applied to Llama 3.1 8B, LLaMaFlex produces compressed models that outperform state-of-the-art pruned models, elastic frameworks, and models trained from scratch. The router can interpolate smoothly between different model sizes, allowing deployment flexibility without accuracy compromises.

strengths:
- The paper introduces several innovative technical components, including an end-to-end learnable router with Gumbel Softmax for architecture selection, policy-aware modulation for enhanced expressiveness, and support for both width and depth pruning dimensions. This comprehensive approach advances the state-of-the-art in elastic LLM architectures.
- The method demonstrates impressive performance across multiple downstream tasks (ARC-E, LAMBADA, PIQA, WinoGrande, HellaSwag, MMLU), consistently outperforming existing approaches including Flextron, Minitron, and models trained from scratch. The results are particularly notable given that LLaMaFlex requires only 60B tokens for training compared to hundreds of billions for traditional approaches.
- LLaMaFlex produces uniform architectures that are compatible with common deployment frameworks like TensorRT-LLM and llama.cpp, addressing a significant limitation of previous elastic approaches that generated heterogeneous architectures. The ability to interpolate between model sizes without additional training also offers practical deployment flexibility.

weaknesses:
- While the paper presents strong technical contributions, there are opportunities for additional context and analysis that could further strengthen the work. The authors could enrich their discussion by connecting their approach to the broader context of supernet-based neural architecture search. For instance, works like HAT and HELP have previously explored supernet-based approaches for generating target-budget subnets in the field of small language models and Meta-NAS or DARTS-EGS have handled Gumbel-Softmax-based supernet. Additionally, a more thorough discussion and comparison of the paper's relationship to MatFormer's subnet extraction approach for Transformer architectures (applied to Llama or ViT) would be valuable.
- The paper would benefit from additional analysis of training dynamics typically associated with supernet approaches. In supernet-based methods, there are well-known challenges to address. First, training imbalances between larger and smaller subnets often occur, where larger architectures receive less training attention and may not fully converge. Second, weight sharing interference can arise when different architectural configurations compete for optimal parameter values. While the authors present strong results, providing insights into how LLaMaFlex addresses these potential training imbalances would be informative. Similarly, a discussion of how weight sharing affects different architectural configurations could offer valuable implementation insights for practitioners.
- The comparative analysis could be more comprehensive, particularly regarding scenarios with fixed target sizes. A direct comparison with knowledge distillation approaches would be valuable, considering factors such as training computational costs, final model performance, and the associated trade-offs. Such analysis would help practitioners better understand when to choose LLaMaFlex over traditional approaches.
- The process for selecting optimal architectures when multiple configurations achieve similar parameter counts needs more detailed explanation.

[HAT] HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
[HELP] HELP: Hardware-Adaptive Efficient Latency Predictor for NAS via Meta-Learning
[MatFormer] MatFormer: Nested Transformer for Elastic Inference
[MetaNAS] Meta-Learning of Neural Architectures for Few-Shot Learning
[DARTS-EGS] Differentiable Architecture Search with Ensemble Gumbel-Softmax

questions:
Please address concerns in Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705378172

thread_type:
official_review

<thread_object>
creator: reviewer_upjk

<content>
summary:
LLAMAFLEX introduces a new approach to efficiently and effectively generate high-accuracy SLMs. LLAMAFLEX employs a nested weight-shared architecture. Starting from a pre-trained model, it uses only 60 billion tokens for a single training phase, enabling zero-shot pruning across both width and depth and a “train once, deploy many” pruning paradigm.

strengths:
+ Introduces a gumbel softmax-based end-to-end learnable router to train the candidate models.
+ Enhances generality by incorporating a policy-aware modulation technique.
+ Enables interpolation-based extrapolation of pruning rates for flexible model compression.
+ Achieves accuracy comparable to state-of-the-art SLMs like minitron.

weaknesses:
- The training process and methodology flow are difficult to understand and a figure like Figure 3 in Flextron paper should be provided.
- The presentation in the main text is somewhat weaker than in the abstract, e.g. the explanation of Equation 3 (L170-L180), and the textual ordering needs to be organized. For another example, paragraphs L199-L210 are intended to be placed in the heading “Router Design” rather than under “Problem Formulation”. Apart from that, parts of 2.3 seem to fit into the experimental chapter. More, the description of the amount of data on line 325 should be placed on line 309 and the source of the data set should be described.
- In Eq. 2, if \lambda_i^j=0, it would be misleading to think that all previous layers have been pruned, and the formula should be reshaped to alleviate the ambiguity. The upper limit of j is not defined, and if j can be taken to infinity, this should be stated.
- Unfair comparison with Flextron, you should reproduce Flextron-Llama3.1-8B to provide a fair comparison. Unfair comparison with structured pruning methods: 1) All of these methods need to be reproduced with llama-3.1-8B as the base model. 2) Structured pruning methods almost always support fine-tuning after pruning, and you should implement training using the same amount of data as your work to achieve a fair comparison.
- While LLAMAFLEX can incubate SLMs similar to Minitron accuracy, the layer-pruning-only version of Minitron-LLaMA-3.1 will embody far more efficiency than an SLM of the same size, and you need to make a full comparison.

questions:
- L094-L096: The value of a heterogeneous architecture should not be dismissed because of unfriendly backend support. I believe that loosening the homogeneous restrictions in LLAMAFLEX will boost the potential to produce stronger models at the same scale. Can you provide some relevant results (or insights if it is resource-intensive to build this experiment)?
- Are open-source SLMs in the Phi series and the LLaMA 3.2 series comparable in effectiveness and worth adding to the evaluation?
- How much computational and time resources were used for this work?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705581318

thread_type:
official_review

<thread_object>
creator: reviewer_7wyq

<content>
summary:
This paper introduces a novel LLM compression framework named LLaMaFlex. With around ~60B tokens for elastic training, the network can be resized both in the depth and the width dimension, producing better results than previous pruning or distillation-based methods. The method composes a Gumbel-softmax router to select the sub-network after the elastic training. For the training, the model is first arranged by the importance of sub-networks and a policy-aware modulation is proposed. Experiments show that it can achieve much better results than previous compression methods.

strengths:
1. The idea of “train once, generate many” for LLM is interesting for compressing large language models. While this paradigm isn’t entirely new, this paper achieves impressive results to apply this on LLM.
2. The proposed method is both well-founded and novel. A particularly interesting discovery is that the learned router can interpolate, enhancing the method’s generalizability.
3. The experimental results demonstrate strong performance.

weaknesses:
I did not find significant weaknesses for this paper

questions:
1. How does the performance compare if elastic training is not applied to retrain the LLM? Specifically, if only the Gumbel router is learned on the pre-trained LLM, while retaining the pre-trained weights of LLaMA-3.1-8B.
2. Have you conducted experiments on other LLMs to demonstrate that this method can be extended across different types of language models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732771652223

thread_type:
official_comment

<thread_object>
creator: area_chair_gxjg

<content>
title:
Reviewers, please kindly respond

comment:
Dear Reviewers,

If you have not responded to author's rebuttal, please kindly do so as soon as possible. The deadline is Dec 2, but the authors can potentially further clarify questions if you respond earlier. Thanks!

Best, AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
(a) Summary: introduces LLaMaFlex, a new weight-shared training architecture with routers, for LLMs. It enables zero-shot pruning for different widths and depths.

(b) Strengths: being able to "train once, deploy many"; strong performance; easy deployment with uniform architectures, etc.

(c) Weaknesses: limited evaluation (LLaMA 3.1 8B); clarity issues in writing/organization; more fair comparisons needed.

(d) Reasons for decision: reviewers' unanimous support; "train once, deploy many" is highly demanded for LLMs.