PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Long-Horizon Interpretability: Efficient and Faithful Multi-Token Attribution for Reasoning LLMs

Wenbo Pan, Zhichao Liu, Xianlong Wang, Haining Yu, Xiaohua Jia

arXiv:2602.01914v4cs.LG

TL;DR

Existing token attribution faces high costs for multi-token targets and difficulty tracing importance through intermediate reasoning to source inputs. FLASHTRACE combines span-wise aggregation with recursive attribution, and experiments report faster attribution with superior faithfulness across long-context retrieval and reasoning tasks.

  • Problem

    Multi-token attribution can require O(M · N) operations, while intermediate reasoning can absorb attribution mass that does not reach source inputs.

    Source: §1 · p. 1, Finding 2: Reasoning chains degrade… · p. 4

  • Method

    FLASHTRACE aggregates target-span contributions and recursively traces importance through intermediate reasoning back to input tokens.

    Source: §1 · p. 2, §4.2 · p. 5

  • Results

    FLASHTRACE is over 130× faster than IFR for a 5k-token target span and demonstrates superior faithfulness across long-context and reasoning tasks.

    Source: §1 · p. 2, §5.1 · p. 7, §6 · p. 9

  • Takeaways & Limitations

    Experiments across long-context retrieval, mathematical reasoning, and multi-hop QA support scalable interpretability for modern agentic workflows.

    Source: §6 · p. 9

Abstract

from arXiv · show

Token attribution methods provide intuitive explanations for language model outputs by identifying causally important input tokens. However, as modern LLMs increasingly rely on extended reasoning chains, existing schemes face two critical challenges: (1) efficiency bottleneck, where attributing a target span of M tokens within a context of length N requires O(M*N) operations, making long-context attribution prohibitively slow; and (2) faithfulness drop, where intermediate reasoning tokens absorb attribution mass, preventing importance from propagating back to the original input. To address these, we introduce FlashTrace, an efficient multi-token attribution method that employs span-wise aggregation to compute attribution over multi-token targets in a single pass, while maintaining faithfulness. Moreover, we design a recursive attribution mechanism that traces importance through intermediate reasoning chains back to source inputs. Extensive experiments on long-context retrieval (RULER) and multi-step reasoning (MATH, MorehopQA) tasks demonstrate that FlashTrace achieves over 130x speedup over existing baselines while maintaining superior faithfulness. We further analyze the dynamics of recursive attribution, showing that even a single recursive hop improves faithfulness by tracing importance through the reasoning chain.

1. Introduction

FLASHTRACE targets two limitations of multi-token attribution: prohibitive cost and importance trapped in intermediate reasoning. It combines span-wise aggregation with recursive attribution and outperforms baselines in efficiency and faithfulness across long-context and reasoning tasks.

1. Introduction

  • Attributing each target token separately makes multi-token explanations prohibitively costly, while intermediate reasoning absorbs attribution mass and obscures important input tokens.Integrated Gradients can take over 10 hours to explain a single 5k-token generation.Source: §1 · pp. 1–2
  • FLASHTRACE aggregates a contiguous generated-token span in one pass per hop and recursively traces importance through intermediate reasoning to the input.The authors report significant speedup while maintaining superior faithfulness across task domains and model architectures.Source: §1 · p. 2
  • FLASHTRACE surpasses baselines in efficiency and faithfulness across long-context and reasoning tasks, with performance gains consistent across models and data distributions.The reported advantage is particularly pronounced on complex reasoning such as multi-hop QA.Source: §1 · p. 2

2. Related Work

Prior token-attribution work spans perturbation-, gradient-, attention-, and relevance-propagation approaches, but multi-token and reasoning-chain attribution remains challenging. CAGE recursively traces reasoning at sentence granularity, while feature-circuit methods study a different interpretability target.

Token Attribution for Transformers

Reasoning LLM and Long-Context Agent

3. Multi-Token Attribution

The section frames attribution as identifying which context tokens causally explain an output, then examines how intermediate reasoning affects attribution quality. It describes both the resulting loss of input recovery and the computational cost of naive recursive attribution.

3.1. Background

  • The context is decomposed into input tokens I, intermediate reasoning tokens T, and final output tokens O; attribution seeks context tokens with causal influence on the generated output.The overview figure depicts span-wise aggregation of causal attention and residual contributions to support multi-hop attribution.Source: §3.1 · p. 3, Figure 2 · p. 3

