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

AI summary

Overview

Research area: Mechanistic interpretability of large language models, specifically token attribution methods for reasoning and agentic LLMs.

Technical level: Intermediate to Advanced. The paper assumes familiarity with Transformer internals (attention heads, residual streams, MLP blocks), causal attribution concepts, and the ALTI/IFR (Aggregation of Layer-wise Token-to-Token Interactions / Information Flow Route) theoretical framework.

Scope: The paper introduces FlashTrace, a method for efficiently attributing a multi-token span of LLM-generated output back to the original input tokens, even when thousands of intermediate reasoning tokens sit in between.

What This Paper Is About

As LLMs increasingly "think" before answering—generating long chains of intermediate reasoning tokens—it becomes hard to tell which parts of the actual user input caused the final answer. Existing attribution methods explain one output token at a time, so explaining a 5,000-token reasoning chain takes hours, and worse, the resulting importance scores pile up on the reasoning tokens rather than the original inputs. FlashTrace solves both problems: it computes attribution for an entire multi-token span in a single pass, then recursively follows the trail of importance backward through the reasoning chain until it reaches the source input.

Key Contributions

1. Formalization of the multi-token attribution problem. The authors define the setting where a contiguous span of model-generated tokens (not a single token) is the explanation target, and identify the two failure modes—an efficiency bottleneck of O(M·N) operations and a faithfulness collapse caused by intermediate reasoning tokens absorbing attribution mass.

2. Span-wise aggregation. FlashTrace exploits the linearity of attention to pre-sum attention weights over the target span, factoring out the expensive value-vector transformation so it is computed once per source token rather than once per (source, target) pair. This reduces complexity from O(M·N) to O(N).

3. Recursive attribution. A mechanism that treats the importance scores on reasoning tokens from one hop as weights for the target span in the next hop, propagating importance backward through the reasoning chain toward the original input. A weighted residual-mass scheme combines scores across hops into one final distribution.

4. Empirical validation across tasks and models. Experiments on long-context retrieval (RULER), mathematical reasoning (MATH), multi-hop QA (MorehopQA, HotpotQA), and code generation (Aider) show over 130× speedup with better faithfulness, validated on both Qwen-3 8B Instruct and LLaMA-3.1-8B-It.

Main Findings

• Reasoning tokens absorb attribution mass. In a controlled study on RULER, attribution to reasoning tokens rose from roughly 80% to over 90% as reasoning chains lengthened, while the recovery rate of ground-truth input tokens dropped from 26% to below 10%.

• Large faithfulness gains on multi-hop reasoning. On HotpotQA expanded to long context, FlashTrace reached recovery rates of 0.698, 0.755, 0.659, and 0.514 across four difficulty settings, versus 0.136, 0.253, 0.202, and 0.155 for the strongest baseline (IFR)—a margin of roughly two to three times.

• Strong speed and memory scaling. At a 5k-token target span, FlashTrace finished in under 20 seconds versus over 38 minutes for IFR (more than 130× faster). It also avoids the memory growth that causes gradient-based methods such as Integrated Gradients to run out of memory on long contexts.

• Matches an exhaustive brute-force baseline. A "Exhaustive Token-Level Rollout" that recursively attributes every individual reasoning token took 11.2 seconds with RISE 0.116 and MAS 0.193; FlashTrace achieved comparable scores (RISE 0.128, MAS 0.205) in 0.72 seconds—a 93.6% runtime reduction.

• A single recursive hop suffices. The authors use K=1 in all main experiments and show that even one hop substantially improves faithfulness, suggesting reasoning-chain dependencies are largely resolved within one step.

• Generalizes beyond reasoning to code. On 133 Python exercises from the Aider dataset, FlashTrace achieved RISE 0.013 and MAS 0.173, compared to RISE 0.710 and MAS 0.782 for IFR attributing only the last line—a dramatic improvement when the intermediate content is code rather than prose reasoning.

