Pense Certo: Aprendendo a Mitigar o Pensamento Sub-Superestimado por meio de Compressão Adaptativa e Atenta

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

October 2, 2025

Autores: Joykirat Singh, Justin Chih-Yao Chen, Archiki Prasad, Elias Stengel-Eskin, Akshay Nambi, Mohit Bansal

cs.AI

Resumo

Modelos de pensamento recentes resolvem tarefas complexas de raciocínio escalonando o cálculo em tempo de teste, mas essa escalonamento deve ser alocado de acordo com a dificuldade da tarefa. Por um lado, o raciocínio curto (subpensamento) leva a erros em problemas mais difíceis que exigem etapas de raciocínio prolongadas; por outro, o raciocínio excessivamente longo (sobrepensamento) pode ser ineficiente em termos de tokens, gerando etapas desnecessárias mesmo após alcançar uma solução intermediária correta. Referimo-nos a isso como subadaptatividade, onde o modelo falha em modular adequadamente o comprimento de sua resposta diante de problemas de dificuldade variável. Para abordar a subadaptatividade e encontrar um equilíbrio entre subpensamento e sobrepensamento, propomos o TRAAC (Think Right with Adaptive, Attentive Compression), um método de RL pós-treinamento online que aproveita a auto-atenção do modelo ao longo de uma trajetória de raciocínio extensa para identificar etapas importantes e podar as redundantes. O TRAAC também estima a dificuldade e a incorpora nas recompensas de treinamento, aprendendo assim a alocar o orçamento de raciocínio proporcionalmente à dificuldade do exemplo. Nossa abordagem melhora a precisão, reduz as etapas de raciocínio e permite um pensamento adaptativo em comparação com modelos base e outras linhas de base de RL. Em uma variedade de tarefas (AIME, AMC, GPQA-D, BBEH), o TRAAC (Qwen3-4B) alcança um ganho médio absoluto de precisão de 8,4% com uma redução relativa no comprimento do raciocínio de 36,8% em comparação com o modelo base, e um ganho de precisão de 7,9% combinado com uma redução de 29,4% no comprimento em relação à melhor linha de base de RL. O TRAAC também mostra forte generalização: embora nossos modelos sejam treinados em conjuntos de dados matemáticos, eles apresentam ganhos de precisão e eficiência em conjuntos de dados fora da distribuição, como GPQA-D, BBEH e OptimalThinkingBench. Nossa análise ainda verifica que o TRAAC fornece ajustes refinados no orçamento de pensamento com base na dificuldade e que a combinação de calibração de dificuldade da tarefa e compressão baseada em atenção produz ganhos em diversas tarefas.

English

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.