Think Right: Learning to Mitigate Under-Over Thinking via Adaptive, Attentive Compression (2510.01581v1)

Abstract: Recent thinking models solve complex reasoning tasks by scaling test-time compute, but this scaling must be allocated in line with task difficulty. On one hand, short reasoning (underthinking) leads to errors on harder problems that require extended reasoning steps; but, excessively long reasoning (overthinking) can be token-inefficient, generating unnecessary steps even after reaching a correct intermediate solution. We refer to this as under-adaptivity, where the model fails to modulate its response length appropriately given problems of varying difficulty. To address under-adaptivity and strike a balance between under- and overthinking, we propose TRAAC (Think Right with Adaptive, Attentive Compression), an online post-training RL method that leverages the model's self-attention over a long reasoning trajectory to identify important steps and prune redundant ones. TRAAC also estimates difficulty and incorporates it into training rewards, thereby learning to allocate reasoning budget commensurate with example difficulty. Our approach improves accuracy, reduces reasoning steps, and enables adaptive thinking compared to base models and other RL baselines. Across a variety of tasks (AIME, AMC, GPQA-D, BBEH), TRAAC (Qwen3-4B) achieves an average absolute accuracy gain of 8.4% with a relative reduction in reasoning length of 36.8% compared to the base model, and a 7.9% accuracy gain paired with a 29.4% length drop compared to the best RL baseline. TRAAC also shows strong generalization: although our models are trained on math datasets, they show accuracy and efficiency gains on out-of-distribution non-math datasets like GPQA-D, BBEH, and OptimalThinkingBench. Our analysis further verifies that TRAAC provides fine-grained adjustments to thinking budget based on difficulty and that a combination of task-difficulty calibration and attention-based compression yields gains across diverse tasks.

• The paper introduces TRAAC, which uses adaptive, attentive compression to balance under- and overthinking in reasoning models.
• It employs Group Reward Policy Optimization and attention-based token scoring to dynamically adjust and compress reasoning steps.
• Experimental results show up to 8.4% accuracy improvement with a 36.8% reduction in reasoning length across diverse tasks.

TRAAC: Adaptive, Attentive Compression for Mitigating Under- and Overthinking in Reasoning Models

Introduction

The paper introduces TRAAC (Think Right with Adaptive, Attentive Compression), a post-training reinforcement learning (RL) framework designed to address the under-adaptivity problem in long-thinking LLMs. Under-adaptivity manifests as overthinking on easy problems (wasting tokens) and underthinking on hard problems (premature termination, loss of accuracy). TRAAC leverages attention-based compression and difficulty estimation to dynamically allocate reasoning effort, improving both accuracy and efficiency across diverse reasoning tasks.

Figure 1: TRAAC adaptively compresses reasoning trajectories, mitigating overthinking on easy problems and underthinking on hard problems by modulating token usage according to estimated difficulty.

TRAAC Methodology

Group Reward Policy Optimization (GRPO) Foundation

TRAAC builds on GRPO, an online RL algorithm that samples multiple rollouts per query and estimates the baseline from the group, eliminating the need for a critic. For each input $q$, the model generates $N$ rollouts, each consisting of a reasoning trajectory $r$ and a final answer $a$, separated by a delimiter (</think>). The pass rate among rollouts is used to estimate problem difficulty, which then informs both the compression ratio and reward assignment.

Attention-Based Compression

The core innovation is the attention-based compression module. After generating reasoning trajectories, the model computes the attention score for each reasoning token from the </think> delimiter, aggregating across all layers and heads:

$$s_j = \frac{1}{LH} \sum_{\ell=1}^L \sum_{h=1}^H \alpha_{</think> \to t_j}^{(\ell, h)}$$

Reasoning steps are segmented using control tokens (e.g., "Wait", "Alternatively", "First", "Then") and assigned importance scores as the mean of their constituent token scores. Steps with lower scores are pruned, yielding a compressed trajectory. The degree of pruning is modulated by the estimated difficulty: easy problems undergo aggressive compression, while hard problems retain more steps.

Figure 2: TRAAC pipeline: (1) generate rollouts, (2) estimate difficulty via pass rate, (3) compute attention scores for reasoning steps, (4) compress trajectory adaptively, (5) update policy using correctness and length rewards.

