TLDR: A new method called Reasoning-Aware Compression (RAC) improves the pruning of large language models that perform step-by-step reasoning. Unlike standard pruning that focuses on input, RAC incorporates the model’s self-generated reasoning steps (chain-of-thought) into the compression process. This leads to significantly better accuracy and more stable reasoning traces in pruned models, especially at high sparsity levels, making them more efficient to deploy without substantial performance loss.
Large language models (LLMs) with advanced reasoning capabilities, such as DeepSeek-R1, have become indispensable for tackling complex, multi-step problems in areas like mathematics, coding, and logical reasoning. These models achieve high accuracy by generating explicit “chain-of-thought” (CoT) traces, which are intermediate reasoning steps. However, this benefit comes at a significant cost: these models produce very long outputs, making them resource-intensive and expensive to deploy at scale.
To address this, researchers often turn to model compression techniques like pruning and quantization. Pruning involves removing redundant weights or neurons from the network, while quantization reduces the precision of model weights. Both methods aim to shrink model size, reduce memory footprint, and lower inference latency. However, applying standard compression methods to reasoning LLMs has proven challenging. These methods are typically optimized to preserve general language modeling performance and often lead to substantial accuracy drops on complex reasoning tasks. In some cases, pruning can even make models slower by causing them to generate more “thinking tokens” that are less accurate.
A new approach called “Reasoning-Aware Compression” (RAC) has been introduced to overcome these limitations. The core insight behind RAC is that standard LLM pruning methods primarily focus on reconstructing activations from the input (or prompt), whereas reasoning is a “decode-dominated” task where the model generates a long sequence of tokens. This means that the activations observed during inference are heavily influenced by the model’s self-generated chain-of-thought tokens, not just the initial prompt.
RAC offers a simple, drop-in solution: during the pruning process, it jointly reconstructs activations from both the input and the model’s on-policy chain-of-thought traces. This means that during calibration, the model generates its own reasoning steps, and these generated activations are included in the data used to determine which parts of the model to prune. This procedure effectively simulates the activations the model will encounter during actual inference, aligning the compression process with the model’s real-world behavior. RAC integrates seamlessly into existing pruning workflows like SparseGPT without requiring any retraining or distillation.
Experiments were conducted using DeepSeek-R1-Distill-Qwen architectures of various sizes (1.5B, 7B, 14B, 32B parameters) pruned with SparseGPT at different sparsity levels (20% to 50%). Three calibration methods were compared: standard C4 corpus, task-specific prompts only, and RAC. The models were evaluated on mathematical reasoning (MATH500 benchmark) and code generation (LiveCodeBench).
The results demonstrate that RAC consistently outperforms both standard C4 calibration and prompt-only calibration, especially as pruning becomes more aggressive. For instance, at 50% sparsity on the 1.5B model, RAC achieved 66.4% accuracy on MATH500, significantly higher than the 35.6% from C4 calibration. Similar improvements were observed for code generation tasks. Even using task-specific prompts without CoT traces showed better performance than generic C4, highlighting the importance of domain-relevant calibration data. However, RAC, which incorporates full reasoning traces, provided substantial additional benefits. Larger models generally showed greater resilience to compression, but RAC still offered meaningful improvements across all model sizes.
Further analysis revealed that RAC effectively mitigates reconstruction errors during the decoding phase, particularly for the long reasoning chains. While prompt-only calibration might slightly outperform on input tokens, RAC consistently yielded smaller errors throughout the much longer decode phase, confirming that reconstructing on CoT activations during calibration reduces decoding error on unseen problems.
Ablation studies explored whether the CoT traces used for calibration needed to come from the same model being compressed (on-policy) or could be from a different model (off-policy). On-policy calibration generally performed better, especially at higher sparsity levels, suggesting that activation patterns during CoT generation are model-specific. The research also looked into throughput gains with semi-structured 2:4 sparsity patterns, which can lead to actual speedups on modern hardware. RAC maintained its effectiveness with structured pruning, and combining RAC with FP8 quantization (RAC+FP8) showed promising results, achieving both high accuracy and substantial throughput gains.
Interestingly, the runtime penalties often associated with pruned reasoning models, which can cause them to generate excessively long and rambling CoT traces, largely disappeared at more practical token limits. When constrained to reasonable decoding lengths (e.g., 8,192 or 4,096 tokens), RAC-pruned models performed similarly to, or even slightly better than, their dense counterparts in terms of both accuracy and runtime.
Also Read:
- Enhancing LLM Reasoning with Latent Thought Optimization
- The Evolving Logic of Language Models: A Cost-Benefit Analysis
In conclusion, Reasoning-Aware Compression (RAC) provides a simple yet powerful modification to existing pruning workflows, significantly improving the accuracy and stability of compressed reasoning language models. By aligning the calibration process with the decode-dominated nature of reasoning tasks, RAC helps maintain performance even under aggressive compression. While RAC does not directly optimize for runtime increases and adds some calibration cost, its indirect benefits in reducing unnecessary decoding steps and preserving accuracy make it a valuable advancement for deploying reasoning LLMs efficiently. For more details, you can read the full paper here.
