Pruning via Causal Attribution Preserves Reasoning Performance in Large Language Models

Technical Summary: Causal Attribution Pruning (CAP)

Problem Statement
Large Language Models (LLMs) demonstrate strong multi-step reasoning capabilities but incur substantial inference costs, hindering real-world deployment. While model pruning offers a solution by removing redundant parameters without retraining, existing methods often degrade performance on complex reasoning tasks. Current approaches rely on correlational signals, such as weight magnitude (e.g., Wanda) or local gradient statistics. These methods assume small weights are unimportant; however, low-magnitude parameters can be critical for multi-step reasoning circuits. Furthermore, purely correlational criteria fail to capture cross-layer interactions essential for chain-of-thought (CoT) inference, frequently leading to "layer collapse" where entire modules lose function.

Methodology: Causal Attribution Pruning (CAP)
The authors introduce Causal Attribution Pruning (CAP), a training-free, one-shot pruning method designed to preserve reasoning performance. Unlike methods that rank weights based on static properties, CAP measures the causal impact of specific model components on task performance. The framework operates in three stages:

1. Calibration and Causal Scoring (Stage 1):

   • A small calibration dataset ($D_{cal}$) is constructed by sampling from reasoning-centric benchmarks: GSM8K (math), StrategyQA (multi-hop commonsense), and ARC-Challenge (science).
   • For each attention head $h_\ell$ in the model, CAP performs an interventional masking experiment. The head is temporarily masked during the forward pass, and the resulting increase in token-level cross-entropy loss ($\Delta(h_\ell)$) is measured relative to the intact model.
   • To ensure robustness against variance from small sample sizes, the dataset is partitioned into $K$ subsets, and the median loss increase is computed.
   • This process identifies heads that are functionally critical for reasoning, distinguishing them from those that are merely numerically large but redundant.

2. Weight-Level Importance Conversion (Stage 2):

   • The head-level causal scores are converted into normalized importance factors $I(h_\ell)$, mapped to a range of $[1.0, 10.0]$. High-impact heads receive higher factors.
   • These factors are propagated to individual weights within the corresponding projection matrices ($W_q, W_k, W_v, W_o$). The importance score for a specific weight $W_{ij}$ is calculated as the product of its magnitude and the importance factor of its associated head region: $S_{ij} = |W_{ij}| \cdot I(h_{region})$.
   • For MLP layers, which lack a natural head-partitioned structure, the method assigns a per-layer average importance derived from the co-located attention heads. This is noted as an approximation.

3. Global Pruning (Stage 3):

   • All weights across the model are ranked by their importance-weighted magnitude scores ($S_{ij}$).
   • A global threshold is determined via binary search to achieve the target sparsity (e.g., 10%, 20%, 50%). Weights below the threshold are zeroed out, creating an unstructured sparsity pattern compatible with standard inference infrastructure.

Key Contributions

• Formalization of Causal Attribution for Pruning: The authors formalize a method to measure the expected increase in task-specific loss when individual attention heads are masked, moving beyond correlational proxies to quantify functional criticality.
• Budget-Aware Weight-Level Pruning: They introduce a procedure that converts head-level causal scores into weight-level importance values. This allows for fine-grained magnitude pruning under a global sparsity constraint while preserving the structure of critical reasoning circuits.
• Empirical Validation: The method is evaluated on Llama-3-8B-Instruct and Mistral-7B-Instruct across GSM8K, StrategyQA, and ARC-Challenge benchmarks.

Results

• Moderate Sparsity (10–20%): CAP generally outperforms the Wanda baseline in preserving reasoning performance.
  • On Llama-3-8B-Instruct at 20% sparsity, CAP achieves a 61% relative improvement over Wanda on ARC-Challenge (70.8% vs. 43.9% accuracy).
  • On GSM8K at 10% sparsity, CAP shows a 20% relative improvement over Wanda (71.5% vs. 59.5%).
  • The authors note that while CAP sometimes trails Wanda on specific tasks (e.g., Llama-3 GSM8K at 20%), it consistently preserves performance on multi-step inference tasks where Wanda's correlational criteria fail.
• High Sparsity (50%): Performance degrades significantly for both methods. CAP suffers catastrophic failure on Llama-3 (perplexity jumps to 428.2), attributed to its coarse-grained treatment of MLP layers. At this level, the global threshold removes essential MLP weights that were assigned low importance based on the attention head approximation.
• Perplexity vs. Reasoning: At moderate sparsity, CAP often achieves better reasoning accuracy despite slightly higher perplexity compared to Wanda, suggesting a decoupling between general language modeling quality and task-specific reasoning preservation.
• Architecture Limitations: Preliminary experiments on Mixture-of-Experts (MoE) models suggest CAP is less effective, as dynamic routing makes head importance non-stationary.

Significance and Claims
The paper claims that attention-head-level causal attribution is a superior criterion for preserving reasoning performance compared to correlational pruning methods at equivalent sparsity levels. By explicitly measuring the causal contribution of heads to reasoning tasks, CAP can identify and protect "numerically small but functionally critical" parameters that magnitude-based methods discard.

The authors emphasize that their work is modest in scope:

• It is a training-free method, making it computationally efficient compared to approaches requiring weight updates (like SparseGPT).
• It is task-aligned, calibrated specifically on reasoning data rather than generic language modeling.
• It acknowledges limitations, particularly the lack of fine-grained causal attribution for MLP neurons and the difficulty of applying the method to dynamically routed architectures (MoE).

The core contribution is the demonstration that interventional, task-aligned pruning can better preserve specific capabilities (reasoning) than methods optimized for general perplexity or static weight statistics, provided the sparsity level remains moderate.