AmoebaLLM: Constructing Any-Shape Large Language Models for Efficient and Instant Deployment
Authors: Yonggan Fu
Zhongzhi Yu
Junwei Li
Jiayi Qian
Yongan Zhang
Xiangchi Yuan
Dachuan Shi
Roman Yakunin
Yingyan Celine Lin
Venue: NeurIPS 2024 poster (2024)
Keywords: Efficient Large Language Models
Model Compression
Primary area: natural_language_processing
TL;DR: We propose AmoebaLLM, a novel framework designed to enable the instant derivation of LLM subnets of arbitrary shapes, which achieve the accuracy-efficiency frontier and can be extracted immediately after a one-time fine-tuning.

Abstract: Motivated by the transformative capabilities of large language models (LLMs) across various natural language tasks, there has been a growing demand to deploy these models effectively across diverse real-world applications and platforms. However, the challenge of efficiently deploying LLMs has become increasingly pronounced due to the varying application-specific performance requirements and the rapid evolution of computational platforms, which feature diverse resource constraints and deployment flows. These varying requirements necessitate LLMs that can adapt their structures (depth and width) for optimal efficiency across different platforms and application specifications. To address this critical gap, we propose AmoebaLLM, a novel framework designed to enable the instant derivation of LLM subnets of arbitrary shapes, which achieve the accuracy-efficiency frontier and can be extracted immediately after a one-time fine-tuning. In this way, AmoebaLLM significantly facilitates rapid deployment tailored to various platforms and applications. Specifically, AmoebaLLM integrates three innovative components: (1) a knowledge-preserving subnet selection strategy that features a dynamic-programming approach for depth shrinking and an importance-driven method for width shrinking; (2) a shape-aware mixture of LoRAs to mitigate gradient conflicts among subnets during fine-tuning; and (3) an in-place distillation scheme with loss-magnitude balancing as the fine-tuning objective. Extensive experiments validate that AmoebaLLM not only sets new standards in LLM adaptability but also successfully delivers subnets that achieve state-of-the-art trade-offs between accuracy and efficiency.

