PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention as a Compass: Efficient Exploration for Process-Supervised RL in Reasoning Models

Runze Liu, Jiakang Wang, Yuling Shi, Zhihui Xie, Chenxin An, Kaiyan Zhang, Jian Zhao, Xiaodong Gu, Lei Lin, Wenping Hu, Xiu Li, Fuzheng Zhang, Guorui Zhou, Kun Gai

arXiv:2509.26628v1cs.LGcs.CL

TL;DR

Existing PSRL approaches have limited exploration efficiency in branching and sampling, and commonly incur the cost of two-step sampling. AttnRL uses attention-based branching and adaptive sampling, and reports better performance and efficiency than prior methods on mathematical reasoning benchmarks.

  • Problem

    Existing PSRL approaches have limited exploration efficiency in branching positions and sampling, and typically use two sampling steps per update.

    Source: §1 · p. 2

  • Method

    AttnRL branches at high-attention positions and adaptively samples problems and batches to maintain non-zero advantage values.

    Source: §1 · p. 2

  • Results

    AttnRL surpasses GRPO and TreeRL by 1.9% and 1.8%, respectively, on average across benchmarks at 1.5B scale, and outperforms DeepScaleR with 500 rather than 1,750 training steps.

    Source: §4.2 · p. 9

  • Takeaways & Limitations

    Experiments on six mathematical benchmarks support AttnRL's effectiveness and efficiency for process-supervised RL in reasoning models.

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

Abstract

from arXiv · show

Reinforcement Learning (RL) has shown remarkable success in enhancing the reasoning capabilities of Large Language Models (LLMs). Process-Supervised RL (PSRL) has emerged as a more effective paradigm compared to outcome-based RL. However, existing PSRL approaches suffer from limited exploration efficiency, both in terms of branching positions and sampling. In this paper, we introduce a novel PSRL framework (AttnRL), which enables efficient exploration for reasoning models. Motivated by preliminary observations that steps exhibiting high attention scores correlate with reasoning behaviors, we propose to branch from positions with high values. Furthermore, we develop an adaptive sampling strategy that accounts for problem difficulty and historical batch size, ensuring that the whole training batch maintains non-zero advantage values. To further improve sampling efficiency, we design a one-step off-policy training pipeline for PSRL. Extensive experiments on multiple challenging mathematical reasoning benchmarks demonstrate that our method consistently outperforms prior approaches in terms of performance and sampling and training efficiency.

1. Introduction

Existing PSRL methods use semantically coarse segmentation, uniform sampling, and costly two-step sampling, limiting exploration efficiency. AttnRL uses attention-guided branching and adaptive sampling; experiments on six mathematical benchmarks report gains in performance and efficiency.

  • Existing PSRL methods use token-length or entropy segmentation, uniform sampling, and typically two sampling steps per update.Source: §1 · p. 2
  • AttnRL branches at positions with high attention values, which the authors associate with reasoning-related behavior.Source: §1 · p. 2
  • Adaptive sampling prioritizes difficult problems and adjusts batch sampling to maintain non-zero advantage values.Source: §1 · p. 2
  • AttnRL outperforms baseline methods in performance and efficiency across six mathematical benchmarks.Source: §1 · p. 2

2. Preliminaries

The preliminaries frame LLM reasoning as a step-level decision process, distinguish outcome- from process-supervised RL, and introduce Transformer attention.

2.1. LLM Reasoning as a Step-Level Markov Decision Process

  • The step-level MDP treats each action as a semantically coherent segment, such as a sentence or paragraph, and each next state as the current state concatenated with that action.Source: §2.1 · p. 2

2.2. Outcome-Supervised and Process-Supervised RL

  • GRPO estimates advantages from sampled rollout rewards without an explicit critic, while PSRL uses Monte Carlo sampling to estimate step-level advantages.Source: §2.2 · p. 3

2.3. Attention Mechanism

  • Transformer attention computes scores from query, key, and value representations, with a causal mask blocking access to future tokens.Source: §2.3 · p. 3

3. Method

AttnRL targets inefficient exploration in process-supervised RL by combining attention-based tree branching, adaptive sampling, and one-step off-policy training. Its method investigates whether high-attention steps mark reasoning behavior and how they can guide exploration.

3.1. Branching at Massive Attention Values

  • Prior work finds massive self-attention values highlight tokens most relevant to answering questions and play a critical role in contextual knowledge understanding, while some reasoning-model heads focus on sentences related to reasoning roles (Bogdan et al., 2025).Source: §3.1 · p. 4

3.1.1. Massive Attention Values in LLMs

  • After segmenting responses into steps, FCI sums attention to sufficiently distant later steps and takes the layer/head maximum; high scores often mark planning or self-verification.The distance parameter Δ is set to 4.Source: §3.1.1 · p. 4

3.1.3. Attention-based Tree Branching

  • ATB selects the top 20% FCI steps, then branches from the two earliest among them to limit misleading early steps and diversify paths.The method notes that earlier steps have more influence on final results and calls the misleading-initial-step phenomenon “Tunnel Vision” (Wen et al., 2025).Source: §3.1.3 · pp. 5–6

3.1.2. The Effects of Steps with High FCI Scores

  • On AIME24 and AIME25, disrupting top-20% FCI steps causes the largest accuracy drop, while earlier disruptions hurt more than later ones.The top-20% step-entropy disruption falls between the other FCI disruption conditions in accuracy impact.Source: §3.1.2 · p. 5, Figure 3 · p. 6

3.2.1. Difficulty-aware Exploration

  • Filtering removes problems below average FCI, which tend to have zero advantages; difficulty-aware expansion assigns different tree counts to retained problems.The original tree count is 6, and the count is adjusted according to problem difficulty.Source: §3.2.1 · p. 7

