🤖 AI

Reasoning Models Can be Accurately Pruned Via Chain-of-Thought Reconstruction

This paper proposes Reasoning-Aware Compression (RAC), a pruning technique that jointly reconstructs input activations and on-policy chain-of-thought traces to effectively reduce the deployment costs of reasoning models while avoiding the performance degradation and increased latency caused by standard pruning methods.

Original authors: Ryan Lucas, Kayhan Behdin, Zhipeng Wang, Qingquan Song, Shao Tang, Rahul Mazumder

Published 2026-05-06

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ryan Lucas, Kayhan Behdin, Zhipeng Wang, Qingquan Song, Shao Tang, Rahul Mazumder

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, over-achieving student named "Reasoning AI." This student is amazing at solving complex math problems and writing code, but they have a quirky habit: before giving you the final answer, they write out a massive, detailed diary of their thought process. They don't just say "The answer is 42"; they write 50 pages of "Let's see, if I try this, then that, but wait, maybe not..." before finally landing on the answer.

While this makes them very accurate, it's also incredibly slow and expensive to run. You want to hire a smaller, cheaper version of this student (a "pruned" model) to do the same job.

The Problem: The "Wrong Homework" Mistake
Usually, when engineers try to shrink these big AI models, they use a method called pruning. Think of pruning like editing a book to make it shorter by cutting out words you don't think are important.

To decide which words to cut, the editors (the pruning algorithms) usually look at a sample of the student's input—the questions they are asked. They assume, "Okay, the student is good at answering these questions, so if we keep the parts of their brain that handle the questions, we'll be fine."

The paper argues this is a huge mistake for reasoning models. It's like trying to train a marathon runner by only watching them tie their shoes. You're ignoring the part where they actually run.

For reasoning models, the "running" part is the long chain of thoughts (the CoT). If you only train your pruning algorithm on the questions (the input) and ignore the long thinking process, the shrunken model gets confused. It starts rambling even more than before, writing 100 pages of confused thoughts instead of 50, and it still gets the answer wrong. It becomes slower and dumber at the same time.

The Solution: "Reasoning-Aware Compression" (RAC)
The authors propose a simple fix called Reasoning-Aware Compression (RAC).

Instead of just showing the pruning algorithm the questions, they say: "Let's make the model answer the questions first, write out its full thought process, and then use that entire process to decide what to cut."

The Analogy: The Rehearsal
Imagine you are editing a play.

Old Method: You read the script's opening lines and decide which actors to fire based on how they look in the first scene.
RAC Method: You watch the entire rehearsal, including the long, messy middle scenes where the actors are figuring out the plot. You only fire the actors who mess up during the actual performance.

By letting the model "rehearse" its own thinking process during the editing phase, the pruning algorithm learns exactly which parts of the brain are needed for the long, complex reasoning steps.

What Happened When They Tried It?
The team tested this on several powerful AI models (like DeepSeek-R1 and Qwen) using math and coding tests.

Accuracy: When they used the old method, cutting 50% of the model's brain made it fail almost everything. With RAC, the shrunken model kept almost all of its smarts, even with 50% of its weights removed.
Speed: The old method made the models so confused they would ramble on for hours, taking forever to answer. RAC kept the models focused. They answered just as fast as the big models, or even faster, because they didn't get stuck in loops of confused thinking.
Versatility: This trick worked whether they were cutting the model randomly or in specific patterns, and it worked on different types of models.

The Bottom Line
If you want to make a smart, reasoning AI smaller and cheaper, you can't just look at the questions it's asked. You have to watch how it thinks while answering them. By including the model's own "thinking diary" in the editing process, you can shrink the model without losing its brain or making it slower. It's a simple tweak that stops the model from getting lost in its own thoughts.

Technical Summary: Reasoning Models Can Be Accurately Pruned via Chain-of-Thought Reconstruction

1. Problem Statement

Large Language Models (LLMs) with step-by-step reasoning capabilities, such as DeepSeek-R1, generate explicit chains-of-thought (CoT) before producing a final answer. While this significantly improves accuracy on complex tasks like mathematics and coding, it results in extremely long output sequences, making these models costly to deploy at scale due to high inference latency and memory requirements.

Model pruning is a standard technique for reducing model size and inference cost. However, the authors observe that applying standard pruning methods (e.g., SparseGPT) to reasoning models yields counter-intuitive and detrimental results:

Accuracy Degradation: Standard pruning causes significant drops in task accuracy as sparsity increases.
Runtime Inflation: Contrary to the goal of compression, heavily pruned reasoning models often become slower. The models produce longer, more rambling, and less accurate CoT traces, increasing the total token count and inference time.

The paper hypothesizes that this failure stems from a distribution mismatch. Standard pruning calibrates weights using only input prompts (where $|x| \gg |y|$ ), whereas reasoning models are decode-dominated tasks where the generated CoT tokens ( $|c| + |y| \gg |x|$ ) constitute the majority of the inference computation. Pruning based solely on prompt activations fails to preserve the model's ability to generate coherent reasoning traces.

2. Methodology: Reasoning-Aware Compression (RAC)

The authors propose Reasoning-Aware Compression (RAC), a "drop-in" modification to existing one-shot pruning workflows that aligns the calibration process with the actual inference-time distribution of reasoning models.

