ATTNPO: Attention-Guided Process Supervision for Efficient Reasoning

# AttnPO: Attention-Guided Process Supervision for Efficient Reasoning

Source: https://arxiv.org/html/2602.09953

Shuaiyi Nie¹,², Siyu Ding³‡, Wenyuan Zhang¹,², Linhao Yu⁴, Tianmeng Yang³, Yao Chen¹,², Weichong Yin³, Yu Sun³, Hua Wu³, Tingwen Liu¹,²†

¹Institute of Information Engineering, Chinese Academy of Sciences
²School of Cyber Security, University of Chinese Academy of Sciences
³Baidu Inc.
⁴Tianjin University

{nieshuaiyi, liutingwen}@iie.ac.cn

## Abstract

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. [Code link: https://github.com/NieSYsc20/AttnPO](https://github.com/NieSYsc20/AttnPO).

## 1 Introduction

Recent advances in Large Reasoning Models (LRMs) have demonstrated remarkable performance on complex reasoning tasks. This success is largely driven by outcome-supervised RL methods, represented by GRPO, which omits the critic model in PPO by utilizing verifiable outcome rewards to estimate group advantages. A hallmark of LRMs is their ability to produce deliberative long Chain-of-Thought that employ advanced strategies such as self-reflection and multi-path exploration. However, long CoT fosters overthinking: LRMs indiscriminately apply verbose reasoning, wasting computation on even trivial operations. Integrating length penalties in outcome-supervised RL is widely adopted to mitigate overthinking. The core idea is to assign higher rewards to responses that are both correct and concise. However, this coarse-grained feedback assigns uniform credit to all reasoning steps, failing to provide discriminative learning signals to distinguish between necessary and redundant steps.

Therefore, recent studies focus on process supervision to down-weight the credit assigned to redundant steps. Sample-based methods estimate the marginal performance gain of intermediate steps via Monte Carlo sampling of immediate final solutions, while model-based methods train a reward model to pinpoint the first correct answer position, treating only the subsequent steps as redundant. However, these methods face two key limitations: (1) **High Overhead**, which requires additional sampling or a reward model; (2) **Inaccurate Credit Assignment**, which fails to provide independent step-wise signals and may lead to erroneous credit allocation.

A key question thus emerges: **Can we, with almost no additional resource costs, rely solely on the model's intrinsic signals to achieve fine-grained, step-wise supervision of which steps are essential or redundant?**

In this study, we delve into the model's intrinsic attention mechanism and identify a striking phenomenon: during the final answer generation (following the `<answer>` token), a subset of heads naturally focus on essential steps by assigning them higher attention weights while suppressing redundant ones. We term these **Key-Focus Heads (KFHs)**. Building on this finding, our experiments further reveal: (1) **Positional Distribution**, KFHs are mainly located in middle-to-late layers; (2) **Training Stability**, where the behavior of KFHs varies little under RL training with length penalties; and (3) **Sparse Sufficiency**, as a small set of KFHs suffice for distinction, with performance saturating rapidly as more heads are added.

Based on these insights, we propose **AttnPO** (Attention-guided Policy Optimization), a novel RL framework that enables step-level supervision by leveraging KFHs' attention scores, achieving superior performance with negligible training overhead. Specifically, **AttnPO** scales the outcome-level advantage in a stepwise manner via two strategies on correct responses: (1) for correct responses with positive outcome advantages, **AttnPO** attenuates the positive advantage assigned to redundant steps, preventing over-encouragement of them and mitigating overthinking; and (2) for correct responses with negative outcome advantages, **AttnPO** attenuates the negative advantage assigned to essential steps, avoiding excessive penalization of valid reasoning and alleviating performance degradation.

Our contributions are threefold: (1) We are the first to reveal the existence of KFHs in LRMs, which naturally focus on essential reasoning steps while suppressing redundant ones; (2) Based on this insight, we propose **AttnPO**, a low-overhead process-supervised RL framework that leverages intrinsic attention signals for step-level credit assignment; (3) Empirical results demonstrate that **AttnPO** substantially shortens reasoning length while improving accuracy. For example, on DeepSeek-R1-Distill-Qwen-1.5B, **AttnPO** achieves an average +7.3-point absolute accuracy gain with a 60% reduction in reasoning length across six math benchmarks.

## 2 Preliminary

### Two-stage Generation of LRMs

Given an input question $q$, an LRM $\mathcal{M}$ generates a response $o=(\mathcal{T}, \mathcal{F})$ in two sequential stages, explicitly separated by a special `<answer>` token. First, the model produces a deliberate thinking process $\mathcal{T}$, which may contain redundant reasoning structures, such as unnecessary reflection or verification. Then, it outputs a concise final solution $\mathcal{F}$ by summarizing $\mathcal{T}$ and retaining only the essential steps required to solve the problem.

### Step Segment

The thinking process $\mathcal{T}$ can be decomposed into meaningful intermediate steps, whose boundaries are typically marked by specific special phrases that capture key points in reasoning, such as pauses, reflections or shifts in direction. Thus we segment the thinking process with special phrases, which fall into three categories:

- **Confused phrases** (e.g., "Wait") indicate the model is entering a reflection, verification, or branching stage
- **Progression phrases** (e.g., "First") signal the continuation of the current reasoning
- **Summary phrases** (e.g., "Therefore") summarize the preceding reasoning

To avoid overly short segments, we adaptively merge segments shorter than a threshold. During segmentation and merging, we separate segments that begin with confusion phrases from the preceding context, since they often trigger redundant verification or reflection. Isolating them makes such redundancy easier to detect and penalize. See Appendix A.1 for details.

## 3 Key-Focus Heads

Following the two-stage generation paradigm described in Section 2, LRMs transform a deliberative thinking process—often containing redundant reasoning—into a concise final solution, which requires non-trivial selection and aggregation of essential information. In Transformer architecture, attention serves as the primary mechanism for information selection, and prior work has shown that different attention heads specialize in distinct functions. Based on this observation, we hypothesize that **during final-answer generation, there exist a specific subset of attention heads that selectively focus on essential reasoning steps while ignoring redundant ones, thereby summarizing and integrating the thinking process.**

To test this hypothesis, we examine the attention distributions over essential and redundant steps across different attention heads.

### 3.1 Probing Setup

#### Probe Data Construction

To analyze how different attention heads attend to essential versus redundant steps, we require a probing dataset with reliable step-level annotations. We therefore sample 300 questions from DeepScaler that are not difficult for DeepSeek-R1-Distill-Qwen-1.5B (average accuracy ≥ 0.875 over eight samples). This restriction is important because, on harder problems, LRMs often exhibit exploratory reasoning. Such reasoning may appear redundant, but still contributes to correct solutions, making reliable step annotation difficult and thereby affecting the head identification.

For each question, we sample eight responses, select the median-length correct one, and segment its thinking process following Section A.1. We then annotate each reasoning step using three state-of-the-art LLMs, and retain only the steps on which all three models agree on the label, thereby reducing annotation noise. Thus, for each thinking process $\mathcal{T}$, we obtain a set of essential steps $E_s$ and a set of redundant steps $R_s$. Details are provided in Appendix A.2.1.

#### Attention-Based Step Scores

For each head at layer $l$ and head $h$, we compute an attention-based step score $\mathcal{S}^{l,h}_{s_k}$ for each reasoning step $s_k$, where $k$ indexes the step index, based on the attention weights used during final-answer generation. Specifically, let $a^{l,h}_{m \to n}$ denote the attention weight from token $m$ in the final solution $\mathcal{F}$ to token $n$ in step $s_k$. We define the step score $\mathcal{S}^{l,h}_{s_k}$ as:

$$\mathcal{S}^{l,h}_{s_k} = \frac{1}{|s_k|} \sum_{m \in \mathcal{F}} \sum_{n \in s_k} a^{l,h}_{m \to n} \quad (1)$$

This score measures the average attention received by each token in step $s_k$ from the final solution.

#### Evaluation Metrics: Step Ranking Accuracy

We evaluate how well each attention head distinguishes between essential and redundant steps. For each thinking process $\mathcal{T}$, we form all essential–redundant step pairs $\mathcal{P}_{\mathcal{T}} = \{(s^+, s^-) \mid s^+ \in E_s, \, s^- \in R_s\}$. For head $h$ at layer $l$, we define Step Ranking Accuracy (SRA) as:

$$\mathrm{SRA}_{l,h} = \frac{|\{(s^+, s^-) \in \mathcal{P}_{\mathcal{T}} \mid \mathcal{S}^{l,h}_{s^+} > \mathcal{S}^{l,h}_{s^-}\}|}{|\mathcal{P}_{\mathcal{T}}|} \quad (2)$$

SRA measures the fraction of step pairs in which the essential step receives a higher score than the redundant one. We report the final SRA by averaging across all examples in the probing dataset.

### 3.2 Probing Analysis

We conduct experiments on DeepSeek-R1-Distill-Qwen-1.5B / 7B and make the following observations; more detailed results are in Appendix A.2.2:

- **During final solution generation, some attention heads selectively focus on essential reasoning steps while ignoring redundant ones, which are mainly located in the middle-to-late layers.** We refer to these specialized heads as Key-Focus Heads. As shown in Figure 2 (a, b), across both evaluated LRMs, a small number of heads achieve SRA above 0.9, with the best-performing head reaching 95% in the 1.5B model and 96% in the 7B model. Most heads, by contrast, perform at or below the random baseline (SRA ≈ 0.5).

- **Ensembling step scores from multiple attention heads provides only limited gains on the probing dataset.** In Figure 2 (c, d), under both top-k selection by SRA and greedy SRA-maximizing selection, performance improves slightly when a small number of heads are combined, but quickly saturates or degrades as more heads are added.

- **Under RL training with length penalties, the behavior of KFHs varies little, suggesting that their functional roles are robust.** We train the model using the reward function described in Eq. 5 and re-run the same probing process on checkpoints at steps 500 and 1000. The resulting SRA distributions show high consistency across checkpoints, with Pearson correlation coefficients exceeding 0.85 in all cases.

- **Even when using non-difficult questions in probing, KFHs exhibit a degree of generalization to challenging problems.** We sample responses on the challenging AIME24 dataset and obtain step scores from the head with the highest SRA. From each thinking process $\mathcal{T}$ with a correct final answer, we remove either the bottom 30% (low-scoring) or the top 30% (high-scoring) steps, and then append the prompt "I think I have finished thinking. Now give the final solution step by step." to generate the final solution. As shown in Figure 2 (e, f), removing high-scoring steps significantly degrades performance, while removing low-scoring steps has minimal impact, confirming that KFHs generalize to harder problems.