What Breaks and What Comes Back: Capability Damage and Recovery in Depth-Pruned Small Language Models

Piyush Dhoka; Vedant Pandhare; Varun Inamdar; Aadarsh Pathre; Pavitha Nooji

cs.CLOct 31, 2026

What Breaks and What Comes Back: Capability Damage and Recovery in Depth-Pruned Small Language Models

Authors: Piyush Dhoka, Vedant Pandhare, Varun Inamdar, Aadarsh Pathre, Pavitha Nooji

Organizations: Department of Artificial Intelligence, Vishwakarma University, Pune, India

Abstract

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.

Fixed budget (1M tokens each; 28 tests)
		Capped mean			Worst case
Family	Comparison	$\Delta$	$p_{\mathrm{H}}$	$p^{\mathrm{seed}}_{\mathrm{H}}$	$\Delta$	$p_{\mathrm{H}}$	$p^{\mathrm{seed}}_{\mathrm{H}}$
Llama-3.2-1B	RAH vs uniform	+0.021	1.000	1.000	+0.168	$\leq$ 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	$\leq$ 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.

\Delta

is the first method minus the second, from per-item scores averaged over seeds.

p_{\mathrm{H}}

: Holm-adjusted

p

-value from the item-level bootstrap;

p^{\mathrm{seed}}_{\mathrm{H}}

: Holm-adjusted

p

-value from the seed-aware bootstrap;

\leq

marks a bound set by the bootstrap resolution.

Figure 4: Fixed-budget versus token-matched comparison (three seeds; bars show

\pm 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.

LLM Pruning

Large language model pruning removes less useful weights or components to reduce memory and inference costs while preserving model capabilities.

See all 188 papers

SHIFT-LLM: Distribution Shift Correction in Depth-Pruned LLMs
Aug 25, 2026Ali Bahri, Hang Li, Hongliang Li +1LLM Pruning LLM Compression
Huawei Noah’s Ark Lab, Canada · Huawei Noah’s Ark Lab, Hong Kong SAR, China
Small LLMs: Pruning vs. Training from Scratch
Jun 12, 2026Yufeng Xu, Taiming Lu, Kunjun Li +3Language Model Pretraining LLM Pruning
Princeton University · New York University · Carnegie Mellon University
Rethinking Layer Redundancy: Calibration Matters More Than Search in LLM Depth Pruning
Apr 27, 2026Minkyu Kim, Vincent-Daniel Yun, Youngrae Kim +3LLM Pruning LLM Inference Acceleration
Neural Superintelligence Lab, MODULABS, Republic of Korea · University of Southern California, United States · Boston University, United States +2

Aug 25, 2026J/K move · Enter open · S save

SHIFT-LLM: Distribution Shift Correction in Depth-Pruned LLMs

Ali Bahri, Hang Li, Hongliang Li, Zhitang Chen

Huawei Noah’s Ark Lab, Canada · Huawei Noah’s Ark Lab, Hong Kong SAR, China

LLM Pruning LLM Compression Distribution Shift Transformer Pruning LLM Inference Acceleration

Open paper AskPDF arXiv

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.

Abstract

LLM Pruning

Explore similar work

SHIFT-LLM: Distribution Shift Correction in Depth-Pruned LLMs

Small LLMs: Pruning vs. Training from Scratch

Rethinking Layer Redundancy: Calibration Matters More Than Search in LLM Depth Pruning