AmoebaLLM: Constructing Any-Shape Large Language Models for Efficient and Instant Deployment
Yonggan Fu, Zhongzhi Yu,, Junwei Li", Jiayi Qian', Yongan Zhang, Xiangchi Yuan,
Dachuan Shi, Roman Yakunin, Yingyan (Celine) Lin
Georgia Institute of Technology {yonggan. fu, celine. (aa.eeelinnngaeeeee lin)@gatechhedu
Abstract
Motivated by the transformative capabilities of large language models (LLMs) across various natural language tasks, there has been growing demand to de- ploy these models effectively across diverse real-world applications and platforms. However, the challenge of efficiently deploying LLMs has become increasingly pronounced due to the varying application-specin performance requirements and the rapid evolution of computational platforms, which feature diverse re- source constraints and deployment flows. These varying requirements necessitate LLMs that can adapt their structures (depth and width) for optimal efficiency across different platforms and application specifications. To address this criti- cal gap. we propose Amoebal.MM, a novel framework designed to enable the instant derivation of LLM subnets of arbitrary shapes, which achieve the accuracy- efficiency frontier and can be extracted immediately after a one-time fine-tuning. In this way, AmoebaLLM significantly facilitates rapid deployment tailored various platforms and applications. Specifically, AmoebaLLM integrates three innovative components: (1) knowledge-preservinn subnet selection strategy that features dyyaamic-poorraminng approach for depth shrinking and an importance- driven method for width shrinking: (2) shape-aware mixture of LoRAs mit- igate gradient conflicts among subnets during fine-tuning: and (3) an in-place distillation scheme with loss-magnitude balancing as the fine-tuning objective. Extensive experiments validate that AmoebaLLM not only sets new standards in LLM adaptability but also successfully delivers subnets that achieve state- oh-the-art trade-offs between accuracy and efficiency. Our code is available at https: //gi httgggttuub... cOw/CATECH- AmoebaLLM.
Introduction
The remarkable abilities and transformative impacts of large language models (LLMs) [1, 2, 3, 4] have been paralleled by a growing interest in deploying them across wide range of real-world applications and diverse platforms. However, given the rapid evolution of computational platforms and varying application-speeifee requirements, the challenge of deploying L.MMs efficiently on various platforms with differing specifications has become more pronounced. This is because diverse platforms often feature different resource constraints and deployment flows, necessitating LLMs with varying structures and shapes (i.e., depth and width) to achieve maximized execution efficiency, as affirmed by our profiling Sec. 2. Moreover, even the same platform may have varying requirements for LLMs' execution efficiency depending These varying requirements demand a flexible framework capable of adapting to both the intrinsic hardware constraints and the extrinsic demands of diverse application scenarios.
factors such on-device battery status.
"Contributed equally.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Existing efficient LLM solutions [5, 6, 7, 8, 9, 10), which primarily use model compression to bridge the gap between the resource constraints of the target device and the prohibitive complexity of L.LMs, fail to fully address the above needs. This is because these solutions either focus on single dimension of compression, resulting in limited efficiency improvements, or require a costly fine-tuning process for each target platform with its unique specifications. This strategy is particularly unscalable and inefficient for deploying widely used public LLMs like LLaMA [1. 2], where each user must compress and fine-tune the LLMs for their specific platform and application needs.
In light of this, is highly desirable to develop suite of LLMs designed such that compressed subnets of arbitrary shapes, which can achieve the eecuracy-efficiene frontier without the necessity of individual fine-tuning, can be instantly extracted, thus allowing for immediate adaptation to the diverse needs of various platforms and applications. To achieve this, previous one-for-all training techniques [11, 12, 13, 14, 15], which strategically sample subnets for joint training to deliver models with switchable complexity, are promising candidates. However, directly applying them to pre-trained LLMs would lead to failure due to the following challenges: (1) their adopted subnet sampling strategies, which are dedicated to models trained from scratch, are not applicable to extensively pre-trained LLMs as informative and critical components that store useful knowledge are highly likely to be skipped; (2) jointly fine-tuning different subnets on commonly adopted small-scale tuning datasets can easily cause severe gradient conflict [16, 17] for LLM weights pre-trained on a large corpus, thus leading to poor performance of all subnets.
To address these challenges, we propose a framework called AmocbaLLM, which endows a given LLM with the ability to instantly derive compressed subnets of arbitrary shapes that can achieve the frontier. We achieve this through the development of three key components of A.oebalLMMss ss one-for-all fine-tuning scheme: subnet selection strategy, a trainable adapter design, and a fine-tuning objective. Specifically, we summarize our contributions as follows:
We develop a framework called AmoebaLLM,, which grants a given LLM the capability to deliver subnets of arbitrary shapes that achieve state-of-the-art (SOTA) accuracy-efficient trade-off's after a one-time fine-tuning. In this way, AMM can greatly facilitate rapid deployment across varying platforms and applications. This is achieved by integrating the following three key components to enable one-for-all fine-tuning.
For extracting high-quality subnets with diverse shapes, we propose a www.ledge-peeeerviie subnet selection strategy that features dynamic programming (DP)-based depth shrinking and importance-driven width shrinking. This addresses the aforementioned challenge (1) by preserving the encoded factual knowledge and reasoning capabilities of pre-trained LLMs.
For the trainable adapter during one-for-all fine-tuning, we propose a shape-aware mixture of LoRAs (SMoL), which selects and combines sparse set of LoRAs [18] using a gating function that takes the subnet shape as input. This technique addresses the aforementioned challenge (2) by mitigating gradient conflicts among subnets. More importantly, once the target subnet shape is determined based on the target platform/application at deployment time, all the selected LoRAs can be merged into the LLM weights, thus climinating overhead.
For the fine-tuning objective, we enhance the in-place distillation strategy [11, 13, 14] by integrating a loss-magnitude balancing scheme. This scheme is based the observation that the loss magnitudes of subnets with different shapes in pre-trained LLMs are unbalanced, leading bias toward specific subnets and thus poor overall performance. Our proposed technique effectively addresses this issue and improves the performance of all subnets.
Extensive experiments validate that our AmoebaLLM can deliver LLMs with instantly serviceable subnets of any shape, each performing better or on par with SOTA efficient LLM solutions. Additionally, when leveraging Amoebal.IMM as a pure L.M compression framework to compress a given L.LM to the target parameter, it can achieve new SOTA compression effectiveness, thanks to its subnet selection strategy.
2 Motivation and Profiling
Before introducing our framework, we first conduct a profiling of generation latency across different devices, deployment flows, and use cases to examine the demand for L.MS with adaptable structures that meet the diverse needs of various platforms and applications.
2

---

Pyflrrh @ NVIDIA A5000
MLC-LLM jetson Orin NX
TensorRT NVIDIA A5000
MLC-LLM NVIDIA A5000
Latency
15.0
3.5
75
1.6
3.0
of Parameters
Figure 1: The latency of LLaMA2 7B with scaled depth/width various devices/deployment flows.
Profiling setup. We study the efficiency of different LLM shapes on different devices by uniformly scaling either the depth or width of LLaMA2 7B [1] to the same model size. Here, depth is defined as the number of self-attention blocks, each including both a multi-head attention module and a
feed-forward network, while width is defined the hidden dimensions. We profile these workloads using (1) two devices, including an NVIDIA A5000 consumer-level GPU and an NVIDIA Jetson Orin NX edge GPU; and (2) three deployment flows, including TensorRT-LLM [19], MLC-LLM [20). and vanilla PyTorch [21].
Observation and analysis. As shown in Fig. 1, we can observe that (1) first of all, the same workload on two GPUs with different resources exhibits large latency gaps, indicating the need for LLMs with adaptable structures to adapt to different devices when aiming to ensure comparable latency to satisfy user needs; (2) the preference regarding LLM shapes differs across deployment flows. Specifically, under the same model size, reducing model depth and width have comparable impact on the measured latency using TensorRT-LLM on A5000 and MLC-LLM on Orin NX, while reducing depth using PyTorch or MLC-LLM on A5000 can achieve notably lower latency than reducing width. This implies that for emerging platforms with limited compatible deployment flows, proper customization of LLM shapes for maximized efficiency is needed.
The Proposed AmoebaLLM Framework
3.1 AmoebaLLM: Methodology Overview
To address the challenges associated with traditional one-for-all network training, as mentioned in Sec. I, our AmoebaLLM is equipped with three components: a knowledge-preservinng subnet selection strategy, an SMoL adapter, and in-place distillation fine-tuning objective with loss- magnituee balancing. We illustrate our overall framework in Fig. 2: given a target LLM, our AmoebaLLM endows it with the capability of instantly deriving capable subnets via two-stage process.
In the first stage, L..ebal.LM generates the subnet selection strategy. Specifically, given the target depth/width remaining ratios, this step decides which layers/neurons to maintain, respectively. To maximally preserve the knowledge and language modeling capabilities of pre-trained LLMs, we propose employing dynamic programming [22] to determine the retained layers under different remaining ratios and leverage neuron importance metrics [8] to retain important neurons in structured manner, as detailed in Sec. 3.2. After this stage, the subnet selection strategy is determined and fixed.
In the second stage, we insert our proposed SMoL adapter into the target LLM for one-time, one-for-all fine-tuning. Specifically. SMoL is composed of a set of LoRAs [18] and employs a gating function to sparsely activate different subsets of LoRAs for different subnets, thus mitigating their gradient conflicts. This process is elaborated upon in Sec. 3.3. In addition, our fine-tuning objective enhances the sandwich sampling and in- in-place distillation methods described in [11, 13, 14) by adding a loss- magnitude balancing scheme. This prevents bias towards specific subnets in pre-trained LLMs, as detailed in Sec. 3.4. At deployment time, the selected LoRAs, which are determined by the extracted subnet shape favorable to the target platform, can be merged into the LLM weights.
3.2 AmoebaLLM: The Proposed Knowledge-Preservine Subnet Selection Strategy
Motivation. As detailed in Sec. 5, previous one-for-all training techniques [11, 12, 13, 14, 15] often select the first layers of model or the first channels of a layer. These techniques, intended for models trained from scratch, are unsuitable for pre-trained LLMs with rich knowledge encoded in their weights. Considering the difficulty of recovering lost knowledge through fine-tuning or our one-for-all fine-tuning relatively small corpus [23]. it is crucial to identify informative and critical
3

---

Stage 1: Knowledg--Preeevving Subnet Selection
Stage 2: One-for-all Fine-tuning
Depth Shrinking
Insplace
Magnitude Balancing
4
B Bandwich Sampling
SMol Adapters
Locaoo Offfrent Subnets
Selnciinn Stratgy
I
DP Table
Importance-de Wisth Shriking
Balnnodd
Figure 2: An overview of our AmoebaLLM framework: (a) Stage 1: Generate the subnet selection strategy: (b) Stage 2: One-for-all fine-tuning. Zoom in for better view.
layers/neurons that store useful knowledge during the subnet selection process instead of relying solely on fine-tuning. To this end, we propose our snowledge-preservee subnet selection strategy to select the most informative layers/neurons under given remaining ratio, detailed below.
DP-based depth shrinking. Previous works have made diverse observations regarding the layer locations that store knowledge in different series of language models [24, 25, 26, 27, 28, 29]. As such, it is highly desirable to have a principled strategy to evaluate the joint contributions different layer combinations in target LLM to derive the optimal layer selection strategy for each remaining ratio. To achieve this goal, we propose a DP-based depth shrinking strategy.
Problem formulation. Given a target LLM with N decoder layers, we define the layer selection indicates that the n-th layer is retained; otherwise, the layer is removed. The objective is to determine the selection strategy s that achieves the optimal target metric, such as maximal accuracy or minimal perplexity (PPL), on a calibration dataset C, subject to the constraint that M layers are removed.
strategy by vector E 10. Here, s s(n]
Key hypothesis. Thanks to the residual structure [30] of common LLMs [1, 2. 3] and the observations that LLMs' knowledge is compositional across layers [24, 31, 321, we hypothesize that the layer selection problem described above can be divided into smaller and approximately independent sub- problems. Consequently, we can employ dynamic programming [22] to effectively and efficiently solve the layer selection problem.
Our DP-based methodology. We define DP table [[]][m] (where n E N1 and m E [1, MJ, which stores the best target metric on the calibration dataset when exactly m layers are removed from the first layers of the target LLM. The corresponding layer selection strategy is denoted as S[n]m]] E {0, 1] Consequently, [[N][M] represents the final strategy derived for removing M layers out of all N layers. Next, we elaborate on how to obtain D[n] [m] and [[n][m]] where we assume the target metric such that larger values are better, without losing generality.
As illustrated in Fig. (a), similar to general DP problems [22], D D(ni][m can be derived by a recurrence relationship. Specifically, derive each D m the following two cases: (1) removing m layers from the first n layers, and (2) removing m layers from the first Tayers and removing the n-th layer. The strategy yielding better metrics is adopted. More formally, this process can be formulated as follows:
[m], []]m we compare the metrics achieved by
D n m]] max (Dín 1[[m], P(n, m))
where (nn, m) is the metric obtained by removing the n-th layer on top of the best-known strategy S In 1[ 1] for removing m layers from the first n layers. This is computed as follows:
P(n, m) = Ceculuate(remvveeee
1][m 1), C)
where remove(s, r)) is function that sets the n--hh layer to in a strategy 8. Leveraging this recurrence relationship. after initializing the DP table with the base cases, i.e., D (Vi E [1, N),, since no layer is removed, the full DP table can established with complexity of (MMN). In practice, set M as the maximum number of layers allowed to be removed in our
and [ = (1)

---

design space. Therefore, after obtaining the corresponding DP table D, the layer configuration of any subnet with m layers removed (Am [1, M)) can be directly obtained from D [[M]mm] Note that constructing the DP table is one-time effort for given L.M..
Differences from and connections with previous methods. The most relevant direction is layer pruning for LLMs. Pioneering works along this direction either utilize single-layer importance to determine which layers to prune [9], failing measure the layers' joint contribution and thus aggressively losing factual knowledge as shown in Sec. 4, or rely on pre-defined rules [33]. i.e., removing consecutive layers among the last ones, which is suboptimal and may not be generalizable to future LLMs. In contrast, our DP-based strategy can measure the joint impact of different layer combinations and is principled and fully automated without relying on human priors for specific LLMs. More importantly. the two aforementioned works [9, 33] are subsets of our DP-based strategy's 's solution space. Thus, it can serve as a new SOTA layer pruning method, as demonstrated in Sec. 4.3, in addition to serving as a component (i.e., the subnet selection strategy) in one-for-all fine-tuning. Previous works [34, 35] have also applied dynamic programming for pruning at different granularities with varied formulations, and our work is the first to introduce this classical approach to LLMs.
Importance-driven width shrinking. Compared layer selection, removing neurons (and all corresponding weight connections) for width shrinking is more fine-grained and involves much larger design space. Therefore, instead of using dynamic programming, we directly employ existing structured pruning metrics for LLMs to obtain importance scores for each neuron and select the most important ones during one-for-all training, given width remaining ratio.
To achieve this, we employ the metric in [8] due to its outstanding performance compared to other alternatives. Specifically, adopt the same width remaining ratio for linear layers within the same block (either the self-attention block or the feed-forward network), where the importance scores of input neurons of the last linear layer in each block are used to determine the width configuration of this block. In other words, when subset of input neurons to the last linear layer is removed, all associated neurons and weights [7] within this block will be removed in structured manner.
Importance metric. The importance score F of the i-th input neuron in the last linear layer W' of the f-th block computed as F A,
(X x))2 I where j is the index of the output
neurons, X a is the input features of the 1-th token in the k-th batch, and X; is the averaged input features across these two dimensions, both received by the i-th input neuron. To better maintain the LLMs' capability under a given remaining ratio: (1) the importance score is further normalized over all input neurons in each layer and globally sorted for non-uniform width shrinking; (2) pre-computed bias term B is added to the output neurons to compensate for the removed input neurons, i.e., B ('1 M X'), where M is the binary mask indicating whether the input neurons are retained. We refer the readers to [8] for more details.
3.3 AmoebaLL.M: The Proposed Shape-aware Mixture of LoRAs
Motivation. As demonstrated in Sec. 4.3, joint weight fine-tuning of different subnets on small- scale datasets can lead to severe gradient conflicts [16, 17], resulting in poor performance across all subnets. One potential solution is to adopt parameter-efficient adapters like LoRA [18]. However, accumulating gradients from all subnets onto the same LoRA still suffers from gradient conflicts, making fine-tuning unstable and causing some subnets to underperform. On the other hand, tuning a separate LoRA for each subnet configuration is infeasible due the large design space. To address this, we propose an intermediate solution to balance performance and efficiency: the SMoL. adapter that features a set of LoRAs, which are sparsely activated based on the subnet shape.
SMoL adapter design. Our SMoL adapter consists of a set of T LoRAs (AW, = B,A,)
which are sparsely activated for each subnet shape using gating function G. Specifically, we employ a one-hot mask M to indicate the shape, i.e., the layer/width configuration, of the subnet. This mask is fed into the gating function G to calculate score for each LoRA. Only the top k LoRAs are activated and weightedly averaged for each subnet shape during fine-tuning, thus mitigating the gradient conflicts among different subnets.
Implementation. We extend the noisy top-K gating mechanism from [36] to implement our SMoL.. Specifically, G(M) = Softmax Softmax(KeepToopp (KeepTopK (H(M), k)), where KeepTopK (v, k), is is if v is in the
5

---

top k of v and -oo otherwise, and :(x); = (x- W Wg) + std Softplus(xx W Woiise):). Here, Wg and W noise are learnable, with the latter used to control the noise std for load balancing [36]. The final
composited weight W base + E == G(M), A W, where W base is the pre-trained weight.
Key difference from previous works. Previous Mexxture-of-Lofe designs [37, 38, 39] are input- adaptive at inference time, and thus the weights of different LoRAs cannot be merged into the original model. In contrast, our SMoL depends only on the subnet shape and is independent of model inputs. This implies that once the target subnet is extracted based on the target platform, all activated LoRAs can be merged into the model weights, thus enhancing parameter efficiency at deployment time.
3.4 AmoebalLLM: The Fine-tuning Objective
Motivation. Strategically sampling and jointly fine-tuning different subnets is necessary to ensure segh-pefformance subnets within the same LLM, where the sandwich sampling and in-place distilla- tion mechanisms [11. 13, 14] can serve as promising sampling and fine-tuning schemes, respectively, However, naively doing so can cause the larger subnets to gradually underperform during fine-tuning. We identify that this is due to notable differences in the loss magnitudes of different subnets, with smaller subnets, which diverge more from the well pre-trained full model, exhibiting much higher losses than larger subnets. To address this, we propose simple but effective solution: equipping the in-place distillation with loss-magnitude balancing mechanism, which we elaborate on as follows.
In-place distillation with loss-magnitude balancing. During each fine-tuning iteration, we employ the sandwich sampling [11, 13, 14] to sample K subnets {T:)== with different layhr/width remaining ratios, including the largest/smallest ones and K - 2 random ones from our design space. Detailed layer/width configurations of sampled subsets can be obtained from the strategies derived in Sec. 3.2. We fine-tune our SMoL adapter as detailed in Sec. 3.3 by accumulating the gradients from all sampled subnets using in- in-place distillation, where only the loss of the largest subnet Ti is calculated using ground truth, while those of other subnets [T use distillation from the largest one [11]. To balance the loss magnitude from different subnets, we normalize all subnets' loss magnitudes that of the largest subnet, as visualized on the rightmost side of Fig. 2 (b). In this way, the final loss direction is jointly determined by all subnets' loss directions, falling on a unisphere without being severely impacted by their unbalanced magnitudes. We formulate the fine-tuning objective as follows:
K Σ [ vl)l (T<))
Luasl = CCI((TTTT)))))))) +
(CE(T(()), (z))
where r and y are the input and ground truth, respectively, and LCE is cross-entropp loss function. Final subnet search after fine-tuning. We adopt simple hierarchical search strategy select sub- nets from the fine-tuned LLM to satisfy the target efficiency constraint while maximizing achievable accuracy. Specifically,
first perform coarse grid search across uniformly spaced depth and width settings based on a small calibration set (e.g., 40 samples from the MMLU dataset) to identify subnets that meet the efficiency constraint with maximized accuracy. Next, conduct a more fine-grained grid search within depth and width ranges surrounding the optimal subnet identified in the coarse grid search stage. More advanced search strategies, such as evolutionary search [11, 12, 13, 14), left for future work.
4 Experimental Results
4.1 Experiment Setup
Baselines. Our baselines include two SOTA structured width pruning methods: LLM-Pruner [7] and FLAP [8], and one layer pruning method Shortened LLaMA [9]. All these baselines are open-sourced and we apply their official code to different LLMs. All baselines, including FLAP which were not fine-tuned in their original paper, are fine-tuned using the settings below for fair comparison.
Fine-tuning setting. Following [7, 9], we adopt 50K samples from Alpaca [40] for our one-for-all fine-tuning as well as for fine-tuning all baselines. For both our method and the baselines, we adopt constant learning rate of 2e-4 with an Adam optimizer and a LoRA rank of 64, and fine-tune for 10K iterations. It takes 40 GPU hours an NVIDIA A5000 GPU for our one-for-all ll fine-tuning.

---

Table 1: Compare with baseline methods under varying remaining ratios on LLaMA2 7B.
Ratio
Method &
MMLU1
Average 1
BoolQ
PIQA
HellaSwag
WinoGrande
ARC-e 1
ARC-c I
OBQA i
LLM.Prumet [7]
29.63
56.95
58.53
76.39
65.80
60.38
64.60
34.56
38.40
FLAP
40.21
60.98
65.43
80%
Shortened LLaMA 19]
26.45
58.72
62.17
76.01
68.22
58.88
68.98
38.10
38,10
Amoebal (Ours)
10.70
62.29
72.70
76.80
70.60
67.60
68.30
38,80
41.20
AmoobaL.LM (Ours)
72.50
76.30
70,80
70.30
40.20
39.60
LLM-Prumer [7]
23.15
54,09
60.73
74.97
58.57
57.85
55.68
33.02
37,80
FLAP [8)
33.87
65%
Shoriened LLaMA [9]
24.89
52.57
62.32
72.03
55.10
52:41
59.47
30.63
36.00
Amocbal (Ours)
36.00
$6.96
72.10
70.70
$9.70
63.20
62.20
34,00
36,80
Amochal.LMM (Ours)
57.26
70.50
70.90
61.50
63.50
34.50
37,20
LLM-Pruner [7]
22.90
47.52
61.83
67.79
43.31
$1.22
46.13
28.16
34,20
FLAP [8]
27.67
59.45
51.33
56.75
36,00
Shortened LMA [9]
24.76
47.35
62.23
66.00
43.60
51.54
50.63
26.45
31.00
Amochal J (Chur))
30.60
$2.19
65.70
66. 10
51.30
60:10
56,60
31.50
34.00
Amochal.LM" (Ours)
32
$2.63
64.70
66.70
53,00
60.30
58.00
30.10
35,60
Table 2: Compare with baseline methods under varying remaining ratios on Vicuna 7B v1.5.
Ratio
Method
MMLU1
Average
BoolQ
PIQA
HellaSwag | WinoGrande
ARC-
ARC-c 1
OBQA 1
LLM-Pruner
64.28
64.35
35.07
39,40
FLAP
13.50
60.54
72/45
73:67
66.59
63.98
68.14
37.51
39,10
809
Shortened LLaMA
35.27
60.93
67.58
75.68
68.12
64.48
70.20
40.44
20,00
A. (Ours)
47.40
60.77
71.30
72.70
68.80
66.10
66.50
38.60
41,40
Amocbal I.M (Curr)
48.30
73.10
73.33
68.
69:30
40:20
L.M.-Prnner
24.07
$4.24
61.19
73.50
$7.67
$7.85
58.67
32.42
38.40
FLAP
56.47
70.78
58.97
63.43
34.21
37.60
6.5%
Shortened aMA
25.59
$2,50
64.28
70.62
56.78
$7.46
52.74
31.40
34.20
Amochal (Ours)
40:30
55.93
68.20
69. 69.80
58
62.40
62.20
34.10
36. 36.20
ArsocaoLMM (Ours)
71,50
35.60
LLM-Pruner
23.24
47,98
59.08
68.55
44.24
$2.17
49.03
28.41
34,40
FLAP
67.52
$1.81
$7.69
56.57
31.06
34,40
50%
Shortened LLaMA
25.03
45.56
51.71
65.67
43:28
51.38
09.96
26.51
30.10
M (Ours)
34.00
51.11
61.20
65.00
50.90
59.10
56.20
30.90
33.60
MrxxbalLLM (Ours)
64.70
35.40
Models. We apply our method to LLaMA2 7B [1] and Vicuna 7B v1.5 [41].
Evaluation. Following our baselines [7, 9, 8, 6], we leverage Im-evaluation- [42] to measure
the zero-shot accuracy on 7 commonsense reasoning datasets. including BoolQ [43], PIQA [44]. HellaSwag [45], WinoGrande [46], eaCy [47], ARC-challenge [47], and OpenbookQA [48].
We also benchmark on the factual knowledge dataset MMLU [49].
AmoebaLLM'ss setting. We set the depth choices to range from 20 32 and the width remaining ratios as (1, 7/8, 3/4, 5/8, 1/2) by default unless specifically stated. In each iteration, we sample subnets for joint training and select LoRAs out of total of LoRAs in SMoL for each subnet.
4.2 Benchmark with SOTA LLM Compression Methods
We benchmark our Amoebal.LM against SOTA LLM width/layer pruning methods on LLaMA2 7B/Vicuna 7B v1.5 in Tab. 1/Tab. 2, respectively. Note that the subnets produced by Amoebal.LM are instantly extracted from the same one-for-all fine-tuned LLM, where the (depth, width scale) settings for 50%//65%/505 remaining ratios, determined by the final subnet search in Sec. 3.4, are (30, 0,875)(28, 0,7 0.75), respectively. We also provide the per-subnet fine- tunettned counterparts of our delivered subnets, denoted as AmoebaLLM', to compare one-for-all and individual fine-tuning.
Benchmark under comparable model sizes. As shown in Tab. 1 and Tab. 2, we observe that (1) the subnets instantly extracted by AmoebaLLM can achieve higher MMLU accuracy compared to all baselines, suggesting that our method better preserves the factual knowledge acquired during pre- training, as further analyzed in Sec. 4.3 and Sec. 4.4; (2) AmocbaLLM's delivered subnets, extracted from the same model, can also achieve better or comparable average commonsense reasoning accuracy compared the strongest baselines, each trained separately; (3) Amoebal.MM achieves the best
performance across all metrics and tasks compared to the baselines, indicating that Amoebal.L.M can serve as new SOTA L.M compression framework in addition to its one-for-all functionality, thus advancing the achievable accuracy-efficiency trade-off: (4) compared to our per-subnet fine- tuned variant AmoebaLLM the instantly delivered subnets achieve comparable performance. This demonstrates the effectiveness of our one-for-all fine-tuning scheme, as further ablated in Sec. 4.3.
7

