Restoring Pruned Large Language Models via Lost Component Compensation

Conference: NeurIPS 2025 arXiv: 2510.21834 Code: GitHub Area: Model Compression / LLM Pruning Recovery Keywords: LLM pruning, performance restoration, attention activation, component compensation, PEFT

TL;DR

RestoreLCC proposes a targeted recovery strategy for pruned LLMs: it uses contrastive probing to identify critical attention heads, applies SVD decomposition to extract activation components lost during pruning, and injects them back into the pruned model as learnable bias vectors — significantly restoring performance without compromising sparsity or inference speed.

Background & Motivation

LLM pruning methods (e.g., Wanda, SparseGPT, SlimGPT) are essential for reducing model size and accelerating inference, yet they inevitably cause performance degradation. Existing recovery methods primarily rely on parameter-efficient fine-tuning (PEFT), such as LoRA, to restore pruned model performance.

Key Challenge: PEFT methods like LoRA were originally designed for dense models to adapt to downstream tasks. When applied to pruned models, they overlook a key property specific to pruned models — the need to compensate for lost information — resulting in inefficient parameter search and suboptimal recovery.

Key Insight: Information loss caused by pruning is reflected in attention head activations. By contrastively analyzing the activation differences between dense and pruned models, three observations emerge: 1. Directly injecting the lost components back into pruned heads can substantially recover logit differences and final accuracy. 2. Different attention heads vary greatly in importance and recovery behavior. 3. Discriminative information may reside in minor components rather than principal ones.

Key Insight: Rather than performing generic, undirected fine-tuning, the paper explicitly compensates for the critical information directions lost during pruning.

Method

Overall Architecture

RestoreLCC comprises two core modules: (1) Contrastive Probing, which identifies the attention heads most critical for recovery; and (2) Lost Component Compensation (LCC), which extracts and optimizes directional components of lost information and injects them back into the pruned model.

Key Designs

1. Contrastive Probing:

   • Contrastive sample construction: For each sample \((q, r^+)\), a sentence encoder (e.g., MiniLM-L6) retrieves the semantically most similar negative sample \(r^-\), forming a triplet \((q, r^+, r^-)\).
   • Activation editing: The recovered query activation is modeled as \(z_c^q = z_p^q + c^q\), where \(c^q\) denotes the lost principal component.
   • Attention head probing: The problem is cast as a natural language inference task — if an attention head is important and its corresponding component is useful, the recovered activation should "entail" the correct sequence activation and "contradict" the incorrect one.
   • A linear probe classifier is trained to assess the discriminative capability of each head, ranking heads by classification accuracy.

2. Lost Component Compensation (LCC):

   • SVD decomposition is applied to the pruning-induced activation difference \(\Delta\mathbf{Z}^{(l,h)} = \mathbf{Z}_d^{(l,h)} - \mathbf{Z}_p^{(l,h)}\).
   • All directional vectors \(v_i\) (orthonormal) are fixed; only a scalar magnitude \(\beta_i\) is learned for each direction.
   • The lost component is modeled as: \(c_{\text{learned}} = \sum_{i=1}^{d_h} \beta_i v_i + b\)
   • \(b\) is a trainable bias vector that captures information beyond the predefined directions.
   • The recovered activation is: \(\tilde{z}_p = z_p + c_{\text{learned}}\)

3. Plug-and-Play Bias Injection:

   • The learned \(c_{\text{learned}}\) is a constant bias vector that captures the key information commonly lost across all samples.
   • It is directly absorbed into the multi-head attention module as a constant bias, adding negligible computational overhead at inference.
   • Parameter overhead is minimal: at most \(1/(2d_l)\) additional parameters per layer, representing less than 0.05% for models with hidden dimension > 1000.

Loss & Training

• General recovery training uses the Alpaca instruction dataset.
• Task-specific recovery requires only 100 training samples.
• The top 10%–25% of attention heads by importance are selected for compensation.
• Component count \(K=1\) is used for head importance identification.

Key Experimental Results

Main Results (LLaMA-7B, General Recovery)

