> ICML 2026 · Shanghai Jiao Tong University × Alibaba × Shanghai AI Laboratory
The old problem: who gets the credit?
Imagine a student solving a geometry problem, filling two pages of scratch work, then making an arithmetic error on the final step and getting zero. In RL training of large models, the same dilemma arises: the model generates hundreds of tokens, but only the final answer gets a score. How should that score be distributed?
GRPO's answer is everyone gets an equal share — near-uniform credit per token, like a communal pot. The obvious problem: tokens that genuinely support the correct answer get diluted signals, while filler or meandering tokens soak up reward.
FlowTracer starts from one observation:
> "Attention is not decoration — it is itself the roadmap of reasoning."
Existing credit assignment methods are either uniform (GRPO) or isolated heuristics (per-token entropy, attention statistics), both ignoring the global structure of information propagation: which tokens communicate with each other, and which are isolated islands?
Turning water ripples into a map
The core intuition is clean: attention weights tell you where a token pulls information from. Stacked across layers and heads, they form a directed information-transport graph — tokens as nodes, attention weights as pipe capacities.
But raw attention maps have two flaws:
1. Too many dead ends. Strong local connections that never reach the answer — energy circulates locally, a structural waste. 2. Flow is not conserved. Attention weights don't guarantee inflow = outflow, so it isn't a true flow network and max-flow algorithms can't run on it.
The authors apply a Doob h-transform (fancy name, simple idea).
The filter: make every drop flow toward the answer
Assume you know the answer region (e.g., the final numeric tokens of a math solution). Define h(i): the potential of "being able to reach the answer" from token i, computed by propagating an "answer gravity" backward from answer nodes across the graph. Then reweight all pipes:
This step enforces local flow conservation (Theorem 3.1). Mathematically it is equivalent to superimposing an "answer-biased potential field" on the original graph. Why conservation? Because without it, "throughput" is undefined — you can't tell how much information a node processes. In plain language: you add a gravity field pointing toward the answer; water not flowing that way is canceled out. What remains is a clean, directed, dead-end-free reasoning backbone. The name "Doob" is deliberate: in probability theory, Doob's h-transform modifies transition probabilities of a stochastic process to favor reaching a target state — exactly what's done here to attention weights.
Installing water meters: who relays at the critical junctions?
With a clean network, the authors add a super-source before the prompt and a super-sink after the answer region, inject unit flow, and run max-flow. Each token's throughput f(k) measures how much flow passes through it. High-throughput tokens are key relay hubs of the reasoning process — like parcel-sorting centers: they don't deliver the package themselves, but without them nothing arrives.
Experiments: not icing on the cake, but coal in a snowstorm
FlowTracer costs only 2.2%-4.5% extra compute (one additional forward pass).
Main results
| Model | Baseline (GRPO) | FlowTracer | Gain | |------|------------|-----------|------| | Qwen3-8B | 39.4% | 43.4% | +4.0% | | Qwen3-4B | 37.1% | 39.4% | +2.2% | | Llama-3.1-8B | 7.7% | 9.1% | +1.4% | | Llama-3.2-3B | 4.8% | 5.9% | +1.1% |
A +4.0% absolute gain on Qwen3-8B is significant on reasoning benchmarks. Notably, at 8K long context the advantage grows: Qwen3-4B gains +5.8% on AIME25. The longer the sequence, the more uniform credit assignment blunders, and the more "knowing which tokens actually carry information" pays off.
On Countdown logic puzzles the gain is a striking +10.6% — tasks demanding precise multi-step arithmetic hit FlowTracer's sweet spot: critical intermediate steps must be precisely identified. CrossThinkQA (cross-domain reasoning) gains +2.2%, showing the method is not overfit to one task type.
Ablations: 40% of hubs is enough
- Top-20%: insufficient coverage, missing key relays — worse
- Top-40%: optimal balance of information vs. noise
- Top-60%: too many low-flow tokens, noise drowns signal
- Paper: How Does Reasoning Flow? Tracing Attention-Induced Information Flow for Targeted RL in LLMs
- Authors: Zhichen Dong, Yang Li, Yuhan Sun, Weixun Wang, Yijia Luo, Zinian Peng, Taiheng Ye, Chao Yang, Wenbo Su, Yu Cheng, Bo Zheng, Junchi Yan
- Institutions: Shanghai Jiao Tong University · Alibaba · Shanghai AI Laboratory
- Venue: ICML 2026 (Accepted)
- arXiv: 2606.10646
Mid-layer attention (layers 15–25) yields the most effective networks — too shallow lacks semantic structure, too deep is highly compressed. γ_flow = 1.5 is optimal. An interesting detail: hard masking beats continuous reweighting, because the flow distribution is extremely skewed — a few tokens carry most of the flow, and continuous weighting lets low-flow tokens freeload on credit. Cut it clean: reward only the hubs.
Structural separators are the hidden MVP
Counter-intuitively, the highest-flow nodes are not semantic content tokens (nouns, verbs) but structural separators — commas, periods, newlines, indentation — acting like traffic lights controlling where information moves between semantic blocks. This may explain why prompting models to wrap reasoning in <think> tags works: those tags are artificially created high-flow structural nodes.
Causal validation: not correlation, causation
| Perturbed target | Answer change rate | Correct→wrong rate | |---------|----------|-----------| | High-flow tokens | 45.9% | 14.9% | | Random tokens | 29.5% | 4.5% | | Low-flow tokens | 14.9% | 0.5% |
Perturbing high-flow tokens changes nearly half the answers; low-flow tokens barely matter. This is causal verification through active intervention, not post-hoc attribution.
Honest limitations
1. Requires localizing the answer region. Open-ended generation lacks a clear answer token set; extending h is future work (last-turn responses could serve as the answer region in dialogue). 2. Outcome reward alone isn't enough. Valid intermediate reasoning with a wrong final answer still gets low scores — FlowTracer is complementary to PRMs (process reward models), not a competitor: it finds key tokens, PRMs judge whether the reasoning is correct. 3. 16K+ contexts may produce noisier attention graphs and less accurate flow estimates — though the 8K gains suggest the issue is robustness of graph estimation, not method failure.
A question worth pondering
FlowTracer answers: "given the answer, which tokens are irreplaceable?" But human thinking often works the other way — a token can be critical because it eliminates a wrong branch, not because it leads to the right answer. The h-function captures only "positive reachability," not the credit for "negative exclusion." How do we distinguish absent information transfer from important-but-weak transfer in attention flow?