LLaMaFlex: Many-in-one LLMs via Generalized Pruning and Weight Sharing
Authors: Ruisi Cai
Saurav Muralidharan
Hongxu Yin
Zhangyang Wang
Jan Kautz
Pavlo Molchanov
Venue: ICLR 2025 Poster (2025)
Keywords: large language models
elastic networks
training efficiency
inference efficiency
Primary area: foundation or frontier models, including LLMs

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.

Published as a conference paper ICLR 2025
LLAMAFLEX: MANY-IN-ONE LLMS VIA GENERALIZED PRUNING AND WEIGHT SHARING
Ruisi Cai Saurav Muralidharan', Hongxu Yin', Zhangyang Wang', Jan Kautz', Pavlo Molchanov' NVIDIA "The University of Texas at Austin
ABSTRACT
Large Language Model (LLM) providers typically train a family of models, each of a dif- ferent size targeting a specific deployment scenario. Models in the family are all trained from scratch, making the process extremely resource intensive. Recent work has suc- cessfully 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 effec- tive, 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 pruned across both width and depth dimensions in zero-shot manner to instantly yield large number of highly accurate compressed models. LLAMAFLEX starts from a pre- trained model, and only requires 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 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 of aht pruned, elastic/flexible, and trained-frr coimed-from-seeee from-scratch models.
1 INTRODUCTION
Large Language Models (LLMs) have demonstrated remarkable effectiveness across diverse range of tasks (Hendrycks et al. 2020 Chiang et al.] 2024 Zheng et al.] 2024). However, the generality and ro- bustness of L.Mss are largely attributed their vast scale, with parameter counts ranging from one billion to several hundred billion (Tovvron et al.) 2023 Zhang et al.) 2022). This substantial model size, in turn, makes them extremely resource-intensive train and deploy. This problem further exacerbated in model deploy- ment scenarios characterized by varying compute and memory requirements, requiring model providers to train entire LLM "family", containing different-sized models for each scenario; for instance, the Llama 3.1 model family includes three different variants with 8, 70 and 405 billion parameters, each trained from scratch on 15 trillion tokens (Dubey al. 2024).
Recent work has attempted to reduce the training cost of LLM families via a combination of structured pruning and knowledge distillation (Muralidharan et al., 2024]; in this setting, only the largest model in the family is trained from scratch, with smaller models obtained via compression. Specifically, the original pretrained model is pruned along both depth (layers) and width (attention heads, hidden dimension, MLP in- termediate size) dimensions to target particular parameter budget; this pruned network is then distilled with
"Correspondence sauravm@nvi dia.

---

