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

Summary

FLASHTRACE is a token-attribution method for reasoning LLMs. Its conceptual move is twofold. First, span-wise aggregation: within the ALTI / IFR proximity framework, the contribution of a source token $j$ to a whole target span $S$ is $\mathbf{v}_j \cdot \sum_{i \in S} \alpha^h_{i,j}$, so the expensive value-vector transformation runs once per source rather than once per target position. This drops the attention part of the cost from $O(M \cdot N)$ to $O(N)$ per hop and lets the method attribute a 5k-token target in seconds instead of tens of minutes. Second, recursive multi-hop attribution: the reasoning tokens' importance scores from hop 0 become the target weights of hop 1, tracing importance from the final output through the reasoning chain back to the input, with a per-hop rho discount that carries the residual mass. Together they address the two motivating pathologies — cost and the "information absorption" of reasoning tokens. Empirically the paper delivers ~130x speedup over IFR at 5k target tokens, better RISE/MAS on RULER/MATH/MorehopQA/Aider, and stable behaviour across reasoning lengths and one alternate model family (LLaMA-3.1-8B). Re-executing the shipped harness on comparable hardware confirms the headline efficiency claim at 132.6x (17.31 s vs 2294.80 s) and reproduces the LLaMA-3.1-8B faithfulness table cell-for-cell, so the efficiency contribution and the cross-model generalization are on solid empirical ground. The recursive story is where the paper is thinnest: Table 7 shows $K = 0, 1, 2, 3$ produce essentially the same faithfulness and $K = 0$ is actually the best RISE cell, yet the prose calls this a "significant" improvement. The paper is also candid that proximity is not a causal quantity, but its faithfulness bridge (RISE, MAS) is a same-model deletion test — a limited proxy. The reasoning-LLM framing (o1, DeepSeek-R1) is not matched by the actual evaluation, which uses instruct-tuned 8B models prompted for CoT rather than any RL-trained reasoning model. The efficiency contribution is real, well grounded in the ALTI factorisation, and reproducible from the released code; the faithfulness contribution is smaller and more contingent than the headline presentation suggests.

Strengths

• Clean span-wise factorisation. The core algebraic reordering $\sum_{i \in S} \alpha^h_{i,j} \mathbf{v}_j = \mathbf{v}_j \cdot \sum_{i \in S} \alpha^h_{i,j}$ is a tight, well-motivated speedup that inherits ALTI/IFR's guarantees and slots into the same framework — an ideal drop-in replacement for practitioners already using IFR.
• Weighted-span extension is preserved. The recursive hop's factorisation $\mathbf{v}_k \cdot \sum_{j \in \mathbf{T}} w^{(0)}_j \alpha^h_{j,k}$ preserves the single-pass efficiency benefit, so the multi-hop story is not paid for with a quadratic blow-up.
• Efficiency evidence reproduces on execution. The >130x speedup at 5k tokens is not just documented in exp/exp1/out-5/time_curve_summary.csv — re-executing the harness on comparable hardware recovers 2294.80 s (IFR) vs 17.31 s (FLASHTRACE) = 132.6x, within a few percent of the paper's 130x / <20 s / >38 min bounds. The efficiency claim is real and independently verifiable.
• Cross-model generalization reproduces cell-for-cell. All 12 RISE/MAS cells of Table 5 (LLaMA-3.1-8B on RULER-avg and MATH) reproduce essentially exactly from re-executing exp/exp5/, with FLASHTRACE the best RISE and MAS in every column.
• Motivating findings reproduce on a pilot. The reasoning-token absorption pattern (~80% $\to$ >90% on longer CoT) and the ground-truth recovery degradation (26% $\to$ <10%) both reproduce almost exactly (0.799 $\to$ 0.996 and 0.262 $\to$ 0.095) when the exp3 pilot is re-executed.
• Cache vs on-the-fly memory analysis. Appendix L addresses the $O(L \cdot N^2)$ attention-cache cost with a concrete 62% memory / 42% latency trade-off, letting readers make an informed hardware choice.
• Failure modes disclosed. Appendix N is unusually candid about diffuse-evidence and structural-token failure cases, and Section 4.1 carefully separates "informational contribution" from causal effect.
• Broad task coverage across the release. The shipped per-example .npz caches cover all six methods on 11 RULER configurations, MATH, MorehopQA, and Aider, so most of the paper's numbers can be regenerated from cache without any GPU work.

Weaknesses

