Source-linked AI summary
Attention Illuminates LLM Reasoning: The Preplan-and-Anchor Rhythm Enables Fine-Grained Policy Optimization
Yang Li, Zhichen Dong, Yuhan Sun, Weixun Wang, Shaopan Xiong, Yijia Luo, Jiashun Liu, Han Lu, Jiamang Wang, Wenbo Su, Bo Zheng, Junchi Yan
TL;DR
LLM reasoning is difficult to interpret, and uniform sequence-level RL credit does not distinguish pivotal tokens from routine steps. The paper derives attention-based preplan and anchor signals, uses them for targeted credit assignment, and reports consistent gains across reasoning tasks while framing the signals as a useful regularity rather than a complete theory of reasoning.
Problem
LLM reasoning remains opaque, while sequence-level reinforcement learning commonly distributes credit uniformly across generations instead of distinguishing pivotal from routine tokens.
Method
The paper analyzes local and global attention with WAAD and FAI, then uses preplan, anchor, and coupling signals to reweight token-level RL advantages.
Results
The proposed attention-derived credit assignment delivers consistent gains across reasoning tasks and benchmarks.
Takeaways & Limitations
Aligning optimization with the model’s internal preplan-and-anchor rhythm offers a path toward more transparent and effective reasoning training.
Takeaways & Limitations
WAAD and FAI are filtered internal signals for locating influential positions, not unbiased explanations, complete causal decompositions, or assessments of each intermediate step’s correctness.
Abstract
from arXiv · showhide
The reasoning pattern of Large language models (LLMs) remains opaque, and reinforcement learning (RL) typically applies uniform credit across an entire generation, blurring the distinction between pivotal and routine steps. This work positions attention as a privileged substrate that renders the internal logic of LLMs legible, not merely as a byproduct of computation, but as a mechanistic blueprint of reasoning itself. We first distinguish attention heads between locally and globally focused information processing and reveal that locally focused heads produce a sawtooth pattern near the diagonal indicating phrasal chunks, while globally focused heads expose tokens that exert broad downstream influence over future tokens. We formalize these with two metrics: 1) Windowed Average Attention Distance, which measures the extent of backward attention within a clipped window; 2) Future Attention Influence, which quantifies a token's global importance as the average attention it receives from subsequent tokens. Taken together, these signals reveal a recurring preplan-and-anchor mechanism, where the model first performs a long-range contextual reference to generate an introductory token, which is immediately followed by or coincides with a semantic anchor token that organizes subsequent reasoning. Leveraging these insights, we introduce three novel RL strategies that dynamically perform targeted credit assignment to critical nodes (preplan tokens, anchor tokens, and their temporal coupling) and show consistent performance gains across various reasoning tasks. By aligning optimization with the model's intrinsic reasoning rhythm, we aim to transform opaque optimization into an actionable structure-aware process, hoping to offer a potential step toward more transparent and effective optimization of LLM reasoning.
1 Introduction
The paper uses attention dynamics to expose recurring local and global structures in LLM reasoning, then aligns reinforcement-learning credit with tokens that appear structurally decisive.
- 1 Introduction: LLM reasoning remains opaque, while sequence-level rewards commonly distribute credit across all tokens and blur pivotal steps with routine elaboration.The paper frames fine-grained credit assignment as a response to this mismatch between apparent reasoning structure and optimization practice.
- 1 Introduction: Attention analysis reveals locally focused heads tracking phrasal chunks and globally focused heads highlighting anchor tokens that steer downstream reasoning.Local heads show a sawtooth pattern near the diagonal, while global heads revisit a small set of semantically influential tokens.
- 1 Introduction: WAAD measures backward attention distance within a clipped window, while FAI measures the average attention a token receives from later tokens.Together, the metrics operationalize local preplan signals and global anchor importance from an auxiliary forward pass.
- 1 Introduction: The proposed policy-optimization strategies reweight token-level credit toward chunk-onset introductory tokens, high-FAI anchors, and their temporal coupling.This replaces uniform trajectory-wide advantages with attention-based, process-aware credit assignment.
- 1 Introduction: The paper reports stable empirical gains across various reasoning benchmarks while maintaining compatibility with existing RLVR workflows.The contribution combines attention-based reasoning analysis with three reinforcement-learning methods targeting structurally important tokens.
2 Related Work
Related work spans reinforcement-learning methods for LLM post-training, analyses of internal reasoning mechanisms, and approaches that target critical or high-entropy tokens.
- 2 Related Work: LLM reinforcement learning includes on-policy policy-gradient methods and offline preference optimization without on-policy sampling.The related-work discussion situates the paper within both major strands of RL-based post-training.
- 2 Related Work: White-box reasoning studies examine information-propagating components, critical attention heads, and mechanisms for planning or backtracking.These analyses provide precedents for investigating internal structures involved in reasoning execution.
- 2 Related Work: Within RLVR, prior work links gains to explicit-reasoning macro-structure and targets decision points or high-entropy tokens that govern divergent paths.The paper distinguishes its approach by deriving credit signals from model-internal attention patterns.
3 Preliminaries
The preliminaries define autoregressive attention maps and reinforcement-learning objectives used to model and optimize LLM-generated sequences.
- 3 Preliminaries: A decoder-only LLM policy πθ autoregressively generates an output sequence o for a sampled question q, with o_t denoting its t-th token.This establishes the sequence notation for the later attention and policy-optimization formulations.
- 3 Preliminaries: Each self-attention head produces a lower-triangular T×T map whose weights represent attention over past positions under causal masking.The maps provide the head-specific dependencies analyzed by the paper.
- 3 Preliminaries: RL trains the LLM policy to maximize cumulative verifier rewards received from generated outputs.This frames verifiable-reward optimization as the objective underlying the methods.
- 3 Preliminaries: PPO uses frozen-old-policy data, importance sampling, and a clipped surrogate objective, with A_t as the advantage and ε as the clipping hyperparameter.The advantage estimates are commonly obtained with GAE.
- 3 Preliminaries: GRPO removes the value-function critic by normalizing rewards among grouped responses to the same prompt and adds a KL penalty to the clipped objective.This preserves comparative gradient information in sparse-reward settings.
4 Dissecting Attention Dynamics to Expose the Reasoning Rhythm
The analysis identifies complementary local and global attention patterns and shows that their coordination forms a recurring preplan-and-anchor rhythm in reasoning.
- 4.1 Head Grouping and the Local/Global Patterns: Local-focused heads produce near-diagonal sawtooth attention patterns: chunk onsets reach farther back, then subsequent tokens rely on immediate context.Low WAAD values correspond to tight local continuation, while peaks mark long-range consultation at semantic boundaries.
- 4.1 Head Grouping and the Local/Global Patterns: WAAD measures backward attention distance within a clipped window, whereas FAI measures the average attention a token receives from later positions.Together, the metrics connect local contextual retrieval with global token influence across the reasoning trajectory.
- 4.2 Joint Dynamics: A Preplan-and-Anchor Pattern: High-FAI tokens function as semantic anchors that receive sustained attention from future positions and steer unfolding reasoning toward a stable frame.Perturbing these positions frequently shifts downstream reasoning logic, whereas low-FAI perturbations more often change only local phrasing.
- 4.1 Head Grouping and the Local/Global Patterns: Global-focused heads highlight sparse tokens with broad downstream influence, while local-focused heads emphasize nearby context within phrasal or semantic chunks.Heads are grouped by their attention-weighted mean backward distance, separating recent-context focus from frequent long-range dependencies.
- 4.2 Joint Dynamics: A Preplan-and-Anchor Pattern: WAAD peaks coincide with higher token entropy, and FAI profiles from receiver heads and global-focused heads are strongly correlated.These couplings support coordinated local-global attention dynamics rather than isolated single-metric patterns.
- 4.2 Joint Dynamics: A Preplan-and-Anchor Pattern: FAI peaks typically follow or coincide with WAAD peaks, forming a two-beat process in which long-range retrieval precedes or accompanies an influential anchor token.The preplan step prepares an upcoming concept or reasoning step, while the anchor organizes subsequent reasoning.
5 Fine-Grained Policy Optimization Driven by Attention Signals
The paper replaces uniform sequence-level credit with attention-derived, token-level reweighting that targets preplan tokens, semantic anchors, and their coupling. These strategies consistently improve reasoning performance across puzzle, QA, and mathematical benchmarks while adding little computational overhead.
- Attention Calculation in the RL Framework: Attention maps are collected with one auxiliary forward pass over the prompt and response, avoiding the thousands of generation passes and introducing little additional latency.Five evenly spaced layers in the middle third of the network provide representative attention snapshots.
- Targeted Credit Assignment to Critical Nodes: The method rescales each token’s on-policy advantage as ˜At = Atγt, using attention-derived weights for local preplan, global anchor, or joint rhythm signals.The auxiliary attention pass supplies the signals while leaving the RL objective unchanged.
- Targeted Credit Assignment to Critical Nodes: Local-chunk credit selects preplan tokens from WAAD distance variations, while global-anchor credit selects top-FAI tokens to emphasize positions organizing subsequent reasoning.The local strategy targets peak-valley chunk boundaries; the global strategy amplifies tokens with high future influence.
- Results on Simple Logical Puzzles and QA: Coupled rhythm credit reaches 63.1% on Countdown versus 52.6% for GRPO, and 50.1% on CrossThink-QA versus 48.0%, with faster convergence and higher plateaus.The same section reports coupled credit as the earliest-improving and highest-scoring variant on the training curves.
- Results on Mathematical Reasoning: On mathematical reasoning, coupled rhythm credit consistently outperforms GRPO and naive alternatives, including gains of +5.0pt on AIME25 and +6.3pt on AMC23 for Qwen3-8B.With an extended 8K context, it still improves AIME25 by +4.3pt over GRPO.
- Robustness across model families and scales: The method improves average scores across model families and scales, including +18.8% on Llama3.2-3B, +13.0% on Llama3.1-8B, and +5.78 on Qwen3-0.6B.It also raises DeepSeek-R1-Distill-Llama-8B from 47.28 to 49.34 and improves Qwen3-1.7B by +3.22.
- Scope of attention-derived signals: The attention-derived signals locate structurally influential positions for credit assignment but do not provide unbiased explanations, causal decompositions, or local correctness judgments.The preplan-and-anchor rhythm is presented as a useful regularity for reallocating sequence-level feedback, not a complete theory of reasoning.
6 Conclusion
The paper presents attention dynamics as a lens for understanding and directing LLM reasoning. It uses WAAD and FAI to identify critical tokens and applies targeted credit assignment, reporting consistent gains on reasoning tasks.
- Attention dynamics reveal a recurring preplan-and-anchor rhythm in which long-range consultation precedes anchor tokens organizing downstream inference.
- WAAD and FAI formalize signals for identifying preplan and anchor tokens, respectively.
- Targeted reinforcement of critical nodes delivers consistent gains on reasoning tasks and offers a path toward more transparent and effective training.
A.1 Analysis in Sec. 4
The analysis evaluates attention-based metrics and their alignment across reasoning tokens, benchmarks, and computational cost. It finds that the proposed attention computation adds minimal overhead to RL training.
- Experimental setup: The evaluation samples 70 GSM8K problems for quantitative metric analysis and uses a shortest correct response for qualitative visualization.The qualitative example asks for Siobhan’s jewels and produces the answer 23.
- Metric validation: WAAD peaks are compared with all response tokens through entropy, while FAI peaks are tested for co-occurrence and temporal alignment against random baselines.The protocols treat higher observed values as evidence of greater uncertainty or non-random attention alignment.
- Benchmark coverage: The benchmark suite spans Countdown, CrossThink-QA, AIME24, AIME25, AMC, MATH500, and OlympiadBench to test generalization across reasoning modalities and difficulty levels.The tasks include symbolic puzzles, heterogeneous question answering, and advanced mathematical problem solving.
- Computational overhead: 2.0-3.3% additional wall-clock time is required for attention-based credit assignment relative to total generation time.The method performs one parallel forward pass through frozen attention layers after autoregressive generation.
B.2 Ablation on Layer Range Selection for Attention Signal Sources in RL
The layer-range ablation tests where attention signals should be extracted for Coupled Rhythm Credit. Middle-layer signals perform best across the evaluated benchmarks.
- Ablation design: The ablation compares four layer sources: uniformly sampled All Layers, the middle third, the upper third, and the lower third of the network.The middle variant uses five evenly spaced layers within the middle third.
- Ablation result: 39.2 average score is achieved by the Middle Layers variant, which leads in every individual benchmark category.The comparison includes All Layers, Middle Layers, Top Layers, and Bottom Layers, with the final middle-layer setting highlighted in Table 9.
- Interpretation: The middle-layer choice is presented as an architecture-scaled heuristic rather than a finely tuned model-specific configuration.The authors attribute weaker alternatives to dilution or omission of the transition zone where structural reasoning crystallizes.
B.3 Layer-Wise Analysis of Attention Patterns
The layer-wise analysis finds distinct depth profiles for global influence, local attention structure, and their temporal coupling. These profiles support using different layer ranges for different attention signals.
- Global influence: 66.52% FAI-receiver co-occurrence occurs in top layers, representing a +296.4% lift over random.The result suggests global influence signals are strongest in deeper layers and are refined there rather than newly created.
- Local structure: WAAD peak entropy follows a U-shaped depth profile, with a +52.3% lift in bottom layers and a +25.6% lift in middle layers.The middle layers have the lowest entropy and are described as a possible compression stage for local dependencies.
- Temporal coupling: 56.88% of FAI events follow or coincide with WAAD peaks in top layers, a +60.0% lift over random.The coupling remains above baseline in bottom layers, at a +36.5% lift.
- Design implication: FAI is most informative in top-to-mid layers, whereas WAAD gains discriminative power in later layers for fine-grained local refinement.This division of roles motivates the architectural design of the attention-based credit signals.
B.4 Attention Pattern Generalizability to Knowledge Retrieval Tasks
The study tests whether the preplan-and-anchor rhythm extends from procedural mathematics to associative knowledge retrieval. Its signatures remain significant on CrossThink-QA, but their coordination is weaker.
- CrossThink-QA analysis: FAI-receiver co-occurrence, elevated WAAD peak entropy, and FAI-WAAD temporal coupling remain above random baselines on CrossThink-QA.CrossThink-QA requires synthesizing answers from scattered factual evidence rather than following predominantly sequential reasoning.
- Generalizability: The rhythm is detectable and functionally relevant in associative settings, but its coordination is attenuated compared with symbolic reasoning.The findings support adaptation of the rhythm to task structure rather than an identical pattern across tasks.
B.5 Attention Pattern Generalizability to the Llama Model Architecture
Replication on Llama3.1-8B shows that the paper’s attention-based reasoning signatures remain robust across architectures, while targeted head grouping improves structure-aware credit assignment.
- Attention Pattern Generalizability: The preplan-and-anchor signatures remain robust and statistically significant on Llama3.1-8B, supporting generality beyond Qwen.Global influence tokens co-occur with receiver positions 72.25% of the time versus a 24.19% random baseline, while WAAD peaks show a +28.2% entropy lift and temporal coupling a +34.1% alignment lift.
- Attention Head Selection: Targeted attention-head groups consistently outperform all-head aggregation across credit-assignment methods and datasets.Coupled Rhythm Credit reaches an average score of 39.2, 2.6 points above its all-heads counterpart.
- Attention Head Selection: The results indicate that effective credit assignment depends on prioritizing functionally specialized computational pathways, not only selecting network layers.Coupled Rhythm benefits from a union of local and global specialized heads.
B.7 Evaluation on Long-Horizon Planning Task
On ZebraLogic’s long-horizon constraint-satisfaction tasks, structure-aware credit assignment improves performance over GRPO, with the coupled rhythm variant achieving the strongest reported accuracy.
- Evaluation on Long-Horizon Planning Task: 93.5% accuracy on ZebraLogic represents a +9.6% absolute gain over GRPO for the full coupled rhythm variant.The benchmark requires deducing unique global arrangements from interdependent logical clues, testing extended multi-step planning.
- Evaluation on Long-Horizon Planning Task: Structure-aware credit assignment methods yield substantial improvements over standard GRPO and simple baselines on long-horizon planning.The result supports coherent reasoning over extended action sequences under the reported Qwen3-4B, two-epoch, 1,024-token setup.
- Evaluation on Long-Horizon Planning Task: Across the broader baseline comparison, the coupled rhythm variant achieves the highest average performance across datasets.Its average score is 43.2 against newly introduced token-level credit-assignment baselines.
B.9 Quantitative Ablation of Top-k Ratio
The coupled rhythm mechanism performs best when credit is assigned to the top 40% of tokens, while perturbation examples distinguish globally influential tokens from locally stylistic ones.
- Quantitative Ablation of Top-k Ratio: Top-k = 0.4 yields the strongest coupled-rhythm performance, with the highest scores across datasets and a peak average of 43.2.Both k = 0.2 and k ≥0.6 produce consistent performance drops, indicating that overly narrow or broad credit allocation dilutes the signal.
- Representative Perturbation Examples: The perturbation analysis treats FAI as a signal for identifying tokens that receive sustained attention from future reasoning steps.The high-FAI example examines “much” as a global influence anchor, whereas low-FAI examples assess less influential positions.
- Representative Perturbation Examples: High-FAI token perturbations can reorganize the reasoning trajectory despite minor lexical substitutions, and one substitution leads to an incorrect final answer.The “much” perturbation with Jaccard similarity 0.383 changes the reasoning structure, while another high-FAI example has similarity 0.255 and deviates from the correct path.
- Representative Perturbation Examples: Low-FAI token substitutions usually preserve the reasoning strategy and final answer while changing only local phrasing.Replacing “over” with “in” yields Jaccard similarity 1.000, and replacing “spends” with “will” yields 0.946 with the same solution logic.
- Additional Results Across Model Families and Scales: Across additional model families and scales, coupled rhythm credit consistently improves over GRPO under the reported evaluation protocol.The reported settings include Llama-family base models, DeepSeek-R1-Distill-Llama-8B, and smaller Qwen3 models.