3.2. How Reasoning Tokens Impact Attribution

  • In a RULER multi-hop retrieval setup, AttnLRP assigns a growing share of attribution mass to reasoning tokens as the chain length increases, leaving little for the original input.The comparison uses Qwen-3 8B Instruct in direct-answering and chain-of-thought settings, with correct answers in both.Source: §3.2 · p. 3

Finding 2: Reasoning chains degrade attribution quality

  • With reasoning, standard attribution recovers fewer ground-truth input tokens in the top 10% and fails to propagate importance back through the chain.Naive recursive attribution requires M separate runs for a reasoning span of length M, increasing time complexity from O(N) to O(M · N) for context length N.Source: Finding 2: Reasoning chains degrade… · p. 4

4. FLASHTRACE

FLASHTRACE combines span-wise aggregation with recursive attribution to explain multi-token outputs efficiently and trace importance through intermediate reasoning to source inputs.

4.1. Preliminaries

  • The method builds on ALTI and IFR, using L1-based proximity to estimate how much a source token’s representation contributes to a target, rather than a direct causal effect.The layer decomposition includes attention, feed-forward, and residual contributions.Source: §4.1 · p. 4
  • The hop visualization shows attribution shifting from output-near reasoning tokens toward earlier reasoning tokens and input context.The difference map marks gains in green and losses in red.Source: Figure 3 · p. 5, §4.1 · p. 5

4.2. Span-wise Aggregation

  • Span-wise aggregation reduces multi-token attribution complexity from O(M · N) to O(N) by computing proximity for an aggregated target span.The target span’s representations and source contributions are summed before applying the proximity metric.Source: §4.2 · p. 5

4.3. Recursive Attribution

  • Recursive attribution uses prior importance scores on reasoning tokens as weights for the next target, tracing influence from output through reasoning to input.Experiments use a small hop count; one hop is reported sufficient to resolve information absorption while keeping computation comparable to a single forward pass.Source: §4.3 · pp. 5–6
  • Across hops, reasoning-token importance supplies weights for the next hop, while input importance is accumulated across hops using the residual mass carried forward.The process can repeat for K hops and aggregates the resulting input scores into one distribution.Source: §4.3 · p. 6

5. Experiments

Experiments evaluate FLASHTRACE on long-context retrieval, multi-step reasoning, and multi-hop QA. The reported results include faster attribution, stronger faithfulness, and comparable performance to exhaustive token-level rollout on MorehopQA.

5. Experiments

  • The evaluation covers RULER retrieval and tracking, MATH and MorehopQA reasoning, and long-context HotpotQA; Qwen-3 8B Instruct generates responses by default.Experiments use one recursive hop unless otherwise specified.Source: §5 · p. 6

5.1. Main Results

  • FLASHTRACE maintains superior Recovery Rates across RULER context lengths, shows superior faithfulness on MATH and MorehopQA, and has a pronounced margin on HotpotQA.The RULER comparison includes NIAH and Variable Tracking; faithfulness on the reasoning tasks is evaluated with causal metrics.Source: Table 1 · p. 7, §5.1 · pp. 7–8, Table 2 · p. 8
  • At a 5k-token target span, FLASHTRACE is over 130× faster than IFR, taking under 20 seconds versus over 38 minutes.The efficiency comparison also reports that span-wise aggregation avoids memory overhead that scales with target span length.Source: Figure 4 · p. 7, §5.1 · p. 7
  • FLASHTRACE maintains stable MAS across reasoning lengths and consistently outperforms baseline methods.Figure 5 plots MAS against bins of generated reasoning-token length.Source: Figure 5 · p. 8, §5.1 · p. 8

Compare with Exhaustive Token-Level Rollout

6. Conclusion

FLASHTRACE combines span-wise aggregation and recursive attribution to address efficiency and faithfulness challenges in reasoning-LLM interpretation.

6. Conclusion

  • FLASHTRACE combines span-wise aggregation with recursive attribution to trace importance through reasoning chains back to source inputs.The method is presented as addressing both efficiency and faithfulness challenges.Source: §6 · p. 9
  • Experiments across long-context retrieval, mathematical reasoning, and multi-hop QA report significant speedup while maintaining superior faithfulness.Source: §6 · p. 9
  • The authors present the method as enabling scalable interpretability for modern agentic workflows.Source: §6 · p. 9