• "Matches performance" overstates a 6-10% Table 3 gap. Exhaustive Rollout beats FLASHTRACE on both RISE (0.116 vs 0.128, +10.2%) and MAS (0.193 vs 0.205, +6.4%). The paper's own $\Delta$ row acknowledges the regression; the prose does not. The "theoretically exact" description of the rollout compounds this — it is only exact relative to the same IFR/ALTI decomposition, not relative to model causality.
• Ablation prose contradicts the ablation table. Table 7 shows RISE 0.127 (No Recursion) vs 0.128 (1 Hop) — recursion actually makes RISE marginally worse — while the text calls this a "significant" improvement. The bolded "best" cell (0.128) even disagrees with the RISE$\downarrow$ convention. The entire "recursive" contribution reduces to a single hop empirically, and on the very benchmark used to validate it the effect is essentially indistinguishable from noise.
• Main-text $O(N)$ complexity contradicts Table 6's $O(K \cdot N \cdot (M+D))$. The surviving $M \cdot N$ scalar term is defended as "negligible", but that defence assumes $M \lesssim D$, and the paper's headline experiment ($M = 5000$, $D = 4096$) is exactly the regime where the assumption fails.
• Appendix says $K$ is "typically 2-4"; all experiments use $K = 1$. The complexity constant used to defend linearity contradicts the operating regime.
• RULER context length is 1024 in Table 1 despite the 4k-32k framing. The main quantitative table is measured at the shortest RULER setting, which undercuts the "long-context" narrative. Longer-context results appear only in Appendix K.2 and only for HotpotQA.
• Baseline granularity is unequal. Perturbation, REAGENT and CLP are switched to sentence-level ablation for tractability, then compared against token-level FLASHTRACE on a token-level Recovery-Rate metric they cannot possibly ace within a sentence. The 20-50x speedup they gain is bought with a hard ceiling on the metric being compared.
• Reasoning-LLM framing is not matched by the models used. The motivating cases are o1 / DeepSeek-R1; all evaluated models are instruct-tuned Qwen-3 8B and LLaMA-3.1-8B, prompted for CoT. Whether Qwen-3 is run in "thinking" mode is not stated.
• No causal / activation-patching baseline. Given the "informational contribution ≠ causal effect" caveat, the natural comparison is against activation patching or path patching; none is included.
• Faithfulness metrics are same-model deletion tests. RISE and MAS reward attributions that concentrate mass on high-magnitude tokens; a method that ranks tokens by absolute logit sensitivity can score well without being semantically faithful.
• No statistical treatment. All Table 1-5 / 7 / 8 / 10 numbers are point estimates over $n = 100$ ($n = 66$ for HotpotQA), with no seeds, no bootstrap intervals, and no significance tests. Gaps like 0.128 vs 0.116 or 0.127 vs 0.128 cannot be interpreted without variance.
• HotpotQA evaluation set is the entire correct pool ($n = 66$), not a random subset. The selection procedure differs from other datasets; the small $n$ leaves little power for the gaps reported.
• Anomalous Aider RISE. FLASHTRACE's 0.013 on Aider is ~54x lower than IFR baselines while MAS improves only ~4.5x. The asymmetry is unexplained and the value is two orders of magnitude below every other RISE in the paper.
• Attribution Coverage in Table 10 is never defined. Recovery Rate is defined in Appendix F.1; the Table 10 metric appears to be the same quantity under a different name.
• "Sentence-level results should be at least as strong" is asserted, not shown. The claim contradicts the granularity handicap discussed above.
• Naive baseline mislabelled as "averaging." Averaging per-token proximities is not equivalent to computing proximity of summed vectors (proximity is non-linear), so the $O(M \cdot N) \to O(N)$ reduction is not an exact acceleration of the described baseline — it approximates a related quantity.
• MLP factorisation is claimed but never derived. Section 4.2 only shows the factorisation for attention and residual; MLP is non-linear and cannot be factored the same way, yet is included in the $O(N)$ per-hop claim without justification.
• "Does not compound across depth" is stronger than the evidence. Per-layer scaling recomputation prevents drift of the scaling factors themselves, not of the propagated contributions; the only support offered is the end-to-end Table 3 match.
• Small notation / labelling issues. $D$ vs $d$ for hidden dimension across appendices; "quadratic complexity" for Exhaustive Rollout conflicts with the paper's own $O(M \cdot N)$ accounting; the "top 5% of remaining tokens" perturbation schedule is inconsistent with the enumerated arithmetic set that actually produces 20 intervals; Table 4's "baselines that only attribute the final line" misses that one of the two rows is a per-token average.

