Scalable LLM Reasoning Acceleration with Low-rank Distillation

1. Problem Statement

Large Language Models (LLMs) are increasingly used for complex reasoning tasks (e.g., mathematics), which require long Chain-of-Thought (CoT) generations. While existing efficient inference methods (such as sparse activation pruning or quantization) successfully reduce latency and computational costs for standard language tasks, they often severely degrade performance on reasoning tasks.

• The Core Conflict: Reasoning tasks are sensitive to approximation errors. Small errors introduced by efficient inference algorithms (like skipping neurons or thresholding) accumulate over long generation sequences, causing the model to deviate from the correct reasoning path.
• The Gap: Current efficient methods (e.g., GRIFFIN, CATS) can reduce accuracy on math benchmarks (like GSM8K or MATH) by 30–50% while maintaining language task performance. There is a need for a method that preserves reasoning capabilities while maintaining the efficiency gains of sparse inference.

2. Methodology: Caprese

The authors propose Caprese (CAPability REcovery with Scalable Efficiency), a resource-efficient distillation method designed to recover the reasoning capabilities lost by compressed Feed-Forward (FF) blocks.

Key Observations

• Residual Structure: The authors observed that the error (residual) introduced by compressing an FF block (e.g., via Top-k selection or hard thresholding) exhibits a strong low-rank structure.
• Efficiency vs. Accuracy: Adding a small low-rank approximation to the residual significantly reduces error compared to simply increasing the density of the sparse selection.

Architecture & Mechanism

Caprese operates by adding a low-rank linear layer to the compressed FF block to approximate the missing information.

1. Base Compression: The method starts with an existing efficient inference algorithm (e.g., GRIFFIN or CATS) that compresses the FF block weights ($\hat{W}$).
2. Low-Rank Distillation: A low-rank matrix pair ($L \in \mathbb{R}^{D \times r}, R \in \mathbb{R}^{r \times D}$) is trained to model the residual error:
   $$\text{FF}_{\text{full}}(x) \approx \text{FF}_{\text{compressed}}(x) + xLR$$
   Where $r$ is the low rank (e.g., 256), which is negligible compared to the intermediate feature size ($D_{FF} \approx 14,000$).
3. Parallel Integration: The low-rank parameters are concatenated with the existing projection matrices. This allows the computation to run in parallel with the original FF operations, ensuring no additional sequential latency.
4. Training Strategy:
   • Data: Trained on only 20,000 synthetic math samples with CoT instructions.
   • Layer-wise Distillation: Each FF block is trained independently to minimize the residual error.
   • End-to-End (E2E) Distillation: The entire model is fine-tuned jointly to minimize the final embedding error before the output head.
   • Frozen Weights: The original model weights remain frozen; only the low-rank layers are updated.

3. Key Contributions

1. Performance Recovery: Caprese recovers "much if not all" of the reasoning performance lost by efficient inference methods. In some cases (e.g., DS-Qwen 7B with CATS), it even exceeds the original full model's accuracy (78.13% vs. 75.00% on AMC 2023).
2. Parameter Efficiency: It adds only ~1% additional parameters (roughly 256 rank) while enabling the removal of ~2 billion active parameters in 8B/9B models.
3. Latency Reduction: By integrating cleanly into existing layers, it preserves the latency reductions of the underlying sparse method, achieving >16% reduction in time-to-next-token compared to full models.
4. Brevity Induction: Caprese implicitly encourages shorter responses (up to 8.5% fewer tokens) without explicit length regularization, likely due to the model learning to be more precise with the recovered reasoning capabilities.
5. Test-Time Scaling Compatibility: The method works effectively with "Best-of-N" sampling and dynamic reselection strategies, further boosting performance.

4. Experimental Results

The paper evaluates Caprese on various models (Llama 3.1/3.2, Gemma 2, DeepSeek-R1-Distill-Qwen) and benchmarks (MATH-500, GSM8K, AIME, AMC, etc.).

• Reasoning Recovery:
  • DeepSeek-R1-Distill-Qwen 1.5B: GRIFFIN drops MATH-500 accuracy from 79.40% to 42.00%. E2E Caprese recovers it to 60.40%, and with reselection, it reaches 74.80% (near the original 79.40%).
  • DeepSeek-R1-Distill-Qwen 7B: CATS drops AMC 2023 accuracy to 62.50%. E2E Caprese restores it to 78.13%, surpassing the original model.
• Language Preservation: On instruct models (Llama 3.2, Gemma 2), Caprese maintains high performance on language tasks (CoQA, XSum, CNN/DailyMail) while recovering math performance, whereas the base sparse methods often degrade math scores significantly.
• Efficiency Metrics:
  • Active Parameters: Reduces active parameters by ~2B for 8B/9B models.
  • Latency: Reduces end-to-end latency by ~16% on Qwen 2.5 14B compared to the full model.
  • Comparison to LoRA: Unlike LoRA, which adds sequential computation and negates sparsity benefits, Caprese's parallel architecture maintains the speedup of the underlying sparse method.

5. Significance

Caprese addresses a critical bottleneck in deploying LLMs for reasoning: the trade-off between efficiency and accuracy.

• Practical Deployment: It enables the use of highly compressed, low-latency models for complex math and reasoning tasks without the massive computational cost of running full models.
• Scalability: The method is compatible with "test-time scaling" (generating multiple responses or longer CoTs), making it suitable for the next generation of "thinking" models.
• Cost-Effectiveness: The training cost is minimal (20k samples, low-rank parameters), making it accessible for organizations with limited compute budgets.

In summary, Caprese provides a scalable, low-cost solution to "fix" the broken reasoning capabilities of sparse LLMs, paving the way for efficient, high-performance AI agents capable of complex multi-step reasoning.