PaperlayerView exact revision on arXiv

Source-linked AI summary

How Does Reasoning Flow? Tracing Attention-Induced Information Flow for Targeted RL in LLMs

Zhichen Dong, Yang Li, Yuhan Sun, Weixun Wang, Yijia Luo, Zinian Peng, Taiheng Ye, Chao Yang, Wenbo Su, Yu Cheng, Bo Zheng, Junchi Yan

arXiv:2606.10646v1cs.LGcs.CL

TL;DR

Token-level credit assignment in LLM reinforcement learning struggles to distinguish decisive reasoning steps from routine tokens and to capture global information routing. FlowTracer models attention as an answer-targeted flow graph to assign token credit, consistently improving reasoning performance across mathematical benchmarks, including 43.4% versus 39.4% for Qwen3-8B against GRPO.

  • Problem

    Token-level credit assignment in LLM reinforcement learning remains difficult because delayed supervision and point-wise heuristics provide noisy or uniform signals while overlooking global information propagation.

  • Method

    FlowTracer models tokens as an attention-induced DAG, reweights edges for answer-targeted conserved flow, and scores tokens by influence routed toward the answer.

  • Results

    43.4% average accuracy versus 39.4% for GRPO on Qwen3-8B, with FlowTracer consistently improving reasoning performance across standard mathematical benchmarks and model sizes.

  • Takeaways & Limitations

    Flow-based token importances focus learning signals on high-impact reasoning tokens and provide a structural signal for more efficient and interpretable model alignment.

  • Takeaways & Limitations

    FlowTracer assumes a localized answer region, while outcome-only rewards may not distinguish useful intermediate reasoning from reasoning leading to an incorrect final answer.

Abstract

from arXiv · show

Token-level credit assignment remains a key obstacle for reinforcement learning (RL) in large language models (LLMs), where RL recipes typically treat all tokens equally, failing to distinguish decisive reasoning steps from routine formatting or fluent filler. Recent attempts leverage model-internal signals to assign finer-grained credit, but these are often point-wise heuristics that ignore the global structure of information propagation. We propose FlowTracer, an RL framework that traces answer-targeted reasoning flow on an attention-induced directed acyclic graph in which nodes correspond to tokens and edge capacities come from aggregated attention weights and derives token credit from this global structure. The edge capacities are reweighted to retain only the influence that can reach the answer region, while enforcing local flow conservation so intermediate tokens neither lose nor gain effective mass due to path length or irrelevant branches. On this graph, FlowTracer extracts an information-flow backbone connecting the question to the answer and scores tokens by flow throughput, revealing high-impact hubs and aggregation checkpoints that mediate long-range dependencies. These derived importances are used to shape token-level rewards, enabling learning signals to focus precisely on the tokens that route information toward (or away from) correct answers and delivering consistent performance gains across a range of reasoning tasks.

1 Introduction

FlowTracer addresses token-level credit assignment in LLM reinforcement learning by tracing how information flows from prompts to answers through an attention-induced graph. It uses answer-targeted global flow analysis to identify influential reasoning tokens and guide token-level RL updates.

  • Motivation: Token-level credit assignment is difficult because autoregressive trajectories are long and supervision is sparse and delayed, while value-based methods depend on accurate state-value estimates.Correctness may be judged only at the end of generation, making individual-token attribution challenging.
  • Method: FlowTracer models the token sequence as an attention-induced directed acyclic graph whose nodes are tokens and whose edge capacities derive from aggregated attention weights.The capacities represent the strength of potential information transfer between tokens.
  • Method: FlowTracer reweights attention edges to retain only influence that can reach the answer and enforces local flow conservation to prevent dilution along long paths or absorption by irrelevant branches.This counters under-crediting of early decisive premises and over-weighting of late tokens near the answer.
  • Method: Unit flow from prompt tokens to answer tokens yields an information-flow backbone whose token throughputs reveal high-impact transit hubs and aggregation checkpoints supporting the final prediction.A super-source injects flow into prompt tokens, and a super-sink collects it from answer tokens.

2 Related Work

