Prompt-based Depth Pruning of Large Language Models

Dear reviewer 9fAo,

We appreciate your positive evaluation on our paper, acknowledging the practical advantages of our method. We respond to the concerns you raise in what follows.

---

It would be better to show the training convergence analysis of the router.

Following your suggestion, we have attached the loss curves of the router training in the following LINK (training loss), and LINK (test loss). The results confirm that the test loss of the router has successfully converged by the end of the training.

We will add the plot in the revised manuscript.

---

We hope that you find our response reasonable. Please let us know if you have any further questions.

Best regards,
Authors.

Dear reviewer KB85,
Thank you for your constructive and thoughtful feedback. In what follows, we address the concerns you raised one-by-one.

---

Missing paper on dynamic model pruning [1].

Thank you for the pointer to this concurrent work. The paper [1] missed our attention, as it appeared on arXiv after finalizing our paper (Jan 2025). Still, it looks interesting and relevant; we will cite and discuss it in the revised manuscript.

We note that [1] differs from our work in several key aspects, making two approaches complementary to each other.

• Sparsity structure: Our work considers depth pruning (very coarse) which is friendlier to hardware using grouped computations (c.f. Kim et al. (2024)), On the other hand, [1] considers row/col pruning.
• Training footprint: Our work only trains the router and keeps the base LLM untouched (except for Section 6.3). On the other hand, [1] involves training the base LLM much further, conducting both pre-training and fine-tuning. Overall, these differences make two approaches complementary to each other---ours focusing on versatility and [1] focusing on good sparsity-quality tradeoff---making both valuable contributions with distinct scopes.

Kim et al., “Shortened LLaMa: A simple depth pruning for large language models,” arXiv 2024.

---

This paper only achieves 20% sparsity while hurting the performance by more than 10%. (…) The performance degradation is still significant, comparing with dynamic pruning.

As we have explained briefly above, our work imposes (1) coarse sparsity structure, and (2) minimal retraining, which increases the performance degradation given the same sparsity. This is because we intend to make the algorithm easily adoptable by low-budget end users, who might be the ones that want to develop their own input-dependent depth-pruned LLMs specialized for their tasks. In this sense, we are trading the performance for better usability.

---

Missing discussion on contextual sparsity.

Thank you for pointing this out. The reviewer is correct; we will add discussions on [3,4,5] (note that we already discuss [2] in Section 2).

---

Method design 1: Why not introduce a learnable network that predicts a mask for each layer?

We have decided to confine the routing decisions to pre-determined options to make the routing decision easier and simpler (only 10 choices, much less than $2^{32}$), so that we can train a lightweight router with limited amount of training data. In fact, we have also tried a learnable mask approach, but have failed to train a sufficiently light yet performant router. Such difficulty is circumvented in existing works by using the intermediate features to make routing decisions (as in D-LLM) or joint training with base LLM parameters (as in Hou et al. (2025)). We have deliberately avoided these options to enable a one-time memory load and enhance the usability of our technique by GPU-poor end users, respectively.

---

Method design 2: Why not directly train a classifier?

By conducting a regression on the likelihoods, we are basically training using the soft label instead of hard labels. Training with such objectives is known to enjoy better generalization performance, by learning to mimic the “dark knowledge” in these models. In fact, this is a popular technique in knowledge distillation; Hinton et al. (2015) shows that one can approximate knowledge distillation by the $ell_2$ regression of pre-softmax activations.

Empirically, in Table 11, we have compared the regression-trained router against the classification-trained router, where we observe that the regression-trained router performs better.

Hinton et al. “Distilling the knowledge in a neural network,” arXiv 2015

---

We sincerely hope that you find our response reasonable. Please don’t hesitate to let us know if there are any further questions.

Best regards,
Authors

Dear reviewer r1RD,
Thank you for your constructive feedback. In what follows, we respond to the points raised by the reviewer one-by-one.

---

No Real-World Deployment Results (...) Missing mobile/edge device benchmarks weakens the practical relevance of the method.

Thank you for this comment. Following the reviewer’s suggestion, we have measured performance of LLaMA-3.1-8B PuDDing on MacBook Pro, where the inference took place on the M3 Pro chip (ARM-based processor) with 18GBs of RAM (we run using C++).

We will add this result in the revised manuscript.

---

How does the dynamic loading of omission sets affect inference latency?

As shown in Table 7, the initial loading of PuDDing takes ~0.2s using PCIe Gen4 (0.02s when using NVlink). This latency is relatively small comparing with the latency from the running the generation steps, which can take more than 16s for generating 512 tokens using RTX 6000 Ada. We will add the relevant discussion to the revised version.

---

The reported 1.2x inference speedup is not very high compared to other compression techniques.

