Depth pruning removes whole Transformer blocks to make a language model smaller and faster, but the damage is not spread evenly across what the model can do, and a short fine-tuning run, called healing, decides what comes back. We measure this for small models used as agents. Four instruction-tuned models of 1.2-1.7B parameters (Llama-3.2-1B, Qwen2.5-1.5B, SmolLM2-1.7B and OLMo-2-1B) are pruned and healed, and six capabilities are tracked: structured output, function calling, instruction following, mathematical reasoning, safety refusal and knowledge. Three results hold across the families. First, function calling and math collapse in every model (retention 0.00-0.07 of the unpruned score), and the order in which capabilities break agrees closely across the four (Kendall's W = 0.92; 0.62 after correcting for chance floors). Second, which parameters are trained matters as much as which data: fully tuning the last three blocks instead of adding low-rank adapters to all blocks costs up to 0.18 in capped mean retention on three families, and one 300k-token pilot per scope detects this. Third, we test a planner that allocates a 1M-token healing budget from 21 short pilot heals. At a fixed budget it raises the weakest capability on two families, although only one of these gains survives a test that also resamples training seeds. Once its 3.48M pilot tokens are counted, uniform healing with the same total number of training tokens matches or beats every planned variant on every family and retains math significantly better on three of them in an item-level test (on one when seeds are also resampled). Healing longer is a strong and simple baseline that planned recovery must beat. Code, per-item predictions and every result are public.
Figures & tables
Figure 1: Study design. Within a family, every lane uses the same pruned model, data pools, teacher targets, seeds and test items, and one training recipe per adaptation scope. RAH spends 3.48M training tokens on 21 pilot heals before its 1M-token heal. We compare it with the top lane at the same 1M-token healing budget, and with the bottom lane, which receives the same 4.48M total training tokens.
Family
Checkpoint
Params
Blocks
Removed
GPU
Llama-3.2-1B
unsloth/Llama-3.2-1B-Instruct
1.24
16
12, 13 (12.5%)
RTX 4060 (laptop); L4
Qwen2.5-1.5B
Qwen/Qwen2.5-1.5B-Instruct
1.54
28
13–16, 24–26 (25.0%)
RTX 4060 (laptop); L4
SmolLM2-1.7B
HuggingFaceTB/SmolLM2-1.7B-Instruct
1.71
24
6, 20–22 (16.7%)
A100
OLMo-2-1B
allenai/OLMo-2-0425-1B-Instruct
1.48
16
8, 13 (12.5%)
A100
Table 1: Model families, parameters (billions), Block-Influence cuts (original block indices) and the GPUs the runs used.
Capability
Benchmark
Test
Dev
Score
Structured output
JSONSchemaBench (GitHub-easy)
150
60
schema-valid JSON
Function calling
BFCL v3 simple + multiple
200
60
AST match (name, arguments)
Instruction following
IFEval
300
60
strict prompt-level accuracy
Mathematical reasoning
GSM8K test
200
60
exact final answer
Safety refusal
XSTest (100 unsafe, 100 benign)
200
60
refuse unsafe, answer benign
Knowledge
MMLU (all subjects)
300
60
answer-letter likelihood; ECE
Table 2: Evaluation suite. Development items are used only for planning; test items give every reported number.
Figure 2: Retention of each capability after the main pruning cut (test split), as raw retention and corrected for the chance floors of knowledge and safety. Columns are ordered by mean damage rank. Function calling and math collapse in every family.
Family
Quantity
Struct.
Tool
Instr.
Math
Safety
Knowl.
Llama-3.2-1B
unpruned score
0.32
0.54
0.45
0.47
0.83
0.32
retention
0.06
0.00
0.43
0.07
0.75
0.76
chance-corrected
0.06
0.00
0.43
0.07
0.39
-0.05
Qwen2.5-1.5B
unpruned score
0.78
0.86
0.38
0.48
0.73
0.58
retention
0.04
0.00
0.20
0.00
0.80
0.57
chance-corrected
0.04
0.00
0.20
0.00
0.37
0.24
Table 3: Damage after the main cut on the test split: unpruned absolute scores, retention (pruned score divided by unpruned score) and chance-corrected retention.
Figure 3: Adaptation scope under identical data. Both heals use the uniform mixture and 1M tokens; only the trained parameters differ. Labels give the difference (last three blocks minus all-LoRA).
Family
Method
Tokens
Struct.
Tool
Instr.
Math
Safety
Knowl.
Capped
Worst
Over-ref.
Llama-3.2-1B
Uniform
1M
1.49
0.88
0.67
0.38
1.04
1.15
0.82
0.38
0.16
Damage-prop.
1M
2.02
0.90
0.66
0.43
0.99
1.09
0.83
0.43
0.15
RAH
1M + 3.48M
1.49
0.82
0.76
0.54
0.94
1.09
0.84
0.54
0.09
RAH-capped
1M + 3.48M
1.13
0.77
0.74
0.55
0.95
1.11
0.83
0.55
0.09
Uniform
4.48M
1.33
0.90
0.70
0.58
1.02
1.15
0.86
0.58
0.14
Qwen2.5-1.5B
Uniform
1M
0.73
0.54
0.58
0.21
0.96
0.67
0.62
0.21
0.59
Table 4: Retention on the test split, averaged over three seeds. “Tokens” gives the healing tokens and, for RAH, the pilot tokens. Bold marks the best capped mean and worst case per family, not significance (Table 5 ). Over-ref.: fraction of benign XSTest prompts refused. On Llama both RAH variants use the last- k scope chosen by the scope pilots; all other rows use all-LoRA.
Capped mean
Worst case
Family
Comparison
Δ
pH
pHseed
Δ
pH
pHseed
Fixed budget (1M tokens each; 28 tests)
Llama-3.2-1B
RAH vs uniform
+0.021
1.000
1.000
+0.168
≤ 0.028
0.525
Llama-3.2-1B
RAH vs damage-prop.
+0.013
1.000
1.000
+0.116
0.460
1.000
Llama-3.2-1B
RAH-capped vs uniform
+0.012
1.000
1.000
+0.172
≤ 0.028
0.378
Llama-3.2-1B
RAH-capped vs damage-prop.
+0.004
1.000
1.000
+0.119
0.288
1.000
Table 5: Paired bootstrap tests on the primary metrics. Δ is the first method minus the second, from per-item scores averaged over seeds. pH : Holm-adjusted p -value from the item-level bootstrap; pHseed : Holm-adjusted p -value from the seed-aware bootstrap; ≤ marks a bound set by the bootstrap resolution.
Figure 4: Fixed-budget versus token-matched comparison (three seeds; bars show ±1 seed standard deviation). The dark line is uniform healing at 1M and 4.48M tokens. Hollow markers show the RAH variants as they are usually reported, at their 1M healing budget; filled markers place them at the 4.48M training tokens they actually used.
Figure 5: Math recovery curves fitted from the pilots (development split; dashed beyond the largest pilot budget) against test retention of uniform heals at 1M and 4.48M tokens, placed at the effective math tokens of Equation ( 1 ).
Figure 6: Refusal operating points: fraction of unsafe XSTest prompts refused against fraction of benign prompts refused. The ideal is the top-left corner. Pruning lowers refusal of unsafe prompts in every family; healing moves the operating point in family-specific directions.
Depth pruning removes entire Transformer blocks to reduce the inference cost of large language models, but disrupts the hidden-state distributions expected by downstream layers, leading to significant accuracy loss. We introduce SHIFT-LLM, a training-free post-pruning correction framework that inserts a Linear Residual Adapter (LRA) at each pruning site. Each LRA preserves the identity pathway of the original residual block and adds a lightweight affine residual correction. This correction is calibrated via closed-form least-squares regression on a small held-out set, without gradient computation, to approximate the missing residual update produced by the pruned block. Together with the preserved identity pathway, the resulting LRA output approximates the hidden state produced by the original block, thereby mitigating the distributional mismatch introduced by layer removal while avoiding the expensive attention and feed-forward computations of the removed blocks. The resulting LRAs support low-rank factorization and exact merging across consecutive pruned layers for additional compression, and combine naturally with parameter-efficient fine-tuning for further recovery beyond fine-tuning the pruned model alone. Experiments on five model families, six layer-selection criteria, and seven zero-shot benchmarks show that SHIFT-LLM consistently recovers accuracy lost to depth pruning across most configurations, achieving gains up to +15.7 points on Llama-3.1-8B-Instruct while requiring only a few hundred calibration samples and no gradient computation.
Ali Bahri, Hang Li, Hongliang Li +1
Huawei Noah’s Ark Lab, Canada · Huawei Noah’s Ark Lab, Hong Kong SAR, China
Pruning promises a shortcut to strong small language models. In this work, we examine this promise by pruning Llama-3.1-8B at pruning ratios of 0.5--0.8 with six methods spanning depth, width, and sparse granularities, under two controlled token-matched settings. (1) With the same training token budget, pruned initialization consistently outperforms random initialization. This shows that the parent model provides a strong starting point, although the advantage narrows as the training token budget grows and as the pruning ratio rises, nearly vanishing at the highest pruning ratio we study. (2) When training from scratch is instead given the full token budget consumed by the whole pipeline, pruning at finer granularities still retains an advantage, while coarser structured pruning can be matched or surpassed. This suggests that the parent model transfers knowledge that additional training tokens alone cannot fully recover, but only at fine granularity. Taken together, our results yield a clear recommendation: with a large pretrained model in hand and a limited training token budget, pruning is better than training from scratch; when the training budget is not limited, training from scratch can be competitive for coarser pruning, so a large pretrained parent is not always necessary.
Yufeng Xu, Taiming Lu, Kunjun Li +3
Princeton University · New York University · Carnegie Mellon University
Depth pruning improves the inference efficiency of large language models by removing Transformer blocks. Prior work typically treats layer redundancy as an inherent structural property of pretrained networks, emphasizing importance criteria and search algorithms to identify removable layers. In this study, we empirically investigate depth pruning from a functional perspective. Evaluating representative LLM families across diverse calibration configurations and multiple search algorithms, we show that different configurations produce different pruning patterns. Furthermore, under a fixed calibration configuration, complex search algorithms yield marginal performance improvements over simple one-shot methods, converging to similar pruned subsets. Overall, our results suggest that the calibration configuration plays a substantially larger role than the choice of search algorithm in shaping pruning patterns and calibration perplexity, while contributing comparably to variance in downstream reasoning accuracy. This indicates that future pruning efforts may benefit from prioritizing the calibration configuration over search complexity.
Minkyu Kim, Vincent-Daniel Yun, Youngrae Kim +3
Neural Superintelligence Lab, MODULABS, Republic of Korea · University of Southern California, United States · Boston University, United States +2