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

Ali Bahri; Hang Li; Hongliang Li; Zhitang Chen

cs.CVAug 25, 2026

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

Authors: Ali Bahri, Hang Li, Hongliang Li, Zhitang Chen

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

Abstract

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.

Figures & tables

Figure 1: The proposed LRA partially restores the hidden-state distribution toward that of the original model.

Figure 2: Overview of SHIFT-LLM for depth pruning in LLMs . Depth pruning induces a hidden-state distribution shift. At each pruning site, the LRA preserves the identity pathway and adds an affine residual correction estimated in closed form from calibration activations to approximate the missing residual update. The resulting corrected hidden states remain closer to the original distribution, while the LRAs support low-rank compression and exact merging across consecutive pruning sites.

	Qwen2-1.5B			Qwen1.5-7B			Llama-3.1-8B-It			Vicuna-7B
Original	59.87			65.54			71.03			66.80
Pruning Criteria	Base	+LRA	Gain	Base	+LRA	Gain	Base	+LRA	Gain	Base	+LRA	Gain
Block Influence	45.78	46.85	+1.08	52.06	52.82	+0.76	57.18	61.62	+4.44	55.05	55.52	+0.48
Reverse-order*	43.71	45.71	+2.00	49.25	53.79	+4.54	43.33	59.06	+15.74	53.47	52.96	-0.51
Magnitude- $\ell_{1}$	47.44	48.47	+1.03	40.72	47.59	+6.88	37.81	39.42	+1.61	34.86	35.09	+0.22
Magnitude- $\ell_{2}$	43.76	45.78	+2.02	43.83	45.74	+1.91	36.00	39.58	+3.58	35.10	37.03	+1.92

Table 1: Zero-shot average accuracy across seven evaluation benchmarks for different pruning criteria without fine-tuning at a pruning ratio of 25%. “Original” denotes the unpruned model before layer pruning. “Base” denotes the pruned model without post-pruning correction, while “+LRA” denotes the same pruned model equipped with the proposed LRA . “Gain” reports the average improvement obtained by adding the LRA . Full per-benchmark results for each model are provided in the supplementary material.

Dataset	Qwen2-1.5B			Qwen1.5-7B			Llama-3.1-8B-It			Vicuna-7B
Dataset	Orig.	LoRA	LoRA+LRA	Orig.	LoRA	LoRA+LRA	Orig.	LoRA	LoRA+LRA	Orig.	LoRA	LoRA+LRA
BoolQ	73.36	63.15	65.54	81.9	71.28	73.24	83.98	70.31	69.11	80.55	74.34	75.99
PIQA	75.46	64.31	64.25	78.40	68.72	69.91	79.98	71.71	72.25	77.26	70.51	71.16
HellaSwag	65.41	45.64	49.18	76.97	61.70	63.54	79.27	67.95	70.43	73.76	66.14	67.28
WinoGrande	66.22	58.41	57.93	66.14	59.98	58.09	74.19	65.98	66.14	69.61	65.19	63.38
ARC-e	66.16	45.03	46.30	70.71	58.75	58.00	81.78	63.22	66.37	75.63	62.75	63.80

Table 2: Zero-shot results with fine-tuning on C4 under Reverse-order ∗ pruning. “Orig.” denotes the unpruned model, “LoRA” denotes standard post-pruning LoRA fine-tuning, and “LoRA+LRA” denotes LoRA fine-tuning combined with the proposed LRA , whose affine residual correction is kept trainable during fine-tuning.

Model	Criterion	Base	+LRA	Gain
Qwen2-1.5B	BI	87.16	52.98	+34.18
	Reverse ∗	458.44	111.33	+347.11
Qwen1.5-7B	BI	146.65	54.66	+91.99
	Reverse ∗	130.21	64.12	+66.09
Llama-3.1-8B	BI	161.51	33.83	+127.68
	Reverse ∗	1179.23	51.37	+1127.86