Related work increasingly extracts fine-grained optimization signals from LLM internal computations, while credit assignment remains an evolving challenge in RL post-training. Prior studies identify attention heads, receiver heads, functional layers, and representation-space directions as important information-handling components.

  • Deriving Optimization Signals From LLM Internal Dynamics: Recent work exploits LLM internal computational processes to derive fine-grained optimization signals beyond black-box, end-to-end treatment.This line of research examines internal dynamics to support more precise optimization.
  • Deriving Optimization Signals From LLM Internal Dynamics: Studies identify task-specific attention heads, context-aggregating receiver heads, specialized functional layers, and representation-space steering directions as important for information handling.These insights support LLM augmentation through methods including representation editing.
  • Credit Assignment for RL in LLMs: RL is standard for LLM post-training, but effective credit assignment remains an evolving challenge amid sparse outcome rewards and differing off-policy and on-policy approaches.Off-policy RL aligns probabilities, whereas more prevailing on-policy RL relies on sparse outcome rewards.

3 Methodology

FlowTracer models LLM reasoning as conserved, answer-targeted influence flow on an attention-induced DAG, then uses forward flow throughput to assign token-level credit. The framework identifies high-flow hubs that connect question information to the final answer and uses them to focus RL updates.

  • Raw Influence Graph from Attention: FlowTracer constructs a time-ordered DAG whose token nodes are connected by attention-derived influence edges.Edge weights aggregate attention scores across layers and heads, while the graph need not be stochastic because outgoing weights may sum above or below one.
  • Doob-h-Like Reweighting for Effective Influence: Doob-h-like reweighting enforces local conservation and retains only influence paths that ultimately reach the answer.The potential h(i) measures total influence from node i to a virtual answer sink, and reweighting suppresses dead-end branches while reallocating flow to answer-reaching paths.
  • Forward Flow for Token-Level Throughput: Forward propagation injects unit source flow across question tokens and scores each token by throughput in the reweighted graph.Token throughput combines incoming flow with outgoing edge flow, identifying nodes that carry effective question-to-answer influence and are suitable for fine-grained credit assignment.
  • Information-Flow Analysis: High-flow tokens form a sparse reasoning backbone, often appearing as periodic structural delimiters or symbolic anchors.The analysis uses Qwen3-4B-Base on GSM8K and reports that information flow is not uniformly distributed across tokens.
  • Causal Validation: Masking 20% of high-flow tokens strongly influences final answers, providing causal evidence that these hubs drive reasoning.The intervention compares perturbations targeting high-flow, low-flow, and randomly selected tokens on GSM8K.
  • RL Credit Assignment: The derived backbone enables RL to assign higher credit to answer-driving tokens instead of rewarding every token uniformly.The method uses γflow = 1.5 as an emphasis factor so policy updates are more aggressive for tokens that contribute to the answer.

4 Experiments

Experiments evaluate FlowTracer across mathematical, logical, and puzzle reasoning tasks, model families, context lengths, and token-selection or layer-signal ablations. FlowTracer consistently improves reasoning performance, with especially large gains in long-context and combinatorial-planning settings, while adding only marginal computational overhead.

  • Mathematical reasoning: FlowTracer consistently outperforms GRPO and alternative token-level credit heuristics on Qwen3-4B and Qwen3-8B mathematical reasoning at 1K context.The experiments vary only credit assignment criteria while keeping other settings fixed.
  • Mathematical reasoning: 43.4% average accuracy is achieved by FlowTracer on Qwen3-8B in the reported 1K-length mathematical reasoning evaluation.This result is reported across the standard mathematical benchmarks in Table 2 and Fig. 4.
  • Long-context reasoning: FlowTracer’s advantage becomes more pronounced from 1K to 8K context, where sparse reasoning signals make precise credit assignment increasingly important.The passage attributes standard RL degradation in longer chains to credit dilution and increased noise.
  • Cross-domain generalization: +10.6 absolute over GRPO on Countdown and +2.2 on CrossThinkQA demonstrate gains in combinatorial planning and natural-language logical reasoning.Both evaluations use Qwen3-4B-Base with a 1K context window and the same RL recipe as the mathematical experiments.
  • Ablations: Top-k token selection consistently improves over GRPO, while the ablation varies both flow-score direction and the selection ratio.The final chosen setting is Top-k with k = 0.4.
  • Efficiency: FlowTracer adds only 2.2%–4.5% time overhead by using one additional batched forward pass, whose cost is overshadowed by autoregressive sampling.The pass extracts mid-layer attention maps, computes answer-targeted flow scores, and selects Top-40% tokens for policy-gradient reweighting.

