Scalable LLM Math Reasoning Acceleration with Low-rank Distillation

TL;DR

This work tackles the latency of mathematical reasoning in LLMs under efficient inference. It introduces Caprese, a low-rank residual distillation approach that learns compact corrections to FF outputs with a small inner dimension (e.g., $r=256$) and trains on a modest dataset of 20K synthetic math problems while keeping original FF weights intact. Caprese recovers most or all math performance lost due to sparse FF methods and preserves language-task performance, functioning with baselines like GRIFFIN and CATS and delivering substantial latency reductions. The method scales across instruct and thinking models, requires only about 1% additional parameters, and achieves practical gains in coverage and response brevity, advancing degradation-free efficient LLM inference for complex reasoning tasks.

Abstract

Due to long generations, large language model (LLM) math reasoning demands significant computational resources and time. While many existing efficient inference methods have been developed with excellent performance preservation on language tasks, they often severely degrade math performance. In this paper, we propose Caprese, a resource-efficient distillation method to recover lost capabilities from deploying efficient inference methods, focused primarily in feedforward blocks. With original weights unperturbed, roughly 1% of additional parameters, and only 20K synthetic training samples, we are able to recover much if not all of the math capabilities lost from efficient inference for thinking LLMs and without harm to language tasks for instruct LLMs. Moreover, Caprese slashes the number of active parameters (~2B cut for Gemma 2 9B and Llama 3.1 8B) and integrates cleanly into existing model layers to reduce latency (>16% time-to-next-token reduction) while encouraging response brevity (up to 8.5% fewer tokens).

Figures (7)

• Figure 1: A full FF block maximizes accuracy without any benefit to efficiency. Sparse FF algorithms can be very efficient by using subsets of the FF block but harm math performance. Our method, Caprese, uses a sparse FF algorithm and a small distilled low-rank linear layer, which can be merged with existing FF weights, for performative inference in language and math settings while being efficient. Layers are drawn as trapezoids to highlight the relative size of inputs and outputs.
• Figure 2: Average relative FF output error of generated tokens with varying top-$k$ densities and low-rank approximations. The density is the fraction of non-zero intermediate FF features maintained by top-$k$. Samples consist of 16 random MATH generations by the original model. A relatively small low-rank approximation to the top-$k$ residual reduces error more effectively than increasing density.
• Figure 3: Coverage and standard deviation of 140 samples from MATH as the number of generation attempts, $N$, scales. We define $\text{Relative Compute Units} = N \times A$ where $A$ is the fraction of total parameters activated per input. FF sparsity set to 50%. Best viewed zoomed.
• Figure 4: Average number of response tokens for MATH-500 queries with increasing problem difficulty. The global averages are indicated by the dashed lines. Sparsity is set at 50%.
• Figure 5: Llama 3.2 1B Instruct's performance on MATH with varying densities of CATS and ranks in Caprese with end-to-end training.