PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ali Bahri, Hang Li, Hongliang Li, Zhitang Chen

arXiv:2608.25068v1cs.CV

TL;DR

Depth pruning reduces inference cost by removing Transformer blocks but can cause hidden-state distribution shift and accuracy loss. SHIFT-LLM uses locally calibrated Linear Residual Adapters to approximate missing residual updates without gradients, consistently recovering performance across diverse settings, with gains up to +15.74 on Llama-3.1-8B-Instruct.

  • Problem

    Depth pruning can cause significant performance degradation because removing Transformer blocks shifts the hidden-state distributions received by downstream layers.

  • Method

    SHIFT-LLM inserts LRAs that preserve each pruned block’s identity pathway and use closed-form affine residual corrections fitted on a small calibration set without gradient computation.

  • Results

    SHIFT-LLM provides strong training-free recovery across diverse LLMs, pruning strategies, and evaluation settings; gains reach +15.74 on Llama-3.1-8B-Instruct under Reverse-order∗ pruning.

  • Takeaways & Limitations

    The framework preserves depth-pruning efficiency, supports modular use across layer-selection strategies, and can complement post-pruning fine-tuning.

  • Takeaways & Limitations

    Gains are smaller and occasionally negative when selected blocks’ missing residual updates are too complex for token-wise affine correction, as observed on Vicuna-7B.

Abstract

from arXiv · show

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.

1 Introduction

Depth pruning improves efficiency but can shift hidden-state distributions and degrade performance. SHIFT-LLM addresses this shift with lightweight, training-free residual corrections that preserve pruning efficiency and complement fine-tuning.

  • Structured depth pruning removes entire Transformer layers for direct hardware efficiency and inference speedups.
  • Existing depth-pruning methods can suffer significant performance degradation because removed blocks eliminate residual updates and shift downstream hidden-state distributions.
  • Post-pruning fine-tuning can recover performance, but requires additional optimization, compute, and training time.
  • SHIFT-LLM inserts a Linear Residual Adapter that preserves the identity pathway and adds an affine correction approximating each missing residual update.
  • The correction is calibrated by closed-form least-squares regression on a small held-out set without gradient computation, avoiding removed blocks’ attention and feed-forward costs.
  • SHIFT-LLM supports diverse layer-selection strategies and provides strong training-free recovery while remaining complementary to post-pruning fine-tuning.

2 Related Works

Related work develops layer-ranking, replacement, folding, and compact-parameterization approaches for efficient depth pruning. SHIFT-LLM differs from neighboring methods by locally approximating only each removed block’s residual update while preserving the identity pathway.

  • Layer-pruning methods rank blocks using weight statistics, gradient saliency, activations, or loss changes before removing less important layers.
  • Representative approaches combine structured pruning with lightweight tuning, layer folding, non-uniform pruning, or compact replacement modules.
  • ReplaceMe fits a linear transformation through a neighboring layer, whereas SHIFT-LLM performs local correction at each pruning site.
  • SHIFT-LLM estimates only the missing residual update and preserves the original identity pathway, rather than reconstructing a full hidden-state output.

3 Method

SHIFT-LLM corrects pruning-induced hidden-state shifts by inserting lightweight Linear Residual Adapters that preserve identity pathways and approximate missing residual updates. Closed-form calibration estimates these corrections efficiently, while affine structure enables further compression and exact merging.

  • Pruning-Induced Hidden-State Shift: Depth pruning removes a block’s residual update, causing downstream hidden-state distributions to differ from those produced by the original pretrained model.The shift can propagate through subsequent residual blocks because each block’s output becomes the next block’s input.
  • Linear Residual Adapter Formulation: SHIFT-LLM inserts an LRA at each pruning site that preserves the identity pathway and adds an affine correction approximating the missing residual update.The method estimates only the removed contribution rather than reconstructing the entire hidden state.
  • Linear Residual Adapter Formulation: The LRA is designed to remain cheaper than the removed Transformer block, operate locally, and avoid modifying the remaining network.Its correction is intended to preserve depth-pruning efficiency while requiring only local adaptation at each pruning site.
  • Closed-Form Estimation: The affine correction is fitted from calibration activations by stacking inputs and missing residual updates, then solving a regularized regression problem.The original model supplies calibration activations through forward passes, and the parameters are obtained with a closed-form linear solve rather than iterative backpropagation.
  • Experimental Evaluation: Table 1 evaluates Original, Base, and +LRA models using zero-shot average accuracy across seven benchmarks at 25% pruning without fine-tuning.The table defines Gain as the average improvement from adding the LRA and reports results across pruning criteria and models.
  • Extensions and Complexity: LRAs naturally support low-rank factorization and exact merging across consecutive pruning sites because they define affine mappings.Low-rank factorization reduces parameters from d^2 + d to 2dr + d, while affine composition permits exact merging without additional approximation.