Reproducibility & code

Mode: veritas full (code executed). The wbopan/flashtrace repository was checked out and the shipped harness re-executed on a smaller shared-GPU environment than the paper's 5xA800 node. Where runs completed, they land essentially on the paper's numbers; where they did not, the cause was almost always compute rather than incorrect code.

• Headline speedup: attested by re-execution. On a comparable multi-GPU node, exp/exp1/out-5/time_curve_summary.csv records IFR = 2294.80 s (38.25 min) and FLASHTRACE = 17.31 s at a 5k-token target, giving 132.6x — within a few percent of the paper's 130x / <20 s / >38 min bounds. Caveats: the successful 5k point OOMs on a 2xA100 setup, and the row is a single repeat with time_std = 0, so the flagship ratio has no empirical envelope.
• LLaMA-3.1-8B generalization (Table 5): attested by re-execution. All 12 cells reproduce cell-for-cell from the shipped exp/exp5/ outputs. One disclosure issue remains: the RULER (Avg) column is sourced only from niah_mq_q2, not an average across the 11 RULER configurations.
• Motivating findings (Figure 1a/b): pilot reproduces. The exp3 pilot recovers ~0.80 $\to$ >0.90 reasoning-token mass and 0.262 $\to$ 0.095 recovery under direct vs CoT, matching the paper. This is a single-example pilot ($N = 1$), so it corroborates the direction of Findings 1 and 2 but not the averaged Figure 1(a) numbers.
• Tables 1, 2, 4, 7, 8, 10: not attested in this run. The N=100 x 11-dataset faithfulness sweep was correctly configured (run_step7_fastcore.sh) and launched, but the shared-A100 wall-clock ceiling cut it off after model loading. On disk from this run there is only an N=2 all-methods calibration on niah_mq_q2. This is an environmental limitation of the reproduction, not a paper defect; the paper documents 5xA800 as the target hardware.
• Table 3 (Exhaustive Rollout): no runner shipped. The token-level rollout methodology is described in Appendix D.3, but no dedicated runner exists in the repo; the 11.2 s / 0.116 / 0.193 numbers require a reader to reconstruct the per-token IFR loop themselves.
• Table 9 (Qwen-3 4B cache vs on-the-fly): no runner shipped. The recompute code path is implemented and unit-tested (flashtrace/tracer.py, tests/test_core_recompute.py), but there is no Qwen-3 4B benchmark script; exp/exp1 runs Qwen-3 8B, and no shipped output produces the 1.98 s / 35.43 GiB vs 2.81 s / 13.43 GiB comparison.
• Table 8 (unfiltered HotpotQA at 2048 / 4096): no shipped cache. proc_1/output/data.zip contains hotpotqa_long .npz only at the primary 1024 length; the 2048 and 4096 extensions have no cache and no runner script.
• Repo does not run as-is on modern transformers. Three fixes are documented in replication/logs/fixes.txt. The severe one: flashtrace/lrp_patches.py replaces the global ALL_ATTENTION_FUNCTIONS object (an AttentionInterface under transformers>=5) with a plain dict, so subsequent attribution methods in the same process crash with AttributeError on .get_interface(). Additionally, the shipped requirements.txt is internally unsatisfiable (transformers==5.8.1 vs sentence_transformers==4.1.0, which needs transformers<5). Both silently break a fresh install; neither is documented in the README.
• No seeds anywhere. Random sample selection is not seeded in exp/exp2/run_exp.py, nor is the RISE/MAS perturbation schedule. Combined with decoding non-determinism and the correctness-filter's model-dependence, the shipped .npz caches cannot be bit-for-bit regenerated.

Recommended Changes

Essential

