Think Right: Learning to Mitigate Under-Over Thinking via Adaptive, Attentive Compression

TL;DR

TRAAC tackles under-adaptivity in long-horizon reasoning by introducing an online post-training RL framework that uses an attention-based compression module and task-difficulty calibration to dynamically prune the reasoning trajectory. Built on Group Reward Policy Optimization, TRAAC estimates problem difficulty from rollout pass rates and allocates a flexible reasoning budget via an attention-guided pruning mechanism, yielding improved accuracy and significant reductions in token usage across math and non-math benchmarks. Key contributions include the difficulty-aware compression strategy, the attention-based pruning mechanism, and a multi-component reward system that balances correctness, formatting, and length, with strong generalization to out-of-domain tasks. The results demonstrate that adaptive, attentive compression can simultaneously enhance performance and efficiency, offering practical benefits for scalable reasoning in LLMs across diverse tasks and datasets.

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.

Figures (3)

• Figure 1: Overthinking on easy problems wastes tokens by continuing computation after a correct answer has been reached. On the other hand, underthinking on hard problems saves token budgets but fails to maintain accuracy. TRAAC addresses this trade-off by adapting to problem difficulty (estimated during training) through attention-based compression, enabling intelligent resource allocation while improving both accuracy and efficiency.
• Figure 2: Overview of TRAAC. Given a problem, the model first generates $N$ rollouts, and the pass rate of these rollouts is used to estimate the problem's difficulty (easy, medium, or hard). Next, the generated reasoning is fed back into the model, which is asked to compute the attention score of each reasoning token from </think>. During this attention-based compression step, we remove steps with lower scores. The degree of removal is determined by the estimated difficulty: easier problems undergo more aggressive compression. Finally, we compute the correctness and length rewards using the compressed reasoning trajectory, and these rewards are used to update the policy.
• Figure 3: (a) Relative change in compression rate of TRAAC and Qwen3-4B + Compression compared to Qwen3-4B across varying problem difficulty. (b) Absolute accuracy drop of TRAAC and Qwen3-4B + Compression compared to Qwen3-4B across varying problem difficulty.