Impact Statement

The paper presents FlashTrace as an efficient, recursive method for attributing multi-token outputs to their input context. Its impact statement connects improved attribution to understanding, debugging, and auditing model behavior, while noting an attention-cache memory requirement.

Impact Statement

  • The authors say efficient, faithful attribution may help practitioners understand, debug, and audit model behavior in high-stakes applications.They report no direct negative societal consequences specific to this work.Source: Impact Statement · p. 9

A.1. Transformer Architecture Decomposition

  • The normalization linearization is inherited from ALTI and IFR, and the paper reports that it does not compound across layers or recursive hops.The authors report that FlashTrace closely matches exhaustive token-level rollout on RISE and MAS faithfulness in Table 3.Source: §A.1 · p. 13

A.2. Span-wise Attribution (Single Hop)

  • FlashTrace attributes a target span of M tokens simultaneously in O(N) time, rather than processing each target token separately.It aggregates attention contributions across the span so expensive vector transformations need only be computed once per source token.Source: §A.2 · pp. 13–14

A.3. Recursive Multi-Hop Attribution

  • Recursive attribution starts from the final output, traces important reasoning tokens back to context, and combines hop distributions using cumulative flow ratios.The weighted sum captures direct input influence and indirect influence mediated by reasoning.Source: §A.3 · pp. 14–15

B.1. Time Complexity

  • With typically 2–4 recursive hops, FlashTrace remains linear in context length; its remaining O(M·N) scalar aggregation is described as negligible in practice.The method moves the heavy D-dimensional vector arithmetic outside the loop over target span length M.Source: §B.1 · p. 15

B.2. Space Complexity

  • Attribution working memory is independent of target span length M, but explicitly cached attention maps can require O(L·N^2) space.The paper notes that recomputing attention on the fly or using memory-efficient attention can mitigate this cache requirement.Source: §B.2 · p. 17

C. Efficiency Experiment Details

The efficiency experiments vary input and generation lengths and measure peak GPU memory during attribution. Baselines include gradient-, attention-, and perturbation-based methods, though IG and IG-Attn could not be evaluated reliably for faithfulness at the required RULER lengths.

C. Efficiency Experiment Details

  • Efficiency experiments used five NVIDIA A800 GPUs and varied input or generation lengths from 10 to 5,000 tokens while holding the other dimension at 100.Reported memory is peak GPU use during attribution.Source: C. Efficiency Experiment Details · p. 17
  • IG and IG-Attn frequently ran out of memory at the RULER lengths, so their faithfulness scores could not be reliably computed.Their Pareto-plot positions show computational cost, not faithfulness.Source: C. Efficiency Experiment Details · p. 17

D. Baseline Methods

  • The comparison includes gradient-based IG and attention-aware IG, alongside perturbation methods such as Leave-One-Out and REAGENT.The baseline implementations were adapted from Walker & Ewetz (2025), while IFR and AttnLRP followed the original authors' open-source implementations.Source: D. Baseline Methods · p. 18
  • IFR decomposes transformer computation into layer-wise information-flow routes, while AttnLRP propagates relevance through transformer components.The paper also includes Context Length Probing, which varies context length to estimate token importance.Source: D. Baseline Methods · p. 18

D.1. Extending Baseline Methods to Multi-Token Attribution

The evaluation adapts single-token baselines for multi-token outputs, makes perturbation methods tractable for long contexts, and includes exhaustive token-level rollout as a comparison. Across retrieval, reasoning, and code-generation tasks, model-generated intermediate traces and task-appropriate attribution metrics shape the evaluation.

D.1. Extending Baseline Methods to Multi-Token Attribution

  • Single-token baselines attribute each output token separately and average the resulting vectors; evaluation then removes reasoning and output scores and renormalizes over input tokens.This keeps RISE, MAS, and Recovery Rate focused on source information within the input prompt.Source: §D.1 · pp. 18–19

D.2. Optimizing Baseline Methods for Long Contexts Efficiency

  • Sentence- or chunk-level perturbation reduces forward passes by a factor of 20–50×, assigning each sentence's score uniformly to its constituent tokens.The strategy addresses the cost of token-level perturbation on contexts reaching 32k tokens.Source: §D.2 · p. 19

