PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2510.01581v1cs.LGcs.AIcs.CL

TL;DR

Thinking models can misallocate reasoning length, underthinking hard problems and overthinking easy ones. TRAAC combines online reinforcement learning with difficulty-adaptive, attention-based compression. Across benchmarks, it improves accuracy while reducing reasoning length, and also reports gains on out-of-domain tasks.

  • Problem

    Thinking models may underthink difficult problems or overthink simple ones, so their reasoning length does not always match task difficulty.

    Source: §1 · p. 1

  • Method

    TRAAC is an online post-training RL method that uses attention-based compression and adapts reasoning compression to task difficulty.

    Source: §2.1 · p. 3, §2.2 · p. 4, §6 · p. 10

  • Results

    Across AMC, AIME, GPQA-D, and BBEH, TRAAC (Qwen3-4B) gains 8.4% average accuracy and reduces reasoning length by 36.8% versus the base model, with accuracy and efficiency gains also reported on out-of-domain tasks.

    Source: §1 · p. 2, §4.1 · p. 7

  • Takeaways & Limitations

    Removing difficulty calibration reduces average performance by 3.4% and efficiency by 23.8%, while removing both calibration and attention compression further lowers OptimalThinkingBench F1.

    Source: §4.2 · p. 8

Abstract

from arXiv · show

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.

1 INTRODUCTION

Thinking models can underthink hard problems or overthink easy ones, misallocating reasoning budget across tasks. TRAAC uses difficulty-adaptive attention-based compression to improve accuracy and efficiency.

1 INTRODUCTION

  • Under-adaptivity describes misallocated thinking budgets: short reasoning can fail on hard problems, while excessive reasoning on easy tasks wastes test-time computation.Prior approaches largely target overthinking through compressed-CoT fine-tuning, early stopping, length-penalty RL, or binary thinking choices.Source: §1 · pp. 1–2
  • TRAAC estimates task difficulty during online training and uses attention-based compression to prune reasoning, with more compression for easier problems.The model scores reasoning steps using attention from </think>; the estimated difficulty also informs training rewards.Source: §1 · p. 2, Figure 2 · p. 3
  • 8.4% average accuracy gain and 36.8% relative reasoning-length reduction for TRAAC (Qwen3-4B) versus the base model across AMC, AIME, GPQA-D, and BBEH.Against AdaptThink, TRAAC gains 7.9% in average accuracy and reduces length by 29.4%.Source: §1 · p. 2
  • On out-of-domain tasks, TRAAC shows an average 3% improvement on Qwen3-4B and an average 40% response-length reduction across those tasks.The paper also reports that difficulty calibration and attention-based compression together support performance and efficiency gains across diverse difficulty levels.Source: §1 · p. 3

2 TRAAC: THINK RIGHT WITH ADAPTIVE ATTENTIVE COMPRESSION

TRAAC is an online GRPO-based method that calibrates reasoning compression to task difficulty. It uses attention to identify less important reasoning steps and combines correctness, format, and length rewards to train the policy.

2 TRAAC: THINK RIGHT WITH ADAPTIVE ATTENTIVE COMPRESSION

  • TRAAC uses online GRPO training with an attention-based compression module that prunes unnecessary reasoning steps while adapting compression to estimated task difficulty.The design aims to preserve essential information while mitigating resource misallocation at test time.Source: §2 · p. 3

2.1 PROBLEM FORMULATION IN TRAAC

  • Problem difficulty is estimated from the proportion of correct answers among N rollouts, then used to set compression and assign rewards to each rollout.The answer is regenerated from the compressed trajectory, and advantages use both original and compressed rollouts.Source: §2.1 · pp. 3–4

2.2 ADAPTIVE, ATTENTIVE COMPRESSION MODULE

  • TRAAC scores reasoning tokens using attention from </think> across layers and heads, averages token scores within each step, and prunes lower-scoring steps.A prompt appended to the reasoning asks for a summary of key steps; difficulty sets the compression rate, while near-uniform attention scores reduce that rate to avoid removing potentially useful reasoning.Source: §2.2 · p. 4