---

Table 3: Ablation Study on the effectiveness of the DP-based depth shrinking on LLaMA2 7B.
Calib, Data
Method
24
23
22
21
20
17
Unreasonable [33]
12.25
13.72
18.61
31.35
44.28
57.34
80.29
120.57
188.27
Wikitex12
ShortenI.LaMA [9]
11.75
13.18
17.83
29.61
39.42
53.27
71.80
106.63
154.69
Ours
11.65
12.77
17.59
20.06
23.77
28.83
38.70
70.87
95.16
Unreusonable [33]
40.8
39.5
37.5
33.5
34.5
34.0
30.3
30.6
25.8
MMLU (%)
LaMM [9]
42.0
34.7
35.1
32.1
32.0
33.5
33.7
29.3
26.6
Ours
46.2
44.8
44.6
44.1
41.2
41.3
43.1
34.7
28.9
Table 4: Ablation Study on different components in our AmoebaLLM on LLaMA2 7B.
32
20
Method
Wikitext2
MMLU 1
Wikitext2
MMLU (%)
Wikitext2
MMLU
Per-subnet ft.
5.54
46.4
10.57
41.9
15.94
41.7
SMoL (+full model)
5.82
46.6
38.48
32.6
167.74
36.4
-MML (+LoRA)
6.97
40.6
12.71
40.0
19.12
37,9
Loss-mag. Balancing
6.77
42.0
12.63
40.1
18.19
39.3
Full (AmochaLLLM)
6.36
47.2
12.40
45.1
18.15
41.0
Benchmark accuracy-latency trade-offs real devices. further benchmark the achieved trade-off between average common- sense reasoning accuracy and measured latency of LLaMA2 7B using MLC-LLM and Py PyTorch as the deployment flows on an NVIDIA A5000
PyTorch @ A5000
MLC-LLM @ A5000
67.5
AmsrbalLH
AnnrcaLLLM Shartanad LLaNt FLiP
Snarisnad LLaV6
PLAZ
Accuracy (%)
60.0
57.5
57
55.0
35.0
GPU, following the settings in Sec. 2. For our method, we select the subnet shape that favors the hardware characteristics based the pro- filing in Sec. 2 from the one-for-all fine-tuned LLM. As shown in Fig. 3, we observe that (1) our method consistently achieves the best trade- off on both deployment scenarios; and (2) al- though FLAP achieves higher accuracy than Shortened LLaMA under comparable model sizes, its real-device speed is limited when using vanilla PyTorch, aligning with observations in [9]. In contrast, our method can instantly deliver subnets that favor the bardware/deploymenn flow characteristics, thus enjoying both high accuracy and real-device friendliness.
50.0
2.5
3.0
3.5
2.0
2.6
Latency
Figure 3: Benchmark AmoebaL.MM's achieved accuracy-latency trade-offs with SOTA LLM com- pression methods on an NVIDIA A5000 GPU.
4.3 Ablation Study: Effectiveness of Each Component
We perform ablation studies to validate the effectiveness of each component AmocbaLLM.
The DP-based depth shrinking strategy. We benchmark our DP-based strategy against two existing LLM layer pruning methods, Shortened LLaMA [9] and Unreasonable [33], on LLaMA2 7B. Specifically, we employ the three methods to select important layers using Wikitext2/MMLU with PPL/accuracy as calibration metrics, respectively. We directly report the achieved PPL/MMLU accuracy after calibration under various layer remaining ratios without fine-tuning to indicate their effectiveness in identifying important layers
Observations and analysis. As shown in Tab. 3, we observe that our DP-based strategy outperforms the other two strategies both calibration datasets and metries, especially under small remaining ratios, c.g., a +9.4% MMLU accuracy and -33.1 PPL over the strongest baseline when remaining 18 layers. This demonstrates the superiority of method over the two baselines in selecting important layers that optimize the target calibration metric, thus significantly contributing to knowledge preservation.
Remark. This set of experiments supports analysis in Sec. 3.2 that the superiority of our method arises from its consideration of different layers' joint contributions, rather than focusing on single- layer importance [9], and its avoidance of reliance pre-defined rules [33], thus ensuring generality.
The SMoL. adapter. assess the efficacy our SMoL. adapter, we substitute with full model fine-tuning or the standard LoRA [18] and benchmark it against our AmoebaLLM with SMoL as
8

