💬 NLP

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

This paper proposes Dynamic Thinking-Token Selection (DynTS), a method that enhances the efficiency of Large Reasoning Models by identifying and retaining only the Key-Value cache states of decision-critical tokens in reasoning traces while evicting redundant entries.

Original authors: Zhenyuan Guo, Tong Chen, Wenlong Meng, Chen Gong, Xin Yu, Chengkun Wei, Wenzhi Chen

Published 2026-06-05

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhenyuan Guo, Tong Chen, Wenlong Meng, Chen Gong, Xin Yu, Chengkun Wei, Wenzhi Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Over-Thinker" with a Bad Memory

Imagine you are a brilliant detective solving a very complex mystery. To find the answer, you don’t just guess; you write down every single thought, clue, and deduction in a massive notebook. This is how Large Reasoning Models (LRMs) work. They generate long "reasoning traces" (like a chain of thought) before giving you the final answer.

However, there is a catch. As the detective writes more and more, the notebook gets incredibly heavy. In the world of AI, this notebook is called the KV Cache (Key-Value Cache). It stores the memory of everything the model has processed so far.

The Issue: The notebook becomes so huge that it slows the detective down and requires a massive amount of desk space (computer memory).
The Bottleneck: Most existing methods to shrink the notebook are like blindly tearing out random pages or keeping only the most recent ones. This often causes the detective to forget crucial clues, leading to wrong answers.

The Insight: The 80/20 Rule of Thinking

The researchers looked closely at these "notebooks" and discovered something surprising, which they call the Pareto Principle (or the 80/20 rule) in AI reasoning:

The Observation: Out of thousands of words the model writes during its thinking process, only a tiny fraction (about 20–30%) are actually critical for finding the final answer.
The Metaphor: Think of the reasoning trace like a long movie. Most of the movie is background chatter, walking scenes, or setup. But there are only a few specific "plot-twist" moments that actually determine how the story ends. The rest is just "filler."
The Finding: If you keep only those critical "plot-twist" tokens (words) and throw away the filler, the AI can still solve the problem correctly. The filler words are mostly just grammatical glue or redundant steps that don’t change the outcome.

The Solution: DYNTS (The Smart Editor)

To solve the memory problem without losing accuracy, the authors created a method called DYNTS (Dynamic Thinking-Token Selection).

Instead of blindly deleting pages, DYNTS acts like a smart editor who reads the detective’s notes in real-time and decides what to keep.

The Importance Predictor: They attached a small, lightweight "brain" (a simple neural network) to the AI. As the AI writes each word of its reasoning, this little brain predicts: "Is this word crucial for the final answer, or is it just filler?"
Training the Editor: They taught this editor by looking at finished cases. They calculated which words in the past were most important by seeing which ones the final answer "paid attention to." The editor learned to spot these high-value words.
The Dual-Window System: During the solving process, DYNTS manages memory using two zones:
- The Local Window: Keeps the most recent thoughts (like the last few sentences) because context matters.
- The Selection Window: Holds the older thoughts. Here, the "Smart Editor" kicks in. It keeps only the top-scoring "critical" tokens and evicts (deletes) the rest to save space.

The Results: Faster and Smarter

The researchers tested DYNTS on six difficult benchmarks (math and science problems) using two popular AI models. The results were impressive:

Same Accuracy, Less Memory: DYNTS reduced the memory footprint by 3.3 to 5.7 times compared to the standard method. It’s like shrinking a 10-volume encyclopedia down to a single pocket-sized guide without losing the facts.
Faster Speed: Because the notebook is smaller, the AI processes information faster. It reduced inference latency (slowness) by 1.8 to 2.6 times.
Better than Competitors: Compared to other state-of-the-art memory-saving methods, DYNTS was actually 2.6% more accurate while using the same amount of memory budget. Other methods often hurt performance because they accidentally deleted important clues; DYNTS avoided this by specifically targeting the "decision-critical" tokens.

Summary

In short, this paper shows that AI reasoning models talk a lot more than they need to. By using a small, trained helper to identify and keep only the critical thinking steps and discarding the rest, we can make these powerful models run much faster and use significantly less memory, without making them dumber.

Technical Summary: Dynamic Thinking-Token Selection for Efficient Reasoning in Large Reasoning Models