Table 3: WikiText-2 perplexity at 25% depth pruning. Lower is better. Gain denotes the perplexity reduction obtained by adding the LRA .

Dataset	Qwen2-1.5B			Qwen1.5-7B			Llama-3.1-8B-It			Vicuna-7B
Dataset	Orig.	Partial	Partial +LRA	Orig.	Partial	Partial +LRA	Orig.	Partial	Partial +LRA	Orig.	Partial	Partial +LRA
BoolQ	73.36	67.19	65.78	81.90	80.09	79.69	83.98	74.89	63.67	80.55	78.04	77.77
PIQA	75.46	62.68	62.95	78.40	66.05	67.36	79.98	67.63	70.02	77.26	68.44	68.23
HellaSwag	65.41	47.81	47.85	76.97	58.13	61.22	79.27	62.30	67.12	73.76	60.03	63.51
WinoGrande	66.22	54.70	53.99	66.14	58.72	57.38	74.19	60.22	64.01	69.61	60.14	62.19
ARC-e	66.16	39.27	39.10	70.71	43.14	46.89	81.78	56.57	64.44	75.63	55.22	59.01

Table 4: Zero-shot results with fine-tuning on C4 dataset under Reverse-order ∗ pruning. “Orig.” denotes the unpruned model, “Partial” denotes standard post-pruning partial-layer fine-tuning, and “Partial+LRA” denotes partial-layer fine-tuning combined with the proposed LRA , whose affine residual correction is kept trainable during fine-tuning.

Method	PR	BoolQ	PIQA	HSwag	OBQA	ARC-e	ARC-c	MMLU	CMMLU	Wino	Avg.
Original	0%	83.98	79.87	59.10	33.80	81.90	51.71	68.04	55.43	73.72	65.28
ShortGPT (BI)	25%	–	71.76	41.96	20.20	61.07	28.41	24.17	24.94	53.91	40.80
Short. LLaMA (PPL)	25%	–	76.28	49.31	26.40	72.90	38.05	33.67	27.24	57.93	47.72
Short. LLaMA (Taylor)	25%	–	71.38	49.64	27.40	68.48	41.81	28.61	25.04	71.35	47.96
SLEB (FT)	25%	–	75.73	49.73	26.80	69.70	38.65	43.05	33.38	63.85	50.11
LLM-Pruner ∗	20%	–	72.00	54.60	–	–	–	25.30	25.00	–	44.23

Table 5: Comparison with state-of-the-art pruning methods on Llama-3.1-8B-It using Alpaca dataset for fine-tuning. All previous results are taken from the Navigation LLM paper, and ∗ indicates reproduced results reported in that work. “Navigation LLM+LRA” denotes Navigation LLM equipped with the proposed LRA .

Criterion	Base	+Generic	+LRA
BI	45.78	36.96	46.85
Reverse ∗	43.71	37.99	45.71
Magnitude- $\ell_{1}$	47.43	35.74	48.47

Table 6: LRA vs. Generic Affine on Qwen2-1.5B at 25% pruning. Average zero-shot accuracy over seven benchmarks.

Panel of Figure 3 — Figure 3: Ablation studies of the proposed LRA under different calibration sizes and pruning ratios.

Per-block computational cost
Component	Params	FLOPs/token	Relative Cost
Transformer block	218M	436M	100%
Full-rank LRA	16.8M	33.6M	7.7%
Rank-64 LRA	0.5M	1.0M	$<0.25\%$
End-to-end runtime at 25% pruning
Configuration	Params	Layers	Latency (ms)	Tokens/s