---

Table 5: Ablation Study on the selection of calibration datasets on LLaMA2 7B.
32
24
20
Calib, Data
Wikitext2 1
MMLU (%)
Wikitext2 +
MMLU (%)
Wikitext2
MMLU (%)
BonkCorpus
5.64
46.4
16.52
26.8
31.70
27.0
Wikitcxt2
5.64
46.4
11.92
26.4
23.77
24.4
MMLU
5.64
46.4
293.01
45.8
1338.68
41.2
BookCorpus (ft)
6.55
45.4
10.55
29.7
14.36
25.7
Wikitext2 (ft)
6.41
43.4
9.58
32.6
13.26
23.6
MMLU (ft)
6.36
47.2
12.40
45.1
18.15
41.0
well as the per-subnet fine-tuning variant. To more clearly demonstrate its efficacy, this experiment conducted under a depth-shrinking-onk setting, which only enables depth shrinking ranging from 20 to 32 layers during one-for-all fine-tuning. We select three layer configurations, covering both the largest and smallest ones, to report the performance.
Observations and analysis. As illustrated in Tab. 4, we observe that (1) full model fine-tuning on relatively small corpus results in suboptimal performance with large PPL; (2) employing the standard LoRA stabilizes the one-for-all fine-tuning compared to full model fine-tuning. albeit with notable performance reductions in larger subnets, e.g., a 5.8% MMLUU accuracy drop in the largest subnet compared to per-subnet fine-tuning: (3) when equipped with our SMoL., the MMLU accuracy of all subnets is substantially enhanced, e.g., +6.6% improvement on the largest subnet compared to the LoRA case, even exceeding that of the per-subnet counterparts. This demonstrates the essential capability of our AmoebaLLM to deliver high-quality LLM subnets across a wide range of accuracy-efficiecy trade-offs.
The loss- maggituge balancing scheme. As shown in Tab. 4, we further disable the loss mossinag balancing scheme and observe significant MMLU accuracy drop in larger subnets. This supports our analysis in Sec. 3.4 that the larger losses from smaller subnets may dominate the overall objective, thereby impairing the fine-tuning of larger subnets and underscoring the necessity of our method.
4.4 Ablation Study: The Selection of Calibration Datasets
We conduct an ablation study on the choice of calibration datasets for our DP-based depth shrinking introduced in Sec. 3.2. Specifically, we accuracy on the training of MMLU [49] and PPL on the training sets Wikitext2 (50)/BookCorpus [51] as target metrics. We report the evaluation metrics, including MMLU test accuracy and Wikitext2 test PPL, under different layer remaining ratios both after calibration and after one-for-all fine-tuning under a depth-shrinking-onie setting,
Observations. As shown in Tab. 5, observe that (1) after calibration without fine-tuning the subnets perform well on the evaluation metric for which they were calibrated and underperform in terms of the other metric; (2) after fine-tuning, the subnets calibrated using PPL continue to perform poorly on MMLU accuracy, indicating severe loss of factual knowledge. In contrast, the subnets calibrated using MMLU accuracy achieve notably lower PPI. compared before fine-tuning, on par with the subnets calibrated using PPI., while still maintaining high MMLU accuracy.
The key insight. This set of experiments indicates that the loss of factual knowledge during com- pression is hard to restore during fine-tuning, echoing the observations [23], while the language modeling capability is easier to recover through fine-tuning, As such, we adopted MMLU as the calibration dataset throughout the previous experiments, and we believe this insight could inspire future LLM compression frameworks and calibration metrics.
4.5 Limitations and Future Work
One limitation of our work that due to the limited tnn-itning data and resources, our method- ology is applied to parameter-effiiient fine-tuning which mitigates gradient conflicts under small data conditions while limiting the achievable accuracy-efficient trade-off. We anticipate that by leveraging more extensive fine-tuning data beyond our current use of Alpaca [40] and extending our design insights regarding subnet selection and gradient conflict mitigation, more aggressive accuracy-efficiency trade-off's can be achieved, which will be the focus of our future work.

---

5
Related Work
Large language models. Before the advent of LLMs, transformer- language models [52, 53, 54, 55] demonstrated their ability to effectively analyze relationships among tokens in complex input sequences, facilitated by the attention mechanism [52]. These models also exhibit notable scalability [56, 57, 58] with respect to model size and the scale of pre-training datasets. This decent scalability has led to the emergence of LLMs, such as GLM [59], OPT [60], BLOOM [61]. the Llama family [1. 1, 2]. Gemma [3], and GPT-4 [4], which exhibit impressive zero-shot and few-shot in-context learning capabilities. However, these LLMs often feature billions of parameters and prohibitive computation complexity, which limits their widespread use across diverse platforms.
Large language model compression. To facilitate efficient deployment of L.Mss in real-world applications, existing works primarily focus on compressing LLMs by extending traditional compres- sion techniques such as knowledge distillation [62. 63], quantization [64, 65, 66, 67, 68. 69], system acceleration [70, 71J. and pruning [6. 5, 72]. Our work is most closely related to LLM pruning. Along this direction, early works [5, 6] employ unstructured and semi-structured pruning [73] by zeroing out connections among neurons. Despite their plausible performance, these methods require specialized support to achieve real-device speedup. To benefit commodity platforms, structured LLM pruning methods remove more coarse- goasse-raiied components, such as all connections related to single neuron 17, 8, 10] or even entire layers [9. 33). For instance, LLM--runer [7] and FLAP [8] reduce LLM width by eliminating identified redundant neurons, while Sheared-LLLMM [10] learns set of binary masks to reduce both the width and depth of LLMs. However, these methods either focus on single dimension of compression (i.e., depth or width) with limited efficiency improvements, or they require a costly fine-tuning process for each target configuration and platform. In contrast, our AmoebaL can instantly extract subnets of arbitrary shapes that reach the accuracy-efficiency frontier, thus facilitating rapid deployment across devices.
One-for-all networks. Slimmable networks [74, 75] are pioneering works that enable single model to operate at varying widths. Follow-up works [11, 76, 12, 13, 14, 15] further extend this approach to train more general one-for-all networks with switchable depth and width, thus enabling tasks like neural architecture search. In particular, BigNAS [11] builds one-for-all networks using a sandwich sampling strategy that samples random set of subnets and jointly trains them in each iteration through in-place distillation, where the largest model guides the learning of the smaller ones. This approach has been inherited by subsequent works [12, 13, 14]. Additionally, this idea has been extended to any-precision networks [77. 78, 79] that allow switchable precision at runtime. A very recent work [80] has further extended this concept to any-precision LLMs.
Nevertheless, directly applying these methods to the depth and width of pre-trained LLMs would likely fail because their subnet sampling strategies often select the first layers of model or the first channels of a layer, which are intended for models trained from scratch. This approach is unsuitable for pre-trained LLMs, it may omit layers neurons containing crucial knowledge. Our AmoebaLLMM framework addresses these challenges by developing three key components: the subnet selection strategy, the trainable adapter design, and the fine-tuning objective. One concurrent work [81] also aims to train many-in-one LLMs that support instant subnet derivation, targeting full-model continual training setting on 90 billion tokens. In contrast, our method targets parameter-efficient tuning setting with only 8.5 million fine-tuning tokens.
6 Conclusion
In this work, we present a framework called Amoebal LLM that grants a given L.LM the capability frontier and
to instantly deliver subnets of arbitrary shapes, which achieve the accuracy-efficiene can be extracted immediately after a one-time fine-tuning. This is achieved by the development of three dedicated components that enable AmoebaLLM's one-for-all fine-tuning scheme: knowledge- preserving subnet selection strategy, an SMoL adapter, and an in-place distillation objective with boss-mannttude balancing. Extensive experiments validate that our LM framework can deliver efficient LLMs with instantly serviceable subnets of any shape, which outperform SOTA LLM compression techniques in terms of the sccuracy-efficienne trade-off. We believe this work is promising to facilitate the wider use of existing and emerging public LLMs by making them instantly deployable on varying platforms and applications, and by providing new perspective on efficient LLM deployment, thus inspiring future solutions.
10

