Mechanistic Data Attribution: Tracing the Training Origins of Interpretable LLM Units

Summary

This paper introduces Mechanistic Data Attribution (MDA), an influence-function framework that traces the training-time origins of interpretable units — attention heads, MLP neurons, and SAE features — rather than the standard target of global loss. The conceptual move is to swap the usual test loss $L(z_{\text{test}})$ in the influence-function formula with a component-specific probe $f_{\text{probe}}(\theta, D_{\text{probe}})$ and to restrict the EK-FAC-approximated inverse Hessian to a mechanism-localised parameter subspace $\theta_{\text{sub}}$. Instantiated for induction heads and previous-token heads on the Pythia suite (14M–160M), MDA is used both for tracing (Sections 4–5) and for a practical data-augmentation pipeline that distills mechanistic templates from a small proxy model via an LLM (Section 6). The empirical programme is unusually careful in its counterfactual design: retraining with duplication vs. gradient-masking of MDA-selected samples, cross-head transfer tests, probe-variant robustness (QK-only vs. QKVO), a controlled n-gram baseline, and an ablation-logit interventional metric — all pointing in the same direction. The results tell a coherent story: repetitive structural data (LaTeX, XML, Base64, and even 'garbage' repeats) preferentially catalyse induction heads, high-influence samples are approximately uniformly distributed across training time, and interventions that shift induction-head strength move ICL performance in tandem. The main conceptual limitation is that MDA can only intervene on data, not on the head itself; the paper repeatedly upgrades a common-cause data intervention that co-moves induction-head score and ICL score into 'direct causal evidence' for a head→ICL mediation, which is stronger than the design supports. Related tensions run through the necessity/sufficiency framing, the sparse-subset-vs-collective-property duality, and the scalability claims. The core contribution — mechanism-localised influence with a genuine counterfactual retraining validation — is nonetheless novel and useful, and the finding that a 14M-trained proxy discovers structural motifs that accelerate induction on 160M is a striking, actionable result.

Strengths

• Conceptual contribution. Restricting influence functions to a mechanism-specific parameter subspace and swapping loss for a component-specific probe is a clean, well-motivated extension of Grosse et al. (2023), and the Q/K joint-subspace treatment of attention pattern formation is a genuine technical refinement rather than a boilerplate application.
• Counterfactual validation. Unlike most attribution work that stops at ranking correlations, this paper actually retrains Pythia models with data deletion and augmentation and demonstrates that MDA-guided interventions shift emergence-time behaviour while matched random baselines do not. The dispersion/concentration ablation and the QK-only-vs-QKVO probe comparison are the kind of controls this literature has been missing.
• Structural-motif finding. The identification of Base64/LaTeX/repeat-heavy sequences as mechanistic catalysts is a concrete, testable prediction that ties an internal circuit to identifiable properties of the pre-training distribution, and the cross-head transfer results in Figure 4b/e give the effect a mechanism-general character.
• Cross-scale synthetic augmentation. Table 2's finding that 14M-derived synthetic patterns accelerate induction on 160M — and outperform 160M-derived patterns on 160M itself — is a striking practical result that hints at scale-invariant structural drivers.
• Robustness pass. Appendix G's seed-stability, QK-only-probe, controlled n-gram, and ablation-logit-contribution checks together push the argument beyond a single measurement, and the framework is exercised (qualitatively) on billion-parameter OLMo-2 models.

Weaknesses