Pruning Type	Method	PPL↓	Avg. Acc.↑	Gain
Unstructured 50%	Wanda (baseline)	7.26	54.09%	—
Unstructured 50%	LoRA	7.09	56.27%	+2.18
Unstructured 50%	LoFiT	7.35	56.82%	+2.73
Unstructured 50%	RestoreLCC	6.93	58.83%	+4.74
Semi-structured 2:4	SparseGPT (baseline)	11.04	48.99%	—
Semi-structured 2:4	DoRA	9.16	52.35%	+3.36
Semi-structured 2:4	RestoreLCC	8.99	55.00%	+6.01
Structured 20%	SlimGPT (baseline)	7.46	57.54%	—
Structured 20%	DoRA	7.54	58.51%	+0.97
Structured 20%	RestoreLCC	7.53	59.76%	+2.22

Ablation Study

Configuration	Avg. Acc.	Note
Full RestoreLCC	58.83%	—
w/o contrastive probing (random head selection)	57.57% (−1.26)	Contrastive probing effectively identifies critical heads
MSE-based head selection	58.14% (−0.69)	Less precise than contrastive probing
KL-based head selection	57.92% (−0.91)	Less precise than contrastive probing
w/o directional component \(\sum \beta_i v_i\)	57.13% (−1.70)	Directional components are the core contribution
w/o bias \(b\)	58.26% (−0.57)	Bias provides additional flexibility

Key Findings

• RestoreLCC consistently outperforms all PEFT baselines across unstructured, semi-structured, and structured pruning regimes.
• In task-specific recovery at 60% sparsity, RestoreLCC surpasses LoFiT by 3.56% using only 100 training samples.
• Effectiveness generalizes to LLaMA-13B: +1.04% (unstructured), +1.73% (semi-structured), +1.45% (structured).
• The compensating bias vectors impose negligible impact on sparsity and inference speed.

Highlights & Insights

• Shift from "patching" to "restoring": Rather than applying generic PEFT to pruned models without direction, RestoreLCC precisely identifies what information was lost, where it was lost, and how to recover it.
• Efficient design via SVD + scalar learning: Fixing orthogonal directions and learning only scalars drastically reduces parameter count and training cost.
• Finding that minor components may matter more: This challenges the intuition that principal components dominate, suggesting discriminative information may be encoded in low-variance directions.
• Contrastive probing as a general attention head importance estimator: The technique is transferable to other scenarios requiring identification of critical attention heads.

Limitations & Future Work

• Only attention heads are compensated; FFN modules are excluded, potentially leaving critical FFN information unaddressed.
• Contrastive probing requires constructing contrastive sample pairs, introducing some dependence on data distribution.
• SVD directions are currently fixed; future work could explore adaptive direction learning.
• Evaluation is primarily conducted on the LLaMA family; generalizability to architectures such as GPT and Mistral remains to be verified.
• Compensation is limited to a single constant vector, precluding handling of input-dependent activation loss.

Related Work & Insights

• vs. LoRA: LoRA is a general-purpose PEFT method that does not account for pruning-specific properties; RestoreLCC performs targeted compensation for pruning-induced information loss.
• vs. EoRA: EoRA searches for low-rank subspaces in feature space, whereas RestoreLCC leverages all directions from SVD decomposition and learns their importance.
• vs. LoFiT: LoFiT also intervenes in attention activations but does not target the specific information directions lost due to pruning.

Rating

• Novelty: ⭐⭐⭐⭐ — Contrastive probing combined with lost component compensation is a novel framing, though the core mechanism remains bias injection in activation space.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ — Three pruning regimes × multiple model scales × general/task-specific recovery × detailed ablations.
• Writing Quality: ⭐⭐⭐⭐⭐ — The derivation from insight to method is logically clear; the three Findings are well-argued and persuasive.
• Value: ⭐⭐⭐⭐ — Directly applicable to LLM pruning deployment; plug-and-play with no inference efficiency overhead.

Related Papers

• [NeurIPS 2025] A Simple Linear Patch Revives Layer-Pruned Large Language Models
• [NeurIPS 2025] Correlation Dimension of Auto-Regressive Large Language Models
• [NeurIPS 2025] PermLLM: Learnable Channel Permutation for N:M Sparse Large Language Models
• [NeurIPS 2025] The Structure of Relation Decoding Linear Operators in Large Language Models
• [NeurIPS 2025] LayerIF: Estimating Layer Quality for Large Language Models using Influence Functions