Published as conference paper ICLR 2025
the original model acting as teacher. While this approach succeeds in reducing both the compute and number of training tokens required (w.rt. training each model from scratch), it suffers from major drawback: each compressed model targeting a particular parameter or latency budget must be distilled separately; this keeps the training cost for LLM families very high. Instead, we ask the following question: is it possible to create a flexible network architecture which can yield different pruned sub-nelworks (depending on budget) without the need for any additional fine-tuning or distillation in single sho!? Our answer to this question is new elastic LLM architecture and pest-traiinng optimization framework named LLAMAFLEX.
70
LLAMAFLEX utilizes nested weight-shared net- work architecture which can be sliced along both depth (layers) and width (attention heads, hid- den dimension, MLP intermediate size) axes, en- abling on-demand, zero-shot generation of accu-
60-
Zaro
The 45
rate pruned models at inference time. An end-to- end trained router is learned to dynamically decide the optimal network architecture given parame- ter budget. We start from pretrained model, and perform brief elastic pretraining and router train- ing (60.4B training tokens in this work). LLA-- MAFLEX is able to generate large number of pruned networks without the need for any addi- tional fine-tuning or distillation. As shown in Fig- ure the zero-shot pruned models obtained with LLAMAFLEX perform par with better than models obtained via structured pruning and dis- tillation; also notice that LLAMAFLEX mod- els outperform similarly-sized models trained from scratch. Figure 2 provides an overview of how LLAMAFLEX models are created, and compares our approach to (1) training from scratch, and (2) pruning + distillation.
Uame--77
MMLU
30
Pyta+1/B223
Stee: 30278
20-
LlamaFlex
Flextron
Pythia
Minitron
Sheared-LLaMA
LLM-Pruner
10
Non-embedding Params
Figure 1: MMLU vs. model size LLAMAFLEX vs. other similar frameworks/models. LLAMAFLEX en- ables zero-shot generation of a Pareto frontier of models that beats the accuracy of current state-of the-art compressed and trained frondd-fro--see coom-ccatch models.
Elastic nested transformer architectures are not new and have been explored in the past. Specifically, Mat- former (Kudugunta et al.) 2023) and Flextron (Cai al., (024) use a nested weight- architecture that Flextron also includes an input-adaptive router that can select some width axes (MLP intermediate size and number of attention heads) based on a target latency. We note that LLAMAFLEX differs in several key ways: (I) to the best of our knowledge, no existing elastic training
very similar to the one used in LLA.MA
From Scratch
wis maunt
8B
% Pruning
wade Daure
Uaama 48
Architecture
Architecture
Liama 8B
Liama
Liama 2B
Ulama 28
LlamaFlex
4B
3.2B
2.2B
Architecture
Liama
LlamaFlex 88
Pretraining
Bladti Traising
biluts
2B
1.38
1B
Figure 2: High-level overview of LLAMAFLEX. In contrast to training from scratch (top left part of the figure) and pruning (top right), LLAMAFLEX supports the generation of potentially infinite number of pruned models without additional fine-tuiingddisiillaioon (zero-shot) after a brief elastic pretraining phase.
2

---

Published as a conference paper ICLR 2025
framework supports the full suite of width and depth dimensions other than LLAMAFLEX, (2) pure nesting introduces hard constraint on the resulting models which may potentially make them less expressive; we counter this by introducing additional parameters for each nesting level through a technique we term policy- aware modulation, and (3) both Matformer and Flextron produce heterogenous architectures (each layer has potentially different architectural configurations) which are difficult to deploy common frameworks such TensorRT-LLM (NVIDIA 2023) and cpp: LLAMAFLEX always produces uniform architectures which have optimized implementations in such common LLM deployment frameworks.
This paper makes the following key contributions:
1. Introduces novel elastic LLM architecture named LLAMAFLEX which can be resized along both depth and width axes to achieve superior accuracy (compared to other pruning approaches, nested architectures, and training from scratch) for a given deployment scenario with no fine-tuning re- quired. Further, LLAMAFLEX produces pruned architectures that are uniform, making it easy to deploy using existing LLM frameworks such as TensorRT-LLM.
2. Presents a Gumbel Softmax-base end-to-end learnable router that is able to interpolate smoothly from 25% 100% of the original model's 's size zero-shot manner. This enables the "train once, generate many" paradigm.
3. Increases the generality of nested architectures by introducing a policy-aware modulation technique inspired from the diffusion model literature.
4. Produces a state-of-the-ar family of compressed models from Llama 3.1 8B, each of which can be obtained zero-shot after initial elastic pretraining and router training phase consisting of ~ 60B tokens. ALAMMFF.EE achieves better accuracy than state-of-the-art compression techniques such Minitron Muralidharan et al.] (2024), nested weight-shared architectures such Flextron Cai ct al. (2024) and models trained from scratch.
2
METHOD
In the following sections, will cover the design of LLAMAFLEX and go the process of converting a pretrained LLM to an elastic form. We provide a high-level overview LLAMAFLEX's design in Figure
2.1 ELASTIC FRAMEWORK: BACKGROUND AND NOTATION
Large Language Models (LLMs) process diverse inputs and tasks by passing tokens through stacked trans- former blocks, whose shapes are characterized by some key architectural variables: the number of trans- former blocks (N), the hidden dimension size (H), the intermediate dimension of the MLP layer (D), and the number of attention heads (N)). Collectively, these parameters, denoted as (D, N , N), describe the "shape" of the model. Specifically, an LLM with shape (D, N), models language as:
(.I. (x) = NNN
where
(1)
(MHA o MLP (æ;) + i i=,,..,,
L++1
LLAMAFLEX allows the original model to operate diverse model variants, each with unique shape and achieving different performance-efficieney trade-off. Users are able to choose the variant with a suitable size during inference, based on their resource and performance constraints. Specifically, model variant j
with shape (D), N HJ, N can be formalized as follows:
Elastic-LLM (Da NN Ns) (z; X ZN where
E X (MHA (Hi (H))
MLP (H+ (HS.D.) (xx) + Ei, i= 1,..,,,, N - 1,
(2)
3

---

Published as conference paper ICLR 2025
Lager Jaokang
75%
Paasins
Poulls
Rouson
Hester Pogsar
77% MIP
75% vertent
60% Vnaraent
varizn
(100%)
varannt
Figure Overview of LLAMAFLEX Framework. router introduced to determine the sub-network configuration (i.e., the size of the MLP, MHA, hidden dimensions, and the number of layers) that satisfies a given budget target (i.e., the percentage of remaining parameters) while optimizing performance. Although the router is trained on discrete "anchor""budee targets (c.g, 50%, 75%, as shown), generalizes well to unscen budget targets (c.g., 62.5%).
here, X is a binary scaler, controlling whether layer submodel j is skipped (A = 0) or not. X is the i-th item of X; E, X No. The elastic MHA and MLP layers are defined as follows:
MLP (#i, (H),Di) (a:;) = o =0(2: (In/ W m:)") (I:))
(IDs I
MHA (1) (a;) = Connca((eead,...ee NY)) HNy.Woo
(3)
head, Attnnx;; W æ.I W wxI W.).
where, I DDD IH RHH and NNCC RHHH diagonal matrices with the first DJ, HJ. and C diagonal elements equal to and the rest respectively. C denotes the size of single head. These diagonal matrices ensure that the j"h sub-model only utilizes the first HJ hidden features, the first D MLP intermediate neurons, and the first N attention heads. We constrain D D, H. D, H, NA- (1) and (2) are the associated two weight matrices in MLP layers, with (11), (2) E DxH. o()) refers to the non-linear activation function. WQ. WK, W. implementation, the diagonal matrix I can be replaced with a slicing operator. Compared to existing meth- ods (Kudugunta al.) 2023] Kavehzadeh et al. 2023) Cai et al.) 2024), LLAMAFLEX is the first generalized elastic framework that supports learnable layer skipping and flexible hidden dimension.
# and Wo E R.CCC For
2.2 GENERALIZABLE ROUTER DESIGN
Problem Formulation The router is tasked with identifying the optimal configuration of the sub-model (Di,N' HJ, H),N) achieving the best performance while adhering to budget constraints 6,. defined as the percentage of remaining parameters relative to the original model. This process can formalized as a combinatorial optimization problem:
min (Di,N), HH,NN Expp(a)
(4)
s.t. P(D', NO H. NO A,
D ED, N H³ E H, NEE,
4

---

Published as a conference paper ICLR 2025
where P(-) denotes the number of parameters given the network dimensions (i.e., Di, NN HH,, No) and the
binary vector for layer skipping (i.e., X²). PDD, NA, H, N, 1x) denotes the number of parameters in the original pre-trained model. To satisfy the constraint, we introduce an additional loss term, which is detailed
in See 2.3 Each of the dimensions are selected from the pre-defined sets D, NA. H. N,, respectively.
Router Design. Due to the discrete nature of the variables DJ, D HJ, NY, ( [0, N]], traditional gradient descent methods are ineffective they require the continuous differentiability of the objective function with respect to the variables. Direct optimization of these discrete choices would typically in- volve exhaustive search evolutionary algorithms, which can be computationally prohibitive given the large search space. To address this challenge, we utilize the Gumbel-Softmax trick (Jang et al., 2016) as a continuous relaxation of the discrete optimization problem. The Gumbel--Sfttan technique allows proximate categorical distributions over discrete choices with a differentiable softmax function, enabling the use of gradient-based optimization methods. We first represent each architectural variable as a categorical distribution, and instead of directly selecting values, we model the probability of each choice. For example, for DJ, we define a categorical distribution with class probabilities P(D ma) for ma E DD where the probabilities satisfy the constraint En. P(D m.) = 1. During stochastic sampling, the algorithm assigns higher probabilities to values of my that are more likely to produce better performance, reflecting a preference for favorable configurations.
ap-
Router Architecture We model the categorical distribution of each architectural variable using a lightweight router implemented two-layer MLP. As described in Section 2.T architectural variables include the MI.P intermediate dimension D`, hidden dimension HP, number of heads in MHA NA. and the Bernoulli variable X, which controls whether to skip the layer i or not. For each budget by E B, we first encode as a one-hot vector, hj Ono-Hot(b;), where hj E R| For instance, to model the distribution for D we input h, into the MLP and obtain the un-normalized bog-probilities logæ(D) [Da ma) MLP D (hz)) Similarly, Ong=(H-- = ma) = MLP Ir (h2), log=(N)) = m.) = MLPN,, (h;), and logt(X; = m.) MLPx, (hz). This framework allows the router to generalize to unseen budget targets, by & B, crucial capability we will discuss in more detail later in this Section.
compute the logits for the other architectural variables:
Router Optimization To enable differentiable sampling, we apply the Gumbel-Softmax (Jang et al. 2016) trick to these categorical distributions. For each architectural variable, the categorical distribution is repa- rameterized as:
P(D ma) = exp ((Klog ma) + ga)/T) EmED exp ((klog= (Di = -m) + =g/T)²
(5)
where da is sample from the G0mbel(0,1) distribution, and is a temperature parameter that controls the smoothness of the approximation; x the scaling factor to balance the relative magnitude of logits and Gumbel noises. As + O, the distribution approaches one-hot vector, allowing the router to make discrete choices. With the same formulation, we compute P(H) m,), P(N ma). and P(X mx). The optimization problem now focuses on minimizing the following loss function:
Lx () E ((Di, N NA. (V.
(6)
where Og represents the parameters of the router and 0 denotes the parameters of the backbone. 2pp in- dicates that the architectural choices (D), No HP, N),)) are sampled from the probability distributions modeled by the routers, for budget target bj. which is parameterized by Oz.
Router Interpolation LLAMAFLEX allows the router to generalize to unseen budget targets at zero addi- tional cost. This is crucial capability that in turn allows LLAMAFLEX to instantly generate pruned models
5

---

Published as a conference paper ICLR 2025
largeting specific user-defined deployment targets. As previously discussed, when by E B, the embedding vector is obtained through one-hot encoding, h; = One-Hot(bj). For an unseen budget target by & B, its embedding vector can be derived by linearly interpolate between its nearest neighbors in the budget set. Specifically, Assume that b, and b are the nearest known budget targets such that b,, I by I bnt1. The embedding vector hg for the unseen budget target by. can be interpolated linearly as follows:
bn+1
by bn bn
h
hn
+
ha++1
bn+1
by
ba+1
This linear interpolation ensures a smooth transition between the embedding vectors of the known budgets, enabling the router to generalize effectively the unseen budget target by- With this interpolation capability, LLAMAFLEX can produce a spectrum of model variants that adapt to continuously changing budget targets.
2.3 MODEL PREPARATION AND ROUTER TRAINING
Model Preparation Given the pre-trained model, we first prepare it following the approach in prior work (Cai et al., 2024, Muralidharan et al., 2024). Specifically, using small set of data samples, we compute the importance of each MLP neuron, attention head, and hidden feature based on the accumulated magnitude of activations. Once the importance is determined, we reorder the corresponding weight matrices such that neurons, heads, and hidden features are arranged in decreasing order of importance for each layer. Sub-networks can then be constructed by simply selecting the first several neurons heads in each layer, thereby preserving the essential knowledge encoded in the most important channels.
Router Training After the model preparation step, we conduct end-to-end training jointly optimize the backbone parameters 0 and router parameters On. using the loss function below:
L= Σ (Cr(æz; (R. + LB(OR.)))
(8)
Here, we note that each sub-network. j 'uses different data batch Tz. Since the sub-networks are created by slicing and therefore share weights, the gradient propagated by one sub-network can implicitly affect the others. This allows all sub-networks to benefit from the information across different data batches, even though they are directly trained on separate inputs. Unlike traditional pruning (Xia et al. 2023 Ma et al.,
2023 Muralidharan et al. 2024), which generate pruned variants separately, weight sharing helps training efficiency by allowing sub-networks to collectively leverage knowledge from various inputs. addition to the router loss (-))), introduce additional loss term (.) ensure that the sampled sub-network adheres to the parameter budget constraint defined in Equ. Here, P(-) represents the number of parameters based on the network dimensions (i.e., DJ, N3, HJ, N) and the binary vector for layer skipping (i.e., X)). Meanwhile, Piun refers to the total number of parameters the original pre-trained model.
(1D) P(D), NN,, H., No x²) - bj Pran. 0): where
(O,,n) = max HUD..N
Piul = NA, N, 1)),
Following prior approaches (Fang et al. 2024). we apply exponential decay to the value of temperature, while linearly increasing the scaling factor of router outputs. We provide more details on hyper-paranee parameee choice in Sec 3.1 In the initial phase, when the temperature is high and the scaling factor is small, signif- icant randomness is introduced, allowing for optimization of the router. As training progresses, the router transitions to making deterministic decisions.
2.4 POLICY-AWARE MODULATION
Due to its simplicity, a nested structure has been widely adopted in most existing elastic LLM work (Kusu- pati et al. 2022 Kudugunta et al. 2023 Kavehzadeh al. 2023). In this approach, each sub-netwee network
6

---

Published as a conference paper ICLR 2025
generated by simply slicing the weight matrices. However, this constraint can potentially limit the represen- tational capacity of elastic networks, as the same set of weights must accommodate inputs across all possible sub-networks.
To address this limitation, we propose the use of policy-aware modulation, a technique in- spired by methods in the literature on dif- fusion models (Ho et al.) 2020, Peebles & Xie 2023). In these models, the time step is encoded in the Transformer layers through learnable normalization, allowing the same parameters to process inputs with varying noise scales. We adopt similar approach in LLAMAFLEX by introducing lightweight non- linear modulation heads, which are applied after the elastic components (i.e., elastic MLP/clastic MHA, as detailed in Equ. 3). These heads modulate the outputs of clastic operations based on the elastic choice. Taking clastic MLP as an example, the current sub-network chooses to use e₁ hidden neurons, we condi- tion the output of the elastic MLP on с₁-- Concretely, we use sinusoidal embedding of e₁ followed by a learnable MLP layer to generate modulation vectors for scaling and shifting. These modulation vectors then transform the output of the elastic MLP y as follows:
NLM-100%
Rouler
Fiddin Featare
Layor
14 P N05
modiatton MVF
Frasón Fealase
Scale Scaeeatero sirtor
Figure 4: The illustration Modulation.
ÿ = y MLP scale (Emb(ex)) + MLP stin (Emb(ex))
(10)
3 RESULTS
3.1 EXPERIMENTAL DETAILS
Pre-trained Model and Dataset We validate our method on Llama 3.1 8B (Dubey et al. 2024), a promi- nent open-source large language model trained 15 trillion tokens. The model comprises N 32 Trans- former blocks, each containing a Multi-head Attention (MHA) layer with 8 attention groups and N = 32 query heads, along with Multilayer Perceptron (MLP) with an intermediate dimension of D 14436. The hidden feature dimension is H 4096. The model has a total of 6.98 billion non-embedding parameters, with 5.64 billion parameters allocated the MLP layers. Since the original training data is not publicly available, we use a proprietary dataset consisting of high-quality pretraining data. With 4 choices of 6,, we sample different tokens for different budget goals (see details in Equ. 8). and use 60.4 billion training tokens in total.
Downstream Tasks We evaluate LLAMAFLEX several downstream tasks including ARC-easy (Clark
ct al. 2018), LAMBADA (Paperno et al. 2016), PIQA (Bisk et al.] 2020), WinoGrande (Sakaguchi et al.) 202T MMLU (Hendryeks al.) 2020). and HellaSwag (Zellers al. 2019). Following the approach of
Xiaa et al., 2023). we report 5-shot performance for MMLU and 10-shot performance for HellaSwag, while presenting zero-shot results for the other tasks.
Elastic Configurations and Training Details As de- scribed in Section 2.2 for every budget target by B, architectural choices, DJ, No and Hi are sampled from the predefined sets D, .A. and H. respectively. We present these sets choices in [able We omit the details of layer skipping, as its indicator variable X is a binary scalar that controls whether layer i in model j is skipped. Unless otherwise specified, we set the sequence length 4096 and the batch size to 128, and
Table 1: Details of clastic configurations.
Budget
8
(25%, 75%, 100%)
MHA MLP inter. dim
25% 50% 75%, 100%
(25%, 37.5%, 50%, 62.3%, 75%, 87.5%, 100%
Hid. dim
H
(500%, 02.5%, 5%, 87.5%, 100%)
7

---

Published as conference paper ICLR 2025
fine-tune the model for 28800 iterations. We set the initial learning rate 4e 5, and use cosine learning rate decay for LLM parameters.
Downstream task evaluation of LLAMAFLEX framework. Here, #Params refers the number of embedding parameters. The prefix Exp. indicates the sub-models are explicitly trained, i.c., b E B, while prefix
Table 2:
Inter. refers to models generated through router interpolation.
Type
Model
#Params
ARC-E
LAMB.
PIQA
Wino.
Hell. (10)
MMLU (5)
Avg.
Exp.22%%
1.8 B
66.7%
61.6%
74.3%
61.9%
67.6%
42.2%
62.4%
1n1er.38.8%%
2.7 B
71.3%
64.6%
76.0%
62.5%
71.1%
53.8%
66.6%
Exp.55%%
3.5 B
76.1%
66.3%
77.0%
67.9%
76.0%
60.3%
70.6%
LLAMAFLEX
Inter-62.5%%
4.4 B
77.4%
68.1%
77.4%
69.4%
76.4%
61.2%
71.7%
Exp--7%%
5.2 B
78.9%
71.7%
78.1%
72.9%
77.2%
62.7%
73.6%
Inter-866.1% 15%
6.0
79.2%
72.1%
78.8%
74.1%
77.9%
63.1%
74.2%
Exp--100%
7.0
81.9%
72.8%
79.7%
75.4%
78.0%
64.7%
75.4%
0.5x
3
65.9%
61.7%
74.8%
61.9%
67.6%
35.9%
61.3%
0.6x
B.0
66.1%
63.8%
75.0%
62.1%
68.0%
37.7%
62.1% 1%
Flextron
0.7x
B.2
65.8%
64.2%
75.6%
62.3%
67.1%
41.9%
62.8%
Full
6.5
75.1%
71.5%
77.5%
69.1%
78.1%
45.1%
69.4%
Minitron- Minitron-Depth Minitron- Width
3.7 B
74.7%
61.5%
74.9%
72.5%
73.2%
58.7%
69.3%
3.7 B
75.1%
63.0%
75.0%
73.1%
76.1%
60.5%
70.5%
Sheared.-1.3B
1.2
61.5%
61.0%
73.4%
57.9%
60.7%
25.7%
56.7%
Sheared.-2.78
2.5 B
67.0%
68.4%
75.8%
64.2%
70.8%
26.4%
62.1%
Compression
NutePrune
3.2
51.7%
71.0%
57.5%
55.9%
LLM-Pruner
4.5
59.2%
73.4%
64.2%
56.5%
23.9%
Compresso
B.5
66.0%
72.9%
63.4%
25.9%
SliceGPT
4.8
66.2%
50.3%
28.9%
LaCo
B.7
69.8%
55.7%
26.5%
-amma-311-BB
7.0
81.8%
72.9%
81. 81.0%
75.7%
81.8%
65.3%
76.4%
Llama2-7B
B
75.2%
68.2%
78.8%
69.2%
78.6%
45.3%
69.2%
OpenI. LaMA3BB22
3.2 B
63.7%
59.1%
78.1%
63.3%
71.6%
25.7%
60.3%
OpenI LMMA-BB22
B
69.5%
63.8%
79.9%
66.0%
76.6%
40.4%
66.0%
Open-Source
GPT3-8B
6.4 B
70.1%
70.5%
79.7%
69.8%
77.7%
40,2%
08.0%
Pythia- 1.4B
1.2
53.9%
46.8%
70.6%
57.1%
52.2%
25.6%
51.0%
Pythia-2.8B
2.5 57.9%
50.1%
73.8%
58.6%
60.0%
26.8%
54.5%
Pythia-6.9B
6.4
60.2%
47.1%
75.2%
59.9%
64.4%
25.5%
55,4%
Router Architecture and Modulation Details For each architectural variable, we use a two-layer MLP followed by a shared embedding layer. The embedding layer converts the scalar by into embedding vector of dimension 128, and is shared across all architectural choices. Each MLP processes an intermediate dimension of 128 and outputs logits vector, with dimension corresponding choices. The routers contains 0.51 million parameters in total. We use Gumbel-Softmax to optimize the routers, following the practice outlined in MaskLLM Fang et al. (2024). Specifically, we exponentially decay the temperature with rate 0.9999, and linearly scale the scaling factor к from to 10. We set the initial learning rate of both soft and hard Gumbel-Softmax techniques (Jang et al.) 2016). For each elastic choice, we modulate the elastic output by a sinusoidal embedding, of size 16, followed by learnable lightweight MLP with an intermediate dimension of 128. Every modulation network only contains 0.02 million parameters. We also set the initial learning rate of 4c 2 for modulation networks.
the number of elastic
for router tuning. During the tuning process, we employ a combination of
3.2 DOWNSTREAM TASK EVALUATION
We present downstream task evaluation results in Table Here, compare LLAMAFLEX against existing flexible inference frameworks such as Flextron (Cai et al. 2024), as well as open-source models, including Llama 3.1 8B (Dubey al. 2024). Llama2 7B (Touvron 2023), Pythia (Biderman al.
8

---

Published as a conference paper ICLR 2025
2023). Ope [Geng & Liu| 2023), and GPT3 8B (Shoeybi et al. 2019). Additionally, we com-
pare with models generated by post-hoc hoc compression methods, including Minitron (Muralidharan et al.
2024). Sheared-LLaMA (Xia al. 2023). Compresso (Guo al., 2023). LLM-Pruner (Ma et al. 2023)
SliceGPT (Ashkboos et al., 2024). and LaCo (Yang et al. 2024).
Since MMLU (Hendrycks et al. 2020) is a comprehensive benchmark for evaluating general knowledge across a wide range of domains, we provide a detailed visualization of the model's performance on bench- mark in Figure [ From Table and Figure [ we notice that LLAMAAF.FX achieves the best performance- efficiency trade-off' compared to bascline methods. In addition, the interpolated LLAMAFLEX models (de- noted by the prefix Inter. .) demonstrate a smooth transition in the performance-efficiene trade-off com- pared to the models explicitly receiving gradients (denoted by the prefix Exp. .), highlighting the generaliz- ability of our routers.
3.3 ROUTER VISUALIZATION
In Table 3 we visualize the 4 sub- networks learned by the LLAMAFLEX router; here, each sub-network satisfies unique budget requirement by. As shown in the table, for higher budgets (e.g., 50% or 75%), the sub-networks gener- ally avoid skipping layers and utilize rel- atively smaller MHA and MLP blocks. However, for lower budgets (e.g., 25%), the router tends to skip nearly half of the layers. Additionally, router finds it hard to reduce the num- ber of attention heads for Llama 3.1, and it remains the same for all different sub- networks. For the 25% variant, layers [1-7, 17-24, 29-30, 32] are retained while the others are skipped. In the 50% variant, the 13-th and 31-th layers are removed.
Table 3: Architecture details of sub-networks selected by the learn- able router. Each sub-network corresponds to different budget by: the architecture of the pre-trained Llama 3.1 8B model is shown in the by 100% row. We present the MLP intermediate dimension D', the number of attention heads (the number of queries) N the hidden dimension H3, the number of layers N and the resulting number of non-emben non-embedding parameters. Note that since the pretrained model uses grouped-query attention and we keep the number of attention groups fixed, modify the number of queries per group.
by
#Param
MLP Dim.
MHA #Heuds
Hid. Dim
#Remain. layers
observe that the
25%
1.848
10732 (75.0%
24 (75%)
2560 (22.5%)
18 (56.25%)
50%
3.48B
12544 (87.5%)
24 (75%6)
2560 (62 5%)
30 (93.75%)
75%
5.20.
12544 (87.5%)
24 (75%)
3584 (87.5%)
32 (100%)
10096
6.988/
14336
32
4096
32
4 ABLATION STUDY
Effect of Policy-Aware Modulation We validate the effectiveness of our modulation technique by comparing per- formance with and without modulation.
Table 4: Results of our ablation study on policy-aware modulation. We run LLAMAFLEX for 800 iterations report the validation loss. All sub-networks show improved performance when modulation is enabled, reducing validation loss by 0.08 on average.
We set = f25%, 50%, 75%, 100%). and perform LLAMAFLEX training with and without modulation. We train Llama 3.1 8B (Dubey et al. (004) for 800 iterations, and show the obtained performance in Table We notice that including modulation consistently improves the performance of all sub-networks, reducing validation loss by 0.08 average. We provide more experimental results in Appendix
by
25%
50%
75%
100%
avg.
Modulation wo Modulation
2.53 (L 0.08)
2.41 (L3.1
2.08 (1 0.09)
1.88
2.22 (L 0.08)
2.6l
2.54
2.17
1.88
2,30
5 RELATED WORK
Structured Pruning A number of recent structured pruning papers specifically target LLMs; we can broadly classify these works into two main categories: (1) ones that prune only depth (layers), and (2)
9

---

Published as a conference paper ICLR 2025
ones that prune width (attention heads, MLP intermediate dimension, etc.) and/or depth. Recent work in the first category (depth pruning) includes ShortGPT Men et al. (2024). LaCo] Yang et al. (2024), and Shortened LLaMa Kim al. (2024). while those in the second category includes Minitron Muralidharan al. (2024). ShearedLlama Xia et al. (2023), Dery et al. (2024), SliceGPT Ashkboos et al. (2024), and LLM-Prrnner Ma jet al. (2023). To the best of our knowledge, all the work this area requires fine-tuning or distillation on a per-compressed-modde basis (not zero-shot). We compare the accuracy of LLAMAFLEX models against a large subset of these frameworks in this paper.
Flexible Architectures Flexible inference has been extensively studied, particularly in the context of con- volutional neural networks (CNNs). Yu et al. (2018); Yu & Huang (2019) introduced slimmable neural networks, which enabled the deployment of the same model with varying numbers of convolutional ker- nels. Building on this, OFA (Cai et al. 2019) generalized pruning techniques to create a single model that can adapt multiple configurations. Recent research has extended slimmable models Transformer architectures. Kusupati et al. (2022) introduced nested weight structure for Transformer networks, focus- ing primarily on hidden features. Matformer (Kudugunta et al. (023) further developed this architecture by applying it to the intermediate dimensions of MLPs. Additionally, SortedNet (Valipour et al., (023) employed a sampling-basee
training strategy to train multiple models via gradient accumulation, while SortedLlama (Kavehzadeh et al., (002)) explored depth-wise sampling in large language models (LLMs). Recently, Flextron (Cai et al.] (024)) introduced elastic multi-head attention (MHA) and search-based rout- ing to enhance flexibility. Similar to Flextron, LLAMAFLEX leverages router mechanism, but supports end-to-end training, interpolation between trained targets, and produces more easily deployable uniform ar- chitectures. Additionally, LLAMAAFLEX introduces novel policy-aware modulation technique, providing additional model expressivity and enhancing adaptability.
SuperNet-based Neural Architecture Search (NAS) HAT (Wang et al., 2020) and HELP (Lee et al.,
(02T) have explored supernet-based approaches for generating sub-networks, primarily targeting relatively small-sized models, while Meta-NAS (Elsken et al., 2020) and DARTS-EGS (Chang et al.] 2019) have utilized wwwnnbel-Softaxx-bll
supernet techniques. Our work focuses on addressing the unique challenges of LLMs, which differ significantly from prior studies targeting smaller models. To avoid costly training processes and huge GPU memory usage, transform pre-trained model into elastic framework, and employ weight-sharing unlike previous attempts. To avoid repeated computation, we enable zero-shot sub- network generation for any parameter budget through router interpolation capabilities not achieved by supernet-based methods.
6 CONCLUSIONS
In this paper, we have presented LLAMAFLEX, a novel elastic LLM architecture that supports zero-shot resizing along both width and depth dimensions, yielding large space of uniform (for ease of deployment) compressed models without any fine-tuning. It utilizes a Gumbel Softmax- based end-to-end learnable router that requires a one-time training phase, and is then able smoothly interpolate across model sizes, ranging from O to 100% of the original model size. We have also introduced a novel policy-aware modulation technique that improves the generality of nested architectures. LLAMAFLEX on Llama 3.1 8B produces a frontier of highly accurate compressed models that outperform similarly-sized compressed, elastic/flexible, and sainned-rrom-scrateee models.
Limitations.
Our method is not training-free, as it requires router tuning and necessitate modification on backbone weights to convert the pre-trained model into an elastic one. There is a potential application of training LLAMAFLEX from scratch, we leave this to future work.
10

---

Published as conference paper ICLR 2025
REFERENCES
Saleh Ashkboos, Maximilian L Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hens- man. Slicegpt: Compress large language models by deleting rows and columns. arXiv preprint arXiv:2401. 15024, 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hal- lahan, Mohammad Aflah Khan. Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430 PMLR, 2023.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference on Artificial Intelligence, 2020.
Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once-for-all: Train network and specialize it for efficient deployment. arXiv preprint arXiv: 19008.9791,, 2019.
Ruisi Cai, Saurav Muralidharan, Greg Heinrich, Hongxu Yin, Zhangyang Wang. Jan Kautz, and Pavlo Molchanov. Flextron: Many-in-one flexible large language model. arXiv preprint arXiv:2406. 10260, 2024.
Jianlong Chang, Xinbang Zhang, Yiwen Guo, Gaofeng Meng, Shiming Xiang, and Chunhong Pan. Differ- entiable architecture search with ensemble gumbel-sofimax. arXiv preprint arXiv:1905.01778 2019.
Wei-Lin Chiang, Lianmin Zheng, Ying Sheng, Anastasios Nikolas Angelopoulos, Tianle Li, Dacheng Li, Hao Zhang, Banghua Zhu, Michael Jordan, Joseph E. Gonzalez, and lon Stoica. Chatbot arena: An open platform for evaluating Iims by human preference, 2024.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 18033:1803.055 2018.
Lucio Dery, Steven Kolawole, Jean-Francois Kagey, Virginia Smith, Graham Neubig, and Ameet Tal- walkar. Everybody prune now: Structured pruning of lims with only forward passes. arXiv preprint arXiv:2402.05406, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama Herd of Models. arXiv preprint arXiv:2407.21783. 2024.
Thomas Elsken, Benedikt Staffler, Jan Hendrik Metzen, and Frank Hutter. Meta-learning neural architec- tures for few-shot learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 12365-12375. 2020.
Gongfan Fang, Hongxu Yin, Saurav Muralidharan, Greg Heinrich, Jeff Pool, Jan Kautz, Pavlo Molchanov, and Xinchao Wang. Maskllm: Learnable semi-structured sparsity for large language models, 2024. URL https://arxiv.. ab/p/////2/2/// 17481
Xinyang Geng and Hao Liu. Openllama: An open reproduction of Ilama, May 2023. //github. com Sro.ttubb.commoppnnn-eeeeesee
URL https:
11ama
Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A Roberts. The unrea- sonable ineffectiveness of the deeper layers. arXiv preprint 4rXiv:240..17888 17887, 2024.
11

---

Published as a conference paper ICLR 2025
Song Guo, Jiahang Xu, Li Lyna Zhang, and Mao Yang. Compresso: Structured pruning with collaborative prompting learns compact large language models. arXiv preprint arXiv:2310.0501., 2023.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Stein- hardt. Measuring massive multitask language understanding. arXiv preprint N.XI::20090033000 2020.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-66 6851, 2020.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax, arXiv preprint arXiv: 144, 2016.
Parsa Kavehzadeh, Mojtaha Valipour, Marzieh Tahaei, Ali Ghodsi, Boxing Chen, and Mehdi Reza- gholizadeh. Sorted llama: Unlocking the potential of intermediate layers of large language models for dynamic inference using sorted fine-tuning (soft). arXiv preprint arXiv:2309.0800 08966, 2023.
Bo-Kyeong Kim, Geonmin Kim, Tae-Ho Kim, Thibault Castells, Shinkook Choi, Junho Shin, and Hyoung- Kyu Song. Shortened LLaMA: A simple depth pruning for large language models. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024. URL https:
7/openreview net. forum? id= BVGXuOdpu
Sneha Kudugunta, Aditya Kusupati, Tim Dettmers, Kaifeng Chen, Inderjit Dhillon, Yulia Tsvetkov, Han- naneh Hajishirzi, Sham Kakade, Ali Farhadi, Prateek Jain, et al. Matformer: Nested transformer for elastic inference. arXiv preprint 10X0::2310.07707. 2023.
2023.
Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanujan, William Howard-Snyder, Kaifeng Chen, Sham Kakade, Prateek Jain, et al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 35:30233-30249, 2022.
Hayeon Lee, Sewoong Lee, Song Chong, and Sung Ju Hwang. Help: Hardware-adaptive efficient latency prediction for nas via meta-learning. arXiv preprint 1rXiv:2106.08631 2021.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models. Advances in neural information processing systems, 36:21702-21720 2023.
Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint wwXiv:2403.03885. 2024.
Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Compact language models via pruning and knowledge distillation. arXiv preprint arXiv:2407. 14679, 2024.
NVIDIA. Tensortt-llm: A tensort toolbox for optimized large language model inference, 2023. URL
wew.oorrTTTToTTRTt- LLM
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernandez. The LAMBADA dataset: Word pre- diction requiring broad discourse context. In Proceedings of the 54th Annual Meeting of the Associa- tion for Computational Linguistics (Volume 1: Long Papers), pp. 1525-1523 Berlin, Germany, August 2016. Association for Computational Linguisties. URL http:/ www. www.wb...........co...........eb org/antholooy/ P16 1144
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3195-4205, 2023.
12

---

Published as a conference paper ICLR 2025
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9)999-106; 2021.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1999008053 2019.
Shoaib Ahmed Siddiqui, Xin Dong, Greg Heinrich, Thomas Breuel, Jan Kautz, David Krueger, and Pavlo Molchanov. A deeper look at depth pruning of llms. arXiv preprint arXiv:2407. /628, 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bash- lykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprini 2xXx:23377.09288 2023.
Mojtaba Valipour, Mehdi Rezagholizadeh, Hossein Rajabzadeh, Marzieh Tahaei, Boxing Chen, and Ali Ghodsi. Sortednet, a place for every network and every network in its place: Towards a generalized solution for training many-in -n neural networks. arXiv preprint arXiv:2309.002. 2023.
Hanrui Wang, Zhanghao Wu, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, and Song Han. Hat: Hardware-aware transformers for efficient natural language processing. arXiv preprint arXiv:2005. 14187, 2020.
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared Ilama: Accelerating language model pre-training via structured pruning. arXiv preprint arXiv:2310.06622 2023.
Yifei Yang, Zouying Cao, and Hai Zhao. Laco: Large language model pruning via layer collapse. arXiv preprint arXiv:2402.1 01.4402.11888 117, 2024.
Jiahui Yu and Thomas S Huang. Universally slimmable networks and improved training techniques. In
Proceedings of the IEEE/CVF international conference on computer vision, pp. 1803-1811, 2019.
Jiahui Yu, Linjie Yang, Ning Xu, Jianchao Yang, and Thomas Huang. Slimmable neural networks. arXivv 2018
preprint
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv: arXiv: 1905007830, 2019.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher De- wan, Mona Diab, Xian Li, Xi Victoria Lin, al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
Lianmin Zheng Wei-Lin Chiang. Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging llm-as-a-judge with mt-bench and chatbot arena. Advances in Neural Information Processing Systems, 36, 2024.
13

---

Published as conference paper ICLR 2025
A MORE EXPERIMENTAL RESULTS
Performance of Learnable Layer Skipping All existing techniques layer skipping that we are aware of use hand- crafted heuristies to measure the importance of layers and subsequently prune the layers with the lowest importance scores (Men et al., 2024 Gromov al. 2024 Siddiqui JMM24MMrrlidee Muralidharan al. (0244). In contrast, our pro- posed method uses Gumbel Softmax-based router op- timize layer skipping. To evaluate how well our approach works, we use the full transformer block and only enable layer skipping: thus, to produce a sub-network with 50% remaining parameters we skip half the transformer layers. We compare against two baselines: (1) random layer skip- ping, and (2) skipping the last half of the layers, except for the final layer (i.e., layers 15 to 31 in 32-layers Llama 3.1 8B model), which follows the depth-pruning approach used by Muralidharan et al. 2024). We train the models for 1600 iterations and present the evaluation loss in Figure 5 We observe that while heuristic skipping converges more quickly than other methods, our approach demonstrates su- perior performance once training has converged.
Heuretic Skipping
Random Setoping
Learablle Skipping
2.6
et at. (
Eowiistion
2.4
600
650
Figure 5: Evaluation of the proposed learnable layer skipping method. Here, former block and only enable layer skipping in the router, setting it to use 50% of layers. Our method (green curve) is compared against random skip- ping (orange curve) and the approach used by Mini- (Muralidharan et al. 2024 (blue curve).
the full Trans-
14


Reviews:
[{"review_id": "pn80zReG7l", "reviewer": "Reviewer_7Wyq", "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.", "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.\n2. Have you conducted experiments on other LLMs to demonstrate that this method can be extended across different types of language models?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. 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.\n3. The experimental results demonstrate strong performance.", "weaknesses": "I did not find significant weaknesses for this paper", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lNlAUFj0xC", "reviewer": "Reviewer_uPJk", "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.", "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)?\n- Are open-source SLMs in the Phi series and the LLaMA 3.2 series comparable in effectiveness and worth adding to the evaluation?\n- How much computational and time resources were used for this work?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "+ Introduces a gumbel softmax-based end-to-end learnable router to train the candidate models.\n+ Enhances generality by incorporating a policy-aware modulation technique.\n+ Enables interpolation-based extrapolation of pruning rates for flexible model compression.\n+ 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.\n- 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.\n- 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.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VmZ1bLV0cZ", "reviewer": "Reviewer_XXXs", "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.", "questions": "Please address concerns in Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n- 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.\n- 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.\n- 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.\n- 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.\n- The process for selecting optimal architectures when multiple configurations achieve similar parameter counts needs more detailed explanation.\n\n[HAT] HAT: Hardware-Aware Transformers for Efficient Natural Language Processing\n[HELP] HELP: Hardware-Adaptive Efficient Latency Predictor for NAS via Meta-Learning\n[MatFormer] MatFormer: Nested Transformer for Elastic Inference\n[MetaNAS] Meta-Learning of Neural Architectures for Few-Shot Learning\n[DARTS-EGS] Differentiable Architecture Search with Ensemble Gumbel-Softmax", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CssC4803TQ", "reviewer": "Reviewer_gENy", "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.", "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}?\n\n2. In Line 196, the author indicates that \\mathcal{D}, \\mathcal{N_A}, \\mathcal{N}, \\mathcal{H} are predefined, how to define them?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\n2. The end-to-end trained router could dynamically determine the optimal network architecture under a given parameter budget.\n\n3. 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.\n\n2. 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.\n\n3. 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?\n\n4. The authors did not explicitly discuss the limitations of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)