Core Mechanism

Standard pruning algorithms (like SparseGPT) minimize the reconstruction error of layer weights ( $W_\ell$ ) using an activation matrix ( $X_\ell$ ) derived solely from input prompts. RAC modifies this by augmenting the calibration dataset with on-policy CoT traces.

On-Policy Rollout: During the calibration phase, for each prompt $x^{(m)}$ , the dense model is allowed to autoregressively generate its own CoT and answer tokens.
Activation Collection: The hidden states ( $x^{(\ell)}_t$ ) corresponding to both the input prompt tokens and the self-generated CoT tokens are collected.
Joint Reconstruction: The pruning algorithm solves the layer-wise optimization problem using a combined activation matrix $X^{RAC}_\ell = [X^P_\ell, X^D_\ell]$ , where $X^P$ represents prompt activations and $X^D$ represents decode-time (CoT) activations.

The objective function becomes:
$\min_{\hat{W}_\ell} \sum_{m} \sum_{t \in P_m \cup D_m} \| (W_\ell - \hat{W}_\ell) x^{(\ell-1,m)}_t \|_2^2$
subject to sparsity constraints. This ensures the compressed weights are optimized to reconstruct activations encountered during the full reasoning process, not just the initial prompt.

Implementation Details

One-Shot Pruning: The method requires no retraining or distillation. It operates in a single pass on pre-trained weights.
Compatibility: RAC integrates seamlessly with existing layer-wise pruning algorithms (e.g., SparseGPT, WANDA, ALPS).
Calibration Cost: While collecting CoT activations is more expensive than prompt-only calibration, it remains significantly cheaper than full model retraining or distillation.

3. Key Contributions

Diagnosis of Pruning Failure: The paper identifies that standard pruning fails on reasoning models because it ignores the decode-dominated nature of CoT generation, leading to distribution shifts that cause models to "ramble" and lose accuracy.
RAC Algorithm: The introduction of a simple calibration strategy that injects on-policy CoT activations into the pruning objective.
Empirical Validation: Extensive experiments demonstrating that RAC allows reasoning models to be pruned to high sparsity levels (up to 50%) while maintaining accuracy and preventing the runtime blow-up associated with standard pruning.

4. Experimental Results

The authors evaluated RAC on multiple reasoning model families (DeepSeek-R1 distilled variants and Qwen3) across various sizes (1.5B to 70B) and benchmarks (MATH-500, AIME-25, LiveCodeBench).

Accuracy and Runtime

Accuracy Preservation: At 50% sparsity, RAC significantly outperforms both standard C4-calibrated pruning and prompt-only pruning.
- Example (DeepSeek-R1-7B on MATH-500): RAC maintains 0.900 accuracy, compared to 0.812 (prompt-only) and 0.744 (C4).
- Example (Qwen3-8B on MATH-500): RAC achieves 0.862 accuracy at 50% sparsity, versus 0.564 (C4) and 0.470 (prompt-only).
Runtime Reduction: RAC prevents the "rambling" behavior.
- Example (DeepSeek-R1-7B): Standard pruning increased runtime from 23.3 minutes (dense) to 135.0 minutes at 50% sparsity. RAC reduced this to 35.3 minutes, effectively maintaining the dense model's speed profile.
Generalization: The benefits hold across different architectures (Qwen, Llama), model sizes, and domains (math and code). RAC also shows robustness when combined with structured pruning (2:4 sparsity) and quantization (FP8).

Ablation Studies

On-Policy vs. Off-Policy: Using CoT traces generated by the same model being pruned (on-policy) yields better results than using traces from a larger, different model, though off-policy data still offers some improvement over prompt-only.
Token Budget Sensitivity: The runtime penalty of pruned models is largely driven by excessive token generation limits. When constrained to reasonable decoding lengths (e.g., 4,096 tokens), RAC-pruned models match or exceed dense model performance in both accuracy and speed.
Algorithm Agnosticism: RAC improves performance when paired with various pruning algorithms (SparseGPT, ALPS, WANDA), not just SparseGPT.

Reconstruction Error Analysis

Token-wise error analysis on held-out test problems reveals that while prompt-only calibration preserves input activations well, it incurs high reconstruction errors during the decoding phase (the CoT). RAC specifically reduces these errors during the long reasoning chains, aligning the pruned model's behavior with the dense model during the critical generation phase.

5. Significance and Claims

The paper claims that Reasoning-Aware Compression is a critical, simple modification required to make pruning viable for reasoning LLMs.

Practical Impact: RAC enables the deployment of smaller, faster reasoning models without the accuracy collapse or latency spikes currently observed with standard compression.
Theoretical Insight: The work highlights that for tasks dominated by generation (like reasoning), calibration data must reflect the entire inference trajectory, not just the input.
Limitations: The authors modestly note that RAC does not explicitly optimize for sequence length reduction (though it indirectly helps by improving generation quality) and that the calibration step incurs a modest computational overhead compared to prompt-only methods. However, they emphasize that this overhead is negligible compared to the cost of retraining.

In summary, the paper demonstrates that by aligning the pruning calibration distribution with the model's on-policy reasoning traces, one can accurately prune reasoning models to 50% sparsity, preserving up to 95% of the dense model's accuracy and preventing the inference slowdowns that typically plague compressed reasoning systems.