5 Additional Analysis and Limitations

Additional analyses find that hard Top-40% flow-based credit is more stable and robust than continuous or cumulative alternatives, with structural delimiters contributing most of the gain. FlowTracer uses attention only to extract answer-directed routing signals and has limitations involving target definitions, outcome-only rewards, and long contexts.

  • Credit-assignment ablations: Hard Top-40% credit is more stable than continuously scaling tokens by raw or transformed flow values.FlowTracer selects the Top-40% tokens and multiplies their GRPO surrogate terms by γflow.
  • Credit-assignment ablations: γflow = 1.5 performs best, while smaller or larger values degrade performance.The analysis also finds fixed token-ratio selection more robust than cumulative-flow-mass thresholds because flow mass is highly concentrated.
  • Robustness analyses: Top-40% high-flow tokens remain nearly unchanged across answer formats, with overlap 1.00 for full-line and explanation-augmented answers and 0.93 for multiple-choice outputs.This tests variations within the answer region despite the main experiments using boxed answers.
  • Robustness analyses: Structural delimiters recover most of the full gain, reaching 38.8% versus 39.4%, while semantic-content tokens reach 37.6%.Structural delimiters therefore dominate the improvement, with semantic-content tokens providing a smaller complementary gain.
  • Limitations: FlowTracer does not claim attention fully explains reasoning; it filters attention into an answer-targeted, flow-conserving multi-hop backbone for analysis and RL credit assignment.Attention is used narrowly as an explicit token-to-token interaction graph from which routing signals are extracted.
  • Limitations: Current limitations include localized answer-region assumptions, outcome-only rewards that cannot distinguish useful intermediate reasoning supporting wrong answers, and noisier graphs in very long contexts.The paper identifies flexible target definitions and PRM-style process supervision as relevant future directions.

6 Conclusion … A.2 Reinforcement Learning with Verifiable Rewards

FlowTracer addresses token-level credit assignment in RL4LLM by reconstructing attention-based global information flow and using it to focus rewards on influential reasoning tokens. The preliminaries establish how decoder-only LLMs route information through causal self-attention and how RLVR methods assign token credit.

  • 6 Conclusion: FlowTracer models reasoning as a directed flow network to identify a reasoning backbone and distinguish decisive steps from routine filler.The framework uses model-internal attention signals to reconstruct global information structure based on tokens’ actual contributions.
  • 6 Conclusion: Flow-based reward shaping focuses learning signals on high-impact tokens and delivers consistent gains across complex reasoning tasks.The work further positions attention geometry as a structural signal for more efficient and interpretable model alignment.
  • A.1 Information Propagation in Decoder-Only LLMs: A decoder-only LLM generates response y from prompt x while maintaining hidden-state information streams over the concatenated sequence S = (x, y<t).At each decoding step, the sequence has length N and hidden states have dimension d.
  • A.1 Information Propagation in Decoder-Only LLMs: Multi-Head Self-Attention routes information by projecting hidden states into queries, keys, and values, with attention matrices determining token-to-token flow.The attention matrix A^(h,l) is computed within each layer l and head h.
  • A.1 Information Propagation in Decoder-Only LLMs: Causal masking restricts each token to preceding context, while attention outputs aggregate contextual values before FFN and residual processing produce globally enriched states.The mask sets M_i,j = −∞ for j > i, enforcing j ≤ i information access.
  • A.1 Information Propagation in Decoder-Only LLMs: Attention entries A^(h,l)_i,j quantify the proportion of information token i retrieves from token j, making attention the natural basis for global-flow topology analysis.This matrix acts as the explicit data-routing mechanism among the model components.
  • A.2 Reinforcement Learning with Verifiable Rewards: RLVR trains an LLM with sparse verifiable rewards, while PPO uses a value-function critic and GRPO replaces the critic with group statistics as a baseline.PPO therefore incurs computational and memory overhead, whereas GRPO samples G outputs for each query from the old policy.
  • A.2 Reinforcement Learning with Verifiable Rewards: Standard GRPO sets γ_i,t = 1, assigning uniform credit to every token despite evidence that fine-grained signals can distinguish critical reasoning steps from trivial tokens.The generalized surrogate objective aggregates gradients over all tokens using token-wise importance coefficients.

