Prompt-based Depth Pruning of Large Language Models
Authors: Juyun Wee
Minjae Park
Jaeho Lee
Venue: ICML 2025 poster (2025)
Keywords: Depth pruning
Model Compression
Primary area: deep_learning->algorithms
TL;DR: We develop a prompt-based depth pruning algorithm.

Abstract: Depth pruning aims to reduce the inference cost of a large language model without any hardware-specific complications, by simply removing several less important transformer blocks. However, our empirical findings suggest that the importance of a transformer block may be highly task-dependent---a block that is crucial for a task can be removed without degrading the accuracy on another task. Based on this observation, we develop a dynamic depth pruning algorithm, coined PuDDing (**P**rompt-ro**u**ted **D**ynamic **D**epth Prun**ing**), which determines which blocks to omit from the model based on the input prompt. PuDDing operates by training a lightweight router to predict the best omission set among a set of options, where this option set has also been constructed in a data-driven manner. Empirical results on commonsense reasoning benchmarks demonstrate that PuDDing effectively accelerates the inference language models, and achieves better on-task performance than static depth pruning baselines.

Prompt-based Depth Pruning of Large Language Models
Juyun Wee Minjae Park
Jaeho Lee'
Abstract
LLaMA 3.1 8B Size:16.06GB
Out of Memory
Depth pruning aims to reduce the inference cost of a large language model without any hardware- specific complications, by simply removing sev- eral less important transformer blocks. However, our empirical findings suggest that the impor- tance of a transformer block may be highly task- dependent- a block that crucial for task can be removed without degrading the accuracy on an- other task. Based on this observation, we develop dynamic depth pruning algorithm, coined PuD- Ding (Prompt-routed Dynamic Depth Pruning), which determines which blocks to omit from the model based on the input prompt. PuDDing oper- ates by training a lightweight router to predict the best omission set among set of options, where this option has also been constructed in data-driven manner. Empirical results on com- monsense reasoning benchmarks demonstrate that PuDDing effectively accelerates the inference lan- guage models, and achieves better on-task perfor- mance than static depth pruning baselines. Project Page:
20% Reduced
Which blocks should be pruned for the question "Has NEL played the Superbowl home?'
Crii broeka lrcated Mt: 18. 10.15, 83, 29
Predictor
Figure I. The general framework of prompt-based depth pruning. Given some query from the user, the goal is to identify which layers from an LLM can omitted, so that one can make accurate prediction low-memory consumer devices.
that is particularly effective for on-device scenarios (Song et al., 2024; Kim et al., 2024). Such methods simply remove several transformer blocks (which we call "omission set") from the pretrained model, based on some measures of block importance computed using a small amount of calibration samples. As everything is identical except for the number of blocks, the pruned model is suitable to be deployed on any hardware without tailored supports on low- now-precision (e.g., integer cores) or fine-grained sparsity (e.g., 2:4 spar- sity). Furthermore, as there is no extensive training involved, depth pruning can be easily done in a device-by-device man- ner for deployment on various devices.
grreOl.gihhubbio/uuDDD gitubb.com/aadd00477PDDDDDn
Code:
1. Introduction
Recent advances in large language models (LLMs) have achieved remarkable success in a wide range of natural lan- guage processing tasks (Brown et al., 2020; Touvron al., 2023; Dubey al., 2024). However, significant computa- tional requirements of LLMs pose challenges in resource- constrained environments, limiting their practicality. For ex- ample, LLaMA-3.3 3.. needs 140GB of RAM to be loaded in bf16, which is often too big for memory-constrainee local devices. Thus, reducing the model size is essential to make LLMs feasible for on-device applications.
A key limitation of typical depth pruning algorithms is that their pruning decision is static, i.e., the same omission set is removed regardless of the query given to the model. While this choice allows one to save storage (e.g., flash drives) by discarding the pruned parameters at the local device, sacrifices the ability to adapt to various downstream tasks. Indeed, our empirical observations show that pruning some transformer blocks in LLM may ineur significant ac- curacy degradation on certain tasks, while being highly unnecessary for other tasks (see Section 3).
Depth pruning is versatile model compression technique
'qqual contribution 'POSTECH. Correspondence to: Jacho Lee <jacho.lee@@ postech.ac.kr>>.
Can we make dynamic depth pruning decisions to improve the performance on various tasks? This question has not been well studied yet, especially in the context of on-device
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Prompt-based Depth Pruning of LLMs
inference. A recent line of work develops effective dynamic token routing mechanisms save training/inference com- putation by processing each token with limited number of transformer blocks (Raposo et al., 2024; Wang et al., 2024). However, such methods require all parameters to be loaded on high-speed memories (e.g., on-GPU memory); thus, the methods are appropriate for large-scale server clusters, not for on-device inference with memory constraints.
underscoring the need for task-dependent depth pruning. We consider the task of prompt- based depth pruning for the first time (to our knowledge), and propose a training- based strategy as a solution.
Comparing with static depth pruning algorithms, our al- gorithm achieves a much higher zero-shot accuracies on various tasks, while being competitive in terms of the computational efficiency.
Contribution. To overcome the limitations, we develop new prompt-based depth pruning approach (Section 4): In the pre-fill stage, based on the prompt given from the user, a limited number of transformer blocks are selected and loaded to the on-device RAM from the storage drive. This approach does not require a large memory to hold all parameters or highly repeated per-token routing, and thus can effectively accelerate inference on low- low-memory devices.
Table 1. A high-level comparison of the proposed prompt- hased depth pruning framework with related depth pruning approaches: Static depth pruning and dynamic token routing.
Tisk Adaptioo
Peak Memory
Routing
Static Pruning they sot 200 Kind al 2004
x
Spurse
Token Rouring RRopo ana a00e (011))
/
Dense
Per token
Prompt based depth pruning (this paper)
<
Spurse
Per prompt
A naïve way to achieve this goal might be to conduct conven- tional static depth pruning at each inference, using the given prompt as calibration samples. However, this approach in- curs a large latency in running static pruning algorithms in every inference. Furthermore, such a method is likely to fail making good pruning decision due to the shortage of calibration data, especially in single-batch inference cases common in on-device scenarios.
2. Related Work
In this section, we provide an in-depth comparison of the proposed framework against existing depth and width spar- sity frameworks. See Table for concise summary.
2.1. Static Depth Pruning
To this end, we propose a training-based method for the prompt-based depth pruning of large langauge models (Sec- tion 5). Our method, coined Prompt-routed Dynamic Depth Pruning (PuDDing), works in two steps.
Static depth pruning methods select and remove unnecessary blocks from a pretrained LLM using various proxy metrics to measure the importance of the blocks. ShortGPT (Men et al., 2024) measures the block importance using the ex- pected cosine similarity between the input and output activa- tions of the block; a block that does not change the direction of the activation is deemed unnecessary. Shortened-L.Laan (Kim et al., 2024) directly measures the perplexity drop after removing each transformer block, and SLEB (Song et al., 2024) combines this idea with an iterative pruning.
1. Candidate omission set generation. We construct small yet diverse and performant family of omission sets. This is done by drawing multiple splits of calibration data from various task dataset, and then finding an omission set which achieves low loss on each split; here, we use a newly developed task-centric loss instead of perplexity.
2. Router training. We train a lightweight router which pre- dicts the appropriate omission set from the given prompt. This is done by generating a training dataset consisting of prompt-loss pairs for each omission set, and training the model to predict the loss from the prompt; routing can be done by choosing the minimum-loss option.
Several recent works also focus layer-level depth pruning, instead of removing an entire transformer block. In partic- ular, Siddiqui et al. (2024), He et al. (2024) discover that pruning out self-attention layers have a much less significant impact than removing the feed-forward layers.
Unlike these works, this paper aims to perform dynamic depth pruning using the prompts for the downstream tasks; to account for this difference, we design and use new saelihood-based metrics to measure the block importance.
Empirically, we find that the proposed PuDDing enjoys clear advantage over static depth pruning algorithms, achiev- ing more than 4%p accuracy increase on zero-shot common- sense reasoning tasks (Section 6). At the same time, as the algorithm uses the router only once per each prompt, PuD- Ding enjoys over 1.2x generation speedup over the dense model, similar to the static depth pruning algorithms.
2.2. Dynamic Token Routing
Inspired by the success of mixture-of-experts -f-eees (Jacobs et al., 1991; Fedus et al., 2022), several recent works have devel- oped mechanisms to route tokens through only fraction of all transformer blocks. Mfxture-of-Depeee (Raposo et al., 2024) adopts the depth sparsity during the training phase with a jointly trained router, to reduce the training cost of
Our key contributions can be summarized as follows:
Our observations reveal that optimal depth pruning de- cisions may be highly depend on the task given at hand,
2

---