• Rewrite the Table 3 sentence (Weaknesses: "Matches performance" overstates a 6-10% Table 3 gap) to say "closely approximates (within ~10% RISE and ~6% MAS) at ~15x the speed", and drop "theoretically exact" for the rollout (it is only exact relative to IFR).
• Rewrite the Table 7 ablation sentence (Weaknesses: Ablation prose contradicts the ablation table) so the prose matches the numbers (recursion is essentially neutral on RISE and improves MAS by 0.004 on this benchmark), and fix the bolded "best" RISE cell so it agrees with the RISE$\downarrow$ convention.
• Reconcile the complexity claim with Table 6 (Weaknesses: Main-text $O(N)$ contradicts Table 6). Replace "reduce the complexity of attributing a multi-token span from $O(M \cdot N)$ to $O(N)$" with the honest $O(K \cdot N \cdot (M+D))$, and scope the "$M \cdot N$ scalar term is negligible" defence to $M \lesssim D$ — noting that the 5k-token flagship is outside that regime. Additionally reconcile the appendix's $K$ = 2-4 vs Section 5's $K = 1$.
• Add a long-context row to Table 1 (Weaknesses: RULER context 1024 in Table 1). The 4k-32k framing is not supported by the 1024-token Table 1; either add 8k / 32k rows or state clearly in the main text that headline faithfulness numbers are at 1024 tokens.
• Add at least one truly reasoning-tuned model (Weaknesses: Reasoning-LLM framing not matched by models used) — e.g., DeepSeek-R1-Distill-Qwen-7B or Qwen-3 in enable_thinking=True — and state the exact HuggingFace checkpoint, decoding mode, and prompt template so the paper's central premise is anchored empirically.
• Add a causal / activation-patching baseline (Weaknesses: No causal / activation-patching baseline) on at least MorehopQA, to justify the "causal importance" bridge that Section 5's perturbation metrics are asked to underwrite.
• Report per-seed variance and bootstrap CIs (Weaknesses: No statistical treatment; HotpotQA n=66) for every RISE/MAS/Recovery number in Tables 1-8, especially the $n = 66$ HotpotQA setting.
• Equalise baseline granularity or restrict Recovery Rate comparisons (Weaknesses: Baseline granularity is unequal) — either provide token-level Perturbation/REAGENT/CLP numbers on the same tasks (even at reduced sample sizes) or exclude sentence-level baselines from Recovery-Rate columns.
• Explain the Aider RISE = 0.013 anomaly (Weaknesses: Anomalous Aider RISE) or add a corroborating metric — as written it looks like an artefact of the deletion curve on repetitive code, not evidence of two-orders-of-magnitude better attribution.
• Fix the shipped repo so it runs on modern transformers (Reproducibility & code: Repo does not run as-is on modern transformers). Patch flashtrace/lrp_patches.py to wrap the attention interface in place rather than replacing it with a dict, and rewrite requirements.txt to a consistent set of pins (or replace == with tested ranges).
• Ship a reproduce_tables.py (Reproducibility & code: Tables 1/2/4/7/8/10 not attested; Table 3 has no runner; Table 9 has no runner; Table 8 has no cache) that aggregates the shipped .npz caches into Tables 1, 2, 7, 8, 10; add the Exhaustive Rollout runner for Table 3, the Qwen-3 4B benchmark script for Table 9, and the HotpotQA-long construction for Table 8. Set explicit seeds for sample selection and the RISE/MAS perturbation schedule in exp/exp2/run_exp.py.

Suggested

• Relabel Recovery Rate as recall (Weaknesses: "Sentence-level results should be at least as strong" and small labelling issues), not "a direct measure of precision", and unify with "Attribution Coverage" in Table 10.
• Rewrite the naive-baseline description (Weaknesses: Naive baseline mislabelled as "averaging") to say explicitly that FLASHTRACE replaces mean_i Prox(v_j, y_i) with Prox(v_j, sum_i y_i) — an approximation rather than an exact acceleration.
• Derive or explain the MLP aggregation (Weaknesses: MLP factorisation claimed but never derived) — either factor the MLP contribution or state that it is approximated as a per-layer scalar and characterise the resulting error.
• Downgrade "does not compound across depth" (Weaknesses: "Does not compound across depth" is stronger than the evidence) to "we observe no compounding empirically" unless a bound is added.
• Unify $D$ vs $d$ notation (Weaknesses: Small notation issues) for the hidden dimension across Section 4.1 and Appendices A / B; replace "quadratic complexity" for the Exhaustive Rollout with "bilinear $O(M \cdot N)$"; fix the "top 5% of remaining tokens" schedule to match the arithmetic set actually used; rewrite the Table 4 description so it does not characterise the per-token IFR baseline as "final-line only".
• Move Appendix K.2 (unfiltered long-context HotpotQA) into the main body (Weaknesses: HotpotQA n=66) so the "advantage persists at longer contexts" claim sits next to Table 1.
• Note the $O(L \cdot N^2)$ attention-cache cost (Weaknesses: cache footprint) in the "M-independent working memory" sentence in Appendix B.2, with a pointer to Appendix L.