A State-Transition Framework for Efficient LLM Reasoning

TL;DR

This work tackles the high computational and memory cost of long Chain-of-Thought reasoning in LLMs by modeling reasoning as a state-transition process and introducing a Mixed Attention Module that combines linear and softmax attention. A persistent reasoning state S_t is maintained via linear attention, enabling each new reasoning step to access history without quadratic attention costs, while preserving CoT interpretability. A state-based reasoning strategy mitigates overthinking from noisy steps through momentum-guided corrections and diversified thinking patterns. Training uses long CoT data with autoregressive and distillation losses, and experiments across seven benchmarks show consistent improvements in both reasoning speed and accuracy, with strong ablations validating the contributions and a demonstrated domain generalization to scientific and code reasoning tasks.

Abstract

While Long Chain-of-Thought (CoT) reasoning significantly improves Large Language Models (LLMs) performance on complex reasoning tasks, the substantial computational and memory costs of generating long CoT sequences limit their efficiency and practicality. Existing studies usually enhance the reasoning efficiency of LLMs by compressing CoT sequences. However, this approach conflicts with test-time scaling, limiting the reasoning capacity of LLMs. In this paper, we propose an efficient reasoning framework that models the reasoning process of LLMs as a state-transition process. Specifically, we first apply a linear attention mechanism to estimate the LLM's reasoning state, which records the historical reasoning information from previous reasoning steps. Then, based on the query prompt and the reasoning state, the LLM can efficiently perform the current reasoning step and update the state. With the linear attention, each token in the current reasoning step can directly retrieve relevant historical reasoning information from the reasoning state, without explicitly attending to tokens in previous reasoning steps. In this way, the computational complexity of attention is reduced from quadratic to linear, significantly improving the reasoning efficiency of LLMs. In addition, we propose a state-based reasoning strategy to mitigate the over-thinking issue caused by noisy reasoning steps. Extensive experiments across multiple datasets and model sizes demonstrate that our framework not only improves the reasoning efficiency of LLMs but also enhances their reasoning performance.

Figures (9)

• Figure 1: A comparison of traditional token‑based reasoning with our state‑based reasoning. $\textbf{think}_{t}$ denotes one reasoning step. In state-based reasoning, LLMs can efficiently generate $\textbf{think}_{t}$ only based on the query prompt and the reasoning state $\boldsymbol{S}_{t{-}1}$.
• Figure 2: In our framework, the original softmax attention module in LLMs is replaced with our mixed attention module (MAM), thus improving reasoning efficiency and reducing memory cost.
• Figure 3: Illustration of our reasoning and training strategies. SAM is the softmax attention module.
• Figure 4: (a) shows the computational and memory efficiency of our model and the base model. (b) and (c) present our model’s performance with different values of hyper-parameters $\beta$ and $\alpha_{\rm max}$, respectively. These experiments are conducted on Qwen2.5-1.5B.
• Figure 5: (a) shows the performance of our model and the baselines on the GPQA Diamond and HumanEval benchmarks. (b) illustrates the average gating scores in the LA submodule for tokens at different positions within each reasoning step. These experiments are conducted on Qwen2.5-7B.