• Data-intervention causality is over-sold as head-mediated causality. The abstract's 'direct causal evidence for the long-standing hypothesis regarding the functional link between induction heads and ICL' and Section 5.4's 'rigorously verify their causal link' both describe a study in which the only intervention is on training data. Concurrent movement of two downstream quantities under a common upstream perturbation is consistent with head→ICL mediation but does not isolate that path from data→other-circuit→ICL. Section 5.4 concedes latent confounders in passing; the abstract and Section 5.4 headers should match that concession, or a direct head-ablation experiment (holding data fixed) should be added.
• Internal tension between 'sparse subset' and 'collective property'. Contribution 3 asserts induction head formation is 'not driven by sparse sample subsets,' while Section 5.1 states the influence distribution 'adheres to a distinct power-law, indicating that the emergence of mechanistic circuits is disproportionately driven by a sparse subset.' The intended rate-vs-saturation nuance is real but the two statements as written read as contradictions; this needs explicit reconciliation.
• Necessity/sufficiency framing is stronger than the design. 'Sufficiency' would require the identified samples to (dominantly) cause the mechanism; the paper's own observation that both regimes 'converge to comparable saturation scores' and that 95% pruning still yields non-trivial emergence supports rate-shifting, not sufficiency or necessity in the strict sense. The bidirectional design is also asymmetric — augmentation duplicates samples (adding gradient steps), while 'deletion' masks gradients on the fixed schedule — and that asymmetry is not called out or controlled for.
• Statistical rigor of the main comparisons. Figure 2, Figure 5, Table 2, and the appendix figures show curves and point estimates without error bars, seed variance, or statistical tests. Given that Table 2's cross-scale gains (+9.84% to +15.8%) are the headline application result, the absence of any seed-variance estimate makes the claim 'no such effect' for random baselines and 'consistently accelerates' for synthetic augmentation hard to defend quantitatively.
• Scalability claims are inconsistent with each other and with Table 6. Figure 3 argues $y \\propto x^{0.57}$ ('sub-linear'), but Appendix E.5 says Table 6 shows 'an approximately linear relationship between computational cost and model size.' Interpreted in parameter count, Table 6 is strongly sub-linear (14M→7B is ~500× params vs. ~33× GPU-hours). It is also unclear whether the Figure 3 fit spans only the four Pythia sizes or also anchors on OLMo-2-1B/7B — a load-bearing distinction for extrapolation.
• Local overclaims and framing issues. 'Fundamental thermodynamic explanation' for the aggregate-sign observation, the 'lottery ticket hypothesis' analogy that seems logically inverted (constant pressure vs. sparse winning subnet), the 'preserves long-term general model capabilities' inference from two loss curves plotted over a short window, and the ranking-stability figure that quietly reports a strong probe-variant anti-correlation are all soft points where the writing outruns the evidence.
• Under-specified controls. The 'Controlled Random' n-gram baseline (Section G.2), the 95% random-mask procedure and the 3σ band around MDA-Repl (Figure 4d), and the Head-Specific Ablation Logit metric (Appendix G.3) are all under-specified in ways that block reproduction and make the effect-size interpretation fragile.
• Presentation errors that impede reading. The parenthetical inserted mid-verb-phrase ('completely ceased to exhibit(the one that shows a decrease of 0.26 in induction scores) induction behavior') obscures the load-bearing 0.26 number in the circuit-reconfiguration argument; Figure 8 is cited for a stability analysis whose subject is not actually plotted in Figure 8 (which shows compensatory behaviour instead); and Table 1's 'LaTeX' row content is PHP followed by LaTeX, suggesting row misalignment.
• Generalisation surface untested. MDA is validated on two attention-head types whose expected catalysts are inherently pattern-repetition-heavy. Whether MDA delivers similar cleanness on mechanisms whose catalysts are not obviously surface repetitions (e.g., knowledge-storage neurons or safety-relevant heads) is left to a qualitative SAE aside in Appendix L.3.

Reproducibility & code

The released materials cover the mathematical framework thoroughly (Appendix A) and the retraining protocol adequately (Appendix E, Tables 4–5). Below are the specific coverage gaps that a careful re-implementer would hit — grounded in the code/data availability that veritas inspected.

• Headline aggregates lack seed variance. Table 2's per-scale deltas and the Figure 2/5 curves are single-seed values; the paper says configurations 'strictly adhere to the original Pythia repository' but does not list the seeds actually used, the number of runs averaged, or any variance. This is a major reproducibility gap for the headline application result.
• Synthetic-augmentation pipeline outputs are not released. The DeepSeek-V3 pattern-extraction step, its temperature/decoding parameters, the deduplication rule that yields '~900 unique patterns,' the resulting JSON pattern registry, and the generated Python data-synthesis scripts appear to be undocumented at the artifact level. Table 2 and Figure 6b both depend on these outputs; without them, exact reproduction is infeasible even for a group with API access.
• Controlled n-gram baseline under-specified. Section G.2's 'Controlled Random' baseline is the paper's defence against the concern that MDA merely recovers repetition. Its construction (which bigrams/trigrams, at what frequency, how entropy and domain are matched, insertion mode and count) is not spelled out; a reader cannot judge whether Figure 10's gap reflects the ranking or an artifact of matching quality.
• 95% pruning and 3σ band procedures are ambiguous. The random-mask procedure that supports the emergence-dynamics conclusion (Figure 4d, 'Pruned 95%') is not enumerated, and the σ used for the 3σ MDA-Repl bound is not defined (over runs? over steps? number of replicates?).
• Cross-head heatmap identities are not enumerated. Figure 4b's IHead-2/3 and NHead-1/2/3 layer/head indices are not tabulated (only L3H3 is named), so the block-structure test cannot be reproduced without independent head-ranking work.
• OLMo-2 configuration missing. Tables 4–5 stop at Pythia-160M; the OLMo-2-1B/7B EK-FAC batch sizes, analysed-sample counts, target heads, and developmental windows behind the Table 6 runtimes are not tabulated, so the scalability numbers cannot be reproduced without inference.
• Δ_h task specification incomplete. Appendix G.3 defines the ablation-logit metric conceptually but not the prefix/gap length distribution, the number of test sequences, or the per-position aggregation, so Figure 11 is not directly reproducible.

Recommended Changes

Essential