D.3. Exhaustive Token-Level Rollout Implementation

  • Exhaustive rollout recursively attributes output tokens to reasoning tokens and then to input, using token-level passes as a costly upper-bound comparison.Unlike CAGE's sentence-level granularity, the rollout maintains token-level precision and requires O(M) passes per hop.Source: §D.3 · p. 19

E. Dataset Construction

  • The evaluation datasets span long-context retrieval, mathematical problem-solving, multi-hop question answering, and code generation, formatted as input, reasoning, and output.RULER includes needle retrieval, variable tracking, and long-context HotpotQA with ground-truth evidence in the input.Source: E. Dataset Construction · p. 19, §E.1 · pp. 19–20

E.2. Complex Reasoning: MoreHopQA and MATH

  • MoreHopQA extends two-hop questions to three or more hops and supplies provenance for each reasoning hop; MATH contains 12,500 competition-level problems.MATH solutions provide step-by-step reasoning, but lack token-level evidence labels, so the study uses RISE/MAS rather than recovery rate.Source: §E.2 · p. 20

E.3. Code Generation: Aider

  • The Aider benchmark uses 133 Python exercises, attributing generated code bodies to natural-language instructions and function definitions.Its intermediate tokens are structured code rather than natural-language thoughts.Source: §E.3 · pp. 20–21

E.4. Generating Intermediate Reasoning Tokens T

  • Qwen-3 8B Instruct generated the reasoning traces; evaluation retained correct answers and randomly selected 100 samples per dataset, except for 66 HotpotQA samples.Incorrect-answer traces were filtered because their reasoning could be flawed or hallucinatory.Source: §E.4 · p. 21

F. Evaluation Details

  • Recovery Rate measures recall of ground-truth evidence among the top 10% of input tokens, while complex reasoning uses perturbation-based faithfulness metrics when token labels are unavailable.The evaluation combines ground-truth recovery for synthetic tasks with perturbation metrics for complex reasoning tasks.Source: F. Evaluation Details · p. 21, §F.1 · p. 21

F.2. Perturbation Metrics: RISE and MAS

This section defines deletion-based faithfulness metrics and adapts their evaluation to long contexts. A case study visualization shows recursive attribution strengthening input-region attribution while attenuating reasoning-region mass.

F.2. Perturbation Metrics: RISE and MAS

  • RISE and MAS deletion evaluate faithfulness by masking high-importance input tokens first and tracking the resulting drop in the model’s target-sequence probability.Lower deletion scores are better; MAS additionally penalizes misalignment between attribution magnitude and response magnitude.Source: §F.2 · pp. 21–22
  • Masking the top 5% of remaining tokens at each step yields 20 evaluation intervals and a fixed budget of 20 forward passes per sample.This proportional schedule adapts the metrics to long contexts, where masking one token per step would require thousands of passes.Source: §F.2 · p. 22

G. Case Study Visualizations

  • On MATH examples, Hop 2 strengthens attribution on key input tokens and attenuates attribution mass on reasoning tokens relative to Hop 1.The visualization shows the change between direct context attribution and recursive attribution through the reasoning chain.Source: G. Case Study Visualizations · p. 22

H. Ablation: Effect of Recursive Hops

The section examines recursive hops, additional faithfulness settings, and a memory–runtime trade-off. Its findings include improved attribution after one hop, retained faithfulness at longer contexts, and a lower-memory recomputation option.

H. Ablation: Effect of Recursive Hops

  • One recursive hop improves faithfulness over direct attribution, with stable performance across hop counts and diminishing returns at higher depths.The authors suggest noise accumulation may explain diminishing returns; Table 7 reports RISE and MAS on MoreHopQA across hop configurations.Source: H. Ablation: Effect of Recursive Hops · p. 22, Table 7 · p. 22
  • From hop 0 to hop 1, attribution mass shifts away from reasoning tokens toward input tokens and, within the reasoning span, toward earlier reasoning tokens.The MoreHopQA visualization aggregates 100 samples; Figure 6 also shows increased input attribution and decreased reasoning attribution between Hop 1 and Hop 2 on MATH examples.Source: H. Ablation: Effect of Recursive Hops · p. 22, Figure 7 · p. 23, Figure 6 · p. 23