2.3 REWARDS

  • The reward combines correctness, format, and length signals, with correctness weighted highly and length penalties adapted to task difficulty.The length reward uses a smooth bonus beyond the median and is granted only when the rollout reaches a correct final answer.Source: §2.3 · p. 5

3 EXPERIMENTAL SETUP

The experiments compare TRAAC with reasoning-model and compression baselines across math and non-math benchmarks, using repeated accuracy and response-length measurements. Training uses math data, while difficulty calibration assigns progressively lower compression scores to harder problems.

3 EXPERIMENTAL SETUP

  • Training uses Qwen3-4B and DeepSeek-R1-Distill-Qwen-7B, with eight rollouts per GRPO step at temperature 1.0.The rollout response-length maximum is 10k; difficulty calibration bins problems as easy, medium, or hard and assigns decreasing compression scores.Source: §3 · pp. 5–6
  • Evaluation covers AIME, AMC, GPQA-D, OverthinkingBench/UnderthinkingBench, and BBEH; only AIME and AMC are math-specific.OverthinkingBench evaluation uses automatically verifiable multiple-choice or numerical-answer problems rather than LLM judging.Source: §3 · p. 5
  • Mean accuracy and response length are reported across five runs at temperature 1.0, with a 10k-token maximum response length.OverthinkingBench also uses AUCOAA, and OptimalThinkingBench combines its AUCOAA with UnderthinkingBench accuracy into an F1 score.Source: §3 · p. 5
  • Table 1 reports that TRAAC achieves the highest average performance while substantially compressing response length.Acc. denotes accuracy (%) and Len. denotes average response length (k).Source: Table 1 · p. 6
  • TRAAC is compared with the base model, TokenSkip, L1-Max, LC-R1, and AdaptThink.The baselines use compressed-CoT supervised fine-tuning, length-constrained RL, externally guided compression, or adaptive thinking/no-thinking modes.Source: §3 · p. 6

4 RESULT AND DISCUSSION

TRAAC improves accuracy and efficiency across benchmarks while adapting reasoning effort to task difficulty. Ablations and scaling analyses examine the contributions of its components and the persistence of its gains under different settings.

4.1 MAIN RESULTS

  • 8.4% average accuracy improvement and 36.8% shorter reasoning length: TRAAC (Qwen3-4B) outperforms the base model across AIME, AMC, GPQA-D, and BBEH.TRAAC also outperforms TokenSkip in performance and efficiency for both evaluated models, and achieves higher accuracy with shorter traces than efficiency-focused RL baselines.Source: §4.1 · p. 6
  • 3% average improvement on Qwen3-4B and 2.8% on Deepseek-Qwen-7B, alongside 40% and 20% fewer reasoning tokens, respectively, on out-of-domain tasks.The improvement reaches 6.8% on UnderthinkingBench; training used math data, while these evaluations include GPQA-D, BBEH, OverthinkingBench, and UnderthinkingBench.Source: §4.1 · p. 7
  • 7.9% higher performance and 29.4% fewer tokens than AdaptThink on Qwen3-4B; TRAAC also improves OptimalThinkingBench F1 over the base models.On OptimalThinkingBench, F1 gains are 7.36% for Qwen3-4B and 12.55% for Deepseek-Qwen-7B; OverthinkingBench AUCOAA gains over the base are 5% and 0.5%, respectively.Source: §4.1 · p. 7, Table 2 · p. 7