4 Experiments

Experiments show that SHIFT-LLM improves depth-pruned models across diverse model families, pruning criteria, evaluation tasks, and recovery settings, both without fine-tuning and alongside parameter-efficient adaptation. Its correction remains effective with limited calibration data, at larger model scale, and across varying pruning ratios, though gains depend on the selected layers.

  • Training-free evaluation: Across four LLMs and seven zero-shot benchmarks, adding the LRA improves pruned baselines across most pruning criteria and models.The affine correction estimates the missing residual update while the identity pathway keeps the corrected hidden state closer to the original block output.
  • Training-free evaluation: +15.74 and +15.71 points are achieved on Llama-3.1-8B-Instruct under Reverse-order∗ and Taylor pruning, respectively.Gains reach +2.03 on Qwen2-1.5B and +6.88 on Qwen1.5-7B; Vicuna-7B gains are smaller and occasionally negative.
  • Training-free evaluation: At 14B scale, average zero-shot accuracy rises from 53.34 to 56.32 under Block Influence and from 46.72 to 52.08 under Reverse-order∗ pruning.These correspond to gains of +2.98 and +5.36 at 25% depth pruning.
  • Language modeling: WikiText-2 perplexity is substantially reduced across both pruning criteria for Qwen2-1.5B, Qwen1.5-7B, and Llama-3.1-8B, and under BI pruning for Vicuna-7B.The exception is Vicuna-7B under Reverse-order∗ pruning.
  • Fine-tuning compatibility: Adding the LRA to post-pruning adaptation improves average zero-shot performance across most evaluated models, including LoRA and partial-layer fine-tuning settings.LoRA gains are +1.41, +0.77, +1.14, and +0.72 points on Qwen2-1.5B, Qwen1.5-7B, Llama-3.1-8B-Instruct, and Vicuna-7B-v1.5; partial-layer gains are +0.98, +2.58, and +1.97 on three models, while Qwen2-1.5B decreases by 0.58.
  • Ablation studies: Calibration performance saturates quickly beyond 256 samples, and the residual LRA parameterization consistently outperforms a Generic Affine baseline.Under varying pruning ratios, SHIFT-LLM typically exceeds the pruning-only baseline by +2 to +3 average accuracy points.

5 Conclusion

SHIFT-LLM mitigates hidden-state distribution shift from depth pruning through training-free Linear Residual Adapters, while experiments show performance recovery and preserved efficiency benefits.

  • 5 Conclusion: SHIFT-LLM’s computational efficiency is reported for Llama-3.1-8B-Instruct.
  • 5 Conclusion: SHIFT-LLM uses LRAs with identity pathways and closed-form affine residual corrections to approximate missing updates after depth pruning.The framework is training-free and targets the hidden-state distribution shift caused by removed blocks.
  • 5 Conclusion: SHIFT-LLM consistently recovers performance across models, pruning criteria, and evaluation settings while preserving depth-pruning efficiency.

Supplementary Material

The supplementary evaluation defines broad model, pruning, calibration, and benchmark settings for testing SHIFT-LLM without using evaluation data for adaptation.

  • Supplementary Material: Experiments cover five decoder-only LLMs, with LRA corrections fitted using 256 C4 validation samples.Data-driven pruning criteria use 2,000 C4 validation samples.
  • Supplementary Material: Compatibility experiments assess LoRA and partial-layer fine-tuning separately from SHIFT-LLM’s training-free procedure.These experiments use C4 training and validation splits with specified optimization settings.
  • Supplementary Material: Zero-shot evaluation uses seven standard benchmarks spanning question answering, commonsense reasoning, sentence completion, coreference, and science question answering.Likelihood-based scoring is used, with length normalization for selected multi-token answer-choice datasets.
  • Supplementary Material: The study evaluates six primary layer-selection strategies plus Reverse-order∗, which preserves the final Transformer layer.The criteria range from heuristics to data-dependent importance measures.
  • Supplementary Material: Per-benchmark results report accuracy for each pruning criterion before and after adding LRA at a 25% pruning ratio without post-pruning fine-tuning.The “Original” row denotes the unpruned model, and “+LRA” denotes the pruned model equipped with LRA.

