SARATHI: Decoupled Data-Free Probing and Adaptive Reconstruction for Structured LLM Pruning

AI summary

Overview

Research area: Machine learning systems and model compression — specifically post-training structured pruning of large language models for efficient inference.

Technical level: Intermediate. Readers should be comfortable with transformer feed-forward layers, matrix factorization, and the general idea of calibration-based pruning, but the paper's core argument (separating selection from reconstruction) is conceptually simple.

Scope: The paper introduces SARATHI, a structured pruning framework that removes FFN neurons using a data-free NMF-based importance probe plus a memory-localized weight reconstruction, evaluated on six models across two architecture families at 15–40% sparsity.

What This Paper Is About

Structured pruning shrinks LLMs by physically deleting neurons and weight rows, which yields smaller dense matrices that run faster on ordinary hardware. Existing post-training methods reuse a single calibration dataset both to decide which neurons to cut and to repair the remaining weights, which tilts the pruning mask toward the calibration domain and blows up GPU memory at high sparsity. SARATHI separates these two jobs: neuron selection becomes completely data-free, driven only by the geometry of the weight matrix itself, while weight repair uses a small localized batch that fits in O(d²) memory per layer.

Key Contributions

1. A decoupled Probe-Reconstruct architecture. Neuron selection (P1) and weight reconstruction (P2) are formulated as two sequentially independent subproblems with disjoint information sources, removing the calibration bias and the O(Ld²) peak-memory bottleneck the authors attribute to coupled methods such as SoBP, SliceGPT, and FLAP.

2. NMF Residual Subspace Probing. A data-free neuron scoring method that applies non-negative matrix factorization to |W^(ℓ)|, treats neurons with large residual entries as structurally irreplaceable, and comes with a formal connection to Robust PCA plus a pruning-safety bound (Proposition 1) of O(τ√d ‖X‖_F) for zero-score neurons.

3. Bias Shift Compensation. A closed-form bias update for post-LayerNorm architectures that corrects the systematic mean shift induced by removed neurons, using the same local batch as the reconstruction stage with no extra forward passes. On OPT-2.7B at 40% sparsity this takes WikiText-2 PPL from 2,082 to 18.98.

4. Robust global budget allocation via MAD thresholding, with a min-keep floor α, plus an evaluation on six models (LLaMA-3-8B, Mistral-7B, OPT-2.7B/6.7B, LLaMA-2-7B/13B) at 15–40% sparsity without retraining.

Main Findings

• Strongest perplexity on gated architectures: SARATHI-NMF achieves the lowest WikiText-2 PPL on LLaMA-3-8B and Mistral-7B at every tested sparsity level. At 40%, LLaMA-3-8B reaches 9.92 PPL versus 20.69 for SliceGPT and 19.11 for Dynamic Slicing; at 30%, Mistral-7B reaches 7.41 versus 8.76 for Dynamic Slicing.

• Survives the sparsity regime where baselines fail: At 40% on OPT-6.7B, the coupled baselines and SARATHI-Wanda hit out-of-memory failures; SARATHI-NMF completes compression at 16.52 PPL and 54.30% zero-shot accuracy. SoBP fails at 40% on OPT-2.7B (marked as a code error), where SARATHI-NMF reaches 18.98 PPL.

• Lower PPL but not always higher accuracy: On LLaMA-2-13B at 30%, SARATHI-NMF reaches 5.85 PPL versus SoBP's 6.27, yet SoBP scores 66.82% zero-shot versus 62.62%. On OPT models at 30%, SARATHI-NMF has higher PPL than SoBP but higher zero-shot accuracy (OPT-2.7B: 52.65% vs. 51.45%; OPT-6.7B: 56.20% vs. 56.02%).

• Activation probing wins accuracy, geometric probing wins perplexity: At 30% on LLaMA-3-8B, SARATHI-NMF gives 8.70 PPL (−3.8% versus Wanda) while SARATHI-Wanda gives 58.45% zero-shot (+4.1 pp versus NMF).

• Bias Shift Compensation is decisive for OPT: Without it, OPT-2.7B PPL is 47.06, 132.00, and 2,082 at 20%, 30%, and 40% sparsity; with it, 13.25, 15.61, and 18.98, reductions of 72%, 88%, and 99.1%.

• Non-negativity, not rank, drives the probe: Replacing NMF with a signed SVD residual raises PPL by more than 8,500 points on OPT-6.7B despite a 98.95% overlap in top-ranked neurons. PPL varies by at most 0.28 across NMF ranks r ∈ {7, 16, 32, 64}.

• Adaptive beats uniform budgets: On LLaMA-2-13B at 30%, adaptive MAD allocation gives 5.85 PPL and 62.62% zero-shot versus 7.48 PPL and 60.15% for a uniform baseline.

• Against data-free PruneNet on LLaMA-2-7B at 30% (5 tasks): SARATHI-Wanda averages 59.11% versus PruneNet's 55.45% (+3.7 pp); SARATHI-NMF is comparable at 54.25% while reporting lower perplexity (8.58 PPL).

• Throughput is hardware-agnostic but uneven: At batch size 16, OPT-6.7B reaches up to 2.05× over dense at 40% sparsity and LLaMA-3-8B 1.66×; OPT-2.7B gains only 1.06–1.11×, and Mistral-7B regresses to 0.82–0.92× because FFN dimensions shift away from Tensor-Core-aligned multiples of 64. Gains at 15–20% sparsity are below 10%.