Table 7: Computational efficiency of SHIFT-LLM on Llama-3.1-8B-Instruct.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Method	PR	BoolQ	PIQA	HellaSwag	WinoGrande	ARC-e	ARC-c	OBQA	Avg.
Qwen2-1.5B	Original	0%	73.36	75.46	65.41	66.22	66.16	36.09	36.40	59.87
Qwen1.5-7B	Original	0%	81.90	78.40	76.97	66.14	70.71	42.83	41.80	65.54
Llama-3.1-8B-It	Original	0%	83.98	79.98	79.27	74.19	81.78	55.03	43.00	71.03
Vicuna-7B	Original	0%	80.55	77.26	73.76	69.61	75.63	45.82	45.00	66.80
Qwen2-1.5B	ShortGPT (BI)	25%	55.57	62.51	43.17	55.01	45.20	28.58	30.40	45.78
Qwen2-1.5B	ShortGPT (BI)+LRA	25%	57.74	64.80	44.88	56.12	48.32	26.71	29.40	46.85

Appendix

Table 8: Per-benchmark zero-shot accuracy corresponding to the average results reported in Table 1 . Results are shown for different pruning criteria at a pruning ratio of 25%, without post-pruning fine-tuning. “+LRA” denotes the pruned model equipped with the proposed LRA .

Method	BoolQ	PIQA	HellaSwag	WinoGrande	ARC-e	ARC-c	OBQA	Avg.
Original	86.27	81.07	82.91	75.14	82.45	58.87	45.20	73.13
BI	66.51	67.03	54.66	63.30	50.97	36.69	34.20	53.34
BI + LRA	70.52	67.46	59.20	71.35	56.06	37.46	32.20	56.32
Reverse ∗	62.23	59.19	42.48	58.01	40.32	31.83	33.00	46.72
Reverse ∗ + LRA	64.92	66.92	51.90	66.61	48.57	34.90	30.80	52.08

Appendix

Table 9: Zero-shot results on Qwen2.5-14B at 25% depth pruning.

LRA Configuration	BoolQ	PIQA	HellaSwag	WinoGrande	ARC-e	ARC-c	OBQA	Avg.
Full-rank	62.23	67.95	43.25	54.38	52.99	25.60	32.20	48.38
Rank 64	62.29	67.68	45.37	53.99	52.57	25.60	31.80	48.47
Rank 128	62.26	68.44	45.65	52.96	53.66	25.51	32.20	48.67
Rank 256	62.26	69.04	45.18	52.57	53.79	25.51	31.80	48.59
Rank 512	62.26	68.39	44.21	53.59	53.87	25.34	32.20	48.55

Appendix

Table 10: Effect of low-rank factorization of the affine residual correction under Magnitude-

\ell_{1}

pruning on Qwen2-1.5B at 25% depth pruning.

Criterion	Pruned	+LRA (C4)	+LRA (WikiText)	+LRA (Alpaca)
BI	55.05	55.66	56.14	54.66
Reverse ∗	53.47	52.96	53.36	52.20
Magnitude- $\ell_{1}$	34.86	35.09	35.80	35.17
Magnitude- $\ell_{2}$	35.10	37.02	38.16	37.32

Appendix

Table 11: Sensitivity to the calibration domain on Vicuna-7B at 25% depth pruning. Results are average zero-shot accuracy over seven benchmarks.

Method	BoolQ	PIQA	HSwag	Wino	ARC-e	ARC-c	OBQA	Avg.
Random	57.83	55.82	27.42	50.59	28.41	22.87	24.60	38.22
Random + LRA	38.84	54.84	27.60	48.70	31.06	22.87	28.40	36.07

Appendix

Table 12: Random pruning with and without the proposed LRA on Qwen2-1.5B at 25% depth pruning.

Criterion	Calibration	BoolQ	PIQA	HellaSwag	WinoGrande	ARC-e	ARC-c	OBQA	Avg.
BI	Independent	82.14	70.84	46.21	71.82	64.60	38.14	22.60	56.62
BI	Sequential	83.21	70.13	45.53	72.30	64.06	37.03	21.80	56.29
Reverse ∗	Independent	62.42	70.84	46.12	71.82	64.23	41.21	25.60	54.61
Reverse ∗	Sequential	62.54	71.87	46.15	72.61	66.62	41.21	24.20	55.03