A.3 Per-Benchmark Results Without Fine-Tuning

Supplementary results examine per-benchmark recovery, scaling, low-rank compression, calibration-domain sensitivity, random-pruning failure, and calibration strategy.

  • A.3 Per-Benchmark Results Without Fine-Tuning: Per-benchmark results show how average improvements are distributed across seven zero-shot benchmarks at 25% depth pruning without fine-tuning.Table 8 reports each criterion before and after adding LRA.
  • A.3 Per-Benchmark Results Without Fine-Tuning: +2.98 points under BI and +5.36 points under Reverse-order∗ are achieved on Qwen2.5-14B at 25% depth pruning.These results indicate effectiveness at the 14B scale.
  • A.3 Per-Benchmark Results Without Fine-Tuning: Low-rank LRA variants match or slightly outperform the full-rank LRA under Magnitude-ℓ1 pruning on Qwen2-1.5B.Rank 128 achieves the highest reported result in the supplied passage.
  • A.3 Per-Benchmark Results Without Fine-Tuning: Changing calibration domains does not eliminate the negative gain under Vicuna-7B Reverse-order∗ pruning.Differences across C4, WikiText, and Alpaca remain relatively small in most settings.
  • A.3 Per-Benchmark Results Without Fine-Tuning: With random layer removal, LRA decreases average accuracy from 38.22 to 36.07.The result is attributed to removed blocks whose residual updates may contain stronger nonlinear or cross-token structure.
  • A.3 Per-Benchmark Results Without Fine-Tuning: Independent and sequential LRA calibration produce nearly identical performance.This supports using independent calibration when corrected hidden states remain sufficiently close to original representations.

B.5 Robustness to Calibration Samples

The supplementary analyses test calibration robustness, fine-tuning compatibility, baseline differences, and the design rationale for SHIFT-LLM’s local identity-plus-correction formulation.

  • B.5 Robustness to Calibration Samples: Across two random calibration seeds, Llama-3.1-8B-Instruct gains remain consistent: BI improves by +4.44 and +5.09 points, while Reverse-order∗ improves by +15.74 and +16.21 points.Variation remains below one point.
  • B.5 Robustness to Calibration Samples: Training affine corrections during LoRA fine-tuning provides additional gains over freezing fitted LRAs.The closed-form solution can therefore be refined when gradient-based recovery is available.
  • B.5 Robustness to Calibration Samples: Under Reverse-order∗, preserving the final block yields 45.71 versus 43.71 accuracy, a +2.00 gain; original Reverse-order instead changes 44.36 to 43.80, a −0.56 change.The comparison is reported on Qwen2-1.5B.
  • B.5 Robustness to Calibration Samples: SHIFT-LLM regresses an affine correction from the full hidden state at each pruning site, preserving the identity pathway and estimating only the missing residual update.ReplaceMe instead transforms the retained predecessor’s MLP contribution to approximate the output after a removed span.
  • B.5 Robustness to Calibration Samples: A generic affine mapping performs substantially worse than the identity-plus-correction parameterization despite equal representational capacity.This supports explicitly preserving the identity pathway.
  • B.5 Robustness to Calibration Samples: SHIFT-LLM achieves higher average zero-shot accuracy than ReplaceMe (L2) across all evaluated model–pruning combinations using the same 256-sample C4 budget.The difference is particularly large for Vicuna-7B under Reverse-order∗ pruning.

C Additional Computational Analysis

SHIFT-LLM adds minimal computational overhead compared with the Transformer blocks removed by depth pruning. The overhead scales linearly with pruning sites, and consecutive affine adapters can be merged exactly.

  • 0.5M parameters and 1.0M FLOPs per token are required by a rank-64 LRA, versus 218M parameters and 436M FLOPs for the original block.The comparison excludes the original block’s quadratic attention term.
  • LRA overhead grows linearly with the number of pruning sites, while depth pruning removes ⌊ρL⌋ Transformer blocks for pruning ratio ρ.
  • At 25% pruning, rank-64 LRAs add approximately 4.2M parameters and 8.4M FLOPs per token after removing eight blocks.The removed blocks account for approximately 1.74B parameters and 3.49G FLOPs per token.
  • The added LRA computation is less than 0.25% of the computational cost removed by pruning.
  • Consecutive pruned layers can have their affine mappings merged exactly into a single transformation.
  • Runtime comparisons cover the original, depth-pruned, rank-64 LRA, and exactly merged models under fixed V100 generation settings.Measurements use 128-token inputs and greedy generation of 128 output tokens.
Loading 2608.25068v1…