Source-linked AI summary
ATTNPO: Attention-Guided Process Supervision for Efficient Reasoning
Shuaiyi Nie, Siyu Ding, Wenyuan Zhang, Linhao Yu, Tianmeng Yang, Yao Chen, Weichong Yin, Yu Sun, Hua Wu, Tingwen Liu
TL;DR
Long reasoning models can overthink, while existing length penalties and process-supervised methods lack fine-grained, low-overhead step-level credit assignment. ATTNPO uses intrinsic attention signals from Key-Focus Heads to discourage redundant steps while preserving essential reasoning, substantially shortening reasoning and improving accuracy across benchmarks.
Problem
Existing methods overthink or incur high overhead because they treat reasoning steps uniformly and lack accurate independent step-wise credit assignment.
Method
ATTNPO uses Key-Focus Heads’ intrinsic attention signals for step-level credit assignment, scaling penalties to suppress redundant steps while preserving essential reasoning.
Results
ATTNPO achieves an average +7.3-point absolute accuracy gain with a 60% reduction in reasoning length across six math benchmarks.
Takeaways & Limitations
KFH-guided process supervision improves reasoning efficiency without sacrificing performance.
Takeaways & Limitations
Experiments use only 1.5B and 7B models and train exclusively on mathematical datasets, although evaluations indicate generalization to OOD scenarios.
Abstract
from arXiv · showhide
Large reasoning models trained with reinforcement learning and verifiable rewards (RLVR) achieve strong performance on complex reasoning tasks, yet often overthink, generating redundant reasoning without performance gains. Existing trajectory-level length penalties often fail to effectively shorten reasoning length and degrade accuracy, as they uniformly treat all reasoning steps and lack fine-grained signals to distinguish redundancy from necessity. Meanwhile, process-supervised methods are typically resource-intensive and suffer from inaccurate credit assignment. To address these issues, we propose ATTNPO, a low-overhead process-supervised RL framework that leverages the model's intrinsic attention signals for step-level credit assignment. We first identify a set of special attention heads that naturally focus on essential steps while suppressing redundant ones. By leveraging the attention scores of these heads, We then employ two sub-strategies to mitigate overthinking by discouraging redundant steps while preserving accuracy by reducing penalties on essential steps. Experimental results show that ATTNPO substantially reduces reasoning length while significantly improving performance across 9 benchmarks.
1 Introduction
Large reasoning models achieve strong complex-task performance but often overthink, while existing length penalties and process-supervised methods struggle with fine-grained, accurate, low-cost credit assignment. ATTNPO addresses this gap with attention-guided step supervision, reducing reasoning length while improving accuracy.
- Long chain-of-thought reasoning can waste computation by applying verbose reasoning to trivial operations.
- Trajectory-level length penalties assign uniform credit across reasoning steps, failing to distinguish necessary from redundant steps.
- Existing process-supervised methods require additional sampling or reward models and may allocate step-wise credit inaccurately.
- Key-Focus Heads naturally emphasize essential steps and suppress redundant ones, with stable, sparse, and mainly middle-to-late-layer behavior.
- ATTNPO uses KFH attention signals for low-overhead step-level supervision, attenuating rewards for redundant steps while reducing penalties on essential steps.
- +7.3-point absolute accuracy gain accompanies a 60% reasoning-length reduction across six math benchmarks on DeepSeek-R1-Distill-Qwen-1.5B.
2 Preliminary
LRMs generate responses in two stages: a deliberative thinking process followed by a concise final solution. The thinking process is segmented into meaningful intermediate steps using phrases that mark reflections, verifications, or shifts in direction.
- An LRM first generates a potentially redundant thinking process T, then produces a concise final solution F.
- Thinking processes are decomposed into intermediate steps whose boundaries are marked by special phrases.
- These phrases signal reflections, verifications, branching, pauses, or shifts in reasoning direction.
3 Key-Focus Heads
The paper probes whether attention heads distinguish essential from redundant reasoning steps during final-answer generation. A small set of Key-Focus Heads achieves strong ranking accuracy, concentrates in later layers, remains stable under training, and generalizes beyond easy probing examples.
- Attention heads are hypothesized to select essential reasoning steps while ignoring redundant ones during final-answer generation.
- The probing dataset samples 300 relatively easy Deepscaler questions so step annotations are less confounded by exploratory reasoning.
- Attention-based step scores measure the average attention that final-solution tokens give to tokens within each reasoning step.
- Step Ranking Accuracy measures the fraction of essential–redundant pairs in which the essential step receives the higher score.
- A small number of middle-to-late-layer heads selectively focus on essential steps; these specialized heads are called Key-Focus Heads.
- The best head reaches 95% SRA in the 1.5B model and 96% in the 7B model, while most heads remain at or below the random baseline.
- Combining more heads yields limited gains because performance quickly saturates or declines after a small number of heads.
- KFH behavior remains consistent across RL checkpoints and shows some generalization from non-difficult probing questions to challenging AIME24 problems.
4 ATTNPO
ATTNPO rescales outcome-level advantages at each step using attention-derived redundancy signals, enabling finer-grained credit assignment without additional resources. It uses complementary strategies to reduce reinforcement of redundant steps and avoid over-penalizing essential steps.
- Stepwise Advantage Rescaling: ATTNPO rescales each correct response’s outcome-level advantage at the step level using a non-negative factor derived from Key-Focus Head attention scores.The scaling factor reflects relative step redundancy while preserving the advantage’s sign.
- Stepwise Advantage Rescaling: Advantage scaling is applied only to correct responses because identifying essential and redundant steps in incorrect responses can produce noisy credit assignment.
- Pos-Adv Attenuation for Redundant-Step: For positive advantages, ATTNPO attenuates relatively redundant steps so they receive less reinforcement, mitigating overthinking.The method applies this attenuation to steps identified as potentially redundant.
- Pos-Adv Attenuation for Redundant-Step: Redundant steps are identified by comparing attention-based step scores with a response-specific, difficulty-aware baseline.The baseline uses average attention received by response tokens from the final solution as a reference for under-attended steps.
- Pos-Adv Attenuation for Redundant-Step: The positive-advantage attenuation schedule delays stronger attenuation on harder problems and uses a lower bound on the stepwise scaling factor.The schedule allows early exploration, while the attenuation strength depends on problem difficulty.
- Neg-Adv Attenuation for Essential-Step: For negative advantages, ATTNPO clips penalties on potentially essential steps to zero, concentrating the remaining penalty relatively on redundant steps.This design is intended to avoid performance degradation from penalizing essential reasoning.
5 Experiments
Across math, code, and science evaluations, ATTNPO generally improves the efficiency–performance trade-off, substantially shortening reasoning while maintaining or improving accuracy. Ablations and analyses show that its gains depend on stepwise advantage rescaling, conservative compression on difficult problems, and retention of exploratory capacity.
- Main Results: ATTNPO achieves the best or second-best average efficiency-adjusted score across all six math benchmarks.It attains the highest average AES across all datasets.
- Main Results: 61% shorter reasoning and +7.3 points in pass@1 accuracy are achieved on the 1.5B model.On 7B, ATTNPO cuts tokens by 55% and gains +2.9 accuracy points.
- Main Results: On AIME2024, ATTNPO gains +9.6 accuracy points while reducing reasoning length by 54%.Its training dynamics initially lag TLMRE on harder tasks but ultimately match TLMRE in accuracy after rapid compression.
- Ablation Studies: Positive-Advantage Attenuation substantially reduces length, while adding Negative-Advantage Attenuation leaves length nearly unchanged but improves accuracy.This supports weakening penalties on necessary reasoning to preserve performance.
- Ablation Studies: Increasing the number of Key-Focus Heads yields only marginal gains with rapidly diminishing returns.A small top-N set therefore provides sufficient learning signals.
- Ablation Studies: Removing difficulty-based modulation shortens responses further but increasingly degrades accuracy, showing that difficult problems require conservative compression.The ablation varies the difficulty-aware baseline β and magnitude λ.
- More Analyses: ATTNPO matches or outperforms the original model across nearly all pass@k sampling budgets and generalizes to code and science with reduced length and maintained or slightly improved accuracy.Under tight token budgets it outperforms the base model and TLMRE, while at higher budgets it matches TLMRE with shorter outputs.
- More Analyses: The frequency of confused phrases drops sharply during training, indicating reduced redundant self-reflection.The same trend appears in cross-model comparisons.
6 Related Work
Related work addresses overthinking through reward design, concise training data, latent reasoning compression, dynamic strategy selection, and inference-time intervention. Outcome-supervised methods use length-aware response rewards but cannot distinguish efficient from inefficient reasoning segments, while process-supervised methods add finer credit assignment at higher resource cost.
- Overthinking and Its Mitigation: Large reasoning models can overthink by generating redundant reasoning, which may degrade performance on some tasks.Existing mitigation approaches span reward design, concise CoT data, latent compression, dynamic strategy selection, and inference-time intervention.
- Outcome-supervised RL: Outcome-supervised RL combines correctness with length-aware penalties, sometimes modeling response length jointly with question difficulty.These methods operate on the response as a whole.
- Outcome-supervised RL: Outcome-supervised methods fail to distinguish efficient from inefficient reasoning segments because they treat the output as a whole.This limits their ability to provide fine-grained step-level signals.
- Process-supervised RL: Process-supervised RL introduces intermediate-step supervision, including methods that sample final solutions after early termination at intermediate steps.These approaches seek finer-grained credit assignment than response-level rewards.
7 Conclusion
ATTNPO is a low-overhead process-supervised RL framework that uses Key-Focus Heads’ attention signals for stepwise credit assignment. It suppresses redundant reasoning while preserving necessary reasoning, improving reasoning efficiency without sacrificing performance.
- ATTNPO leverages Key-Focus Heads’ attention signals to assign credit at the reasoning-step level.
- Two targeted strategies suppress redundant steps while preserving necessary reasoning.
- The framework improves reasoning efficiency without sacrificing performance.
Limitations
The experiments are limited to relatively small models and mathematical training data, although evaluations indicate generalization to several out-of-distribution scenarios. The section also describes step segmentation and adaptive merging procedures used to identify and penalize redundant reasoning.
- Limitations: Experiments cover only 1.5B and 7B models because of limited computational resources.
- Limitations: Training uses exclusively mathematical datasets because they provide readily available, verifiable reward signals.
- Limitations: Evaluations on LiveCodeBench, GPQA, and MMLU indicate generalization to out-of-distribution scenarios.
- Step Segmentation: Short segments are adaptively merged according to the initiating phrase category to avoid semantically weak fragments.
- Step Segmentation: Confusion-initiated segments remain separate from preceding context to facilitate redundant-reasoning detection and penalization.
A.2.1 Details of Probe Data Construction
The probe dataset contains 300 relatively easy Deepscaler questions selected to support reliable essential-versus-redundant step annotations. Agreement is moderate, and trajectories are retained only when they contain both step types for within-context attention comparisons.
- The probe dataset contains 300 Deepscaler questions on which the 1.5B model achieves at least 0.875 average accuracy over eight samples.
- Questions are restricted to relatively easy problems because exploratory reasoning on harder problems can appear redundant while remaining necessary.
- Krippendorff’s α is 0.64 for the 7B model and 0.59 for the 1.5B model, indicating moderate inter-annotator agreement.
- Trajectories containing both essential and redundant steps are retained to enable attention comparisons within the same context.
A.2.2 Details of Probe Results
Probe results are evaluated across checkpoints using SRA, with top-k and greedy selection procedures ranking attention heads by their contribution to this score.
- SRA is compared across checkpoints at 0, 500, and 1000 training steps using heatmaps and Pearson correlation coefficients.
- Top-k selection ranks attention heads by average SRA and chooses the k highest-scoring heads.
- Greedy SRA-maximizing selection iteratively adds the remaining head that yields the largest increase in overall SRA.
B.1 Training Hyperparameters
The evaluation and training setup specifies benchmark coverage, sampling and scoring procedures, model-specific hyperparameters, and the additional overhead of attention-based scoring.
- Implementation: The 1.5B and 7B models use different α values and KFH head selections, while ATTNPO's attention-score computation is integrated into rollout log-probability calculation.The 1.5B model uses α = 0.2 and heads 16-2, 20-9, and 23-2; the 7B model uses α = 0.1 and heads 22-4, 22-7, and 22-22.
- Evaluation: 128 samples are used for AIME2024 and AIME2025, while all other benchmarks use four samples with a 16K context and temperature 0.6.These settings follow the described evaluation protocol.
- Metrics: AES jointly measures Pass@1 performance and token cost relative to the base policy, rewarding improvements and penalizing performance reductions.The metric captures both efficiency and performance effects.
- Metrics: For Pass@k, n is the total number of samples and c the number correct; when k = 1, the metric becomes average accuracy.This defines the interpretation of Pass@1 in the evaluation.
C.3 Ablation for Difficulty-Aware Schedule
The difficulty-aware schedule trades slower early compression for sustained later shortening, while the accompanying examples illustrate how ATTNPO avoids redundant reflection and reasoning detours.
- Difficulty-Aware Schedule: A difficulty-aware schedule slows early reasoning-length reduction but sustains a more consistent compression trend throughout training.Without the schedule, length falls faster initially but reaches a clear later plateau.
- Difficulty-Aware Schedule: Delaying advantage attenuation on difficult problems prioritizes early accuracy improvement and prevents premature, excessive compression.After performance improves sufficiently, the model can reduce reasoning length more flexibly.
- Step Annotation: The annotation protocol labels steps as necessary, redundant, or uncertain, treating error-correcting verification as necessary and repeated verification as redundant.Mixed or borderline steps are classified as uncertain.
- Case Studies: In a MATH500 example, TLMRE-1.5B uses unnecessarily reflective reasoning whereas AttnPO-1.5B produces a concise solution.The comparison is presented as a case study of reasoning efficiency.
- Case Studies: In an AIME 2024 example, TLMRE introduces unnecessary variables and an incorrect modeling idea before correction, while ATTNPO reaches the correct formulation directly.The direct formulation yields more concise reasoning in the described case.
- Case Studies: MMLU examples show TLMRE adding redundant self-reflection or unnecessary information that ATTNPO omits.These examples cover both a simple calculation and a simple knowledge-based question.