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.