I. Pilot Study Details

  • The pilot study uses one RULER VT validation example with Qwen-3 8B Instruct, comparing direct answering against chain-of-thought prompting.The baseline in Figure 1(b) is the perturbation-based attribution approach described in Section D.Source: §I · p. 23

J. Additional Related Work

  • FLASHTRACE localizes influential input tokens, while feature-circuit methods analyze internal components; the paper presents these approaches as complementary.The suggested workflow is to localize relevant input regions first, then use circuit analysis within them for a mechanistic account.Source: J. Additional Related Work · p. 24

K. Additional Faithfulness Experiments

K.1. Faithfulness without Correctness Filtering

  • RISE and MAS are evaluated on HotpotQA without correctness filtering because these perturbation metrics do not require correct answers.The experiment tests whether the main-text conclusions depend on filtering generations for correctness.Source: §K.1 · p. 24

K.2. Faithfulness at Longer Contexts

  • RISE falls from 0.036 at 2048 tokens to 0.025 at 4096 tokens, while FLASHTRACE retains an advantage over IFR and AttnLRP without correctness filtering.Table 8 reports RISE and MAS on HotpotQA at both context lengths.Source: §K.2 · p. 24, Table 8 · p. 24

L. Memory and Runtime Trade-offs

  • On Qwen-3 4B with a 2048-token context, on-the-fly attention recomputation cuts peak GPU memory by 62% but increases latency by 42%.The comparison averages five runs; recomputation trades added compute for a smaller memory footprint.Source: L. Memory and Runtime Trade-offs · p. 25, Table 9 · p. 25

M.1. Aggregation Operators

The aggregation analysis compares pooling and targeting alternatives, notes a potential cancellation caveat, and distinguishes token- from sentence-level spans by their information resolution.

M.1. Aggregation Operators

  • Mean pooling loses span-wise single-pass efficiency and recovers fewer ground-truth tokens than FLASHTRACE, because aggregation after the nonlinear L1 proximity operation breaks the exploited linearity.Mean pooling corresponds to the IFR baseline applied to each target token and then averaged.Source: §M.1 · p. 25
  • Last-token targeting reduces attribution coverage from 0.384 to 0.161 and degrades both RISE and MAS compared with span-wise aggregation on HotpotQA.Table 10 compares last-token targeting with FLASHTRACE; the result supports attributing the full multi-token response span.Source: §M.1 · p. 25, Table 10 · p. 25
  • Positional weighting would require task-specific knowledge about which response tokens deserve more weight, so the paper leaves a principled general instantiation for future work.The target span includes all response tokens after the reasoning chain, with no task-agnostic signal for assigning nonuniform weights.Source: §M.1 · p. 25
  • Summing token vectors could let opposing semantics partially cancel before proximity is computed, though the authors did not observe this harming results in practice.The authors identify this as a caveat motivating consideration of alternative operators.Source: §M.1 · p. 25

M.2. Span Granularity

  • Token- and sentence-level spans have equivalent single-pass cost, but sentence-level aggregation removes within-sentence detail for coarser, more stable attributions.Experiments use token-level granularity, described as the more demanding case.Source: §M.2 · p. 26

N. Qualitative Failure Analysis

The failure analysis identifies localized evidence as a favorable setting and describes how diffuse evidence or structural tokens can reduce attribution precision. It frames these limits as an efficiency–precision trade-off.

N. Qualitative Failure Analysis

  • FLASHTRACE performs best when output evidence is clear and localized, as in named-entity multi-hop reasoning with a small, well-separated set of relevant input tokens.In that setting, span-wise aggregation concentrates attribution mass on the relevant tokens.Source: N. Qualitative Failure Analysis · p. 26
  • Diffuse evidence can spread attribution thinly so key tokens fall outside the top-ranked set, while high-frequency formatting or connective tokens can receive salience beyond their semantic role.The latter can occur because structural tokens participate in the residual stream feeding the target.Source: N. Qualitative Failure Analysis · p. 26
  • These failure modes reflect an efficiency–precision trade-off: span-wise aggregation is fast but can obscure fine-grained contributions, so outputs may be coarse localizations rather than precise token evidence.For fine-grained precision, the paper suggests complementing FLASHTRACE with exhaustive token-level rollout.Source: N. Qualitative Failure Analysis · p. 26

Related papers

Loading 2602.01914v4…