4.2 ABLATIONS AND ANALYSIS

  • Removing difficulty calibration reduces average performance by 3.4% and efficiency by 23.8%; on OptimalThinkingBench, removing calibration lowers F1 by 5.2%, and removing compression as well lowers it by 7.6%.The ablations compare successive component additions on the four benchmark datasets and OptimalThinkingBench.Source: §4.2 · p. 8, Table 3 · p. 8, Table 4 · p. 8
  • As task difficulty rises, TRAAC's compression rate falls, while the uncalibrated compression variant remains roughly uniform; TRAAC maintains higher accuracy than that variant.The plotted tasks are ordered by increasing difficulty, and the reported accuracy gains mainly come from harder problems.Source: §4.2 · p. 8, Figure 3 · p. 9
  • 3.5% average accuracy improvement and 23.4% efficiency gains persist when TRAAC's maximum training and test response length is increased to 15k tokens.The reported evaluation covers AIME, AMC, and GPQA-D.Source: §4.2 · p. 9, Table 5 · p. 9
  • 11% average accuracy drop with random pruning and 7.25% with least-confidence pruning, compared with TRAAC (Qwen3-4B).The alternatives prune reasoning steps randomly or by least confidence instead of using attention-based compression.Source: §4.2 · p. 9, Table 6 · p. 9
  • 7.28% higher accuracy and 27.5% higher efficiency for fully trained TRAAC than inference-only compression, averaged across AIME, AMC, and GPQA-D.The inference-only variant still improves over the base model by 3.11% in accuracy and 12.65% in efficiency on average.Source: Table 7 · p. 9, §4.2 · p. 10

5 RELATED WORK

Prior efficient-reasoning methods compress or shorten reasoning but generally do not explicitly account for problem difficulty. TRAAC instead targets how much to think and uses attention-based supervision through online RL.

5 RELATED WORK

  • Prior work identifies sequential reasoning-token scaling as effective, while efficient reasoning seeks to maintain performance with minimal token use.Approaches include early exiting, budgeted training, latent thoughts, supervised fine-tuning, and length penalties.Source: §5 · p. 10
  • Earlier approaches generally leave budget allocation to the model, whereas TRAAC explicitly uses difficulty-based supervision and addresses both overthinking and underthinking.The passage contrasts this with methods that mainly shorten outputs, sometimes at the cost of performance drops.Source: §5 · p. 10
  • Unlike binary think-or-not decisions or post-hoc pruning, TRAAC chooses how much to think and adapts compression during online RL training.The passage reports that TRAAC empirically outperforms a binary adaptivity baseline.Source: §5 · p. 10

6 CONCLUSION

The conclusion presents TRAAC as online post-training RL with difficulty-adaptive, attention-based compression. It reports adaptive pruning, generalization to OOD tasks, and analysis supporting combined difficulty calibration and attention-based compression.

6 CONCLUSION

  • TRAAC is an online post-training RL method using a difficulty-adaptive, attention-based compression module to prune reasoning steps.The pruning is adapted to task difficulty.Source: §6 · p. 10
  • The method aims to improve performance and efficiency by thinking longer on harder problems and shorter on easier ones.The conclusion links longer thinking on harder problems with better exploration and shorter thinking on easier problems with less wasted test-time compute.Source: §6 · p. 10
  • The conclusion reports OOD generalizability and says analysis and ablations support fine-grained budget adjustment and gains from combining difficulty calibration with attention-based compression.Source: §6 · p. 10

A.1 DATASET DETAILS

The appendix describes benchmark datasets and an accuracy-evaluation method. It lists source, scope, or size information for several math, science, and reasoning benchmarks.

A.1 DATASET DETAILS

  • AMC draws on AMC12 2022 and 2023 questions (83 total), while AIME uses AIME 22–24 questions (90 total).Both question sets were extracted from the AOPS wiki page.Source: §A.1 · p. 14
  • GPQA-D covers physics, biology, and chemistry (198 questions), and BBEH contains 460 harder tasks probing capabilities similar to BBH.BBEH replaces each BBH task with a novel task of significantly increased difficulty.Source: §A.1 · p. 14
  • OptimalThinkingBench includes OverthinkingBench (607 questions) and UnderthinkingBench (550 questions), while BBH has 2,115 questions and SuperGPQA has 1,620.SuperGPQA covers 285 graduate-level disciplines, with 540 questions at each of three difficulty levels.Source: §A.1 · p. 14
  • Accuracy is calculated with Math-Verify, except UnderthinkingBench uses the Reasoning-Gym evaluation scripts.Source: §A.1 · p. 14