Prompt-based Depth Pruning of LLMS
LLMs. Here, the trained router can also be used infer- ence. D-LLM (Wang et al., 2024) trains router that can be applied on pretrained LLMs to reduce their inference cost.
2:19
Set. 3:
25.
27.
30.30
Set 4: 9, 19, 25, 27. 27.
Our approach differs from both of these works in the sense that it needs only a limited number of transformer blocks active for single input query (or prompt); the routing is conducted once per imput prompt, not per token.
Accurrea (%)
Zero-shat Accuracy
2.3. Contextual Sparsity
Our work is most closely related to the idea of contex- tual sparsity, where lightweight router selects an input- dependent subnetwork at inference time without updating the base weights. In the context of width pruning. prior works-Dejp Vu (Liu et al., 2023), ShadowLLM (Akhauri et al., 2024), Sirius (Zhou et al., 2024), and CATS (Lee et al., 2024)-avv demonstrated that context-aware routing can be done with minimal or no degradation in task performance. PuDDing extends this paradigm depth pruning for the first time: instead of skipping neurons or channels, our router decides which entire transformer blocks to omit. This preserves the original matrix shapes and avoids hardware mismatches often caused by width pruning.
BoalQ
PIDA
Wintronde
Tasks
Figure 2. The impact of pruning the transformer block 29 vs. block 30. On the BoolQ dataset, pruning the block 29 instead of block 30 incurs dramatic performance degradation, with over 20%p drop. On the other hand, PIQA and WinoGrande, the accuracy does not change much, or even increases.
inference of task information from the prompt.
3. A Motivating Observation
4. Problem Description
Before describing the proposed framework, we briefly de- scribe motivating observation which demonstrate that:
Inspired by the observations in Section 3, now formalize the problem of prompt- baset depth pruning.
The importance of a transformer block in a language model may be highly task-dependent.
In nutshell, given some pretrained LLM and a prompt, the goal of the prompt-based depth pruning is to desig- nate which transformer blocks should be removed from the model to generate the most accurate response to the prompt.
Setup. To show this point, we have compared the zero-shot accuracies of the LLMs whose omission sets differ by a single transformer block. More concretely, we compare the performance of an omission set (b1, b2, ... ,,, .a- 1. bk) to another omission set (b,, b2........ b b,). on the LLaMA 3.1-8B model. Here, we have used the SLEB (Song al., 2024) to generate omission set, and then replaced single block to get another one. Then, we observe the impact of such replacement on three commonsense reasoning tasks: BoolQ, PIQA, and WinoGrande.
More concretely, let be the prompt given the model, and let W (W1..... Wa) be the weight parameters of pretrained LLM consisting of d transformer blocks, with W; indicating the weights of the ith block. The prediction quality of this language model is measured by the expected loss between the model output and the ground-truth, i.e.,
L(W) = -[(((x,,))) W)],
Result. Figure illustrates our findings. We observe that pruning out block 29 instead of block 30 has two-sided impact: On BoolQ, the change makes dramatic drop in accuracy (62.2%- + 38.0%, 62.5% 37.9%). However, on PIQA and WinoGrande, we observe slight accuracy boost. This phenomenon suggests that the block 29 may contain more knowledge relevant to answering BoolQ ques- tions, while 30 may be more knowledgeable about PIQA and WinoGrande. This observation highlights the need to consider task variability during the selection of the omission set. To formally address such need, this paper considers
where E((:) :); W) ) is some loss function which also encap- sulates the generative procedure of language model with parameter W (e.g., perplexity). In static depth pruning, the goal is to find which blocks to prune from the given LLM. More formally, define omission set as a(n unordered) set of transformer block indices
b = (b
S dJ.
(2)
which designates which blocks will be omitted from the target LLM. Then, let Wib be a sequence of d weights, with b,th block eliminated from the W. Then, the static
3

---

Prompt-based Depth Pruning of LLMs
depth pruning aims solve the minimization
Coverage:
For any realistic prompt-answer pair (x; y) from wide range of tasks, the set B should contain least one b, with a small loss {(y, f(x: W)b))).
min L(Wib), b: b|2k
(3)
Cardinality: The number of omission sets m should be sufficiently small, that one can train nice predictor for B with limited number of samples.
given the depth constraint & designated by the operational constraints, such as the desired latency or the peak memory.
Prompt-based Depth Pruning. The problem of prompt- based depth pruning can be described as optimizing the omission set as function b(x), L.e., solving
To obtain these properties, we adopt the following strategy: First, we collect calibration datasets Di...., ,, D, on a di- verse set of downstream tasks. Then, on each calibration dataset, we select the omission set that minimizes some loss criterion, i..., solve
min E [((((x,)) y); W ). b(-)
(4)
subject to Pr[[b(x) 2 k| 1.
by arg min Ed, [[ f(x; W.b...
(6)
Note that we are constraining the omission set to have the cardinality greater than k for all In other words, the pruned model should always have d - k or less blocks. This is because we mainly consider the peak memory constraint, ie., the RAM cannot hold more than d - k blocks. Other- wise, one can consider slightly modified version of the problem (4) with a probabilistic constraint.
Here, the minimization done in a greedy manner, similar to Song et al. (2024). We apply different loss criteria on each calibration dataset to get m = x omission sets.
new task-focused vari- ants of the perplexity loss, which we call the task likelihood losses. The perplexity measures the fluency of the generated sentences by measuring the average log-likelihood losses over the whole sequence. That is, for a sample sentence z = (Z :2. the perplexity
Losses. As the loss function,
5. Method
We now formally describe the proposed PuDDing (Prompt- routed Dynamic Depth Pruning)- an algorithm to train a router b(.) for the prompt-hased depth pruning.
ppl(z; W) = =exp (- 1 Σ logri ()): (7)
(-T
In nutshell, PuDDing operates in two steps:
where (: W) denotes the conditional generative prob- ability of the target language model with parameters W, at the ith token. We modify this loss to measure the likeli- hood only the sequence that matters for on-task performance. That is, if the given datum z can be broken down into the prompt and answer pair:
1. Generating candidate omission sets using the prompt- answer dataset collected from various tasks (Section 5.1)
2. Training router to predict the best option among the candidate omission sets (Section 5.2)
During the inference phase, the given prompt is fed to the router, which predicts which omission set (among the candi- dates) one should use for the given prompt. Then, the model parameters are loaded from the storage to the high-speed memory to constitute the depth-pruned LLM (see Figure 3).
z (x.y) = (21....: 25. 28-1....,
(8)
then we can define the task likelihood (tl) loss as:
We note that this classification- based approach is in contrast with the approach of dynamic token routing (Wang et al., 2024), where one makes yes/no decisions for omitting each block in a sequential manner; this change is to make the router training easier and generalizable.
tl(z: w) T (=5+1
log W).
(9)
In addition, also consider the task likelihood difference (tld) loss, which is defined as follows: In many tasks, the answer choices are limited (e.g. "true" or "false"). In such cases, we can also use the likelihood difference of the correct and wrong answers, i.c.,
5.1. Candidate Omission Set Generation
The first step is to generate a candidate pool of omission sets. That is, we generate a family of omission sets
tid(z; (z; W) ti((x, y): W) tl((x, y wrong): W), (10)
where y wrong denotes the wrong version of the answer. We use both (l(-) and tld(.) as our loss criteria.
B = (b....., b
(5)
We note that the task likelihood losses Equations (9) and (10) is different from the perplexity (Equation (7)), in the sense that we do not exponentiate the values. We use this version as it empirically works better than the exponentiated one.
which will be used as the codomain of the router b(-), which will simply be an m-class classifier.
Desirable properties of the candidate set B are as follows:
4

---

Prompt-based Depth Pruning of LLMS
Step 2. Router Training (Section 5.2): Train router using the soft labes of candicate omission sets.
Step 1. Candidate Omission Set Generation (Section 5.1): Apply static degth pruning on diverse caliraaion sets, using a new task spealfi lass function.
Storage [Low-speed Memoryy
Candidate Ortisaion Sets
Trained
Transformar Encoder
27.48
.....O
5.10 11,11
Dense Model
IA 25
Holhoperd manrrry
User Input
"How do ready guinea pig for it's new cocupants""
Pruned Model
Figure 3. A visual overview of the proposed pipeline. Whenever the prompt is given from the user, a trained router predicts which set of blocks can be omitted with minimal loss, among the small number of candidate omission sets. Then, LLM transformer blocks are loaded from the storage (e.g., flash drive) to the high-speed memory (e.g., GPU RAM) except for the omitted blocks, saving the time and energy for the data communication. Finally, the depth-pruned model operates and generates the text.
5.2. Router Training
6. Experiment
After generating the candidate omission set B, we train a router that maps the given prompt to the best omission set. Roughly, this done by first constructing soft-labeled dataset with task-specific datasets and then training BERT- based router the constructed dataset (Devlin et al., 2019)
We now empirically validate the performance of the pro- posed algorithm, PuDDing, on zero-shot tasks against the static depth and width pruning baselines.
6.1. Experimental Setup
Dataset Construction. To construct the training dataset, we first collect various prompt-answer pairs from the task datasets, similarly to the calibration datasets in Section 5.1. Then, for each sample, we compute the task likelihood losses on all omission sets, and store them label vector. That is, each datum inside the dataset takes the form (xi; 8;), where is the prompt and the is length-m vector with (xI((x,)))) y:); w W'b'),... (l((x;, y:); W (Y))) (11)
Models. We evaluate the proposed method on compressing three popular open-weight language models. As the main model, we use the LLaMA-3. model with 8B parameters (Dubey al., 2024). In addition, we evaluate on two lan- guage models: Vicuna 1.5 with 7B (Chiang al., 2023), and OPT with 6.7B parameters (Zhang et al., 2022). We use these models for two reasons. First, the models have an appropriate scale for on-device deployments. Second, all three models consist of 32 transformer blocks, and thus can compared with the same sparsity criterion.
Note that we no longer need to store the correct answers y;-
Router training, We train a router to accurately predict the label vector given the input prompt x, for all samples in this dataset. That is, we train function fxx such that
Baselines.
We mainly compare against four recent static depth and width pruning baselines with open source.
SLEB (Song et al., 2024): A depth pruning baseline that iteratively selects the omission set based on perplexity.
S 2 holds. We the MSE loss
MSE(s, s) ||s $|
(12)
Shortend LLaMA (Kim al., 2024): A depth pruning algorithm which selects the omission set one-shot; here, we compare with the version that uses perplexity as the loss criterion and does not apply LoRA.
to train the router. At inference, will select the omission set with the minimum siainnum--cciiiinng index of the predicted s.
Router architecture. We use lightweight transformer- based encoder as our router. More specifically, we insert a single linear layer on pretrained BERT-base (Devlin al., 2019), and jointly fine-tune during the training. While this router has more parameters (~110M) than typical routers that are used for dynamic token routing- surhing--uchh D-LLM (Wang et al., 2024) which uses 2-layer MLP-the computa- tional cost bearable as we route only once per prompt. In our experiments, the routing cost typically takes up around 2 4% of the total fre-fill cost.
FLAP (An et al., 2024): A retraining ffree width pruning algorithm based structural fluctuation metric.
SliceGPT (Ashkboos et al., 2024): Another width pruning algorithm based principal component analysis.
In addition, we also compare with "SLEB per prompt," which is simply SLEB which is conducted by using each given prompt as the calibration data. As this option does not work well general, and requires a long inference time, we
5

---

Prompt-based Depth Pruning of LLMS
Table Zero-shot accuracy comparisons of PuDDing against baseline compression algorithms I...M-3.1 8B commonsense Reaoning tasks. The best performances are marked in bold, and the runner-up is marked with underline.
Pruned Blocks
Per-task Accuracies
Average Acc. (%)
Method
Structure
(Sparsity)
Are-C
Arc-E
BoolQ
HellaSwag
PIQA
WinoGrande
Donse
0
53.50
81.52
82.20
78.81
79.98
73.40
74.90
FLAP
Width
(20%)
26.54
46.80
62.32
46.93
64.58
58.56
50.96
SliceGPT
Width
(20%)
34.30
65.15
44.52
60.55
73.67
56.43
50.28
SLEB
Depth
7 (>>21%)
34.90
66.25
49.11
61.60
74.37
57.22
57.24
SLEB per prompt
Depth
7 (>-21%)
33.44
50.59
57.95
53.57
63.44
56.51
52.58
Shortened LJ aMA
Depth
7 (>=21%)
34.30
65.15
44.52
60.55
73.67
56.43
55.77
PuDDing (Ours)
Depth
(>>21%)
41.47
67.09
62.02
62,92
73.94
64.16
61.93 (+4.69)
FLAP
Width
-155%
33.19
60.02
69.45
58.18
71.16
61.88
58.98
SliceGPT
Width
(1.5%)
32.59
59.60
49.82
58.59
67.14
64.56
55.38
SLEB
Depth
5 (>115)
39.59
70.58
58.17
67.16
75.63
63.77
62.48
Shortened LLaMA
Depth
5 (>-15)
40.78
69.11
60.67
67.46
76.28
64.09
63.07
PuDDing (Ours)
Depth
5 (>115%
42.32
72.39
65.11
67.28
75.79
65.35
64.71 (41.64)
FLAP
Width
(10%)
36.43
66.20
69.69
63,29
74.10
66.61
62.72
SliceGPT
Width
(10%)
38.14
68.90
63.67
65.47
70.78
62.21
SLEB
Depth
3 (>9%))
45.73
76.01
68.93
71.96
77.53
68.98
68.19
Shortened LLaMA
Depth
3 (x99%)
38.57
69.91
69.72
71.28
77.31
67.48
65.71
PuDDing (Ours)
Depth
(>-9%)
48.98
77.02
70.18
73.26
77.20
68.11
69.13
tasks: ARC-Challenge, ARC-Easy, HellaSwag, PIQA, and WinoGrande. That is. we use total 10 omission sets (as we use two different losses). For training the router, we have used the full training splits. BoolQ dataset has been left out in order to evaluate the generalization to unseen sets. The router has been trained with Adam with learning rate 10-5, weight decay 0.01, and batch size 32 for 10 epochs, with 500 warm-up steps. Also, for WinoGrande dataset, use tailored data pre-processing procedure; describe this in detail in Appendix A.
Table Zero-shot task accuracy comparison LLaMA 3.1 8B, OPT 6.7B, and Vicuna 1.5 7B. The best performances are marked in bold, and the runner-up marked with underline. We have applied 20% sparsity (iee, pruned seven blocks).
Method
LLaMA 3.1
OPT 6.7B
Vicuna
Dense
74.90
62.51
70.49
FLAP
50.96
46.68
51.45
SliceGPT
50.28
55.45
59.11
SLEB
57.24
56.55
58.68
Shortened LLaMA
55.77
54.58
59.78
Hardware. We have mainly used NVIDIA RTX 6000 Ada for evaluation and training. addition, we have used cloud instances of NVIDIA A100 for evaluation.
PuDDing (Ours)
61.93
58.37
60.01
only compare on limited number of scenarios.
6.2. Main Experiment
Dataset: Evaluation. We evaluate the test splits of six zero-shot commonsense reasoning tasks: ARC-Challenge and ARC-Easy (Clark et al., 2018), HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2020). WinoGrande (Sakaguchi et al., 2020), and BoolQ (Clark et al., 2019).
Table provides comparison of zero-shot accuracies of the model compression methods, on LLaMA-3.1 8B model. From the table, we observe that PuDDing achieves the high- est average accuracy on all sparsity levels tested. Especially when 7 blocks have been pruned (over 20% sparsity), the improvement over the best baselines is almost 3%p.
Dataset: Calibration for Baselines. For the baseline algo- rithms, we have used the calibration data designated in the have
An interesting observation is the poor performance of "SLEB per prompt," which measures which block to remove on the fly, by using the given prompt as a calibration dataset. In fact, the performance is worse than the vanilla SLEB. We hypothesize that this is because single prompt usually does not contain enough information to work as good cali- bration data. Our training based strategy circumvents such
original papers. For SLEB, FLAP, and SliceGPT, used the WikiText-2 (Merity et al., 2022). For Shortened LLLMA, we have used the BookCorpus (7hu et al., 2015).
Training. To generate the candidate omission set for our algorithm, we have used 128 randomly drawn samples from the training splits of five zero-shot commonsense reasoning
6

---

Prompt-based Depth Pruning of LLMs
Table Zero-shot accuracy comparisons of PuDDing vs. other depth pruning methods I.LaMA-3.1 8B, with LoRA finetuning. The best performances are marked in bold, and the runner-up is marked with underline.
Pruned Blocks (Sparsity)
Per-task Accuracies
Average
Method
Acc. (%)
Are-C
Arc-E
BoolQ
HellaSwag
PIQA
WinoGrande
Dense
53.50
81.52
82.20
78,81
79.98
73.40
74,90
SLEB LoRA
7 (>221%)
45.39
74.92
69.05
70.92
78.35
64.40
67.17
Shortened LLaMA LoRA
(>-21%)
43.52
74.07
63.88
71.74
78.35
63.85
65,90
LLM-Streamline (w/ fine-tune)
(2>21%)
44,80
70.12
70.06
67.15
72.63
71.74
66.08
PuDDing + LoRA (Ours)
(>-21%%
45.39
75.34
71.96
71.58
77.26
66.54
68.01 (+0.84)
Table Accuracy comparisons of PuDDing vs, other depth prun- ing methods LLaMA-3.1 8B in the unseen tasks that require more complicated reasoning. The best performances are marked in bold, and the runner-up is marked with underline.
Table 6. Wall clock inference speed of the
LLaMA-3.1 8B evaluated on NVIDIA A100 and RTX 6000 Ada.
A
Pre-fil (TIFT)
Pre-fill Generation
Prompt Length
128
512
128
128
128
Pruned Blocks
Gen. Length
1
128
256
512
Metbod
OpenBookQA
MatQA
MMLU
PubMedQA
SciQ
Donse
0.137s
0.251s
0,505s
3.2968
6.634s
13.595
Dense
44.60
39,53
63,49
75.80
9600
PuDDDng
0.109
0.201s
0,393s
2.694s
5,375
11.024s
SLEB
25.19
23.76
56.40
20
-Router
+0.004s
0.000s
+0.008s
+0.004s
+0.004s
+0.004s
Shorlened LLaMA
7
31,20
25.76
26.78
$2,60
89.20
Speedup
1.21x
1.22x
1.23x
1.22x
1.23x
x
PuDDing
27.20
60.00
RIX 6000 Ada
Pre-fil (TTFT)
Pre-fil Generation
Prompt Length
128
256
512
128
128
128
Gen. Length
1
128
256
512
Dense
0.008s
0.171s
0.323s
4.923s
9.877s
19.973s
difficulty by training a router from the data.
PuDDing
0.069s
0.134s
0.260s
3.946s
7.926s
16.039s
-Rouler
+0.005
+0.005
+0.005
+0.005
+0.005s
+0.005
Regarding the of-distributton generalization, we observe that PuDDing also works well on unseen dataset (BoolQ). PuDDing outperforms all other baselines except for FLAP, which works extraordinarily well this specific dataset.
Speedup
1.19x
1,23x
1.22x
1,25x
1,25x
1.25x
6.4. More Complicated Tasks
In Table 3, we provide comparisons on other language mod- els: Vicuna and OPT. We confirm that our algorithm works better than other baselines under this setup as well.
In Table 5, we compare the performance of various depth pruning algorithms on more complicated tasks, including OpenBookQA (Mihaylov et al., 2018), MathQA (Amini et al., 2019), and MMLU (Hendrycks al., 2021). From the results, observe that PuDDing continues to perform better than the baselines, even though these tasks have not been observed during the training of the router.
6.3. LoRA Fine-tuning
Next, we compare the performance where we assume that we can recover the accuracies using LoRA (Hu et al., 2022). For PuDDing, we generate LoRA updates for each omis- sion set (thus total 10 for these experiments). This requires additional storage space for storing 10 separate copies of LoRA weights for each omission set. However, this increase only incurs ~2.5% increase in the total storage space. For training LoRA weights, we have followed the setup and hy- perparameters used for LoRA training in shortened LLaMA (Kim et al., 2024); we have used Alpaca dataset (Taori et al., 2023) for training, as in the paper.
7. Analysis
We now provide further analyses on PuDDDngg In particular, we provide the following analyses: Wall clock speedup (Section 7.1), and visualization of omission sets for tasks (Section 7.2). In Appendix B, we conduct ablation studies.
7.1. Wall-clock Speedup
We now provide wall-clock analyses and estimates on the latency and throughput of the PuDDing-comprest
Table 4 provides LoRA fine-tuned results of depth pruning algorithms on zero-shot commonsense reasoning tasks, for ALaMA-3.1-8P pruned to 20% sparsity. We observe that PuDDing continues to achieve the best performance among all options evaluated. That is, the advantages of prompt- adaptivity also exists after fine-tuning.
models.
Inference. Table 6 presents the average wall-clock inference time comparison between the dense and PuDDing-prund version of the LLaMA 3.1 8B, evaluated on NVIDIA A100 and RTX 6000 Ada. For PuDDing, we have pruned seven
7

---

Prompt-based Depth Pruning of LLMs
1.0
Arc-C-
Arc-E
0.8
BoolQ
0.6
HellaSwag
0.4
PIQA
0.2
WinoGrande
0.0
16
21
31
Figure 4. A visual illustration of the PuDDing's 's pruning rate of each transformer block, given the prompts drawn from various zero-shot tasks. The results are for the LLaMA 3.1 8B model, pruned to 20% sparsity (seven blocks removed), The color red indicates that the blocks are likely to be pruned, and the color green indicates that the blocks are likely be retained. We provide additional visualizations on the other language models (OPT 6.7B and Vicuna 1.5 7B) in the Appendix C.
7.2. Pruned Block vs. Task
Table Wall clock inference speed of the PuDDing-comprestee LLaMA-3.1 8B evaluated edge devices (M3 pro, Apple).
Figure 4 depicts the distribution of the pruned transformer blocks LLaMA-3.1 model, given the prompts from different tasks. Again, we consider the case where we drop seven transformer blocks for each prompt.
M3 Pro (Apple)
Pre-fill (TTFT)
Pre-fill Generation
Prompt Length
128
256
512
128
128
128
Gen. Length
1
128
256
512
Dense
0.177s
0.300s
0.480s
7,890s
15,970s
32.520s
PuDDing
0,138s
0.235s
0,376s
6.174s
12.497s
25,447s
From the figure, we make two intriguing observations: First, several blocks are considered almost universally unneces- sary. In particular, the blocks 20, 26, 27 are removed with over 80% probability in all tasks. Similarly, there are certain block which are almost never pruned, e.g., blocks 1-3 and 5-8. Second, regarding some blocks, the importance of the block highly varies over task. For instance, transformer block 4 is pruned with over 80% for ARC-Easy and ARC- Challenge. On the other hand, for PIQA and WinoGrande. the pruning rate is less than 40%; in these tasks, the blocks and are likelier to be less important.
Router
0,009s
0.016s
0.029s
0.009s
0.009s
0.009s
Speed Up
1.20x
1.20x
1.19x
1.28x
1.28x
1.28x
Table The estimated time required to transfer the weight param- eters of LI LLLMA-3.1 8B and PuDDing (with seven blacks pruned) to NVIDIA A100 GPU through various communication channels.
Bandwidth
Dense
PuDDing
PCIe Gen4 x4 NVIDIA NVlink
64GB/s
0.250s
0.198s
600GB/s
0.027s
0.021s
We note that similar patterns can be observed for OPT and Vieuna; Appendix C for visualizations on these models.
layers (21.8%% sparsity). We observe that PuDDing pro- vides a consistent 1.19-1.23x speedup during the pillfill stage, and 1.22-1.25x speedup including the generation stage. The total routing time takes up to 4-8ms, which can be deemed negligible comparing with the overall latency. Also, Table 7 presents results on edge devices (e.g., Apple M3 Pro), showing consistent speedup. This outcome shows that the proposed method is well-suited for both server-like and eige-like hardwares.
8. Conclusion
In this paper, we have developed a new paradigm for the depth pruning of large language models, where we dynami- cally determine which blocks should be utilized for process- ing the prompt given from the user. By doing so, we can save both the memory access cost and the inference com- putation, thus suitable for on-device deployment of large language models. We have proposed PuDDing, an algo- rithm to train a router using various task data. Through our experiments, we have confirmed that such framework is quite effective, clearly outperforming existing static depth pruning algorithms consistently over multiple LLMs.
Parameter loading, Table 8 presents the estimated time required for loading the model parameters of LLaMA-3.1 8B (16GB in FP32) from the storage to the GPU. PuDDing can save around 52ms on PCIe and 6ms on NVLink, which nonnegligibly large comparing with the computational scale of running these models. However, pitfall is that, for repeated inference, PuDDing may require loading additional weights to account for different prompts. This additional cost can be minimized by loading only the previously un- loaded blocks from the storage; in fact, many blocks overlap, as we will demonstrate in Section 7.2.
Limitations and future work. A notable limitation of the proposed method is that we assume that we have access to various task datasets. In particular, we have focused on the case where we use LLMs for commonsense reasoning tasks, instead of an open-ended language generation. promising future direction will be to develop new techniques
8

---

Prompt-based Depth Pruning of LLMs
to harness unlabeled text corpus, such as Alpaca or C4, to generate diverse clusters of calibration data for attaining corresponding omission sets.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
Another limitation is general lack of mechanisms to ac- count for the different difficulties of the tasks. For some tasks, it may be necessary to utilize all layers to generate an answer with sufficiently high quality; on the other hand, some tasks can be simply handled with very few layers. While our decision consider fixed number of trans- former blocks is motivated by the practical constraints of on-device inference, we believe that utilizzng variable-depth depth can be even more effective whenever the on-device memory is spacious but can be preempted to other processes.
Chiang. W.-L.,, Li, 7.,, Lin, 7.,, Sheng, Y., Wu, Z., Zhang, H., Zheng, L.,, Zhuang, S., Zhuang, Y., Gonzalez, J. E., Stoica, I., and Xing, E. P. Vicuna: An open-source chatbot impressing GPT-4 with 90%* ChatGPT qual- ity, March 2023. URL uttps:::: www.ss.rggy..l.......... 2023 03-30----0000000000 30-vicuna/.
Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In North American Chapter of the Association for Computational Linguistics, 2019.
Acknowledgements
This work has been supported by National Research Founda- tion of Korea (NRF) grants funded by the Korea government (MSIT) (Nos. RS- RS-2023-00213770,, RS-2024-00453301..
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. In arXiv preprint, 2018.
Impact Statement
Our paper targets for advancing the general field of machine learning and LLM compression. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for lan- guage understanding. In NAACL-HLT, 2019.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle,
A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Ilama 3 herd of models. In arXiv preprint, 2024.
References
Akhauri, Y., AbouElhamayed, A., Dotzel, J., Zhang, Z., Rush, A. M., Huda, S., and Abdelfattah, M. Shadowl. M.MM: Predictor- based contextual sparsity for large language models. In Empirical Methods in Natural Language Processing, 2024.
Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research, 23 23(12))1139, 2022.
Amini, A., Gabriel, S., Lin, S., Koncel- R., Choi, Y., and Hajishirzi, H. MathQA: Towards interpretable math word problem solving with Operation-Based for- malisms. In North American Chapter of the Association for Computational Linguistics, 2019.
He, S., Sun, G., Shen, Z., and Li, A. What matters in trans- formers? not all attention is needed. In arXiv preprint, 2024.
Hendrycks, D.,, Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D.,, and Steinhardt, J. Measuring massive multitask language understanding, In International Conference on Learning Representations, 2021.
An, Y., Zhao, X., Yu, T., Tang, M., and Wang, J. Fluctuation- based adaptive structured pruning for large language mod- els. In Association for the Advancement of Artificial In- telligence, 2024.
Hu, E. J., Wallis, P., Allen Z., Li. Y., Wang, S., Wang. L.,, Chen, W., et al. LoRA: Low-Rank adaptation of large language models. In International Conference on Learning Representations, 2022.
Ashkboos, S., Croci, M. L., do Nascimento, M. G., Hoefler, T., and Hensman, SliceGPT: Compress large language models by deleting rows and columns. In International Conference on Learning Representations, 2024.
Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E. Adaptive mixtures of local experts. Neural computation, 3(1):79-88) 87, 1991.
Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In Association for the Advancement of Artificial Intelligence, 2020.
Kim, B.-K., Kim, G., Kim, T.-H., Castells, T., Choi, S., Shin, J., and Song, H.-K. Shortened llaaa: simple
9

---

Prompt-based Depth Pruning of LLMs
depth pruning for large language models. In Proceedings of the ICLR Workshop on Memory-Ffficien Foundation Models (ME-FoMo), 2024.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. In arXiv preprint, 2023.
Lee, D., Lee, J.-Y., Zhang, G., Tiwari, M., and Mirhoseini, A. Cats: Contextually- aware thresholding for sparsity in large language models. In Conference On Language Modeling, 2024.
Wang, H., Xie, L.,, Zhao, H., Zhang, C., Qian, H., Lui, C., al. D-LLM: A token adaptive computing resource allo- cation strategy for large language models. In Advances in Neural Information Processing Systems, 2024.
Lin, J., Tang, J., Tang, H., Yang, S., Chen, W.-M., Wang, W.-C., Xiao, G., Dang, X., Gan, C., and Han, Awq: Activation-aware weight quantization for on-device Ilm compression and acceleration. In Proceedings of Machine Learning and Systems, 2024.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. HellaSwag: Can machine really finish your sentence? In Association for Computational Linguistics, 2019.
Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M.,
Liu, Z., Wang. J., Dao, T., Zhou, T., Yuan, B., Song, Z Shrivastava, A., Zhang, C., Tian, Y., Ré, C., al. Deja contextual sparsity for efficient LLMs at inference time. In International Conference Machine Learning, 2023.
Chen, S., Dewan, C., Diab, M., Li, X., Lin, V., et al. Opt: Open pre-trained transformer language models. In arXiv preprint, 2022.
Zhou, Y., Chen, Z., Xu, Z., Lin, X. V., and Chen, B. SIRIUS: Contexual sparisty with correction for efficient LLMs. In Advances in Neural Information Processing Systems, 2024.
Men, X., Xu, M., Zhang, Q. Wang, B., Lin, H., Lu, Y., Han, X., and Chen, W. Shortgpt: Layers in large language models are more redundant than you expect. In arXiv preprint, 2024.
Zhu, Y., Kiros, R., Zemel, R., Salakhutdinov, R., Urta- sun, R., Torralba, A., and Fidler, S. Aligning books and movies: Towards Story-Like visual explanations by watching movies and reading books. In International Conference on Computer Vision, 2015.
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. In International Conference Learning Representations, 2022.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can suit of armor conduct electricity? new dataset for open book question answering. In Empirical Methods in Natural Language Processing, 2018.
Raposo, D., Ritter, S., Richards, B., Lillicrap, T., Humphreys. P. C., and Santoro, A. -fxtureeof-Depthss Dynamically allocating compute in transformer-based language models. In arXiv preprint, 2024.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. In Association for the Advancement of Artificial Intelligence, 2020.
Siddiqui, S. A., Dong, X., Heinrich, G., Breuel, T., Kautz, deeper look at
J., Krueger, D., and Molchanov, depth pruning of LLMs. ICML 2024 Workshop on Theoretical Foundations of Foundation Models, 2024.
Song. J., Oh, K., Kim, T., Kim, H., Kim, Y., et al. Sleb: Streamlining Iims through redundancy verification and elimination of transformer blocks. In International Con- ference Machine Learning, 2024.
Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang. P., and Hashimoto, T. B. Stanford alpaca: An cnstructioo-followiiii llama model, 2023.
10

---

Prompt-based Depth Pruning of LLMs
A. Pre-processing for the WinoGrande Dataset
The WinoGrande dataset, originally consisting of fill-in-the-blank sentences, was initially computed using the sentence-level likelihood (sl) follows:
sl( W)
+ T
Σ log pi (xilzei W).
(13)
By reformulating the dataset into a Question-Answer format and evaluating the task likelihood (tl) score for the answer part using Equation (9), performance improved significantly, from 61. 11.09% 64.16%.
B. Ablation Studies
We have conducted various ablation studies on the proposed algorithm, PuDDing. Below, we provide summary of our key findings, with corresponding pointers to the relevant section.
Number of candidate omission sets (Appendix I): We have varied the number of candidate omission sets inside the set B, and find that having 10 classes is sufficient for handling zero-shot tasks; the gain from adding omission sets quickly saturates.
Proposed task likelihood score (Appendix B.2): We compare the performance of the task likelihood- basedd routing and the perplexity-basee based routing under both static and dynamic setups. We find that the using the task likelihood score leads to a clear advantage in both scenarios.
MSE loss for training (Appendix B.3): We have used the mean-squared error (MSE) loss to train the router using the soft labels. Our experiments show that this leads to a slightly better performance than using the classification loss, namely the cross-entropy loss.
B.1. Number of Candidate Omission Sets
Table Results of zero-shot task accuracy with varying omission set sizes.
Number of Omission Sets
Per-task Accuracies
Average
Acc. (%)
Arc-C
Arc-E
BoolQ
HellaSwag
PIQA
WinoGrande
5
40.27
67.85
61.01
62.60
73.83
61.56
61.19
10
41.47
67.09
62.02
62.92
73.94
64.16
61.93
30
38.57
66.88
63.70
64.23
72.85
63.93
61.69
Table 9 shows the accuracy of zero-shot task with varying omission set sizes, comparing the impact of using 5, 10, and 30 omission sets across common-sense reasoning tasks.
Using only 5 omission results in lower accuracy, with average of 61.19%, it shows insufficient performance for optimal results. 30-set configuration, despite some task-specific advantages, does not lead to consistently higher performance. In contrast, the 10-set configuration provides an improvement across multiple tasks, with highest average accuracy of 61.93%, indicating that this size offers better balance between performance and model efficiency.
B.2. Effectiveness of the Proposed Task Likelihood Score
Table 10. Zero-shot accuracy performance of static pruning methods.
Table 11. Zero-shot accuracy performance of dynamic routing pruning
methods.
Mathed
Merre
Tel--wis
Averge Acc. (%)
Mahool
Marrc
Omiscion Set Soletion
Roster Training
Average Are. (8)
SLED
Batch upl Batch-ppl
57.24
St
Prompt-ppl Batch ppl
per-pemppt
59.32
pre
(Oum)
likelibood
61.02
PaDDing (Ouns)
task libelhoodd ol
selected 10 sess
61.93
Both dynamic and static pruning methods were set up on LLaMA-3.1 8B, with a sparsity of 20%, corresponding to the pruning of seven blocks.

---

Prompt-based Depth Pruning of LLMS
Static pruning experiments in Table 10 show the experimental validation of task-adaptive pruning and our proposed / scoring method. First, we compare two pruning strategies using batch-PPL.: one that applies fixed omission set based Wikitext-calibraaee batch (57.24% accuracy) and another that dynamically selects omission sets per task (59.32% accuracy). The improvement confirms the claim that different tasks require different layer sets. Next, keeping the task-wise adaptive setting constant, we replace batch-PPI. with our task likelihood (il) loss for omission set selection. This further improves accuracy from 59.32% to 61 .0.02 demonstrating that our method is more effective at identifying layers that impact quality of task-specific inference.
By comparing Batch-PPL Router (58.19%) with PuDDing (Task Likelihood Loss) (61.93%) Table 11, we observe that our tl metric results in omission sets that generalize better across tasks, contributing additional performance gain. When comparing these Table 11 results with Table 10, Batch-PPL static pruning (57.24%) improves with router training (58. 188.1 and our method (60.12% also benefits from router training, increasing to 61.93%. This validates that even within the same task, different prompts favor slightly different omission sets, and adapting omission sets dynamically through router training is essential for optimal pruning.
B.3. Using MSE Loss for Training
Table 12. sero-shot accuracy comparison between three different training strategy the router. The best performances are marked in bold, and the runner-up is marked with underline.
Per-task Accuracies
Average Acc. (%)
Label
Loss
Arc-C
Are-E
BoolQ
HellaSwag
PIQA
WinoGrande
One-hot Vector
CB
41.21
66.29
59.66
61.44
72.96
59.35
60.15
Log-likelihood
CE
39.67
67.80
61.77
60.80
73.56
59.04
60.44
Log-likelihood
MSE
41.47
67.09
62.02
62.92
73.94
64.16
61.93
In Table 12, we present the result of zero-shot task accuracies in the different router training settings. For the label, one-hot vector signifies that the router learns only from the highest confidence value within the omission block set derived from the training dataset. In contrast, the log- likelihhoood label allows the router incorporate all confidence values during training. Our findings show that training with log-likeli likkliheeood label leads to improved average accuracy (from 60.16% to 61 .93%). Hence, we observe that the mean-squared error (MSE) loss function outperforms cross entropy (CE). As result, in this case for routers to train for finding optimal omission sets by given prompts, richer information in the label (i.e., un-chosen labels are not assigned to zero values), and employing MSE loss enhances better performance.
B.4. Training the Router with Multi-Domaie Calibration Datasets
Table 13, Zero-shot accuracy comparison of dense, PuDDing, and PuDDing-MultiDomnn using diverse calibration datasets.
Average Acc. (%)
Per-task Accuracies
Methad
Arc-C
Arc-E
BoolQ
HellaSwag
PIQA
WinoGrande
MathQA
PubMedQA
SciQ
Dense
73.42
53.50
81.52
82.20
78.81
79.98
73.40
39.53
75.80
96.00
PuDDing
61.28
41.47
67.09
62.02
62.92
73.94
64.16
27.27
60.00
92.70
LuDDing--ultiionng 62.37
41.38
67.26
67.37
63.68
73.07
64.56
29.58
62.00
92.40
In Section 6, we construct omission sets using common-sense reasoning datasets that are widely used to evaluate language models' reasoning ability-ARC. PIQA, HellaSwag, and WinoGrande- to reflect popular usage scenarios. For the new variant,
we conduct an additional experiment by incorporating calibration datasets from diverse domains: MathQA (mathematics), PubMedQA (biomedical), and SciQ (science). To maintain consistency in sample size with the original setting, include only subset of the original common-sense datasets specifically, ARC-Easy and WinoGrande-in this experiment. All other experimental settings follow the same setup described in Section 6.1.
We observe that with noticeable gains in the newly introduced datasets (e.g.. +2.31 in MathQA, +2.00%p in PubMedQA). This indicates that the router can benefit from more diverse calibration dataset, especially when deployed in tasks requiring domain -eeaii- knowledge. Overall, our method demonstrates consistent performance across both in-domain and doman tasks, and
achieves slightly improved average accuracy compared to the original PoDDing,
12

---

Prompt-based Depth Pruning of LLMs
be further customized to specific applications by adjusting the coverage of the omission pool.
B.5. Quantized PuDDing
Table combined with other compression techniques
LLaMA-3,1-8B
Average
Arc-C
Arc-E
BoolQ
HellaSwag
PIQA
WinoGramde
Dense
74.90
53.50
81.52
82.20
78.81
79.98
73.40
SLEB
57.24
34,90
66.25
49.11
61.60
74.37
$7.22
Shortened LLaMA
55.77
34.30
65.15
44.52
60.55
73.67
$6.43
PuDDing
61.93
41.47
67.09
62.02
62.92
73.94
64.16
PuDDing W8A16 (AWQ)
61.68
41.30
67.00
61.50
62.95
73.72
63.61
PuDDing W4A16 (AWQ)
58.58
37.37
61.45
60.64
57.55
71.71
62.75
In Table 14 presents the performance of the proposed method when combined with a representative compression technique
such as quantization (e.g. Activation-aware Weight Quantization (AWQ) (Lin et al., 2024)). Interestingly, we observe that compressing the weights to 8-bit results no performance degradation. Although 4-bit quantization introduces some degradation, the performance still surpasses that of static pruning methods applied with BF16 precision.
C. Additional Visualizations
Figure 5 illustrates the dynamic block selection process in various tasks, highlighting that this process has also been analyzed with different models to highlight how the block selection strategy varies not only varying tasks but also depending on the specific architectures of the models.
Arc-C
Arc-E
BoolQ
0.6
HellaSwag
0.4
PIQA
0.2
WinoGrande
0.0
16
31
(a) OPT 6.7B
Arc-C-
Arc-E
0.8
BoolQ
0.6
HellaSwag
0.4
PIQA
0.2
WinoGrande
0.0
11
16
21
26
31
(b) Vicuna 1.5 7B
Figure 5. A visual illustration of the PuDDing's pruning rate of each transformer block, given the prompts drawn from various zero-shot tasks. The results are for the OPT 6.7B model vicuna 7B, pruned 20% sparsity (seven blocks removed). The color red indicates that the blocks are likely to be pruned, and the color green indicates that the blocks are likely to be retained.
13


Reviews:
[{"review_id": "dq0zW2XwNk", "reviewer": "Reviewer_2qCi", "summary": "This paper introduces PuDDing, a method to reduce LLM inference costs by skipping transformer layers on a per-input basis. The motivation is that different tasks or queries may not require all layers of a deep model. PuDDing consists of two main components: (1) a procedure to generate a small set of candidate omission sets, and (2) a lightweight router network that, given a new prompt, predicts the best omission set to use. The authors train the router by creating a training dataset of prompts and the optimal omission decisions for those prompts, found via a data-driven search (they evaluate different layer-drop combinations on a set of training prompts to see which yields minimal loss). Once trained, the router can generalize to new prompts. On several commonsense reasoning benchmarks (ARC, PIQA, WinoGrande, BoolQ, etc.), a PuDDing-pruned LLaMA-3.1-8B model at 20% sparsity outperforms equivalent static-pruned model.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer 2qCi,  \nThank you for your insightful comments and suggestions. In what follows, we respond to the points raised by the reviewer one-by-one.\n\n\n---\n\n### **The scope of the task is too narrow, all being the commonsense reasoning tasks**. \n**TL;DR. We have already evaluated PuDDing on MMLU/OBQA/MathQA (Table 5), and newly added experiments on PubMedQA/SciQ.**  \nFirst, 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.\nTo 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.\n\n|LLaMA-3.1-8B|MMLU|MathQA|PubMedQA|sciQ|OpenbookQA|\n|:-|:-:|:-:|:-:|:-:|:-:|\n|Dense|63.49|39.53|75.80|96.00|44.60|\n|SLEB|23.76|25.19|56.40|89.20|36.00|\n|Shortened LLaMA|26.78|25.76|52.60|89.20|34.20|\n|**PuDDing**|**39.00**|**27.27**|**60.00**|**92.70**|**36.40**|\n\n---\n\n### **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.**  \n\nOur point is twofold:\nFirst, 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](https://ibb.co/zTVMKrBL)**, 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.  \nSecond, 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.  \n\n| Method     | Average | Arc-C | Arc-E | BoolQ | HellaSwag | PIQA | WinoGrande | MathQA | PubMedQA | sciQ |\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n| Dense| 73.42| 53.50 | 81.52 | 82.20 |78.81|79.98|73.40| 39.53|75.80|96.00|\n| PuDDing| 61.28| 41.47 | 67.09 | 62.02 |62.92| 73.94 |64.16|27.27|60.00|92.70|\n| newPuDDing | 62.37  | 41.38 | 67.26 | 67.37 |63.68| 73.07 |64.56| 29.58  |62.00| 92.40|\n\n---\n\n### **It fairly covers most literature, but it might be good to cover a few more baselines such as LaCo, LLM-streamline, and FinerCut.**  \n\nFollowing 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.\nUnfortunately, 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.  \n\n| LLaMA-3.1-8B   | Average | Arc-C | Arc-E | BoolQ | HellaSwag | PIQA  | Winogrande |\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n| Dense| 74.90| 53.50 | 81.52 | 82.20 | 78.81| 79.98 | 73.40|\n| LLM-Streamline (w/ fine-tune)| 66.08| 44.80 | 70.12 | 70.06 | 67.15| 72.63 | **71.74**|\n| PuDDing (w/ LoRA fine-tune)| **68.01**| **45.39** | **75.34** | **71.96** | **71.58** | **77.26** | 66.54|\n---\nWe sincerely hope that you find our response reasonable. Please don’t hesitate to let us know if there are any further questions.\n\nBest regards,  \nAuthors", "claims_and_evidence": "The paper claims that prompt-adaptive layer pruning yields better task performance than static pruning for the same speedup, and that it achieves a meaningful speedup over the dense model. The evidence supports these claims yet the scope of the task is too narrow, all being the same commonsense reasoning tasks. In their experiments, PuDDing consistently achieved the highest mean accuracy among various pruning strategies when tested on zero-shot commonsense QA tasks at 9~20% layer sparsity​. For example, Table 4 shows PuDDing surpassing other baselines on ARC, HellaSwag, etc., after fine-tuning with LoRA. As for speed, they measure actual wall-clock time on GPUs: PuDDing yields about 1.21–1.25× speedup in different settings, with the router’s overhead being minimal. Overall, the evidence is convincing that PuDDing meets its goals: it yields a speedup and better accuracy than static pruning baselines, confirmed by multiple benchmarks and metrics.", "methods_and_evaluation_criteria": "The paper focuses on accuracy on NLP benchmarks and actual inference speed. They evaluate on a suite of commonsense reasoning tasks (ARC-Easy/Challenge, HellaSwag, PIQA, WinoGrande, BoolQ), which are challenging tasks that can benefit from the full depth of a model. \\\nThat said, the main drawback of the proposed method is construction of omission sets, while I acknowledge the authors already discussed this in the final section of the paper. The authors need to devise an experimental setting where omission sets largely differ, so that prompts drop notably different layers. There can be math, coding, medical, law tasks. From Figure 4, commonsense reasoning tasks seem to drop approximately the same layers which degrades the need for the proposed method. In case the omission sets overlap much, I wonder whether we can choose to have a single omission set per task via constructing a task-representation vector instead.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experiments and ablation studies are well-designed to isolate the benefits of PuDDing. The authors also test across different model architectures (LLaMA-based, Vicuna, OPT). The speed analysis is detailed as well – breaking down the time into pre-fill and generation phases along with router inference time. This gives a complete picture of how the method performs in practice.", "relation_to_broader_scientific_literature": "N/A.", "essential_references_not_discussed": "It fairly covers most literature, but it might be good to cover a few more baselines such as LaCo, LLM-streamline, and FinerCut.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A\n\n--------------------After Rebuttal--------------------\n\nI have raised my score from 2 to 3, and lean towards acceptance, only if the authors faithfully include new experimental results in the final manuscript.", "questions_for_authors": "N/A", "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": "XJttLdyHDT", "reviewer": "Reviewer_r1RD", "summary": "This paper introduces Prompt-routed Dynamic Depth Pruning (PuDDing), a method for dynamically pruning transformer blocks in LLMs based on the input prompt. The core motivation is that the importance of transformer layers is task-dependent, making static pruning suboptimal. To address this, PuDDing trains a lightweight router that predicts the optimal omission set for each prompt, reducing inference costs while maintaining accuracy.\n\nThe proposed approach consists of two key steps:\n\n1) Candidate Omission Set Generation: A small, diverse set of omission strategies is precomputed using a new task likelihood (TL) loss, which improves on traditional perplexity-based pruning metrics.\n2) Router Training: A BERT-based classifier is trained to predict the best omission set for a given prompt, ensuring minimal accuracy loss while improving efficiency.\n\nEmpirical results demonstrate that PuDDing outperforms static depth pruning methods (e.g., SLEB, Shortened LLaMA) by up to 4% in zero-shot accuracy on commonsense reasoning tasks (ARC, BoolQ, WinoGrande). It also achieves a 1.2× inference speedup compared to the unpruned model. The method is particularly designed for on-device inference, as it loads only the necessary transformer layers from storage, reducing memory requirements.\n\nOverall, the paper contributes to model compression and efficient inference research by demonstrating that depth pruning can be made adaptive to input prompts, improving efficiency without requiring extensive retraining or hardware-specific optimizations.\n\n## Update after rebuttal\nThe authors have done a good job by adding experiments and analysis during the rebuttal which strengthens the paper but I still believe that the idea is not very novel and might need more ingenuity to increase it's impact. Also, the gains are marginal thereby making this technique not useable on it's own. I would suggest the authors to come up with a truly dynamic approach which might help in improving the speedup gains. Good Luck to the authors", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer r1RD,  \nThank you for your constructive feedback. In what follows, we respond to the points raised by the reviewer one-by-one.\n\n---\n\n### **No Real-World Deployment Results (...) Missing mobile/edge device benchmarks weakens the practical relevance of the method.**  \nThank 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++).  \n||Pre-fill (TTFT)|||Generation|||\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Prompt Len|128|256|512|128|256|512|\n|Gen.Len| 1 | 1 | 1 |128|256|512|\n|Dense|0.177|0.300|0.480|7.890|15.970|32.520|\n|PuDDing|0.138|0.235|0.376|6.174|12.497|25.447|\n|Router|0.009|0.016|0.029|0.009|0.009|0.009|\n|Speedup|1.20×|1.20×|1.19×|1.28×|1.28×|1.28×|\n\nWe will add this result in the revised manuscript. \n\n---\n\n### **How does the dynamic loading of omission sets affect inference latency?**  \nAs 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.\nWe will add the relevant discussion to the revised version.\n\n---\n\n### **The reported 1.2x inference speedup is not very high compared to other compression techniques.**\nFirst, 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).\nSecond, 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.  \n|LLaMA-3.1-8B|AVG|Arc-C|Arc-E|BoolQ|HellaSwag|PIQA|Winogrande|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n|Dense|74.90|53.50|81.52|82.20|78.81|79.98|73.40|\n|SLEB|57.24|34.90|66.25|49.11|61.60|74.37|57.22|\n|Shortened LLaMA|55.77|34.30|65.15|44.52|60.55|73.67|56.43|\n|PuDDing|61.93|41.47|67.09|62.02|62.92|73.94|64.16|\n|PuDDing+ w8a16 (AWQ)|61.68|41.30|67.00|61.50|62.95|73.72|63.61|\n|PuDDing + w4a16 (AWQ)|58.58|37.37|61.45|60.64|57.55|71.71|62.75|\nWe will add the relevant discussion to the revised version.\n\n---\n\n### **The pruning method is not truly dynamic.**\nWe 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.”  \n\n---\n\n### **Fixed omission set reduces flexibility.**\nWhile 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.\nFor a more detailed answer, please refer to our ***answer #4 to the reviewer KB85.***  \n\n---\n\n### **Some of the prior work that author could have included (...) How does PuDDing compare to other dynamic pruning methods like MoD or D-LLM?**\nWe 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.  \n\n---\n\n### **Why TL loss is a better metric.**\nThe 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.  \n\n---\n\n### **How well does the router generalize to unseen tasks or domains.**\nWe 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.***  \n\n---\n\n### **Other comments.**\nUnfortunately, 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.  \n\n---\n\nWe sincerely hope that you find our response reasonable. Please don’t hesitate to let us know if there are any further questions.  \n\nBest regards,  \nAuthors", "claims_and_evidence": "Below is an evaluation of key claims:\n\n1) Claim: Transformer block importance is task-dependent.\n- Evidence: Section 3 presents empirical results showing that pruning different blocks affects different tasks differently.\n- Support: Figure 2 demonstrates how removing specific layers leads to varying accuracy drops across datasets.\n- Weakness: While the empirical evidence is strong, the paper lacks theoretical justification for why this happens at a structural level.\n\n2) Claim: PuDDing achieves inference speedup.\n- Evidence: Table 6 shows a 1.2x speedup in inference over unpruned models.\n- Support: The method reduces the number of active layers, which naturally leads to speed improvements.\n- Weakness: The speedup is modest compared to other pruning methods (e.g., structured sparsity, quantization).\n\n3) Claim: PuDDing is a fully dynamic pruning method.\n- Issue: The method selects from precomputed omission sets rather than dynamically pruning layers on a per-query basis. A true dynamic pruning method would adjust pruning layer-by-layer rather than choosing from a fixed set.\n\n4) Claim: Task likelihood (TL) loss is superior to perplexity for pruning.\n- Issue: The paper provides empirical comparisons (Table 9, Table 10) but does not explain why TL loss is theoretically better. A mathematical argument or formal connection to task complexity would strengthen this claim.", "methods_and_evaluation_criteria": "1) Relevant Benchmark Datasets for Task-Specific Evaluation\n- The paper evaluates PuDDing on six widely used commonsense reasoning datasets (ARC, BoolQ, PIQA, WinoGrande, HellaSwag).\n- Since the goal is to show that different tasks require different pruning strategies, these benchmarks make sense.\n\n2) Fair Comparisons Against Strong Baselines\n- The method is compared against multiple pruning techniques, including: Static Depth Pruning (SLEB, Shortened LLaMA)\nWidth Pruning (FLAP, SliceGPT)\n- Multiple sparsity levels (10%, 15%, 20%) are tested, ensuring fairness in evaluation.\n\n3) Ablation Studies for Router Training\n- The paper evaluates different loss functions (MSE vs. CE) for the router, showing that MSE improves generalization (Table 11).\n- The number of candidate omission sets is also analyzed (Table 8).\n\n4) No Real-World Deployment Tests\n- The paper claims PuDDing is suitable for on-device inference, but does not test it on resource-limited devices.\n- All experiments are done on NVIDIA A100 / RTX 6000 GPUs, which do not reflect real-world constraints of mobile or edge hardware.", "theoretical_claims": "The paper does not present any formal theoretical proofs. Instead, it relies entirely on empirical evidence to support its claims.\n\nTheoretical justification cane be provided for following:\n1) why TL loss is a better metric\n2) how well the router generalizes to unseen tasks, given that it is trained on specific datasets", "experimental_designs_or_analyses": "The experimental design is mostly sound. Below is an analysis of the strengths and few weaknesses of the experiments.\n\n1) Comprehensive Benchmarking on Commonsense Reasoning Tasks\n- The paper evaluates PuDDing on six widely used benchmarks (ARC, BoolQ, PIQA, WinoGrande, HellaSwag).\n- Multiple pruning baselines (SLEB, Shortened LLaMA, FLAP, SliceGPT) are included for fair comparison.\n- Accuracy is measured consistently across different pruning levels (10%, 15%, 20%), ensuring robustness.\n\n2) Fair Comparisons with Static Pruning Methods\n- The experimental setup controls for model size and number of pruned layers, ensuring a fair comparison between static and dynamic pruning.\n- The study includes ablation experiments on router training methods (Table 11) and omission set selection (Table 8), adding depth to the analysis.\n\n3) Inference Speed Evaluation\n- Table 6 reports wall-clock speed improvements, showing a 1.2× speedup over unpruned models.\n- The memory efficiency claim is supported by parameter transfer time comparisons (Table 7).\n\n4) No Real-World Deployment Results\n- The paper claims PuDDing is well-suited for on-device inference, but all experiments are conducted on high-end GPUs (A100, RTX 6000).\n- Missing mobile/edge device benchmarks (e.g., ARM-based chips, Jetson devices) weakens the practical relevance of the method.", "relation_to_broader_scientific_literature": "The paper is related to prior work in model compression, pruning techniques, and adaptive computation for LLMs.\n\n1) Static Depth Pruning: SLEB (Song et al., 2024), Shortened LLaMA (Kim et al., 2024) perform static depth pruning, removing less important transformer blocks based on perplexity or activation-based importance metrics. PuDDing extends these methods by introducing prompt-conditioned pruning decisions, making depth pruning task-adaptive rather than fixed.\n\n2) Width Pruning & Sparsity Methods: FLAP (An et al., 2024), SliceGPT (Ashkboos et al., 2024) apply structured width pruning, reducing parameters in weight matrices. PuDDing focuses on depth pruning instead of width pruning, making it hardware-agnostic.", "essential_references_not_discussed": "Some of the prior work that author could have included:\n\n- Mixture-of-Depths (MoD) (Raposo et al., 2024) introduces adaptive layer skipping, where transformer blocks are selectively skipped based on token-level routing decisions.\n- D-LLM (Wang et al., 2024) also uses a router for adaptive pruning, but dynamically decides layer usage per token.\n\nMoD and D-LLM perform token-level routing (more flexible, but computationally expensive). PuDDing precomputes a small set of pruning strategies and selects one per prompt (faster but less flexible).", "other_strengths_and_weaknesses": "Strengths\n\n1) Practical and Computationally Efficient Approach\n- PuDDing introduces a lightweight routing mechanism that selects omission sets once per prompt, reducing computational cost compared to token-level routing methods like MoD or D-LLM.\n- Unlike fine-grained dynamic pruning, PuDDing does not require per-token decisions, making it more efficient for real-time inference.\n\n2) Strong Empirical Performance on Task-Specific Pruning\n- The experiments show that task-dependent pruning improves accuracy over static pruning methods by up to 4% in zero-shot commonsense reasoning tasks.\n- This reinforces the idea that layer importance varies by task, a key insight for adaptive pruning.\n\n3) Hardware-Agnostic Design\n- Unlike width pruning (which often requires hardware-specific optimizations), PuDDing’s depth pruning approach can be used on any hardware without requiring changes to matrix sparsity patterns.\n\n4) Good Paper Organization and Clarity\n- The paper is well-structured, with clear motivation, experimental design, and results.\n- The figures (e.g., Figure 2 on transformer block importance) effectively communicate key insights.\n\nWeaknesses\n\n1) Limited Novelty – More of an Incremental Contribution\n- While the idea of prompt-aware pruning is new, it builds upon static pruning (SLEB) and dynamic routing (Mixture-of-Depths, D-LLM) but does not introduce a fundamentally novel pruning algorithm.\n\n2) Fixed Omission Set Reduces Flexibility\n- The pruning method is not truly dynamic—it selects from a small precomputed set of omission strategies, instead of pruning layer-by-layer in real-time. A more adaptive method would dynamically select layers per query, instead of relying on precomputed omission sets.\n\n3) No Real-World Deployment or Mobile Testing\nThe paper claims that PuDDing is suitable for on-device inference, but all experiments are done on high-end GPUs (A100, RTX 6000).\nMissing evaluation on real constrained devices (e.g., ARM CPUs, mobile GPUs) makes this claim untested.\n\n4) Relatively Modest Speedup (1.2×)\nThe reported 1.2× inference speedup is not very high compared to other compression techniques (e.g., structured sparsity, quantization, MoE models).", "other_comments_or_suggestions": "The paper is well-organized and well-written. There are minor typos:\n\"langauge\" --> \"language\" (Section 5)\n\"calibraion” --> “calibration\" (multiple instances)", "questions_for_authors": "1) How does PuDDing compare to other dynamic pruning methods like Mixture-of-Depths (MoD) or D-LLM?\n2) How well does the router generalize to unseen tasks or domains?\n3) The method selects a pruning strategy per prompt, but this means that different layers may need to be loaded dynamically from storage. How does the dynamic loading of omission sets affect inference latency?\n4) Can PuDDing be applied beyond transformer models?", "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": "gp41qXqaiI", "reviewer": "Reviewer_KB85", "summary": "This paper propose a input-dependent depth pruning method for LLMs. Unlike existing static model pruning methods where the LLM is pruned into a small subnetwork and apply it for all testing examples, this paper considers select differently pruned sub-networks for inference depends on the task that the testing example belongs to.\n\nThe algorithm is very straightforward and contains the following parts: 1. depth-pruned network candidates (\"candidate omission set generation\" called in thiss paper). The authors first collect a series of sub-networks by perform depth-pruning of the LLMs on a set of different datasets. During testing, a router will select one of these sub-networks (each represented as a set of indices of the pruned layers) for each input example. 2. router training. Selecting sub-networks is formulated into a regression problem followed by argmax operation. The authors will pair each example in the training dataset with the loss values given each sub-network. The router is trained to predict these loss values given the input example. 3. Inference. During inference, the router first predict the losses of each sub-network, then the sub-network with the minimum predicted loss is selected for inference.\n\nThe author demonstrate the proposed method outperforms other depth pruning methods. By pruning ~20% parameters, the proposed method brings 13% performance degradation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer KB85,  \nThank you for your constructive and thoughtful feedback.  In what follows, we address the concerns you raised one-by-one.  \n\n---\n\n### **Missing paper on dynamic model pruning [1]**.    \n\nThank 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.\n\nWe note that [1] differs from our work in several key aspects, making two approaches complementary to each other.\n- **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.\n- **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.\nOverall, 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.\n\nKim et al., “Shortened LLaMa: A simple depth pruning for large language models,” arXiv 2024.\n\n---\n\n### **This paper only achieves 20% sparsity while hurting the performance by more than 10%. (…) The performance degradation is still significant, comparing with dynamic pruning.**\n\nAs 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.\n\n---\n\n### **Missing discussion on contextual sparsity**.\n\nThank 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).\n\n---\n\n### **Method design 1: Why not introduce a learnable network that predicts a mask for each layer?**\n\nWe 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.\n\n---\n\n### **Method design 2: Why not directly train a classifier?**\n\nBy 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.\n\nEmpirically, 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.\n\nHinton et al. “Distilling the knowledge in a neural network,” arXiv 2015\n\n---\n\nWe sincerely hope that you find our response reasonable. Please don’t hesitate to let us know if there are any further questions.\n\nBest regards,  \nAuthors", "claims_and_evidence": "The claims made in the submission is well supported by empirical results. The authors first present an empirical test regarding the importances of transformer blocks given different inputs in Figure 2. This verify the correctness of this paper's focus: the model parameters could be input-dependent, and we can only select the necessary parameters (layers) for different inputs.", "methods_and_evaluation_criteria": "The method is generally reasonable and straightforward. The evaluation datasets are also proper. I only have several questions regarding on the method design, where I think the following methods could be better than the proposed method:\n\n- Method design 1: Why not introduce a learnable network that predicts a mask $m_i\\in {0,1}$ for each layer $i$ ? This is a typial method for structured model pruning. You can directly train the parameters of this learnable network while fixing the model to be pruned by optimization the following objective function:\n\n$$\\ell = \\mathcal{L}_{\\text{CE}}[f(\\boldsymbol x; \\boldsymbol m \\cdot \\boldsymbol\\theta), \\boldsymbol y]$$,\n\nwhere $\\boldsymbol m = g(x, \\boldsymbol \\phi)$ is the output of the learnable network with parameter $\\boldsymbol \\phi$, $m\\cdot \\theta$ means we mask the layers of the LLM's parameter $\\theta$ if $m_i = 0$ (otherwise we keep it). This is more straightforward and allows you to select sub-networks with more freedom, instead of only have 10 different subnetworks for selection.\n\n- Method design 2: why not directly train a classifier? The authors propose to train the router with the regression objective and predict the loss values with each sub-network. I wonder the rationale behind such a design.", "theoretical_claims": "N/A. This paper is mainly about empirical findings.", "experimental_designs_or_analyses": "The experiment design is reasonable. I do not find significant issues.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "Are there related works that are essential to understanding the (context for) key contributions of the paper, but are not currently cited/discussed in the paper? Be specific in terms of prior related findings/results/ideas/etc.\n\nThe key contribution of this paper is a dynamic pruning method. However, a very similar paper is missing. Also, although the authors formulate the story of this paper as \"model pruning\", this paper actually closely aligns with the research of **contextual sparsity**.\n\n1. Missing paper on dynamic model pruning [1]. This paper also focuses on dynamic pruning and select the sub-network given different inputs. Furthermore, the authors in [1] demonstrate their dynamic pruning methods only brings less than 2% performance degradation on the same datasets evaluated in this paper's experiments while achieves up to 75% sparsity. This paper only achieves 20% sparsity while hurting the performance by more than 10%.\n2. Missing discussion on **contextual sparsity**. This paper's \"dynamic pruning\" can be interpreted as contextual sparisty, whether the sub-network is selected based on the input during inference. Typical contextual sparsity methods [2-5] also do not fine-tune the LLM. They collect the parameter activation patterns on training dataset and train a router to predict these activated parameters, which is very similar to the proposed method in this paper.\n\nIt is necessary to include these papers into the related work section and discuss the differences.\n\n[1] Hou, Bairu, et al. \"Instruction-Following Pruning for Large Language Models.\" *arXiv preprint arXiv:2501.02086* (2025).\n\n[2] Liu, Zichang, et al. \"Deja vu: Contextual sparsity for efficient llms at inference time.\" International Conference on Machine Learning. PMLR, 2023.\n\n[3] Akhauri, Yash, et al. \"Shadowllm: Predictor-based contextual sparsity for large language models.\" arXiv preprint arXiv:2406.16635 (2024).\n\n[4] Zhou, Yang, et al. \"Sirius: Contextual sparsity with correction for efficient llms.\" arXiv preprint arXiv:2409.03856 (2024).\n\n[5] Lee, Donghyun, et al. \"Cats: Contextually-aware thresholding for sparsity in large language models.\" arXiv preprint arXiv:2404.08763 (2024).", "other_strengths_and_weaknesses": "Please refer to the **Questions For Authors**. The main weaknesses include the poor performance, missing references, similarity to contextual sparisty methods, and method design.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. The discussion with missing relatedd work (see **Essential References Not Discussed**)\n2. The performance degradation is still signficiant. From my perspective, it is easy to propose a method that brings a better tradeoff between sparisty and performance than baselines, but it is hard (and more important) to propose a method that brings a **usable** tradeoff. The current method occurs more than 10% performance degradation while achieves 20% sparsity, which is far away from usable for modern LLMs. The poor performance (although still better than baselines) makes this paper less competitive.\n3. Why do not adopt the method design 1 (see **Methods And Evaluation Criteria**) which is very straightforward and familiar to people in the research area of model pruning?  I think that method could bring better performance given the more freedom and expressivity. It is also a commonly used method for structured model pruning [1]. \n\n[1] Xia, Mengzhou, et al. \"Sheared llama: Accelerating language model pre-training via structured pruning.\" *arXiv preprint arXiv:2310.06694* (2023).", "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": "UOSap4DURc", "reviewer": "Reviewer_9fAo", "summary": "This paper propose prompt-based depth pruning of Large Language Models, where given a prompt, a router is trained to select a best set of LLMs layers, and other layers will be pruned.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer 9fAo,\n\nWe appreciate your positive evaluation on our paper, acknowledging the practical advantages of our method. We respond to the concerns you raise in what follows.\n\n---\n\n### **It would be better to show the training convergence analysis of the router.**\n\nFollowing your suggestion, we have attached the loss curves of the router training in the following [LINK (training loss)](https://ibb.co/VYsX2gYf ), and [LINK (test loss)](https://ibb.co/5XkTWZpq). The results confirm that the test loss of the router has successfully converged by the end of the training.\n\nWe will add the plot in the revised manuscript.\n\n---\n\nWe hope that you find our response reasonable. Please let us know if you have any further questions.\n\n\nBest regards,  \nAuthors.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "NA", "experimental_designs_or_analyses": "No issues", "relation_to_broader_scientific_literature": "This work contributes to efficient LLM deployment.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strength: This paper first show supportive observations that the effectiveness of  LLMs internal layers are prompt dependent, which give empirical evidence that some layers can be disregard and abandoned for a specific prompt.\n\nWeakness: It would be better to show the training convergence analysis of the router", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "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: The paper introduces PuDDing, that skips some of the layers during inference. The key idea is that the importance of a specific layer can be prompt dependent. Hence, PuDDing introduces a small and quick router model to decide which pre-selected layers to use based on the input. The simple and well-designed method is able to provide real speedups without needing special setups. The experiments also show that it outperforms removing the same layers all the time. However, the speedup (1.2-1.25x) was considered modest. Initial concerns about the narrowness of the tasks tested and lack of on-device benchmarks were addressed by the authors during rebuttal.
Decision: Accept (poster)