attnrl-attention-compass-reasoning-exploration

name	attnrl-attention-compass-reasoning-exploration
title	Attention as a Compass: Efficient Exploration for Process-Supervised RL in Reasoning Models
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2509.26628
keywords	["process-supervision","RL-reasoning","attention-mechanism","efficient-exploration","math-reasoning"]
description	Guide LLM exploration in reasoning tasks using attention patterns as navigation signals. This technique branches exploration from high-attention tokens (likely reasoning steps) and applies adaptive sampling to maintain effective gradients, significantly improving training efficiency for mathematical reasoning.

AttnRL: Attention as a Compass for Efficient Reasoning Exploration

Current reinforcement learning approaches to reasoning in LLMs sample from a broad action space—every possible continuation at every step—resulting in high variance and wasted computation. Blind exploration treats all tokens equally, but reasoning models naturally develop attention patterns that highlight critical decision points. AttnRL exploits this signal.

The core observation is that attention weights correlate with reasoning behavior. When an LLM attends strongly to certain tokens during reasoning, those positions are likely decision points worth exploring. By branching exploration from high-attention regions, we concentrate computational effort on promising paths rather than sampling uniformly.

Core Concept

AttnRL uses a three-component approach:

1. Attention-guided branching: Identify tokens with high attention scores (above a percentile threshold, e.g., 75th)
2. Adaptive sampling: Adjust rollout counts based on problem difficulty and batch statistics
3. One-step off-policy updates: Use importance weighting to train from non-uniformly sampled trajectories

The result is a focused exploration strategy that reduces computational cost while improving convergence on math benchmarks like AIME and Olympiad-style problems.

Architecture Overview

• Attention extractor: Collect attention patterns from intermediate transformer layers
• Branching policy: Decide which tokens to explore from (threshold-based or learned)
• Sampler: Generate multiple rollouts, concentrating samples near high-attention positions
• Off-policy corrector: Re-weight samples to account for non-uniform exploration distribution
• Reward signal: Process supervision (step-level correctness signals) or outcome supervision

Implementation Steps

First, extract and aggregate attention patterns from the model. High attention scores signal tokens worth exploring from:

def extract_attention_branches(model, prompt, attention_percentile=75):
    """
    Identify high-attention tokens for exploration branching.

    Args:
        model: Language model with attention hooks
        prompt: Input prompt to analyze
        attention_percentile: Threshold for high-attention identification

    Returns:
        branch_positions: List of (step, token_id) tuples for exploration
    """
    attention_weights = []  

    
     torch.no_grad():
        _ = model(prompt)
        attention_weights = model.get_attention_weights()

    
    avg_attention = attention_weights.mean(dim=(, ))  

    
    max_attention_per_token = avg_attention.(dim=).values

    
    threshold = torch.quantile(max_attention_per_token, attention_percentile / )
    branch_positions = torch.where(max_attention_per_token > threshold)[].tolist()

     branch_positions