Dynamic Thinking-Token Selection for Efficient Reasoning in Large Reasoning Models

TL;DR

This work tackles the heavy memory and computation costs of extended reasoning traces in Large Reasoning Models by revealing a Pareto principle: only a small subset (~20–30%) of thinking tokens critically influence the final answer. It introduces Dynamic Thinking-Token Selection (DynTS), which combines a lightweight Importance Predictor with a decoding-time KV-cache selection strategy to retain only the most impactful thinking tokens and relevant context, bounded by a budget. Empirical results across two LRMs and six math-reasoning benchmarks show DynTS achieves up to 3.32–5.73× KV-cache memory reduction, up to 4.51× peak throughput gains, and a 2.6% accuracy improvement over strong baselines, closely matching full-cache performance. The approach enables memory- and compute-efficient inference for LRMs, with practical implications for deploying powerful reasoning systems in resource-constrained environments.

Abstract

Large Reasoning Models (LRMs) excel at solving complex problems by explicitly generating a reasoning trace before deriving the final answer. However, these extended generations incur substantial memory footprint and computational overhead, bottlenecking LRMs' efficiency. This work uses attention maps to analyze the influence of reasoning traces and uncover an interesting phenomenon: only some decision-critical tokens in a reasoning trace steer the model toward the final answer, while the remaining tokens contribute negligibly. Building on this observation, we propose Dynamic Thinking-Token Selection (DynTS). This method identifies decision-critical tokens and retains only their associated Key-Value (KV) cache states during inference, evicting the remaining redundant entries to optimize efficiency.

Figures (17)

• Figure 1: (Left) Comparison of token selection strategies across different KV cache eviction methods. In each row, colored blocks denote the retained high-importance tokens, while grey blocks represent the evicted tokens during LRM inference. (Right) The average reasoning performance and KV cache memory footprint on DeepSeek-R1-Distall-Llama-8B and DeepSeek-R1-Distall-Qwen-7B across six reasoning benchmarks.
• Figure 2: Importance scores of question tokens and thinking tokens in a reasoning trace, computed based on attention contributions to the answer. Darker colors indicate higher importance. The red dashed line shows the mean importance score, and the annotated ratio indicates the fraction of tokens with importance above the mean.
• Figure 3: (Left) Reasoning performance trends as a function of thinking token retention ratio, where the $x$-axis indicates the retention percentage and the $y$-axis is the accuracy. (Right) Accuracy across all datasets when retaining $30\%$ of the thinking tokens.
• Figure 4: Overview of DynTS. (Left) Importance Predictor Training. The upper heatmap visualizes attention weights, where orange intensity represents the importance of thinking tokens to the answer. The lower part shows a LRM integrated with an Importance Predictor (IP) to learn these importance scores. (Right) Inference with KV Cache Selection. The model outputs the next token and a predicted importance score of the current token. When the cache budget is reached, the selection strategy retains the KV cache of question tokens, local tokens, and top-k thinking tokens based on the predicted importance score.
• Figure 5: Real-time throughput, memory, and compute overhead tracking over total decoding step. The inflection points in the sawtooth correspond to the steps where DynTS executes KV Cache Selection.

Theorems & Definitions (6)

• Theorem 5.1: Computational Gain
• Corollary 5.2: Break-even Condition
• Definition 1.1: Predictor Overhead
• Definition 1.2: Effective KV Cache Length
• Definition 1.3: LLM Overhead
• proof : Proof: Computational Gain