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

2025-10-01

AI Summary

Existing reasoning models lack adaptability in adjusting inference steps according to task difficulty: they tend to under-reason on complex problems—leading to errors—while over-reasoning on simple or already-solved ones, thereby degrading efficiency—a critical limitation in inference-length adaptability. This paper proposes TRAAC, the first method enabling fine-grained, dynamic inference-step control during online post-training. TRAAC introduces a task difficulty assessment module to quantify problem complexity, leverages attention mechanisms to identify and prune redundant reasoning steps, and designs a difficulty-aware reinforcement learning reward function to optimize inference budget allocation. Evaluated across multiple tasks, TRAAC achieves an average accuracy improvement of 8.4% and reduces average inference steps by 36.8%, while generalizing effectively beyond mathematical domains. It thus significantly enhances the trade-off between reasoning efficiency and accuracy.

Technology Category

• Knowledge Representation and Reasoning: Computational Complexity of Reasoning
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning
• Search and Optimization: Metareasoning and Metaheuristics

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Mitigating underthinking and overthinking in reasoning models
• Adaptively allocating reasoning steps based on task difficulty
• Compressing reasoning trajectories by pruning redundant thinking steps

Innovation

• Adaptive attentive compression prunes redundant reasoning steps
• Difficulty estimation allocates reasoning budget appropriately
• Online RL training balances accuracy and efficiency gains