• Holds across model families. On LLaMA-3.1-8B-It, FlashTrace again beat IFR and AttnLRP on both RULER and MATH faithfulness metrics.

• Stable across reasoning lengths. Faithfulness (MAS) stayed flat as the number of generated reasoning tokens increased, whereas baseline methods degrade.

Methodology in Plain English

The approach builds on the observation that a Transformer layer's output is just a sum of contributions: a residual stream term, the outputs of each attention head, and the MLP output. Each attention head's output, in turn, is a weighted sum of transformed vectors from every preceding token—where the transformation depends only on the source token, not the target.

Existing methods measure the importance of each source token by asking how much a target vector's magnitude would shrink if that contribution were removed (an L1-based proximity metric from the ALTI framework).

FlashTrace's first trick is to group the target tokens into a span and ask the same question about the sum of their representations. Because each source token's transformed vector is identical across all target positions, it can be factored out, leaving only the sum of attention weights to be aggregated. So instead of running the attribution algorithm M times for M target tokens, it runs once.

The second trick addresses the "information absorption" problem. The first attribution pass explains the final output in terms of whatever tokens came just before—mostly reasoning tokens. The authors then run a second pass where those reasoning tokens become the target, weighted by how important they were in pass one. This second pass surfaces which earlier tokens (reasoning steps, and eventually the original input) caused the important reasoning. Importance scores on input tokens from each hop are combined into a final distribution, scaled by how much mass remained in the reasoning tokens at each stage.

The authors note that the proximity score measures informational contribution rather than strict causal effect, and they validate causally through perturbation-based faithfulness tests (RISE and MAS), which check whether removing high-attribution tokens actually changes the model's output probability.

Why This Matters

Impact on research. The paper reframes interpretability for the reasoning-model era: attribution is no longer a single-token problem but a multi-hop flow problem. It provides a scalable primitive that other interpretability work—circuit analysis, context optimization, agent auditing—can build on. The finding that reasoning tokens structurally dominate attribution mass is a cautionary result for anyone applying standard explanation methods to chain-of-thought models.

Real-world applications:

• Debugging and auditing LLM agents. When an agent takes a wrong action after a long reasoning trace, FlashTrace can point to the user instruction or retrieved document that actually caused the error.
• Context optimization for RAG systems. Identifying which retrieved passages genuinely drive an answer allows pruning irrelevant context, saving tokens and improving accuracy.
• Regulatory and compliance review. High-stakes deployments (medical, legal, financial) need evidence linking outputs to specific inputs; attribution offers a principled form of provenance.
• Trust and safety investigation. Tracing a harmful or incorrect generated span back to its trigger in the prompt or tool output supports root-cause analysis at scale.

Industry relevance. The efficiency profile is the selling point—attribution that costs roughly a forward pass per hop rather than hours of gradient computation can plausibly run in production monitoring pipelines or as part of agent observability tooling, rather than only in offline research.

Future Directions

• Extending recursive depth and adaptive hop counts. The experiments use K=1; understanding when additional hops help or hurt, and how to decide dynamically, remains open.
• Scaling to truly long agentic workflows. The evaluation covers up to roughly 10K-token contexts; behavior on multi-turn tool-use trajectories spanning hundreds of thousands of tokens is untested.
• Improving the causal interpretation of proximity scores. The proximity metric measures informational contribution, not strict causality; tighter theoretical links between the two, or a causal-corrected scoring variant, would strengthen claims.
• Extending to non-text intermediate content and multimodal agents. The Aider result hints at generalization; whether the method holds for tool-call JSON, image tokens, or mixed-modality chains is an open question.

Target Audience

Interpretability and explainability researchers working on Transformer internals; ML engineers building or debugging long-context reasoning agents and RAG systems; and applied researchers focused on AI transparency, auditing, and compliance for high-stakes LLM deployments. Readers without background in attention decomposition or attribution theory will find the methodology section demanding, though the motivation and results are accessible to a general ML audience.

Authors’ abstract

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.