A.2 DEEPSEEK ABLATION AND ANALYSIS

The appendix describes reward variants for a Deepseek-Qwen-7B ablation evaluated on four datasets. The supplied passages specify the variants and table scope, but not their outcome values.

A.2 DEEPSEEK ABLATION AND ANALYSIS

  • The base-model correctness-reward variant trains with GRPO using only the correctness reward.Source: §A.2 · p. 14
  • A second variant adds a length reward to correctness reward but omits difficulty-level calibration.Source: §A.2 · p. 14
  • Table 8 covers Deepseek-Qwen-7B across AIME, AMC, GPQA-D, and BBEH, with each component addition building on the previous method.Source: Table 8 · p. 14

A.3.1 AUXILIARY PROMPT

The appendix specifies TRAAC’s trajectory segmentation and compression procedures, training rewards, evaluation metrics, and experimental settings. It also lists prompts and implementation details for baselines and training.

A.3.1 AUXILIARY PROMPT

  • Table 9 compares incremental TRAAC component additions on Deepseek-Qwen-7B using OptimalThinkingBench, but the supplied caption gives no result values.The auxiliary prompt appended to each trajectory asks the model to stop and summarize the key solution steps.Source: §A.3.1 · p. 14, Table 9 · p. 15

A.3.2 SPECIAL TOKENS TO SPLIT TRAJECTORY TO CHUNKS

  • A list of control phrases, including “Wait” and “Therefore,” splits each reasoning trajectory into multiple reasoning steps.The full list is provided in the appendix.Source: §A.3.2 · p. 15

A.3.3 UNIFORMITY SCORE

  • The eviction percentage is calculated from an attention-uniformity score, with high uniformity preserving all steps and nonuniformity scaling eviction up to 80%.The algorithm’s target reduction defaults to 0.25; the score is normalized entropy over step-importance scores.Source: §A.3.3 · p. 15

A.4 GRPO DETAILS; A.5 EXPERIMENTAL DETAILS; A.5.1 HYPERPARAMETERS

  • For each question, GRPO samples a response group from the old policy and optimizes the policy using relative verifiable-outcome advantages.The appendix identifies verl as the training framework and provides a hyperparameter table covering training, evaluation, and difficulty calibration.Source: §A.4 · p. 16, §A.5 · p. 16, Table 10 · p. 16

A.5.2 TRAINING REWARD

  • The total training reward sums correctness, format, and length rewards, weighting correctness at +4 and assigning the length component a range of 0 to 2.The format component ranges from 0 to 1 and rewards the presence and proper enclosure of reasoning with <think> and </think> tokens.Source: §A.5.2 · p. 16

A.5.3 EVALUATION METRICS

  • OverthinkingBench evaluation includes OAA, correctness under a reasoning-token threshold, and AUCOAA; OptimalThinkingBench combines AUCOAA with UnderthinkingBench accuracy into an F1 score.The authors report computing AUCOAA and F1 with their own implementation because Cheng et al. (2025)’s evaluation scripts were not open-sourced.Source: §A.5.3 · pp. 16–17

A.5.4 TRAINING PROMPT; A.6.1 BASELINE PROMPTS; A.7 COMPUTE USED

  • For each question in the training set, the instruction was “Let’s think step by step and output the final answer within \\boxed{}”; all training was done on 4*A100 (80GB).Source: §A.5.4 · p. 17, §A.7 · p. 18

A.6 CONFIDENCE BASED COMPRESSION

  • Confidence-based compression scores tokens from their log-probability distribution and removes reasoning steps with the lowest average scores.The algorithm defines confidence as normalized inverse entropy of the top-k token probabilities.Source: §A.6 · p. 17
Loading 2510.01581v1…