3.2.2. Adaptive Batch Sampling

  • Adaptive batch sampling updates prompt sampling using historical and current valid batch sizes, then discards zero-advantage responses so every final-batch sample has non-zero advantage.Compared with DAPO's dynamic sampling, it uses one prompt-sampling and generation round per training step and avoids discarding valid responses when their number exceeds the target batch size.Source: §3.2.2 · p. 7

3.3. Efficient Training with One-Step Off-Policy

  • One-step off-policy training performs initial sampling for the next problem batch alongside Monte Carlo sampling for the current batch, eliminating redundant sampling.The pipeline uses a single sampling operation at each training step, reducing sampling cost and improving training efficiency.Source: §3.3 · pp. 7–8, Figure 5 · p. 8

4. Experiments

The experiments compare AttnRL with outcome-based and process-based baselines across six mathematical benchmarks, then assess its components through ablations.

4.1. Setup

  • Experiments use two DS-R1-Distill-Qwen base models and compare against GRPO, TreeRL, and DeepScaleR-Preview-1.5B.Source: §4.1 · p. 8
  • Evaluation covers six mathematical benchmarks and reports Pass@1 and Pass@K, with K set to 32 for AIME24, AIME25, and AMC23, and 4 for the other benchmarks.Source: §4.1 · p. 8

4.2. Main Results

  • Across six benchmarks, AttnRL improves 7.5% over the 1.5B base model and exceeds GRPO and TreeRL by averages of 1.9% and 1.8%, respectively.It also outperforms DeepScaleR-Preview-1.5B with 500 training steps versus 1,750.Source: Table 1 · p. 9, §4.2 · p. 9
  • AttnRL improves over the base model on all six benchmarks; its AIME24 gains are 11.4% at 1.5B and 5.3% at 7B.Source: §4.2 · p. 9, Table 1 · p. 9

4.3. Ablation Study

  • ATB alone improves over TreeRL by an average of 1.2%, and adding adaptive sampling yields the highest performance; filtering all-correct prompts slightly reduces performance.The authors note that initially easy problems can still produce incorrect Monte Carlo samples that provide valuable training signals.Source: §4.3 · p. 9, Table 2 · p. 10

5. Analysis

The analysis examines how ATB and adaptive sampling shape Monte Carlo sampling, then compares training dynamics and efficiency across methods.

5.1. Sampling

  • ATB outperforms entropy-based branching and enables more effective sampling on both easy and hard problems, with the analysis tracking solve-all, solve-none, and valid-sample ratios.The valid ratio counts problems whose outputs are neither all correct nor all wrong.Source: Figure 6 · p. 10, §5.1 · p. 10
  • Adaptive sampling filters prompts with low FCI scores and dynamically adjusts prompt batch size, yielding training batches in which all tokens have non-zero advantages.The method reduces the ratio of cases where both sampling steps are correct after initially correct sampling.Source: §5.1 · p. 10

5.2. Training Dynamics and Efficiency

  • Compared with TreeRL, AttnRL maintains higher entropy, learns faster with fewer training steps, and produces shorter responses.GRPO entropy decreases during training, whereas PSRL entropy first decreases and then increases.Source: §5.2 · pp. 10–11, Figure 7 · p. 11, Figure 8 · p. 11
  • The one-step off-policy pipeline reduces TreeRL training time by 8%; AttnRL also achieves better performance with less wall-clock time and more valid training tokens under the same compute.The authors attribute these efficiency gains especially to adaptive sampling's dynamic batches and filtering of some low-value easy problems.Source: Table 3 · p. 11, §5.2 · p. 11

6. Related Work

The section distinguishes outcome-based RLVR from process-supervised approaches and situates AttnRL among Monte Carlo methods for process rewards.

6.1. Reinforcement Learning for LLM

  • The section describes RLVR work as focusing on outcome-based rewards, in contrast to AttnRL's focus on process rewards.Source: §6.1 · p. 12

6.2. Process Supervision for LLM

  • Prior process-supervision work includes token-level rewards, discriminative and generative PRMs, and Monte Carlo reward estimation; the authors characterize earlier methods as inefficient for reasoning models in branching, sampling, and generation.Source: §6.2 · p. 12

7. Conclusion

AttnRL uses attention information to locate reasoning-related steps for branching and combines this with adaptive sampling for mathematical reasoning benchmarks.

7. Conclusion

  • AttnRL uses attention information to identify reasoning-related steps and branch at those positions for efficient exploration.Source: §7 · p. 12
  • The framework adds adaptive sampling based on problem difficulty while maintaining valid training batch size.Source: §7 · p. 12
  • Experiments on mathematical reasoning benchmarks support the effectiveness and efficiency of AttnRL.Source: §7 · p. 12

A. Experimental Details

Evaluation follows Luo et al. (2025) and uses a prompt asking for step-by-step reasoning and a boxed final answer.

A. Experimental Details

B.1. Full Test Curves

Figure 9 presents test curves for GRPO, TreeRL, and AttnRL across six mathematical benchmarks.

B.1. Full Test Curves

C. Cases

The cases include a full response with high-FCI steps marked in red and a worked count of three-digit integers divisible by both 11 and 5.

C. Cases

  • Table 5 is identified as the full response, with steps having high FCI scores shown in red.Source: C. Cases · p. 19, Table 5 · p. 19
  • The example counts 17 positive three-digit integers divisible by both 11 and 5 by enumerating multiples of 55 from 110 to 990.The trace uses common difference 55 and verifies the 17 terms by listing them.Source: C. Cases · p. 20

Related papers

Loading 2509.26628v1…