---

Acknowledgement
The work is supported by an National Science Foundation (NSF) CAREER award (Award number: 2345577) and CoCoSys, one of the seven centers in JUMP a Semiconductor Research Corporation (SRC) program sponsored by DARPA.
References
[1] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.0928. 2023.
[2] Meta AL. Introducing Meta Llama 3: The most capable openly available LLM to date. 2024.
[3] Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, al. Gemma 2: Improving open language models practical size. arXiv preprint arXiv:2408.00 000:24088001166 18, 2024.
[4] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint 2rXiv:2303.08774, 2023.
[5] Elias Frantar and Dan Alistarh. SparseGPT: Massive language models can be accurately pruned in one-shot, 2023.
[6] Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. A simple and effective pruning approach for large language models, 2023.
[7] Xinyin Ma, Gongfan Fang, and Xinchao Wang. language models. Advances in neural information processing systems, 36:21702-21720, 2023.
On the structural pruning of large
[8] Yongqi An, Xu Zhao, Tao Yu, Ming Tang, and Jinqiao Wang. Fluctuation- -baeed adaptive structured pruning for large language models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 10865-10873, 2024.
[9] Bo-Kyeong Kim, Geonmin Kim, Tae-Ho Kim, Thibault Castells, Shinkook Choi, Junho Shin, and Kyoung-Kyu Song. Shortened laama: A simple depth pruning for large language models. arXiv preprint arXiv:2402.00 02834, 2024.
[10] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared Ilama: Accelerating language model pre-training via structured pruning. arXiv preprint arXiv:2310.06694 2023.
[11] Jiahui Yu, Pengehong Jin, Hanxiao Liu, Gabriel Bender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song, Ruoming Pang. and Quoc Le. Bignas: Scaling up neural architecture search with big single-stage models. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK. August 23-28, 2020, Proceedings, Part VII 16, pages 702-717. Springer, 2020.
[12) Dilin Wang, Meng Li, Chengyue Gong, and Vikas Chandra. Attentivenas: Improving neural architecture search via attentive sampling. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 6418 6427, 2021.
[13] Dilin Wang, Chengyue Gong. Meng Li, Qiang Liu, and Vikas Chandra. Alphanet: Improved training of supernets with alpha-divergence. In International Conference on Machine Learning, pages 10760-10771. PMLR, 2021.
[14] Chengyue Gong and Dilin Wang. Nasvit: Neural architecture search for efficient vision transformers with gradient conflict-aware supernet training. ICLR Proceedings 2022, 2022.
[15] Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once-for-all: Train one network and specialize it for efficient deployment. arXiv preprint 1907999908009991,220

---

[16] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 2020.
[17] Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Conflict-averse gradient descent for multi-tast halti-taskk learning. Advances in Neural Information Processing Systems, 11118888-18990 2021.
[18] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.0900 09688, 2021.
[19] NVIDIA. Tensorrt-llm, 2024.
[20] ML.C team. MI.C-LLM, 2023.
[21] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[22] Richard Bellman. The theory of dynamic programming. Bulletin of the American Mathematical Society, 5036)5503-515, 1954.
[23] Zorik Gekhman, Gal Yona, Roee Aharoni, Matan Eyal, Amir Feder, Roi Reichart, and Jonathan Herzig. Does fine-tuning llms on new knowledge encourage hallucinations? arXiv preprint 24ii::2005059004 2024.
[24] Ningyu Zhang, Yunzhi Yao, Bozhong Tian, Peng Wang, Shumin Deng, Mengru Wang. Zekun Xi, Shengyu Mao, Jintian Zhang, Yuansheng Ni, et al. A comprehensive study of knowledge editing for large language models. arXiv preprint arXivv2401.01226, 2024.
(25) Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. arXiv preprint arXiv:2104.08690 2021.
[26] Mor Geva. Roci Schuster, Jonathan Berant, and Omer Levy. Transformer -eew-forward layers are key-value memories. arXiv preprint arXiv:2012.. 14913, 2020.
[27] Ganesh Jawahar, Benoît Sagot, and Djamé Seddah. What does bert learn about the structure of language? In ACL 2019-57th Annual Meeting of the Association for Computational Linguistics, 2019.
[28] Kevin Meng. David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35:17359-17372 2022.
[29] Tianjie Ju, Weiwei Sun, Wei Du, Xinwei Yuan, Zhaochun Ren, and Gongshen Liu. How large language models encode context knowledge? layer-wise probing study. arXiv preprint arXiv:2402. 16061, 2024.
[30] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[31] Yunzhi Yao, Peng Wang. Bozhong Tian, Siyuan Cheng, Zhoubo Li, Shumin Deng, Huajun Chen, and Ningyu Zhang. Editing large language models: Problems, methods, and opportunities. arXiv preprint arXiv:2305. 13172, 2023.
[32] Peng Wang. Ningyu Zhang, Xin Xie, Yunzhi Yao, Bozhong Tian, Mengru Wang, Zekun Xi, Siyuan Cheng, Kangwei Liu, Guozhou Zheng, et al. Easyedit: An easy-to-use knowledge editing framework for large language models. arXiv preprint arXiv:2308.00 07266, 2023.
[33] Andrey Gromoy, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A Roberts. The unreasonable ineffectiveness of the deeper layers. arXiv preprint 1rXiv22403.17887, 2024.
12

---

[34] Kwangbae Lee, Hoseung Kim, Hayun Lee, and Dongkun Shin. Flexible group-level pruning of deep neural networks for on-device machine learning. In 2020 Design, Automation & Test in Europe Conference & Exhibition (DATE), pages 79-84. IEEE, 2020.
[35] Kaixin Xu, Zhe Wang, Xue Geng, Min Wu, Xiaoli Li, and Weisi Lin. Efficient joint optimization of layer-adaptive weight pruning in deep neural networks. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 17447-17457, 2023.
[36] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprini arXiv: 170Nvv:700.0000 77770065065 2017.
[37] Xun Wu, Shaohan Huang, and Furu Wei. Mole: Mixture of lora experts. In The Twelfth International Conference on Learning Representations, 2023.
[38] Shaoxiang Chen, Zequn Jie, and Lin Ma. Llava-mole: Sparse mixture of lora experts for mitigating data conflicts in instruction finetuning milmss. arXiv preprint arXiv:2401.1616. 2024.
[39] Dengchun Li, Yingzi Ma, Naizheng Wang, Zhiyuan Cheng, Lei Duan, Jie Zuo, Cal Yang, and Mingjie Tang. Mixlora: Enhancing large language models fine-tuning with lora based mixture of experts. arXiv preprint arXiv:2404. 15199, 2024.
[40] Rohan Taori, Ishaan Gulrajani. Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, et al. Stanford alpaca: An citructtoo-following llama model. https://gitae coom/tatsu-laa stanford_alpaca, 2023.
[41] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, et al. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, 2023.
[42] EleutherAL Language model evaluation harness (package version caaf9ab). https: / htops:://gtthue /ittbb...... com./ELeutherrA/Im-ovvve www.tmmmeeaaaaannnnaaaaa harness, 2023.
[43] Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In NAACL, 2019.
[44] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In AAAI, 2020.
[45] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In ACL, 2019.
[46] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. arXiv preprint 1 2019.
[47] 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 arXiv:1803.05455 2018.
[48] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? new dataset for open book question answering. In EMNLP, 2018
[49] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 2020.
[50] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In ICLR, 2017.
[51] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pages 19-27, 2015.
13

---

[52] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[53] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 010004 2018.
[54] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21 (1(140): 2020.
[55] Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor. Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorme, Aitor Lewkowycz, Alex Salcianu, Mare van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bulian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee, Dan Garrette, James Lee-Thorp, Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Maitin -Satiir,, Noah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, and Andrea Gesmundo. Scaling up models and data with t5x and seqio. arXiv preprint 2rXiv22003.17188 2022.
[56] Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Fei Yuan, Xiao Luo, et al. Scaling transnormer to 175 billion parameters. arXiv preprint arXiv: 14999, 2023.
[57] Jared Kaplan. Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.088 008611, 2020.
[58] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' Brienn, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397-2430 PMLR, 2023.
[59] Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang, and Jie Tang. Glm: General language model pretraining with autoregressive blank infilling. arXiv preprint arXiv:2103. 10360, 2021.
[60] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint wrXiv:2205.01063 2022.
[61] BigScience Workshop, Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, et al. Bloom: 176b- cabamaareeeer open-access multilingual language model. arXiv preprint arXiv:22/1.05// 1.05/00, 2022.
[62] Yao Fu, Hao Peng, Litu Ou, Ashish Sabharwal, and Tushar Khot. Specializing smaller language models towards multi-step reasoning. In International Conference on Machine Learning, pages 10421-10430. PMLR, 2023.
[63] Cheng- Yu Hsieh, Chun-Liang Li, Chih-Kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alexander Ratner, Ranjay Krishna, Chen-Y Yu Lee, and Tomas Pfister. Distilling step-by-step! outperform- ing larger language models with less training data and smaller model sizes. arXiv preprint arXiv:2305.0230) 2023.
[64] Elias Frantar, Salch Ashkboos, Torsten Hoefler, and Dan Alistarh. Opty: Accurate quantization for generative pre-trained transformers. In The Eleventh International Conference on Learning Representations, 2022.
[65] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. LLM.int80): 8-bit matrix multiplication for transformers at scale. In Advances in Neural Information Processing Systems, 2022.
14