First, we highlight that this 1.2x speedup is the real speedup, meaning that it can be achieved on almost all hardwares, without any implementational tricks. This is in stark contrast with many structured pruning algorithms (which often fails to provide speedup on hardwares that use grouped operations, e.g., GPUs with tensor cores) or quantization methods (which needs low-precision hardware and/or needs frequent dequantization). Second, the proposed PuDDing can be combined with other compression techniques, providing orthogonal benefits. For instance, we can apply weight quantization on the PuDDing to reduce the computational cost even further; W8 quantization can be done almost for free on this model, and applying W4 quantization is still better than the baselines.

---

The pruning method is not truly dynamic.

We agree with the reviewer’s point. To avoid any confusion, we will revise the manuscript to avoid calling our method “dynamic” and replace it with other options, such as “contextual.”

---

Fixed omission set reduces flexibility.

While this is true, using the fixed omission set was a necessary choice to reduce the complexity of the routing task, thereby minimizing the size of the router and its associated training costs (both compute and data). In fact, we have tried a “truly dynamic” approach similar to D-LLM, but this option worked worse than the current version. For a more detailed answer, please refer to our answer #4 to the reviewer KB85.

---

Some of the prior work that author could have included (...) How does PuDDing compare to other dynamic pruning methods like MoD or D-LLM?

We clarify that we have made conceptual comparisons against MoD and D-LLM in the sections 1 and 2.2. The works critically differ from our work in that they require loading full models to memory to conduct token-level routing.

---

Why TL loss is a better metric.

The task likelihood loss measures the perplexity (PPL) of the sample, but measured on the “answer” part of the given sample conditioned on the “question,” i.e., $P(x_{\mathrm{answer}}|x_{\mathrm{question}})$. This improves the depth pruning decisions, by letting us optimize how well we answer to the questions (of specific type), not focusing on the fluency in general.

---

How well does the router generalize to unseen tasks or domains.

We empirically confirm that it has strong generalizability to unseen and specialized tasks such as MMLU/MathQA/OBQA (Table 5), and MathQA/SciQA (newly added). Due to the character limit, we refer the reviewer to our answer#1 to the reviewer 2qci.

---

Other comments.

Unfortunately, we cannot give detailed answers to all points in this round due to the space limit. We deeply appreciate these, and will incorporate these sincerely in the revised manuscript.

---

We sincerely hope that you find our response reasonable. Please don’t hesitate to let us know if there are any further questions.

Best regards,
Authors

Dear reviewer 2qCi,
Thank you for your insightful comments and suggestions. In what follows, we respond to the points raised by the reviewer one-by-one.

---

The scope of the task is too narrow, all being the commonsense reasoning tasks.

TL;DR. We have already evaluated PuDDing on MMLU/OBQA/MathQA (Table 5), and newly added experiments on PubMedQA/SciQ.
First, we clarify that we have already evaluated PuDDing (trained on commonsense tasks) on specialized tasks. In particular, see Table 5 for results on MMLU/MathQA/OpenBookQA. From the table, we confirm that PuDDing generalizes well on these tasks as well. To make this point even more concrete, following the reviewer’s suggestion, we have added new experiments on PubMedQA/SciQ (see table below). Again, the results suggest that training PuDDing on commonsense tasks can give a competent router for specialized tasks. We hypothesize that this is because the optimal routing decision is not solely determined by the knowledge domain. Instead, there may be other notions of diversity inside commonsense reasoning tasks that affect how we should route, which also generalizes to specialized tasks.

---

The main drawback of the proposed method is construction of omission sets (...) need to devise an experimental setting where omission sets largely differ, so that prompts drop notably different layers.

Our point is twofold: First, we clarify that despite the deceiving look of Figure 4, the actual omission sets differ quite notably for PuDDing trained with commonsense reasoning dataset. In the LINK, we depict commonly omitted blocks for various tasks—including MMLU/MathQA/(…)—with explicit omission rates. We find that the omission rate of certain blocks differs dramatically over tasks. For instance, block 11 is dropped 99% in PIQA, but is dropped only 34% in MMLU; block 18 is dropped over 40% on PIQA and WinoGrande, but almost never in other tasks.
Second, we have followed your suggestion to train a new version of PuDDing, where we construct omission sets using diverse domain data: math (MathQA), medicine (PubMedQA), science (SciQ), and commonsense reasoning (ARC-Easy, WinoGrande). The results are given in the table below; see newPuDDing. We observe that the average performance slightly increases over PuDDing, driven by the performance boosts in newly added datasets.

---

It fairly covers most literature, but it might be good to cover a few more baselines such as LaCo, LLM-streamline, and FinerCut.

Following the reviewer’s suggestion, we are working towards adding more baselines to our main table. In particular, we have already added LLM-Streamline results (see table below). As the method involves fine-tuning the base model, we compare it with the LoRA fine-tuned version of PuDDing; we find that PuDDing outperforms this baseline as well. Unfortunately, we could not compare against LaCo and FinerCut during the response period, as FinerCut does not provide any code, and LaCo has only released a very limited amount of codes (in ipynb). We are currently implementing the codes for these methods, and will be able to add it in the revised manuscript.

---

We sincerely hope that you find our response reasonable. Please don’t hesitate to let us know if there are any further questions.

Best regards,
Authors