Appendix

Table 13: Independent vs. sequential LRA calibration on Llama-3.1-8B-Instruct at 25% depth pruning.

Model	Frozen	Trainable	Gain
Qwen2-1.5B	47.98	48.94	+0.96
Qwen1.5-7B	56.08	56.28	+0.20
Llama-3.1-8B-It	60.32	61.12	+0.80

Appendix

Table 14: Frozen vs. trainable LRAs during LoRA fine-tuning under Reverse-order ∗ pruning at 25% depth pruning.

Model	Criterion	Pruned	ReplaceMe	SHIFT-LLM
Qwen2-1.5B	BI	45.8	43.9 $(-1.9)$	46.9 $(+1.1)$
Qwen2-1.5B	Reverse ∗	43.7	44.3 $(+0.6)$	45.7 $(+2.0)$
Qwen1.5-7B	BI	52.06	50.34 $(-1.72)$	52.82 $(+0.76)$
Qwen1.5-7B	Reverse ∗	49.25	53.62 $(+4.37)$	53.79 $(+4.54)$
Llama-3.1-8B	BI	57.18	59.46 $(+2.28)$	61.62 $(+4.44)$
Llama-3.1-8B	Reverse ∗	43.33	58.02 $(+14.69)$	59.07 $(+15.74)$

Appendix

Table 15: Comparison with ReplaceMe (L2) at 25% depth pruning. Results are average zero-shot accuracy over seven benchmarks using the same 256-sample C4 calibration budget.

LLM Pruning

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

See all 168 papers

Forward-Free LLM Depth Pruning via Weight Redundancy
Sep 9, 2026Vincent-Daniel Yun, Woosang LimLLM Pruning LLM Inference Acceleration
University of Southern California, United States · 2Seoul National University, Republic of Korea
Ghosted Layers: Unconstrained Activation Alignment for Recovering Layer-Pruned LLMs
May 15, 2026Vincent-Daniel Yun, Junhyuk Jo, Sai Praneeth Karimireddy +1LLM Pruning LLM Compression
University of Southern California, USA · Inha University, Republic of Korea
Locality-Aware Redundancy Pruning for LLM Depth Compression
May 27, 2026Vincent-Daniel Yun, Youngrae Kim, Woosang Lim +3LLM Pruning LLM Compression
University of Southern California, United States · Neural Superintelligence Lab, MODULABS, Republic of Korea · Seoul National University, Republic of Korea +1

Sep 9, 2026J/K move · Enter open · S save

Forward-Free LLM Depth Pruning via Weight Redundancy

Vincent-Daniel Yun, Woosang Lim

University of Southern California, United States · 2Seoul National University, Republic of Korea

LLM Pruning LLM Inference Acceleration Neural Network Pruning Transformer Pruning

Open paper AskPDF arXiv

Depth pruning reduces large language model (LLM) inference cost by removing complete Transformer blocks. Activation-based methods collect hidden states through forward passes on calibration data, while existing forward-free methods score each Transformer block separately without measuring similarity between blocks. We propose Weight-Redundancy Pruning (WRP), a forward-free depth-pruning method that estimates inter-layer redundancy from checkpoint weights to select blocks without calibration data or model forward passes. WRP compares attention output and MLP down-projection weights across layers and combines their pairwise similarities with relative projection-scale information. The resulting all-pairs similarity matrix guides layer grouping and block selection. Across multiple pruning settings, model families, and downstream tasks, WRP consistently outperforms existing forward-free magnitude pruning and approaches the performance of activation-based methods.

Abstract

LLM Pruning

Explore similar work

Forward-Free LLM Depth Pruning via Weight Redundancy

Ghosted Layers: Unconstrained Activation Alignment for Recovering Layer-Pruned LLMs

Locality-Aware Redundancy Pruning for LLM Depth Compression