B Model and Dataset Specification · C Implementation Details for RL Experiments in Sec. 4 · C.1 RL Framework

FlowTracer is evaluated with Qwen3-4B and Qwen3-8B as primary backbones, supplemented by Llama-3.1-8B and Llama-3.2-3B, across mathematical reasoning, multi-domain question answering, and domain-specific puzzle-solving tasks. RL experiments use the ROLL framework with separate inference and training actors, with actor infer generating experience through vLLM for subsequent reward-model or verifier scoring.

  • B Model and Dataset Specification: Llama-3.1-8B and Llama-3.2-3B provide supplementary experimental results for generalization validation.Detailed model specifications, references, and access links are listed in Table 8.
  • B Model and Dataset Specification: FlowTracer is evaluated on challenging mathematical reasoning, multi-domain question answering, and domain-specific puzzle-solving tasks.The dataset specifications and access links are listed in Table 9.
  • B Model and Dataset Specification: CrossThinkQA tests cross-disciplinary reasoning with multiple-choice questions spanning science, humanities, and social sciences.The task requires synthesizing knowledge across domains to select correct answers.
  • B Model and Dataset Specification: Countdown evaluates search and planning by requiring four randomly selected numbers to be combined with +, −, ×, or ÷ to reach a target exactly.It uses a constrained discrete action space.
  • C.1 RL Framework: ROLL separates rollout generation and policy optimization through specialized actor infer and actor train instances.Actor infer uses vLLM for high-throughput experience generation, after which a reward model or verifier scores the experience.

C.2 Credit Assignment Implementation Details

The credit-assignment mechanism sits between sampling and training: it extracts attention maps, performs flow analysis for token-level credit, and modulates token advantages. High-credit tokens are then prioritized during policy optimization by scaling their advantages.

  • Credit assignment pipeline: Credit assignment bridges sampling and training by quantifying token-level credit from attention flow before modulating token advantages.The mechanism extracts attention maps from sampled responses, analyzes their flow, and passes the resulting credits to advantage processing.
  • Attention extraction: After sampling, each prompt is concatenated with its response, followed by an additional batch forward pass that materializes full attention weights.The eager execution mode of attention is disabled so complete attention weights are accessible; maps are extracted across heads in the model’s middle layers.
  • Advantage modulation: 40% of tokens are selected as high-credit tokens, whose advantages are proportionally scaled before loss calculation and policy optimization.The modified advantages complete the iterative training update loop.

C.3 Training Details

FlowTracer training uses task-specific datasets and standard benchmark evaluation, with Countdown and multi-domain QA following prescribed construction or partition protocols. Training applies fixed rollout, sampling, optimization, and model-specific step settings, while Llama experiments bypass supervised fine-tuning through few-shot prompting.

  • Data Usage and Splits: Mathematical reasoning trains on DAPO-Math-17K and evaluates on standard benchmarks such as AIME and OlympiadBench to test generalization.
  • Data Usage and Splits: Countdown uses 20,000 synthesized training samples and 5,000 held-out test samples, while CrossThinkQA uses its official training and validation splits.
  • Training Hyperparameters: Training runs for 600 steps on 8B models and 500 steps on 4B and 3B models, using 64 prompts per rollout batch and 8 responses per prompt.Responses have maximum lengths of 1,024 or 8,192 tokens; sampling uses temperature 0.99, top-p 1.0, and top-k 100, with learning rate 1.0 × 10^-6 and 20-step linear warmup.
  • Experiments Details for Llama Families: Llama-3.1-8B and Llama-3.2-3B experiments bypass supervised fine-tuning and initiate reinforcement learning through few-shot prompting.

