Restoring Pruned Large Language Models via Lost Component Compensation

TL;DR

This work addresses performance degradation in pruned large language models by proposing RestoreLCC, a targeted restoration method that reintroduces information lost during pruning. It identifies critical attention heads through contrastive probing and reconstructs missing directional information via learned magnitudes of lost components, then injects these back into pruned heads without altering sparsity or inference speed. The approach is compatible with structured, semi-structured, and unstructured pruning, and it yields consistent improvements over state-of-the-art baselines in both general and task-specific recovery across multiple LLMs and configurations. By enabling higher pruning ratios without sacrificing performance, RestoreLCC has practical implications for deploying efficient, scalable LLMs in resource-constrained settings.

Abstract

Pruning is a widely used technique to reduce the size and inference cost of large language models (LLMs), but it often causes performance degradation. To mitigate this, existing restoration methods typically employ parameter-efficient fine-tuning (PEFT), such as LoRA, to recover the pruned model's performance. However, most PEFT methods are designed for dense models and overlook the distinct properties of pruned models, often resulting in suboptimal recovery. In this work, we propose a targeted restoration strategy for pruned models that restores performance while preserving their low cost and high efficiency. We observe that pruning-induced information loss is reflected in attention activations, and selectively reintroducing components of this information can significantly recover model performance. Based on this insight, we introduce RestoreLCC (Restoring Pruned LLMs via Lost Component Compensation), a plug-and-play method that contrastively probes critical attention heads via activation editing, extracts lost components from activation differences, and finally injects them back into the corresponding pruned heads for compensation and recovery. RestoreLCC is compatible with structured, semi-structured, and unstructured pruning schemes. Extensive experiments demonstrate that RestoreLCC consistently outperforms state-of-the-art baselines in both general and task-specific performance recovery, without compromising the sparsity or inference efficiency of pruned models.

Figures (10)

• Figure 1: These figures compare the performance of the original dense model, the pruned model, and two recovery methods applied to the pruned model: LoRA and Lost Component Compensation (LCC, ours), which directly adds back components lost due to pruning. The used dataset is BoolQ, and please see § \ref{['sec:observation']} for experimental details. The left figure shows the logit difference (correct minus incorrect) of attention heads with/without recovery, and the right figure shows the accuracies of final model outputs for each method. l[$\cdot$] and a[$\cdot$] indicate the layer index and the head index, respectively. For instance, l12.a16 means the 16-th head at the 12-th layer.
• Figure 2: Logit gain of different attention heads recovered with principal components.
• Figure 3: Logit gain of attention heads recovered with different components.
• Figure 4: Overview of the RestoreLCC framework. It integrates two key modules, (1) contrastive probing and (2) lost component compensation (LCC), to restore the performance of pruned LLMs.
• Figure 5: Visualization of learned components for different attention heads.