Technical Summary: CORP - Closed-Form One-Shot Representation-Preserving Structured Pruning for Transformers

1. Problem Statement

Transformer-based models achieve state-of-the-art performance but incur prohibitive computational and memory costs, hindering deployment on resource-constrained hardware. While structured pruning (removing entire dimensions or substructures) offers practical speedups, most existing methods rely on retraining, iterative optimization, or access to labeled data and gradients. In strict post-training deployment scenarios, these resources are often unavailable.

Naive one-shot structured pruning fails under these constraints because removing channels or attention dimensions introduces systematic representation errors that accumulate across layers, leading to rapid accuracy collapse even at moderate sparsity levels. The authors argue that this failure is driven not merely by the quality of importance ranking, but by the lack of explicit error recovery for the removed components.

2. Methodology: CORP

The authors propose CORP (Closed-Form One-Shot Representation-Preserving Structured Pruning), a framework that reframes pruning as a representation recovery problem. CORP operates without labels, gradients, or fine-tuning, utilizing only a small set of unlabeled calibration data.

Core Philosophy

Instead of viewing pruning solely as a selection task, CORP models the removed components as affine functions of the retained components. It derives closed-form ridge regression solutions to estimate the missing contributions and folds these compensations directly into the weights of the remaining layers. This preserves the layer's output distribution without requiring iterative optimization or additional inference modules.

Technical Components

The method targets two primary structures in Transformers:

A. MLP Channel Pruning

For an MLP layer defined by $y = Wx + b$, where the input $x$ is partitioned into kept ($x_S$) and pruned ($x_P$) channels:

1. Modeling: The pruned activations $x_P$ are approximated as an affine function of the kept activations: $x_P \approx Bx_S + c$.
2. Closed-Form Solution: The parameters $(B, c)$ are derived via ridge regression on the calibration set $D$:
   $$B = \Sigma_{PS}(\Sigma_{SS} + \lambda I)^{-1}, \quad c = \mu_P - B\mu_S$$
   where $\Sigma$ represents the covariance matrices of the activations.
3. Weight Folding: The compensation is folded into the remaining weights:
   $$\hat{W}_S = W_S + W_P B, \quad \hat{b} = b + W_P c$$
   This results in a standard linear layer with reduced dimensions that minimizes the reconstruction error of the original output.

B. Attention Dimension Pruning (Query/Key)

For self-attention, CORP prunes dimensions in the Query ($Q$) and Key ($K$) projections, leaving the Value ($V$) projection intact to avoid complex coupling issues.

1. Logit Recovery: The original attention logits are $L = QK^\top = Q_S K_S^\top + Q_P K_P^\top$. The term $Q_P K_P^\top$ (missing logits) is approximated using the kept subspaces: $Q_P K_P^\top \approx Q_S M K_S^\top$, where $M$ is a small compensation matrix.
2. Closed-Form Solution: $M$ is estimated by solving a ridge regression problem over the calibration set, formulated as a Sylvester equation or a vectorized linear system:
   $$\left[ \sum_b (K_{S,b}^\top K_{S,b}) \otimes (Q_{S,b}^\top Q_{S,b}) + \lambda I \right] \text{vec}(M) = \sum_b \text{vec}((Q_{S,b}^\top Q_{P,b})(K_{P,b}^\top K_{S,b}))$$
3. Weight Folding: The matrix $I + M$ is factorized via SVD ($U\Sigma V^\top$) and folded into the Query and Key projections:
   $$\hat{W}_{Q,S} = W_{Q,S} U \Sigma^{1/2}, \quad \hat{W}_{K,S} = W_{K,S} V \Sigma^{1/2}$$

Ranking Strategy

Before compensation, CORP ranks channels/dimensions for removal using simple, data-driven metrics:

• MLP: A combined metric of activation energy ($E_i = E[x_i^2]$) and weight magnitude ($|W_{:,i}|^2$).
• Attention: Expected logit energy ($s_j = E[\|q_j\|^2 \|k_j\|^2]$).
  Channels with the lowest scores are pruned.

3. Key Contributions

1. Reframing Pruning: The authors shift the focus of one-shot pruning from purely "ranking" to "representation recovery," demonstrating that compensation is critical for preserving accuracy.
2. Closed-Form Framework: CORP provides a unified, closed-form solution for both MLP channel and Q/K attention dimension pruning that requires no gradients, labels, or fine-tuning.
3. Unified Compensation: Unlike prior work that may only compensate the output of a block, CORP performs layer-wise representation recovery and introduces specific logit-space compensation for attention dimensions, folding all corrections directly into the reduced projections.

4. Experimental Results

Experiments were conducted on ImageNet (DeiT and ViT models), WikiText-2 (OPT-1.3B), and downstream tasks (DINOv2 on NYU Depth V2 and ADE20k).

• Accuracy Preservation:

  • DeiT-Huge: Achieves 83.27% Top-1 accuracy after pruning 50% of both MLP and attention structures (a drop of only ~1.7% from the baseline of 84.97%).
  • Scalability: Larger models (e.g., DeiT-Huge) show stronger redundancy and better accuracy retention under aggressive sparsity compared to smaller models (e.g., DeiT-Tiny).
  • Comparison: CORP outperforms or matches state-of-the-art methods like GRAIL, SNOWS, and VBP. For instance, on DeiT-Base with 50% MLP sparsity, CORP achieves 72.0% Top-1, compared to 68.6% for GRAIL and 66.4% for VBP.
  • No Fine-tuning: CORP matches the performance of DC-ViT (which uses few-shot feature-mimicking fine-tuning) without any gradient updates.

• Efficiency Gains:

  • DeiT-Huge: At 50% sparsity, parameters and FLOPs are reduced by ~40–45%, resulting in a 1.6× throughput speedup on an NVIDIA RTX 3090.
  • OPT-1.3B: Pruning 30% of both MLP and attention structures reduces FLOPs by 27.2% and parameters by 25.0%, with perplexity rising modestly (14.60 to 25.41 for joint pruning).
  • Downstream Tasks: On DINOv2, 50% backbone pruning preserves downstream performance on depth estimation (RMSE improves from 0.3308 to 0.3167 for DINOv2-Huge) and segmentation.

• Calibration Efficiency:

  • The method is robust to calibration set size; larger models remain stable even with small calibration sets (e.g., 1000 samples).
  • Runtime analysis shows that calibration dominates the cost, while ranking and closed-form compensation add negligible overhead.

5. Significance and Claims

The paper claims that CORP demonstrates that structured pruning can be effectively performed in a strict one-shot, post-training setting by treating it as a representation recovery problem. The key significance lies in:

• Practical Deployability: It enables compression on hardware where retraining is impossible due to lack of labels or gradients.
• Theoretical Insight: It highlights that the "systematic error" introduced by removing additive components is the primary cause of accuracy collapse, and that simple affine/logit compensation can recover this error without complex optimization.
• Broad Applicability: The method shows strong results across vision models (DeiT, ViT, DINOv2) and language models (OPT), suggesting strong redundancy in both MLP and attention representations across modalities.

The authors modestly note limitations, including the assumption of independent module compensation (ignoring cross-module interactions), the use of uniform sparsity ratios, and the lack of optimization for quantization-pruning composition. They state that performance may degrade under strong distribution shifts between calibration and deployment data.