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

TL;DR

This work tackles the challenge of interpreting long-horizon reasoning in LLMs by addressing two bottlenecks: the efficiency of attributing multi-token spans and the faithfulness of attribution through reasoning chains. It introduces FlashTrace, which combines span-wise aggregation with recursive attribution to enable efficient, faithful multi-token explanations across long contexts, achieving over $130\times$ speedups and superior faithfulness compared to baselines. The method formalizes multi-token attribution, demonstrates effective propagation of importance from final outputs through intermediate reasoning back to inputs, and validates across long-context retrieval, multi-hop reasoning, and code-generation tasks. The results suggest FlashTrace enables scalable interpretability for agentive LLM workflows, with practical implications for transparency, debugging, and auditing in high-stakes applications.

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.

Figures (7)

• Figure 1: Motivation for FlashTrace. Top: Naive token-by-token attribution requires expensive per-token computation, while FlashTrace performs efficient span-wise recursive attribution. Bottom: (a) With extended reasoning, attribution weight on reasoning tokens increases significantly (from approximately 80% to over 90%); (b) This causes recovery rate of ground-truth input tokens to drop substantially (from 26% to below 10%); (c) Naive multi-hop attribution scales poorly with reasoning length, while FlashTrace remains efficient even for 10K tokens.
• Figure 2: Overview of FlashTrace span-wise attribution. The method pre-aggregates causal attention and residual contributions within target spans at each layer, enabling efficient multi-hop attribution through chains of intermediate reasoning.
• Figure 3: Visualization of recursive attribution across hops.Hop 1: Initial attribution concentrates on reasoning tokens nearest to the output. Hop 2: Attribution shifts toward earlier reasoning tokens and input context. $\Delta$ (Hop2 - Hop1): The difference map shows how importance propagates backward: HTML]82bea6green regions indicate tokens gaining attribution in the second hop (typically input tokens), while HTML]db8887red regions show tokens losing attribution (typically intermediate reasoning tokens).
• Figure 4: Efficiency comparison across methods. (a-b) Time cost vs. input and generation length. (c-d) Memory consumption vs. input and generation length. (e) Pareto front of speed vs. faithfulness. FlashTrace achieves the most efficient scaling in both time and memory, while gradient-based methods (IG, IG-Attn, Perturbation) encounter OOM at longer contexts. Dashed lines indicate OOM.
• Figure 5: Faithfulness (MAS) across reasoning lengths. FlashTrace maintains stable performance and consistently outperforms baseline methods.