From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

Technical Summary: From Layers to Submodules (SUBFIT)

Problem Statement
Post-training compression of Large Language Models (LLMs) aims to reduce deployment costs by removing architectural components. While methods that delete entire layers or contiguous blocks are effective, they disrupt the residual stream, often causing sharp quality degradation. Recent "replacement-based" approaches address this by substituting removed components with lightweight fitted modules. However, existing methods share two restrictive design constraints:

1. Full-layer granularity: They treat the entire Transformer layer as the atomic unit of removal.
2. Contiguous selection: They remove components in consecutive depth ranges.

The authors argue these constraints are overly restrictive. Redundancy in pretrained transformers is not confined to contiguous regions, nor is it evenly distributed between Attention and FeedForward (FFN) outputs. Consequently, different strategies are needed to approximate different submodule types, and removable components need not cluster within contiguous depth ranges.

Methodology: SUBFIT
The paper introduces SUBFIT (Submodule-level Fitted residual replacement), a post-training compression method that operates at the granularity of individual Attention and FFN submodules. The process involves two main stages:

1. Non-Contiguous Selection:

   • SUBFIT selects a target number of Attention and FFN submodules independently across the model depth based on a sparsity budget $s$.
   • Selection is non-contiguous; modules are chosen independently by score without forming a consecutive block.
   • Scoring Strategy:
     • Attention: Modules are ranked by an "Impact score" computed on calibration activations, measuring how much a submodule rotates the residual stream and its relative magnitude.
     • FFN: For base models, a replacement-aware score combining removal effects and surrogate error is used. For instruction-tuned models, a cosine-only score is used to avoid noise from poorly calibrated residual magnitudes.
   • Ordering: Attention submodules are selected and replaced first. FFN submodules are then scored and selected on the partially compressed model to account for the residual-stream shift induced by the Attention replacement.

2. Fitted Residual Replacement:

   • Each selected submodule is replaced by a lightweight parametric surrogate $\hat{F}_\ell$ that approximates the original submodule's residual contribution ($y = F_\ell(x)$) rather than its internal computation.
   • Surrogate Structure: The surrogate uses an affine correction (element-wise gate, bias, and centering) plus a structured interaction term defined by a low-rank operator $U_\ell = V_\ell^\top W_\ell$.
   • Asymmetric Parameterization:
     • Attention: Each removed Attention submodule receives its own independent low-rank basis ($V_\ell$) and parameters, reflecting heterogeneous roles across depths.
     • FFN: A shared input basis $V$ is used across all selected FFN layers. Only the output factors ($W_\ell$) and affine terms remain layer-specific. This cross-layer parameter sharing significantly reduces deployed cost.
   • Fitting: The parameters are recovered analytically via closed-form fitting on calibration statistics (ridge-regularized least squares and eigen-decomposition), requiring no gradient-based optimization or full-model fine-tuning.

Key Contributions

1. Identification of Constraints: The authors identify full-layer granularity and contiguous selection as limiting factors in existing replacement-based compression.
2. SUBFIT Framework: A novel method that performs joint selection and replacement at the submodule level (Attention and FFN) with non-contiguous selection.
3. Asymmetric Design: A specific architecture where Attention bypasses use independent bases, while FFN bypasses share a common input basis to minimize parameter overhead.
4. Comprehensive Evaluation: Extensive testing across 10 LLMs (5 base, 5 instruction-tuned) and 5 sparsity levels (12.5% to 37.5%), demonstrating improved trade-offs compared to four baselines (Streamline, ReplaceMe).

Experimental Results

• Perplexity (Base Models): SUBFIT achieves the lowest perplexity degradation in 9 out of 10 model-sparsity settings. The gap over the strongest baseline widens significantly under aggressive compression (e.g., at 25% sparsity, SUBFIT shows a 2.42× perplexity degradation vs. 4.34× for the best baseline).
• Downstream Accuracy (Instruction-Tuned): At 25% sparsity, SUBFIT retains 84.6% of dense downstream accuracy on average, outperforming the strongest baseline (81.6%). It is the only method to retain above 80% accuracy at 25% sparsity and above 73% at 37.5%.
• Stability: SUBFIT exhibits significantly lower variance in perplexity degradation compared to baselines, which show instability (e.g., standard deviation growing by orders of magnitude) at high sparsity.
• Efficiency: The method yields measurable inference speedups (1.18×–1.40× in time-to-first-token) and KV-cache savings (proportional to sparsity) without requiring sparse kernels or custom hardware.
• Ablation Studies:
  • Residual Compensation: Replacing FFNs is the dominant factor for stability; replacing both Attention and FFNs yields the best results.
  • Shared Basis: The shared FFN basis achieves performance comparable to per-layer bases but with significantly fewer parameters.
  • Rank Sensitivity: FFN quality improves up to rank 4096, while Attention saturates at rank 256, justifying the asymmetric rank settings.

Significance and Claims
The paper claims that granularity is an underexplored axis in post-training compression. By moving from layer-level to submodule-level replacement, SUBFIT recovers more of the dense model's behavior than contiguous full-layer replacement under the same parameter budget. The results suggest that non-contiguous, submodule-specific recovery provides a stronger aggregate quality–efficiency trade-off, particularly under aggressive compression. The method is designed to be cheap to deploy (closed-form fitting) and compatible with existing inference stacks. The authors modestly suggest that this approach may generalize to other removal units, such as attention heads or individual neurons, but do not claim immediate generalization to those specific units in this work.