---

[66] Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, 2023.
[67] Elias Frantar, Saleh Ashkboos, Torsten Hoeferr, and Dan Alistarh, GPTQ; Accurate post-training compression for generative pretrained transformers. In International Conference on Learning Representations, 2023.
[68] Tim Dettmers, Ruslan Svirschevski, Vage Egiazarian, Denis Kuznedelev, Elias Frantar, Saleh Ashkboos, Alexander Borzunov, Torsten Hoefler, and Dan Alistarh. Spqr: A sparse-quantized representation for near-lossless I1m weight compression, 2023.
[69] Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-aware weight quantization for Ilm compression and acceleration. arXiv preprint TrXiv:2306600975. 2023.
[70] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXivv2307.08661, 2023.
[71] Woosuk Kwon, Zhuohan Li. Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and lon Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 61 1-626, 2023.
[72] Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models, 2023. version
[73] Aojun Zhou, Yukun Ma, Jumnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. Learning n: m fine-grained structured sparse neural networks from scratch. arXiv preprint (0Xiv:2002.0000) 2021.
[74] Jiahui Yu, Linjie Yang, Ning Xu, Jianchao Yang, and Thomas Huang. Slimmable neural networks. arXiv preprint arXiv:1812.08928,, 2018.
[75] Jiahui Yu and Thomas S Huang. Universally slimmable networks and improved training techniques. In Proceedings of the IEEE/CVÉ international conference on computer vision, pages 1803-1811, 2019..
[76] Jiahui Yu and Thomas Huang. Autoslim: Towards one-shot architecture search for channel numbers. arXiv preprint arXiv: 117228, 2019.
[77] Qing Jin, Linjie Yang, and Zhenyu Liao. Adabits: Neural network quantization with adaptive bit-widths. In Proceedings of the IEFF/CVE Conference on Computer Vision and Pattern Recognition, pages 2146-2156, 2020.
[78] Luis Guerra, Bohan Zhuang, Ian Reid, and Tom Drummond. Switchable precision neural networks. arXiv preprint arXiv:2002.02881 028155, 2020.
[79] Haichao Yu, Haoxiang Li, Humphrey Shi, Thomas S Huang, and Gang Hua. Any-precision deep neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10763-10771, 2021.
[80] Yeonhong Park, Jake Hyun, SangLyul Cho, Bonggeun Sim, and Jae W Lee. Any-precision 11m: Low-cost deployment of multiple, different-sizer 1lms. arXiv preprint 2:Xi:22402.10517. 2024.
[81] 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.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have accurately summarized our paper's contributions and scope in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer:
[Yes]
Justification:
We discussed this in Sec. 5.5 of our paper.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
16

---

Justification:
We didn't 't provide theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, have elaborated on settings and provided the code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
17

---

Answer: [Yes]
Justification:
We have
our code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
Staiic/guudees/Coeeeeeee ceabbbmissoonPooo wooooooooioo olicyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: cc/ www.//pwlccggpuiaeeeeeeeeeeeeeee / /twpps.ccppuu.ll////
Caleaes/doSSSSubiiceee wwwsiisssspopiiigggo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We offered this information in Sec. 5.1 of our paper and also provided sufficient references.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[No]
Justification: We follow our baselines and do not report the error bars as the results are relatively stable.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
18

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provided this information in Sec. 5.1 of our paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that NeurlPS Code of Ethics https httpp://neett /ati:://neurpppee ctpppps..cpppppeeee httiee/eecc.ceee.ccciiiiieieeeee
didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification: Yes, it conforms with NeurlPS Code of Ethnics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work targets the acceleration of large language models to facilitate its widespread use and thus does not suffer from obvious negative societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
19

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We built on top of public models and datasets and do not suffer from high risks of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes, have added the reference for all used data.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
20

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We did not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We did not perform crowdsourcing experiments.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work didn't involve human participants.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
21


