AmoebaLLM: Constructing Any-Shape Large Language Models for Efficient and Instant Deployment

Thank you for recognizing the value of our work and for your constructive comments! We have addressed all your concerns below.

1. The definition of depth and width in decoder models

Thank you for pointing this out! In this work, depth is defined as a whole self-attention block, including both an MHSA and an FFN.

For 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).

We will clarify this in the final version.

---

2. The metric used as the importance in subnets sampling of our DP algorithm

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

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

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

---

3. Whether the backbone is frozen and only the adapters are tunable

Yes, 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.

---

4. Where the adapters are added

We follow the settings in QLoRA [1] to add adapters to all linear layers in the pre-trained LLM, including both MHSAs and FFNs.

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

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

[1] “QLoRA: Efficient Finetuning of Quantized LLMs”, T. Dettmers et al., NeurIPS’23.

Dear Reviewer,

We sincerely appreciate the time you dedicated to providing valuable feedback on our paper. In this author response, we have addressed all of your initial concerns. If you have any further questions or concerns, we are happy to discuss them with you. Additionally, we welcome any new suggestions or comments you may have!

Best,

The Authors of Paper #13852

5. Benchmark with baselines with the same amount of training time

We first clarify that in Table 1, we follow a standard criterion in the literature, i.e., ensuring the same number of fine-tuning tokens across both the baselines and our method.

Following your suggestion, we further train the baselines with the same amount of training time as our method, which corresponds to 20k training iterations on Alpaca-gpt4. As shown in the table below, we report the MMLU accuracy as well as the average accuracy across 7 tasks, following the task list of Table 1. We observe that (1) the baselines fine-tuned with 20k iterations generally maintain comparable performance with those trained with 10k iterations, potentially because 10k iterations is sufficient for fine-tuning on Alpaca-gpt4; and (2) our method still outperforms all baselines across various tasks.

6. The number of remained depth and width in Table 1

Thank you for the suggestion! The (depth, width scale) for AmoebaLLM with 80%/65%/50% remaining ratios in Table 1 are (30, 0.875)/(28, 0.75)/(22, 0.75), respectively. We will follow your suggestion to add this information to Table 1 to improve readability.

---

7. The detailed subnet selection strategy

The profiling results in Section 3 are purely for motivating the problem. Here is our current subnet selection strategy: we adopt a hierarchical search strategy to deliver subnets from our design space that satisfy the target efficiency constraint, e.g., 50% weight remaining ratios in Table 1 of our manuscript, while maximizing the achievable accuracy. Specifically, we first perform a coarse grid search across uniformly spaced depth and width settings based on a small calibration set, e.g., 20 samples from the MMLU dataset, to identify the subnets that satisfy the given efficiency constraint with maximized accuracy. Next, we perform a more fine-grained grid search within depth/width ranges surrounding the optimal subnet identified in the coarse grid search stage. This process typically evaluates 40 promising subnets and takes no more than 10 minutes on an NVIDIA A5000 GPU.

We empirically find that the above strategy works well at the scale of our target problem and can already deliver high-quality subnets that outperform previous compression methods. We also note that more complex subnet selection strategies, such as the evolutionary search adopted by [11]-[15] cited in our manuscript, can also be employed, which will be our future work.

We will clarify this in the final version.

---

8. The rank of LoRA

We follow QLoRA [1] and adopt a rank of 64. We will add this information to the final version.

[1] “QLoRA: Efficient Finetuning of Quantized LLMs”, T. Dettmers et al., NeurIPS’23.

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.

1. Cite the results from Section 5.3 to motivate the issue of full model fine-tuning

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

---

2. Clarify the inconsistent results of Table 3

Thanks 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:

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

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

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

[1] “Does fine-tuning llms on new knowledge encourage hallucinations?”, Z. Gekhman et al., arXiv’24.

[2] “LIMA: Less Is More for Alignment”, C. Zhou et al., NeurIPS’23.

[3] “R-Tuning: Instructing Large Language Models to Say ‘I Don’t Know’”, H. Zhang et al., NAACL’24.

---

3. Missing citations

Thank you for providing these related works! We will cite and comment on them in the final version.

---

4. The DP complexity and run-time

Thank 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)).

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

We will add the above information in our final version.

Dear Reviewer,

We sincerely appreciate the time you dedicated to providing valuable feedback on our paper. In this author response, we have addressed all of your initial concerns. If you have any further questions or concerns, we are happy to discuss them with you. Additionally, we welcome any new suggestions or comments you may have!

Best,

The Authors of Paper #13852

5. Detailed subnet selection strategy

Thank you for pointing this out! Currently, we adopt a hierarchical search strategy to deliver subnets from our design space that satisfy the target efficiency constraint, e.g., 50% weight remaining ratios in Table 1 of our manuscript, while maximizing the achievable accuracy. Specifically, we first perform a coarse grid search across uniformly spaced depth and width settings based on a small calibration set, e.g., 20 samples from the MMLU dataset, to identify the subnets that satisfy the given efficiency constraint with maximized accuracy. Next, we perform a more fine-grained grid search within depth/width ranges surrounding the optimal subnet identified in the coarse grid search stage. This process typically evaluates 40 promising subnets and takes no more than 10 minutes on an NVIDIA A5000 GPU.

We empirically find that the above strategy works well at the scale of our target problem and can already deliver high-quality subnets that outperform previous compression methods. We also note that more complex subnet selection strategies, such as the evolutionary search adopted by [11]-[15] cited in our manuscript, can also be employed, which will be our future work.

We will follow your suggestion and our promise in the abstract to open source all source code and the delivered subnets upon acceptance.

---

6. Which LLMs are used

As mentioned in Section 5.1 of our manuscript, we employ LLaMA 2 7B and Vicuna 7B v1.5.

Thank you for recognizing our work as practical and interesting, as well as for your constructive comments! We have addressed all your concerns below.

1. The definition of latency and more real-device efficiency measurement

We 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).

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

2. The complexity of our DP algorithm

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

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

---

3. The overall training efficiency of one-for-all fine-tuning

First, 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.

Furthermore, 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.

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

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

---

4. The subnet selection may be error-prone when generalized to various platforms

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

However, 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.

Dear Reviewer,

We sincerely appreciate the time you dedicated to providing valuable feedback on our paper. In this author response, we have addressed all of your initial concerns. If you have any further questions or concerns, we are happy to discuss them with you. Additionally, we welcome any new suggestions or comments you may have!

Best,

The Authors of Paper #13852

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.

1. Whether our DP-based depth shrinking strategy needs to evaluate the performance for each selection strategy and its overhead

Yes, 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.

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

We will add this analysis to the final version.

---

2. Whether the optimal subnets change before and after fine-tuning and the potential of searching subnets with fine-tuning

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

Second, 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.

[1] “Does fine-tuning llms on new knowledge encourage hallucinations?”, Z. Gekhman et al., arXiv’24.

[2] “LIMA: Less Is More for Alignment”, C. Zhou et al., NeurIPS’23.

[3] “R-Tuning: Instructing Large Language Models to Say ‘I Don’t Know’”, H. Zhang et al., NAACL’24.

---

3. Training efficiency analysis

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

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

We will add the above analysis to the final version.

Dear Reviewer,

We sincerely appreciate the time you dedicated to providing valuable feedback on our paper. In this author response, we have addressed all of your initial concerns. If you have any further questions or concerns, we are happy to discuss them with you. Additionally, we welcome any new suggestions or comments you may have!

Best,

The Authors of Paper #13852