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

TL;DR

This paper introduces FlashTrace, a fast and faithful multi-token attribution method for reasoning LLMs, enabling efficient explanations of long-context and multi-step reasoning processes with significant speed improvements.

Contribution

FlashTrace is a novel attribution technique that employs span-wise aggregation and recursive importance tracing to improve efficiency and faithfulness in long-horizon reasoning explanations.

Findings

Achieves over 130x speedup compared to baselines.

Maintains superior faithfulness in attribution.

Recursive attribution enhances importance tracing through reasoning chains.

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.