Problem Statement
Large Reasoning Models (LRMs), such as DeepSeek-R1 and Gemini-3-Pro, enhance reasoning capabilities by generating extensive intermediate reasoning traces before producing a final answer. While effective, this process incurs significant computational overhead and memory footprint due to the accumulation of Key-Value (KV) cache states during the long decoding phase. Existing KV cache compression techniques, designed primarily for long-context prefilling or standard language modeling, are ill-suited for the "short-prefill, long-decoding" paradigm of LRMs. Furthermore, current methods often fail to distinguish between tokens critical for the final answer and redundant syntactic scaffolding, leading to a degradation in reasoning performance when compression is applied.

Methodology
The authors propose Dynamic Thinking-Token Selection (DYNTS), a method that identifies and retains only the decision-critical tokens within a reasoning trace while evicting redundant entries. The approach is built on the observation of a "Pareto Principle" in LRM reasoning: only a small subset (~20–30%) of thinking tokens significantly influences the final answer, as evidenced by their attention weights from the answer tokens.

The methodology consists of two main components:

Importance Predictor: Since attention weights from the answer to thinking tokens are unavailable during the reasoning phase, DYNTS introduces a lightweight, trainable Multi-Layer Perceptron (MLP) head attached to the final layer of the LRM. This predictor dynamically estimates the importance score of each generated thinking token in real-time. It is trained using Mean Squared Error (MSE) loss against ground-truth importance scores derived from full autoregressive inference on a training set (MATH), where importance is calculated by aggregating attention weights from answer tokens to thinking tokens. The backbone model parameters remain frozen during this training.
KV Cache Selection Strategy: During inference, DYNTS manages memory using a dual-window mechanism within a fixed budget $B$ $B$ . The budget is partitioned into a question window ( $W_q$ $W_{q}$ ), a selection window ( $W_s$ $W_{s}$ ), and a local window ( $W_l$ $W_{l}$ ).
- Question Window: Stores KV caches for all prompt tokens, which are deemed critical and never evicted.
- Local Window: Retains the most recent tokens to ensure local contextual coherence.
- Selection Window: Stores historical thinking tokens. When the cache reaches the budget limit, the system retains the KV caches of the top- $k$ tokens with the highest predicted importance scores from this window and evicts the rest.

Key Contributions

Empirical Insight: The paper identifies and validates the sparsity of critical thinking tokens in LRMs, demonstrating that retaining only the top ~30% of tokens based on answer-centric attention scores preserves nearly full reasoning performance.
DYNTS Framework: A novel inference-time compression method that uses a learned predictor to dynamically select critical tokens, avoiding the performance degradation associated with heuristic-based eviction methods.
Theoretical Analysis: The authors derive a "Break-Even Condition" for computational gain, proving that the savings from reduced attention computation outweigh the overhead of the importance predictor provided the eviction volume exceeds a threshold determined by model architecture.

Results
DYNTS was evaluated on two mainstream LRMs (DeepSeek-R1-Distill-Llama-8B and DeepSeek-R1-Distill-Qwen-7B) across six reasoning benchmarks (AIME24, AIME25, AMC23, GPQA-D, GAOKAO2023EN, and MATH500).

Accuracy: DYNTS outperforms state-of-the-art KV cache compression baselines (including R-KV, SnapKV, H2O, and StreamingLLM). Under the same memory budget, DYNTS improves Pass@1 accuracy by 2.6% compared to the best baseline. Its performance is comparable to the full-cache Transformer baseline (e.g., 61.9% vs. 61.6% on R1-Llama).
Efficiency: Compared to vanilla Transformers, DYNTS reduces inference latency by 1.84–2.62× and peak KV-cache memory footprint by 3.32–5.73×. It achieves a peak throughput speedup of up to 4.51× as decoding steps increase, effectively bounding resource consumption.
Ablation: Experiments confirm that all components of DYNTS (question tokens, critical thinking tokens, and local window) are essential. Removing the importance-based selection (replacing it with random or bottom-k selection) leads to significant performance drops, validating the effectiveness of the predictor.

Significance
The paper claims that DYNTS addresses the critical deployment challenge of LRMs in resource-constrained environments by significantly reducing memory and computational costs without compromising reasoning accuracy. By shifting from heuristic-based eviction to a learned, answer-centric importance prediction, DYNTS aligns cache retention with the actual decision-critical steps of the model, offering a more efficient and accurate alternative to existing compression techniques.

The Problem: The "Over-Thinker" with a Bad Memory

The Insight: The 80/20 Rule of Thinking

The Solution: DYNTS (The Smart Editor)

The Results: Faster and Smarter

Summary

More like this