Difficulty-Level Calibration

Difficulty is estimated from the pass rate of rollouts and categorized into easy, medium, or hard. Each category is assigned a compression rate, with higher rates for easier problems. Additionally, the uniformity of attention scores is used to avoid excessive pruning when no step stands out as important. This calibration ensures fine-grained, instance-level adaptivity.

Reward Structure

TRAAC employs three reward signals:

• Correctness Reward (CR): High weight for correct final answers.
• Format Reward: Ensures proper use of delimiter tokens for trajectory/answer separation.
• Length Reward (LR): Penalizes unnecessarily long reasoning, with a sigmoid-based smoothing to avoid sharp cutoffs and maintain stability.

Rewards are combined and used to update the policy via GRPO.

Experimental Results

Performance and Efficiency Gains

TRAAC was evaluated on Qwen3-4B and Deepseek-Qwen-7B across AIME, AMC, GPQA-D, BBEH, and OptimalThinkingBench. TRAAC (Qwen3-4B) achieved an average absolute accuracy gain of 8.4% and a 36.8% reduction in reasoning length compared to the base model. Against the best RL baseline (AdaptThink), TRAAC delivered a 7.9% accuracy gain and a 29.4% length reduction.

Generalization to Out-of-Distribution Tasks

Despite training on math datasets, TRAAC generalized to non-math benchmarks (GPQA-D, BBEH, OverthinkingBench, UnderthinkingBench), showing an average 3% accuracy improvement and 40% reduction in response length on Qwen3-4B. Maximum improvement on UnderthinkBench was 6.8%.

Adaptive Token Allocation

TRAAC outperformed AdaptThink, which uses a binary "think/no-think" mode, by 7.9% in accuracy and 29.4% in efficiency. On OptimalThinkingBench, TRAAC improved the F1 metric by 7.36 points (Qwen3-4B) and 12.55 points (Deepseek-Qwen-7B) over the base model.

Figure 3: (a) TRAAC's compression rate decreases as problem difficulty increases, demonstrating adaptive allocation. (b) TRAAC maintains higher accuracy than static compression baselines across difficulty levels.

Ablation Studies

Ablations revealed that both difficulty-level calibration and attention-based compression are essential. Removing either component led to significant drops in both accuracy and efficiency. Attention-based compression outperformed random and confidence-based pruning by 7–11% in accuracy.

Scalability and Test-Time Compression

TRAAC scaled effectively with increased token budgets (up to 15k tokens), maintaining accuracy and efficiency gains. Applying attention-based compression only at test time yielded moderate improvements, but full training integration was necessary for optimal results.

Implementation Considerations

• Computational Requirements: Training TRAAC requires sampling multiple rollouts per query and computing attention scores across all layers/heads, necessitating substantial GPU memory (e.g., 4×A100 80GB).
• Integration: TRAAC is implemented as a post-training RL method and can be integrated with existing reasoning models supporting attention score extraction.
• Deployment: For production, attention-based compression can be applied at inference, but training-time integration yields superior adaptivity and performance.
• Limitations: TRAAC relies on verifiable rewards and may require adaptation for open-ended tasks without automatic answer verification.

Theoretical and Practical Implications

TRAAC demonstrates that explicit difficulty estimation and attention-based compression can overcome the trade-off between reasoning performance and efficiency. The approach provides a principled mechanism for dynamic resource allocation, moving beyond static length penalties or binary decision frameworks. The strong generalization to OOD tasks suggests that attention-based compression captures transferable reasoning patterns.

Future work may explore:

• Extending difficulty estimation to unsupervised or open-ended domains.
• Integrating TRAAC with larger-scale models and multi-modal reasoning.
• Investigating alternative token importance metrics (e.g., attribution, gradient-based scores).
• Applying adaptive compression in agentic or interactive settings.

Conclusion

TRAAC presents a robust framework for mitigating under- and overthinking in reasoning models via adaptive, attentive compression. By leveraging attention scores and difficulty estimation, TRAAC achieves simultaneous improvements in accuracy and efficiency, with strong generalization across domains. The combination of task-difficulty calibration and attention-based compression is critical for fine-grained, instance-level adaptivity, setting a new standard for resource allocation in long-thinking LLMs.