D Additional Experimental Results … E.1 Implementation Details

The additional experiments broaden FlowTracer’s evaluation through stronger baseline comparisons, answer-format and token-type robustness checks, and controlled analyses of token flow and perturbation effects. Implementation details specify datasets, sampling, token selection, perturbation procedures, answer-based metrics, and the analytical prompt.

  • D.1 Additional Baseline Comparisons: Additional comparisons evaluate FlowTracer against generative-PRM, value-uncertainty, confidence-reweighting, and entropy-based token-selection baselines on Qwen3-8B with 1K context.These methods differ in implementation details and supervision assumptions.
  • D.2 Additional Results for Answer-Format and Structural-vs-Semantic Checks: Answer-format robustness compares high-flow-token overlap between the default answer-only format and alternative forms inside the explicitly delimited boxed answer region.The full results correspond to the answer-format check discussed in Sec. 5.2.
  • D.2 Additional Results for Answer-Format and Structural-vs-Semantic Checks: Structural-versus-semantic analysis separates structural delimiters from semantic-content tokens, selects Top-40% high-flow tokens within each subset, and applies the same credit-assignment procedure.Table 11 identifies this as a token-type ablation on Qwen3-4B with 1K context length.
  • E Details for Analytical Experiments in Sec. 3.2: Information-flow analyses use 500 randomly sampled GSM8K instances with the Qwen3-4B model.These experiments analyze flow characteristics and verify its properties.
  • E.1 Implementation Details: Word-cloud analysis generates responses at temperature 0.7 and aggregates the top 20% high-flow and low-flow tokens identified within each sentence.Separate word clouds characterize compositional differences between the two token groups.
  • E.1 Implementation Details: Perturbation experiments use temperature 0.7, randomly truncate sentences, select high-flow, low-flow, or random preceding tokens, and mask them to test implicit multi-hop influence.The procedure uses truncation rather than single-token substitution.
  • E.1 Implementation Details: Perturbation effects are measured by answer change rate and correctness change rate, using extracted original and regenerated answers against GSM8K golden answers.Failed answer extraction is treated as incorrect before aggregating the statistics.
  • E.1 Implementation Details: The analytical experiments use a step-by-step reasoning prompt requiring the final answer inside a boxed region.The prompt template is explicitly provided for these experiments.

E.2 Representative Perturbation Examples

Representative masking experiments show that high-flow tokens are critical for maintaining coherent, correct reasoning, whereas low-flow token masking can preserve correctness by inducing an alternative reasoning path. High-flow tokens therefore function as anchors for intermediate conclusions, numerical values, and structural links across reasoning steps.

  • Experimental setup: The examples target high-flow and low-flow tokens by masking selected positions and observing how the generated response, especially its final answer, changes.The section presents two representative perturbation examples.
  • High-flow perturbations: Masking high-flow tokens causes regenerated responses to lose the task objective and prior conclusions, producing repetitive incoherence and an incorrect answer.In the Julie example, truncation at “remaining” disrupts continuation from the next token onward.
  • Low-flow perturbations: Masking low-flow tokens at the same position switches the model to an alternative reasoning path while preserving coherent reasoning and the correct answer.The perturbed Julie response still computes 42 pages.
  • Interpretation: High-flow tokens act as pivotal anchors carrying intermediate conclusions, key numerical values, and structural cues that bind successive reasoning steps together.Masking these positions disrupts the model’s train of thought and propagates errors.

F Flow Analysis Example

Flow analysis of Qwen3-4B outputs reveals a periodic information-flow pattern for the token “the”: information aggregates from prior occurrences and propagates forward. This supports the findings reported in the main text.

  • F Flow Analysis Example: Flow analysis examines the inflow and outflow of the token “the” in outputs generated by Qwen3-4B.The analysis is illustrated in Fig. 7.
  • F Flow Analysis Example: Information periodically aggregates from previous occurrences of the same token, “the”.The repeated token serves as a source of aggregated information in the observed flow pattern.
  • F Flow Analysis Example: After aggregation, information propagates forward through subsequent occurrences of “the”.This observation corroborates findings presented in the main text.
Loading 2606.10646v1…