Large Language Models (LLMs) are no longer just text generators — they are becoming reasoners, capable of solving mathematical problems, logical puzzles, or planning tasks step by step.
One of the key challenges is how to improve the quality of this reasoning. Traditional Reinforcement Learning (RL) rewards only the final outcome, but in complex reasoning it makes more sense to evaluate each intermediate step. This is called process-supervised RL (PSRL).

The problem: existing PSRL methods are expensive and inefficient, as they explore too many uninformative paths.
The new paper Attention as a Compass: Efficient Exploration for Process-Supervised RL in Reasoning Models proposes AttnRL, which leverages the model’s own attention mechanism as a compass to decide where to branch the reasoning process.

Reasoning with an Attention Compass

Imagine solving a riddle step by step. At each stage, you could go in multiple directions. Traditional RL just tries some random branches and checks which one works — but that wastes time.

AttnRL works differently. While writing, the model naturally “pays attention” to certain previous steps. If a step is strongly referenced, it’s probably important. That’s the signal: “this is a good place to branch and explore an alternative path.”

Analogy

  • Map of the journey = the reasoning process.
  • Crossroads = reasoning steps.
  • Attention compass = signals from the model highlighting important places.

Instead of branching everywhere, we follow the compass — and only explore promising directions.

Adaptive Exploration

AttnRL further improves efficiency by:

  • Giving more exploration to harder problems, less to easy ones.
  • Filtering out trivial prompts that the model always solves correctly.
  • Dynamically adjusting batch sizes to keep useful samples.
  • Cutting training costs with a one-step off-policy trick, requiring just one round of generation per iteration.

Result: the model learns to reason both faster and smarter.

Inside AttnRL

Formalization

Reasoning is modeled as a step-level Markov Decision Process (MDP):

  • Initial state: $( s_1 = q \sim \mathcal{D} )$ (the prompt).
  • Actions: textual segments (e.g., paragraphs).
  • Deterministic transitions:
    $$ s_{k+1} = [s_k, a_k] $$

In Outcome-Supervised RL (OSRL), only the final answer is rewarded.
In Process-Supervised RL (PSRL), each step can be rewarded.

TreeRL estimates the value of node ( s_k ) as:

$$ V(s_k) = \frac{1}{|L(s_k)|} \sum_{\ell \in L(s_k)} \mathbf{1}(\ell \text{ is correct}) $$

And the advantage:

$$ \hat A_{i,k} = \frac{1}{\sqrt{|L(s_k)|}} \Big( (V(s_k) - V(s_1)) + (V(s_k) - V(p(s_k))) \Big) $$

where ( p(s_k) ) is the parent node.

Core Components of AttnRL

  1. Attention-Based Tree Branching (ATB)

    • Defines Forward Context Influence (FCI) for each step:

    $$ y_{l,h,k} = \sum_{j=k+\Delta}^{T_k} \alpha_{l,h}(j,k) $$

    Aggregated influence:

    $$ y_k = \max_{l,h} y_{l,h,k} $$

    Steps with highest FCI are selected as branching points.

  2. Adaptive Sampling (ADS)

    • Filters out trivial prompts.
    • Adjusts branching based on difficulty:

    $$ {tree_num} = \exp(-z_n) \times {original_tree_num} $$

    • Dynamically resizes batch:

    $$ B_m = \mathrm{Round}\big( \lambda B_{m-1} + (1-\lambda),\frac{B’}{B’’} B_{m-1} \big) $$

    where ( B’’ ) = number of samples with non-zero advantage.

  3. One-Step Off-Policy Training
    Combines new prompt generation and rollouts from the previous batch into a single iteration — cutting training costs.

Results

  • Outperforms baselines (GRPO, TreeRL) by +1.8 pp on math benchmarks.
  • Training time reduced by ~8%.
  • Produces more valuable exploration paths (with non-zero advantage).

Conclusion

AttnRL uses the model’s own attention as a compass for efficient exploration in reasoning tasks. This approach:

  • Makes models faster and more effective learners.
  • Avoids wasting effort on uninformative branches.
  • Simplifies and speeds up the training pipeline.

Why it matters

AttnRL is a step towards smarter reasoning models:

  • capable of solving mathematical, physical, and logical problems,
  • assisting in theorem proving or process planning,
  • extending to multimodal reasoning.

It’s not just about making models stronger — it’s about making their thinking process more purposeful.

📎 Links