Reviews:
[{"review_id": "0euS8Q6Z9x", "reviewer": "Reviewer_TKvU", "summary": "To address the problems of diverse resource constraints and deployment flows while using LLM for multiple real-world applications, this paper proposes an AmoebaLLM, featuring a knowledge-preserving subnet selection strategy, a shape-aware mixture of LoRAs and a distillation scheme with loss-magnitude balancing.", "questions": "Please see the weaknesses part.", "limitations": "Yes.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-organized. \n2. The method achieves SOTA performance on most metrics.\n3. The experimental setup is easy to follow.", "weaknesses": "Overall, the motivation and ideas of this paper are commendable. However, I have a few concerns:\n1. DP-based depth shrinking aims to select suitable layers that can achieve the optimal target metric. Does this mean that the proposed method needs to evaluate the performance for each selection strategy? If so, how long is required for each evaluation? What is the relationship between the number of layers and the overall search time?\n2. Another concern is about the optimal subnet. I wonder if the optimal subnet identified in the first phase is truly “optimal.” I believe that the selected subnet is the optimal choice for the initial parameters. However, the proposed method consists of two stages. After the fine-tuning stage, some subnets that were not optimal in the first phase might achieve the best target metrics. I understand that searching subnets with fine-tuning is very expensive, but have the authors considered this issue? Is this correct?\n3. I am particularly interested in the training efficiency comparisons, as the authors mentioned that existing methods are not efficient. However, the experimental section does not show such comparisons. Could the authors provide relevant comparisons in Table 1 and Table 2?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing the idea and performance of our work, as well as for your constructive comments! We have addressed all your concerns below.  \n\n**1. Whether our DP-based depth shrinking strategy needs to evaluate the performance for each selection strategy and its overhead**\n\nYes, you are right. Our DP-based strategy needs to evaluate the performance for each selection strategy, as described in Eq. (2) of our manuscript. Specifically, as mentioned in Section 5.4 of our manuscript, we measure the MMLU accuracy of 20 samples for each evaluation, which takes about 15 seconds on average on an NVIDIA A5000 GPU for LLaMA2 7B.\n\nThe overall search time is proportional to the total number of layers. Specifically, when targeting the removal of M layers out of all N layers, the total number of evaluations is M(N-1)/2. For LLaMA2 7B in our case, where N=32 and M=16, the overall search time is about 1 hour. It is worth noting that constructing the DP table is a one-time effort for each pre-trained LLM model and does not need to be repeated for each target subnet configuration.\n\nWe will add this analysis to the final version.\n\n---\n**2. Whether the optimal subnets change before and after fine-tuning and the potential of searching subnets with fine-tuning**\n\nThanks for the insightful question! First of all, we agree that under different subnet selection criteria, the selection of optimal subnets may vary before and after fine-tuning, making smarter strategies like iterative subnet selection and one-for-all fine-tuning promising.\n\nSecond, under our subnet selection criteria, i.e., the encoded factual knowledge measured by MMLU accuracy as elaborated in Section 5.4 of our manuscript, we observe a high consistency between the optimal subnets before and after fine-tuning. Specifically, this is because, according to observations in recent works [1][2][3], new factual knowledge can hardly be gained during finetuning, and the loss of factual knowledge during compression is hard to restore during finetuning. As such, the selection of optimal subnets stays consistent when using MMLU accuracy as a subnet selection indicator. This motivates us to adopt a simple-yet-effective two-stage strategy in our AmoebaLLM framework.\n\n[1] “Does fine-tuning llms on new knowledge encourage hallucinations?”, Z.  Gekhman et al., arXiv’24.\n\n[2] “LIMA: Less Is More for Alignment”, C. Zhou et al., NeurIPS’23.\n\n[3] “R-Tuning: Instructing Large Language Models to Say ‘I Don’t Know’”, H. Zhang et al., NAACL’24.\n\n---\n**3. Training efficiency analysis**\n\nThanks for constructive comments and questions! Generally, the advantage of our method in training efficiency is that it requires a constant *O(1)* training time with respect to the number (*N*) of subnets with varying efficiency constraints to be delivered, thanks to the one-for-all finetuning. In contrast, previous compression methods require *O(N)* fine-tuning time.\n\nFor example, for delivering 3 subnets with varying weight remaining ratios in Table 1 of our manuscript, when fine-tuned on the same number of tokens, all baseline methods require about 45 GPU hours in total on an NVIDIA A5000 GPU, while our method requires 30 GPU hours. Similarly, when delivering 10 subnets, all baseline methods require about 150 GPU hours, while our method still requires only 30 GPU hours, thanks to its constant training complexity. \n\nWe will add the above analysis to the final version.", "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": "nGEaWyu9eG", "reviewer": "Reviewer_LhJH", "summary": "The paper proposed a new framework, named AmoebaLLM, that adapts any LLMs to achieve optimal efficiency across different platforms and applications. In specific, the framework contains two stages. The first stage denoted as a knowledge-preserving stage, creates a subnet of the LLM by dynamic programming given an arbitrary depth or width. This second stage denoted as the one-for-all finetune stage is to fine-tune the obtained subnet to achieve the optimal performance on the given application (dataset). To finetune the subnet, the paper proposed shape-aware LoRAs to adapt the knowledge and a loss-magnitude balancing to ensure efficient distillation in the second stage. All combined, the method achieves SOTA performance among all the pruning methods and obtains the best accuracy in the same latency in different platforms.", "questions": "1. What is the definition of the latency?\n\n2. Is it possible to provide the details of selecting the subnet shape that favors the hardware characteristics?\n\n3. Which LLM is used across all the experiments?", "limitations": "No negative societal impact is found.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "1. The target of the paper is practical and interesting. It is indeed a good open problem to implement different LLMs within different constraints of devices and platforms while achieving the optimal trade-off between efficiency and accuracy.\n\n2. This method, the DP-based depth shrinking strategy, could construct an any-shape LLM with reasonable performance, which is promising in different applications.\n\n3. The presentation and writing are good and easy to follow.", "weaknesses": "1. About latency in section 5.2:The work claimed that it achieve the accuracy-efficiency frontier. In terms of accuracy, it is convincing as shown in Table 1. However, for efficiency, it seems only Fig.3 is related to this topic yet only shows a trade-off between latency and accuracy. It is not sufficient to evaluate the efficiency purely based on latency. \n\n2. In my understanding, efficiency should also consider the resources consumed in the pruning method. First, the DP algorithm itself has a high computational complexity. Second, the proposed method is somewhat limited to its one-for-all finetune, which is a training process related to a specific dataset, and a manual selection of the subnet, the method is not perfectly efficient. In this case, the author may consider adjusting their claim in efficiency.\n\n3. The selection of the subnet seems to be related to the profiling of the generation latency. It is purely empirical and may be error-prone when it is generalized to other platforms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing our work as practical and interesting, as well as for your constructive comments! We have addressed all your concerns below.\n\n**1. The definition of latency and more real-device efficiency measurement**\n\nWe clarify that the latency used in Fig. 3 of our manuscript is start-up latency, i.e., the end-to-end time used to finish the generation of one sample (batch size = 1, sequence length = 128).\n\nFollowing your suggestion, to provide a more comprehensive evaluation of inference efficiency, we further measure the start-up latency (batch size = 1) and throughput (batch size = 16) for generating/prefilling a sequence with a length of 128, respectively, using the MLC-LLM framework on an NVIDIA A5000 GPU. We provide the achieved accuracy (averaged over the 7 tasks from Table 1 of our manuscript) and efficiency metrics in the table below. We can observe that our method still consistently achieves the best accuracy-efficiency trade-off.\n\n|  |  | Generation |  | Prefilling |  |  |\n|---|---|:---:|:---:|:---:|:---:|---|\n| Remaining Ratio | Method | Start-up latency/s (bs=1) | Throughput (bs=16) | Start-up latency/s (bs=1) | Throughput (bs=16) | Accuracy (%) |\n| 80% | Shortened LLaMA | 2.834 | 5.528 | 0.0405 | 48.810 | 58.72 |\n|  | FLAP | 2.847 | 5.427 | 0.0418 | 45.701 | 60.98 |\n|  | **Ours** | **2.879** | **5.286** | **0.042** | **46.256** | **62.29** |\n| 65% | Shortened LLaMA | 2.308 | 6.785 | 0.0331 | 60.060 | 52.57 |\n|  | FLAP | 2.372 | 6.376 | 0.0349 | 56.437 | 56.12 |\n|  | **Ours** | **2.294** | **6.663** | **0.035** | **56.398** | **57.43** |\n| 50% | Shortened LLaMA | 1.782 | 8.887 | 0.0259 | 77.295 | 47.35 |\n|  | FLAP | 1.883 | 8.056 | 0.0275 | 71.942 | 51.12 |\n|  | **Ours** | **1.860** | **8.124** | **0.0281** | **73.260** | **52.19** |\n\n\n**2. The complexity of our DP algorithm**\n\nTo address your concern, we provide the consumed time of our DP algorithm here. As elaborated in Section 4.2 of our manuscript, when targeting the removal of M layers out of all N layers, the total number of evaluations required by the DP algorithm is M(N-1)/2. For each evaluation, as mentioned in Section 5.4, we measure the MMLU accuracy of 20 samples, which takes about 15 seconds on average on an NVIDIA A5000 GPU for LLaMA2 7B.\n\nAs such, when applying our DP algorithm to LLaMA2 7B in our case, where N=32 and M=16, the total consumed time is about 1 hour, which is about 1/15 of the LoRA fine-tuning time of our baselines in Table 1. It is worth noting that constructing the DP table is a one-time effort for each pre-trained LLM model and does not need to be repeated for each target subnet configuration.\n\n---\n**3. The overall training efficiency of one-for-all fine-tuning**\n\nFirst, we clarify that our one-for-all fine-tuning is not dataset-specific; instead, we perform a generic fine-tuning on alpaca-gpt4, as mentioned in Section 5.1 of our manuscript, and then evaluate the subnets with varying shapes derived from the fine-tuned model across different test datasets and tasks, as shown in Table 1 of our manuscript. In addition, the selection of subnets can be automatically performed within 10 minutes, as elaborated in our response to Question 5.\n\nFurthermore, beyond the inference efficiency analyzed in Section 5.2 of our manuscript as well as in Question 1, the advantage of our method in training efficiency is that it requires a constant *O(1)* training time with respect to the number (*N*) of subnets with varying efficiency constraints to be delivered, thanks to the one-for-all fine-tuning. In contrast, previous compression methods require *O(N)* fine-tuning time.\n\nFor example, for delivering 10 subnets with varying weight remaining ratios in Table 1 of our manuscript, when fine-tuned on the same number of tokens, all baseline methods require about 150 GPU hours in total on an NVIDIA A5000 GPU, while our method requires only 30 GPU hours.\n\nWe will follow your suggestion to emphasize our achieved trade-off between accuracy and inference efficiency, while also clarifying and highlighting our claim regarding our method’s advantage in training efficiency from the above perspective.\n\n---\n**4. The subnet selection may be error-prone when generalized to various platforms**\n\nWe agree that determining the subnet shape is a general problem for hardware-aware compression methods, which can be error-prone and not generalizable across platforms. This issue originates from the fact that theoretical FLOPs cannot always reflect real-device efficiency.\n\nHowever, our method mitigates this issue by allowing for rapid subnet determination for each target platform, thanks to its delivered one-for-all fine-tuned LLM. Specifically, for each target platform, users can select efficient subnets based on platform-specific profiling and directly obtain the accuracy of these selected subnets from the one-for-all fine-tuned LLM, without the need for per-subnet fine-tuning. This improved training efficiency enables users to quickly identify subnets that hit the Pareto frontier of accuracy-efficiency trade-offs, thus obtaining high-quality subnets efficiently.", "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": "hf4CdAA8Ft", "reviewer": "Reviewer_wNCb", "summary": "This paper proposes AmoebaLLM: a one-for-all fine-tuning and compression framework for delivering pruned and accurate subnets from a pre-trained LLM at various pruning (both depth and width) ratios without the need to fine-tune individual subnets. AmoebaLLM consists of two components: 1) a dynamic programming-based approach to layer selection for selecting subnets of certain depth that takes into account various layer interactions, as opposed to greedy-layer based approaches; 2) after the subnets are defined, a one-for-all fine-tuning strategy is proposed based on gating multiple LoRA heads that are subnet-specific. The training strategy uses a hybrid loss where the largest subnet (always sampled) computes its loss with the ground truth, and all subsequently sampled subnets compute a distilled loss from the largest subnet. The final loss is computed by weighing each subnet's loss appropriately. To demonstrate the efficacy of their method, the authors perform experiments on both LLama 2 7B and Vicuna 7b v1.5 and compare against prior LLM pruning techniques.", "questions": "- AmoebaLLM produces different subnets, some may have the same parameter count due to different width/depth shrinkage ratios, but it is not clear how a model is selected for a certain experiment (e.g. Fig 3). Is it simply enumerating all the subnets and picking the fastest  one? Are the profiling results obtained from Section 3 being used for this, or is that purely for motivating the problem?\n- What ranks are used for LoRA?", "limitations": "The authors have addressed the limitations of their proposed work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper tackles a challenging and relevant problem: compressing pre-trained LLMs with various complexities while maintaining their performance\n- The paper is well written and easy to follow\n- The method overall is novel, and performs better than other pruning methods.\n- The paper provides interesting insights into how different LLM metrics are impacted by model compression, e.g., Sec 5.4 shows that MMLU is a better metric for calibration.", "weaknesses": "- The claims that full model fine-tuning suffers from severe gradient conflicts appear rather weak. In line (233) the authors state \"As analyzed in Sec. 2 and demonstrated in Sec. 5.3 ... \", whereas Section 2 simply conjectures that this is an issue as it uses phrases like\" .. would likely fail ..\" and \"as it may omit layers\". I would suggest citing the results from Section 5.3 directly to motivate this issue. As for the results themselves in Table 3, the performance of the full model (second row) does not seem consistent. While it performs poorly on Wikitext2, the MMLU performance remains comparable with the SMoL baseline, and seems to improve when pruning more (24 to 20) as opposed to all other methods. It is not clear how these results lead to the conclusion that there is a \"gradient conflict\". How would the same experiments look when using the Vicuna model? and how would the other evaluation metrics look?\n- Using dynamic programming in the context of neural network pruning is not new (https://arxiv.org/pdf/2308.10438, https://ieeexplore.ieee.org/document/9116287 to name a few). The authors should cite and discuss these works, even if they do not tackle LLMs specifically. \n- In line 199, the authors claim that the DP complexity is O(MN). This is a bit misleading, as it ignores the complexity of computing P(n,m) which I argue is also dependent on M and N. Furthermore, the actual run-time of the DP step is not discussed, is it too negligible? \n- In Table 1, comparing with AmoebaLLM$^\\dagger$ is not a fair as it is trained for longer (on top of the training performed in regular AmoebaLLM). How will the other methods compare with the same amount of training?\n- It would improve reading clarity if the authors explicitly mention how many layers do the models being studied contain, as simply mentioning how many layers are being pruned in depth shrinkage is hard to gauge without that context. This should also be mentioned in the Tables.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing the novelty and interesting insights of our work, as well as for your constructive comments! We have addressed all your concerns below.\n\n**1. Cite the results from Section 5.3 to motivate the issue of full model fine-tuning**\n\nThank you for the suggestion! We believe what you suggested is a more direct way to motivate the adoption of adapters and will follow your recommendation in the final version.\n\n---\n**2. Clarify the inconsistent results of Table 3**\n\nThanks for the interesting question! The reasons for the comparable MMLU performance, which corresponds to the understanding of factual knowledge, between full model fine-tuning and other parameter-efficient settings are two-fold:\n\n*(1)* As highlighted and analyzed in Section 5.4 of our manuscript, as well as echoed by recent observations regarding the role of LLM fine-tuning [1][2][3], new factual knowledge can hardly be gained during fine-tuning, and the loss of factual knowledge during compression is hard to restore during fine-tuning. As such, the MMLU performance is mainly determined by the selection of subnets, which is the same across settings in Table 3 of our manuscript, thus making the MMLU accuracy relatively comparable.\n\n*(2)* Following lm-evaluation-harness ([69] cited in our manuscript), the MMLU accuracy is calculated by picking the largest logit among the four logits corresponding to A/B/C/D, instead of all logits. As such, even if the model performs poorly in language modeling on the MMLU dataset, it still has chances to maintain comparable accuracy. To validate this, we further report the PPL achieved on the MMLU dataset, where the target texts are constructed by question plus correct answer contents. As shown in the table below, which is an extended version from Table 3 of our manuscript, we can observe that: Similar to the PPL on Wikitext2, full model fine-tuning suffers from a notable PPL increase on the MMLU dataset, indicating its failure in language modeling on the MMLU dataset.\n\n| Method | 32 |  |  | 24 |  |  | 20 |  |  |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|  | Wikitext2 (PPL) | MMLU (PPL) | MMLU (Acc %) | Wikitext2 (PPL) | MMLU (PPL) | MMLU (Acc %) | Wikitext2 (PPL) | MMLU (PPL) | MMLU (Acc %) |\n| Per-subnet ft. | 5.54 | 20.62 | 46.4 | 10.57 | 45.59 | 41.9 | 15.94 | 65.9 | 41.7 |\n| - SMoL (+full model) | 5.82 | 27.02 | 46.6 | 38.48 | 291.11 | 32.6 | 167.74 | 1349.17 | 36.5 |\n| - SMoL (+LoRA) | 6.97 | 32.75 | 40.6 | 12.71 | 81.9 | 40.0 | 19.12 | 122.67 | 37.9 |\n| - Loss-mag. Balancing | 6.77 | 26.34 | 42.0 | 12.63 | 59.81 | 40.1 | 18.19 | 91.07 | 39.3 |\n| **Full** | **6.36** | **25.02** | **47.2** | **12.40** | **55.8** | **45.1** | **18.15** | **84.27** | **41.0** |\n\nFollowing your suggestion, we have also provided the benchmark with the full model fine-tuning variant of our method on Vicuna 7B below. The results are consistent with previous observations: (1) full model fine-tuning leads to poor language modeling capabilities across both datasets, and (2) the MMLU accuracy of full model fine-tuning does not suffer a very drastic drop but is notably worse than our method.\n\n|  | 32 |  |  | 24 |  |  | 20 |  |  |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Method | Wikitext2 (PPL) | MMLU (PPL) | MMLU (Acc %) | Wikitext2 (PPL) | MMLU (PPL) | MMLU (Acc %) | Wikitext2 (PPL) | MMLU (PPL) | MMLU (Acc %) |\n| Full Model Finetuning | 7.29 | 25.52 | 46.9 | 40.7 | 164.67 | 39.3 | 134.07 | 882.82 | 37.1 |\n| Ours | 6.85 | 23.27 | 47.9 | 11.87 | 49.51 | 46.1 | 14.77 | 67.73 | 45.4 |\n\n\n[1] “Does fine-tuning llms on new knowledge encourage hallucinations?”, Z.  Gekhman et al., arXiv’24.\n\n[2] “LIMA: Less Is More for Alignment”, C. Zhou et al., NeurIPS’23.\n\n[3] “R-Tuning: Instructing Large Language Models to Say ‘I Don’t Know’”, H. Zhang et al., NAACL’24.\n\n---\n**3. Missing citations**\n\nThank you for providing these related works! We will cite and comment on them in the final version.\n\n---\n**4. The DP complexity and run-time**\n\nThank you for pointing this out! You are right that the DP complexity should take the varying number of layers involved in evaluating *P(n, m)* into consideration. The updated DP complexity is *O(MN(M-N))*.\n\nFor the run-time of our DP algorithm, when applying it to LLaMA2 7B on an NVIDIA A5000 GPU in our case, where N=32 and M=16, the total consumed time is 1 hour.\n\nWe will add the above information in our final version.", "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": "8OpYqHH3kr", "reviewer": "Reviewer_U1ac", "summary": "The authors proposed a novel pipeline for once-for-all training multiple subnets in a supernet LLM under different resource constraints. The  entire pipeline consists of a knowledge preserving subnet selection utilizing DP to sample depth and width and a new LoRA to resolve the gradient conflicts during training multiple subnets jointly. The experimental results on two widely used decoder-only architectures show the efficacy of the pipeline.", "questions": "See above", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-motivated and well-written. To deploy the LLM on multiple resource-constrained devices, fine-tuning multiple subnets jointly is a promising direction to save the computation and training efforts.\n- Resolving the complexity when jointly selecting the removal layers in layer-wise pruning is a remaining question for a long time. Although the authors did not entirely solve the problem, the proposed DP solution may potentially also benefit other layer-wise network architectures, e.g., CNN, encoder model. etc.", "weaknesses": "- Some technical details are not clear to me. For example, how are the depth and width defined in decoder models? Is a self-attention block a layer? Or does the authors consider FFN and MHSA as two individual layers? In terms of the width, it is clear for FFN but is that the number of heads in MHSA?\n- What metric is used as the importance in subnets sampling? I assume that that importance is not derived from gradients, since DP is done layer-by-layer in forward manner. If so, will that metric work better than gradient-based metrics?\n- The authors introduce a new LoRA in the paper. Does the authors freeze the backbone LLM and only tune the adapters? If not, how does the authors avoid the co-adaptation of backbone and adapters? Since the original LoRA is proposed to preserve the learned knowledge and fine-tuning only low-rank adaptation for a small amount of new data. \n- Also are the adapters only added in MHSA? If so, can the authors provide some theoretical or experimental results to show why FFN layers do not suffer from gradient conflicts?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing the value of our work and for your constructive comments! We have addressed all your concerns below.\n\n**1. The definition of depth and width in decoder models**\n\nThank you for pointing this out! In this work, depth is defined as a whole self-attention block, including both an MHSA and an FFN.\n\nFor width, your understanding is correct. Specifically, for FFNs, the width is defined as channels. For MHSA, the width is defined as the number of attention heads, since all channels within one attention head will be kept or pruned as a whole for real-device efficiency, following the definitions and strategies in previous structured LLM pruning works ([6][8] cited in our manuscript). \n\nWe will clarify this in the final version.\n\n---\n**2. The metric used as the importance in subnets sampling of our DP algorithm**\n\nAs elaborated in Section 5.4 of our manuscript, we adopt the MMLU accuracy, which serves as an indicator for the encoded factual knowledge, as the DP metric for subnet selection. This is based on the key insight that the loss of factual knowledge during compression is hard to restore during fine-tuning and thus should be prioritized in subnet selection, while the language modeling capability is easier to recover through fine-tuning.\n\nIn addition to the benchmark between our DP algorithm and existing layer pruning methods in Table 2 of our manuscript, where our method consistently outperforms all baselines, we follow your suggestion to further benchmark our DP algorithm with gradient-based metrics. Specifically, the gradient-based metric we tried is defined as the average gradients received by all MLPs falling within one layer (defined as a whole self-attention block), where layers with lower average gradients are potentially less important and thus pruned. Following the evaluation setting of Table 2 of our manuscript, we summarize the achieved Wikitext2 PPL and MMLU accuracy when remaining different numbers of layers on LLaMA 2 7B in the table below.\n\n| Calib. Data | # Remained Layers | 24 | 23 | 22 | 21 | 20 | 19 | 18 | 17 | 16 |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Wikitext2 (PPL) | Gradient-based | 12.689 | 16.394 | 21.813 | 25.209 | 67.262 | 326.666 | 575.863 | 2295.172 | 2140.926 |\n|  | Ours | 11.66 | 12.77 | 17.59 | 20.06 | 23.77 | 28.83 | 38.70 | 70.87 | 95.16 |\n| MMLU (Acc %) | Gradient-based | 24.4 | 26.2 | 24.9 | 24.6 | 23.9 | 234.0 | 24.8 | 24.8 | 24.2 |\n|  | Ours | 46.2 | 44.8 | 44.6 | 44.1 | 41.2 | 41.3 | 43.1 | 34.7 | 28.9 |\n\nWe can observe that our method consistently outperforms gradient-based metrics. According to our analysis in Section 4.2, this may be because (1) gradient-based metrics cannot effectively evaluate the joint contributions of different layer combinations like our DP method; and (2) gradient-based metrics cannot effectively reflect the encoded factual knowledge.\n\n---\n**3. Whether the backbone is frozen and only the adapters are tunable**\n\nYes, we freeze the backbone and only tune the adapters. This is because we find that one-for-all fine-tuning of the whole backbone on low-resource fine-tuning data results in suboptimal performance of particular subnets, as shown in Table 3 of our manuscript.\n\n---\n**4. Where the adapters are added**\n\nWe follow the settings in QLoRA [1] to add adapters to all linear layers in the pre-trained LLM, including both MHSAs and FFNs.\n\nTo study the relative importance of adapters in MHSA and FFNs, we further explore variants where adapters are added only to one of them on top of our method and summarize the results in the table below. \n\n| Remaining Ratio | Attn only |  | FFN only |  | Attn + FFN |  |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|  | Wikitext2 | MMLU (%) | Wikitext2 | MMLU (%) | Wikitext2 | MMLU (%) |\n| 80% | 13.31 | 44.6 | 12.14 | 45.6 | 10.34 | 47.4 |\n| 65% | 14.11 | 38.3 | 13.63 | 39.1 | 11.76 | 40.3 |\n| 50% | 19.23 | 32.5 | 16.98 | 33.1 | 15.58 | 34.0 |\n\nWe find that (1) adding adapters to both layers achieves the best task performance, and (2) adding adapters to FFNs results in better task performance compared to adding them to MHSAs. This indicates that FFNs may play a more effective role in adjusting the output token distributions after compression.\n\n[1] “QLoRA: Efficient Finetuning of Quantized LLMs”, T. Dettmers et al., NeurIPS’23.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: After the rebuttal, the ratings were: borderline accept, borderline accept, accept, and borderline accept. While none of the reviewers rejected the paper, only one expressed strong support. However, the reviewers did acknowledge the practicality of the problem and the novelty of the proposed method.
Decision: Accept (poster)