• Functional preservation: NTK Cosine Similarity is reported as 1.0000 universally, with NTK Trace > 1 on some OPT and LLaMA-2-13B layers attributed to bias-correction scaling.

Methodology in Plain English

The authors split pruning into two problems that no longer share data.

First, choosing what to cut, without any text at all. For each FFN layer they take the absolute values of the weight matrix and factor it into two non-negative matrices whose product approximates it (NMF, run for 100 iterations with rank r=16 and ε=1e-8, verified to converge within 50 iterations). Whatever the low-rank product cannot explain is the residual. Neurons whose weight rows contain large residual entries are judged structurally irreplaceable and kept; neurons that sit entirely inside the shared low-rank subspace are judged redundant. The justification is that this mirrors Robust PCA: the residual is the sparse outlier component that carries unique information. For gated FFN models, the scores for the gate, up, and down matrices are averaged arithmetically rather than multiplied.

Second, allocating the budget. Scores are pooled across layers and thresholded with a median-absolute-deviation rule, chosen because neuron importance scores are heavy-tailed and mean/std thresholds would overshoot the target sparsity. A binary search on the threshold constant k matches the target rate, and a min-keep floor α stops any single layer from collapsing.

Third, repairing the weights. With the mask already fixed, a standard OBS-style closed-form update recomputes the surviving weights column by column using a regularized Hessian built from only 128 local samples. Because the mask is fixed in advance, the Hessian is used purely for compensation rather than joint search, which keeps peak memory at O(d²) per layer instead of accumulating buffers across all L layers.

Finally, correcting drift. In post-LayerNorm models like OPT, deleting neurons shifts the mean activation entering the next LayerNorm, and that shift compounds through the remaining layers. The authors measure the mean difference on the same 128 samples and fold it directly into the layer's bias term, at zero extra cost.

Why This Matters

Research impact. The paper reframes structured pruning as two separable problems and offers a geometric, calibration-free criterion for neuron importance, backed by a Robust PCA argument and a safety bound. It also documents a failure mode — cascading mean shift in post-LayerNorm architectures — that explains why some models collapse at high sparsity, and shows a closed-form fix. The observation that the sparsity ceiling of prior methods is partly a memory artifact rather than a quality limit is directly actionable for the field.

Real-world applications:

• Serving LLMs on a fixed GPU budget where per-token memory and latency dominate cost.
• On-premise or edge deployment where only dense BLAS kernels are available and specialized sparse kernels are not.
• Domain-specific deployment where the target text distribution differs from generic calibration corpora (code, legal, clinical), and calibration bias is a genuine risk.
• Pruning models for architectures where calibration data is unavailable or restricted for licensing or privacy reasons, since selection needs no data.

Industry relevance. The O(d²)-per-layer memory profile and the reported absence of OOM at 40% sparsity on OPT-6.7B matter for practitioners running compression pipelines on single-GPU nodes. The 16–44 minute pruning time on one A100 with no retraining fits into a standard deployment workflow. The Mistral-7B throughput regression is an honest signal that dimension alignment, not just FLOP count, governs realized speedup.

Future Directions

• Extending the probe to attention heads. The conclusion names attention heads as a target; the current work prunes only FFN intermediate neurons.
• Better probing for non-gated architectures. On ReLU-FFN models like OPT the NMF scores are described as near-uniform, and performance is carried by bias compensation and OBS rather than the probe; activation-conditioned probing for these models is proposed as future work.
• Fixing dimension alignment for throughput. The Mistral-7B regression (0.82–0.92×) and the sub-10% gains at 15–20% sparsity point to a need for sparsity budgets or rounding schemes that land on Tensor-Core-friendly widths.
• Closing the accuracy gap under data-free selection. SoBP's calibration-dependent mask leads on zero-shot accuracy for LLaMA-2-13B at 30% (66.82% vs. 62.62%), leaving open how to recover calibration-domain task alignment without reintroducing coupling.

Target Audience

ML engineers and systems researchers who need to shrink deployed LLMs without retraining or specialized sparse kernels; compression researchers interested in data-free importance criteria and the Robust PCA framing; and practitioners working with post-LayerNorm or ReLU-FFN models, who will find the bias-shift analysis and the OPT-specific results the most directly useful part. Readers focused on gated architectures such as LLaMA and Mistral will get the clearest picture of the probe's benefits.

Authors’ abstract

Structured pruning reduces the inference cost of large language models (LLMs) by permanently removing neurons and weight rows, yielding dense, hardware-agnostic speedup. Existing post-training methods couple neuron selection and weight reconstruction through a shared calibration dataset, introducing domain bias and out-of-memory failures at high sparsity. We propose SARATHI, a framework that decouples these two stages. A data-free \emph{NMF Residual Subspace Probe} scores neurons using intrinsic weight geometry, without accessing any calibration data, and carries a formal connection to Robust PCA. \emph{Adaptive OBS} then reconstructs the pruned weights with a localized calibration batch, operating in $\mathcal{O}(d^2)$ peak memory per layer and eliminating the global memory bottleneck of coupled methods. A \emph{Bias Shift Compensation} step corrects systematic activation drift in post-LayerNorm architectures that would otherwise cause perplexity collapse at high sparsity. Experiments across six models from two architecture families (gated and non-gated FFN) at 15--40\% sparsity show that SARATHI achieves competitive perplexity and matches or exceeds existing structured pruning baselines on zero-shot accuracy.