• Soften the head→ICL causal claim. Rewrite the abstract's 'direct causal evidence' and Section 5.4's 'rigorously verify their causal link' to reflect that the design is a common-cause data intervention that co-moves both quantities; either add a direct head-ablation experiment holding data fixed, or explicitly state that the current evidence is consistent-with, not proof-of, head→ICL mediation. (Weakness: 'Data-intervention causality is over-sold as head-mediated causality.')
• Reconcile the sparse-subset vs. collective-property tension. Rewrite contribution bullet 3 in Section 1 and the framing sentences in Sections 4.2 and 5.1 to spell out that a sparse subset dominates the rate while the saturation is collective — as the results actually show. (Weakness: 'Internal tension between sparse subset and collective property.')
• Retire the strict 'necessity/sufficiency' framing. Reframe Section 4.2 in terms of rate modulation; note explicitly that both regimes converge to comparable saturation, that 95% pruning still yields non-trivial emergence, and that the deletion arm (gradient masking) is not step-count-matched to the augmentation arm. Add a step-count-matched augmentation control if feasible. (Weakness: 'Necessity/sufficiency framing is stronger than the design.')
• Report per-seed variance for the headline results. Add multi-seed runs (at minimum 3) with error bars for Figure 2, Figure 5, and Table 2, and state the exact seeds used. Without this, the 'no such effect' random baseline and the '+9.84% to +15.8%' augmentation gains cannot be defended statistically. (Weakness: 'Statistical rigor of the main comparisons.')
• Release the synthetic-augmentation artifacts and pipeline configuration. Publish (a) the DeepSeek-V3 sampling temperature/top-p/seed, (b) the deduplication rule producing '~900 unique patterns,' (c) the merged JSON pattern registry, and (d) the generated Python data-synthesis scripts. Without these, Table 2 and Figure 6b are not reproducible. (Reproducibility: 'Synthetic-augmentation pipeline outputs are not released.')
• Fix or reconcile the scalability statements. Either revise Appendix E.5's 'approximately linear' to match the Figure 3 sub-linear-in-parameters fit (or clarify the axis interpretation as scaling with model dimension), and state which of the six Table 6 rows anchor the $y \\propto x^{0.57}$ fit and whether OLMo-2 points are inside or outside the fit. (Weakness: 'Scalability claims are inconsistent.')

Suggested

• Spell out the 'Controlled Random' n-gram baseline construction. Describe the bigram/trigram selection rule, the entropy/domain matching procedure, the sample size, and the insertion mode; ideally release the baseline samples. (Reproducibility: 'Controlled n-gram baseline under-specified.')
• Document the 95% pruning and 3σ procedures. Report the mask-selection rule (uniform per step vs. per window), the number of replications, and how σ is computed for the MDA-Repl band. (Reproducibility: '95% pruning and 3σ band procedures are ambiguous.')
• Enumerate the six heads used in Figure 4b and complete Δ_h and OLMo-2 configuration tables. Add a small appendix table listing IHead-1/2/3 and NHead-1/2/3 layer/head indices for Pythia-14M; add prefix/gap distributions and evaluation sample counts for the Δ_h task in G.3; and add an OLMo-2 configuration table analogous to Tables 4–5 covering batch sizes, analysed samples, target heads, and windows. (Reproducibility: multiple items.)
• Rewrite the broken 0.26-decrease sentence and fix the Figure 8 cross-reference. State precisely which Pythia-160M head 'ceased' induction behaviour, its pre/post induction scores, and whether the F.2.1 stability claim really points to Figure 8 or to a figure that needs to be added. (Weakness: 'Presentation errors.')
• Drop or reframe overclaims. Remove 'fundamental thermodynamic explanation'; disambiguate the Nanda-et-al. 'lottery ticket hypothesis' analogy from Frankle–Carbin; scope the 'preserves long-term general model capabilities' claim to the two metrics actually measured over the plotted horizon; and address the Dyn-Prefix anti-correlation in Figure 9 explicitly as a probe-sensitivity finding, not a seed-stability figure. (Weakness: 'Local overclaims and framing issues.')
• Add a diversity measurement for the augmentation crossover. Report a simple structural-diversity statistic (unique-pattern counts, n-gram entropy) for synthetic vs. natural inserts at each $N$, so 'diversity exhaustion' becomes measured rather than asserted. (Weakness: 'Local overclaims.')
• Broaden the mechanism coverage. Add at least one MDA case-study on a mechanism whose expected catalysts are not surface repetition (e.g., a knowledge or safety-relevant unit); the current SAE aside in L.3 is qualitative. (Weakness: 'Generalisation surface untested.')
• Test proxy transfer beyond the 14M→160M pair. Report augmentation results using 31M or 70M as the proxy, so the scale-invariance claim rests on more than one comparison. (Weakness: 'Reliance on Pythia-14M as the mechanism-transfer proxy.')
• Correct minor presentation issues. Fix Table 1 row alignment; clarify the '≤10%' denominator (fraction of total corpus vs. developmental window vs. flagged pool); soften 'reduction in loss' to a signed loss difference for the ICL Score; clarify the temporal offset between intervention step 900 and the 1200–2000 phase transition window. (Weakness: multiple.)