PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jianhui Chen, Yuzhang Luo, Liangming Pan

arXiv:2601.21996v2cs.CLcs.AIcs.LG

TL;DR

Interpretable LLM circuits can be identified, but their causal origins in training data remain insufficiently traced. The paper introduces MDA, which applies influence functions to interpretable-unit behavior and validates targeted sample interventions. Across Pythia models, these interventions modulate head emergence, while structural data catalyzes induction heads and augmentation accelerates convergence across scales.

  • Problem

    Mechanistic interpretability describes circuit computations but lacks tools to identify their causal origins in the training corpus.

  • Method

    MDA uses influence functions and scalable EK-FAC approximations to attribute interpretable-unit behavior to specific training samples.

  • Results

    Removing or repeating at most 10% of high-influence samples significantly modulates targeted induction and previous-token head emergence, unlike random interventions.

  • Takeaways & Limitations

    MDA provides a causal methodology for tracing and steering interpretable mechanisms, including their connection to in-context learning and development across model scales.

  • Takeaways & Limitations

    MDA assumes attribution captures direct effects on a unit, while indirect effects through upstream computational nodes remain systematically unresolved.

Abstract

from arXiv · show

While Mechanistic Interpretability has identified interpretable circuits in LLMs, their causal origins in training data remain elusive. We introduce Mechanistic Data Attribution (MDA), a scalable framework that employs Influence Functions to trace interpretable units back to specific training samples. Through extensive experiments on the Pythia family, we causally validate that targeted intervention--removing or augmenting a small fraction of high-influence samples--significantly modulates the emergence of interpretable heads, whereas random interventions show no effect. Our analysis reveals that repetitive structural data (e.g., LaTeX, XML) acts as a mechanistic catalyst. Furthermore, we observe that interventions targeting induction head formation induce a concurrent change in the model's in-context learning (ICL) capability. This provides direct causal evidence for the long-standing hypothesis regarding the functional link between induction heads and ICL. Finally, we propose a mechanistic data augmentation pipeline that consistently accelerates circuit convergence across model scales, providing a principled methodology for steering the developmental trajectories of LLMs.

1. Introduction

Mechanistic interpretability has characterized what LLM circuits compute but lacks tools to trace their causal origins in training data. MDA addresses this gap by attributing unit behavior to samples and validating targeted interventions that steer circuit formation.

  • Motivation and Contributions: MDA traces the training origins of individual interpretable units rather than attributing only global model behavior.It uses influence functions to quantify how specific samples affect neurons, attention heads, or SAE features.
  • Causal Validation: Across four Pythia scales and two attention-head types, deleting or repeating at most 10% of high-influence samples modulates targeted head emergence.Randomly removing or augmenting the same number of other samples yields no such effect.
  • Mechanistic Insights: Highly repetitive LaTeX and XML patterns significantly accelerate induction-head emergence.The analysis identifies these noisy structural sequences as primary catalysts in high-positive-influence samples.
  • Mechanistic Insights: Induction-head interventions concurrently change in-context learning performance, providing causal evidence for their functional link.The paper frames this bidirectional coupling as evidence that induction heads are foundational to ICL.
  • Practical Application: A data-augmentation pipeline extracts patterns from high-influence samples and synthesizes data that accelerates induction-head formation across model sizes.Synthetic data generated using the smallest model generalizes effectively across various model scales.

2. Related Work

Prior work has separately studied circuit development, training-data influence, and links between data properties and mechanisms. These approaches provide useful evidence but often remain observational or rely on simplified data, motivating MDA’s complementary focus on causal, component-level attribution.

  • Mechanistic Interpretability: Mechanistic interpretability traditionally analyzes where and how circuits operate during inference while treating their training origins as fixed or unexplained.Recent developmental studies examine emergence trajectories, but the related-work discussion distinguishes this from tracing causal data origins.
  • Controlled Interventions: Controlled studies use synthetic data or forward-pass interventions, whereas MDA targets mechanisms in less simplified training-data settings.The related work presents these approaches as complementary rather than equivalent.
  • Training Data Attribution: Training-data attribution applies influence functions to quantify how examples affect model behavior, with scalable approximations addressing Hessian-inverse costs.Prior applications include explaining output likelihood through MLP layers and other behavior-level targets.
  • Data Properties and Development: Observed relationships between data patterns and induction-head emergence had remained primarily correlational before this work’s causal interventions.The paper positions MDA as moving beyond observation toward causal attribution.

3. Mechanistic Data Attribution Framework

MDA extends influence-function attribution from global loss to the behavior of localized interpretable units. Its framework identifies a unit, defines a probe for its function, and computes unit-specific sample influence with scalable curvature approximations.

  • Scalable Influence Estimation: Influence functions estimate training-sample effects, while EK-FAC approximates the inverse-Hessian-vector product needed to scale attribution to LLMs.The curvature approximation is computed layer-wise using Kronecker products of covariance matrices.
  • Framework Overview: MDA attributes fine-grained component behavior to training samples instead of measuring influence only on global model loss.The framework is designed for localized, mechanistic analysis of model development.
  • Stage 1: Localizing Interpretable Units: The framework first localizes an interpretable unit using a monitoring metric, isolates its parameter subspace, and defines a probing function with evaluation data.For induction heads, the monitoring metric can be the prefix-matching score.
  • Causal Validation: Causal validation compares targeted deletion and augmentation of high-influence samples with random baselines using induction- and previous-token-head scores.The figure caption states that head scores use the prefix-matching score metric.
  • Stage 2: Unit-Specific Influence Calculation: MDA replaces standard validation loss with a unit-specific probe objective to measure how a training sample affects the target mechanism.This focuses attribution on functional behavior rather than generic token prediction.
  • Stage 2: Unit-Specific Influence Calculation: The inverse Hessian used for MDA is restricted to the target component’s parameter subspace.This implements component-level rather than full-parameter attribution.

4. Causal Validation: Data Influence on Mechanistic Emergence

Counterfactual retraining tests whether MDA-identified samples causally affect interpretable-head emergence. Targeted deletion or duplication changes emergence dynamics, while larger-scale tracing remains computationally scalable but causal validation is limited to smaller Pythia models.

  • Experimental Setup: MDA interventions remove or duplicate high-influence samples within a developmental window to test necessity and sufficiency.The experiments target induction and previous-token heads across Pythia models, using bidirectional counterfactual retraining.
  • Causal Validation: Removing or duplicating ≤10% of high-influence samples significantly changes targeted-head emergence, unlike random interventions.Targeted deletion hinders emergence, whereas repeating the samples accelerates formation; random removal or augmentation has negligible effect.
  • Causal Validation: Augmentation accelerates head emergence and deletion suppresses or delays it, establishing a causal link between identified samples and circuit development.The intervention effects concern emergence timing and strength rather than only final performance.
  • Causal Validation: Augmentation and deletion regimes eventually reach comparable saturation scores, suggesting specific samples modulate emergence rate rather than determine ultimate head formation.The paper attributes final formation to the broader training distribution while noting early induction-score drops can signal accelerated emergence.
  • Generalizability and Scalability: MDA tracing scales sub-linearly with model parameters and captures linguistic and semantic patterns in OLMo-2 models up to 7B parameters.Full causal retraining was not performed at these larger scales because of computational cost.

5. Mechanistic Insights into Induction Head

MDA reveals that induction-head influence is concentrated in repetitive, mechanism-relevant data distributed across training. Interventions further connect induction-head formation causally with ICL capability.

  • Distributional Patterns of High Influence Data: Influence scores are heavy-tailed, indicating that a sparse subset of high-leverage training signals disproportionately drives mechanistic-circuit emergence.Figure 4 reports that the top 10% of samples contribute up to 50% of cumulative influence.
  • Distributional Patterns of High Influence Data: Highly repetitive structures, including seemingly noisy or garbage sequences, act as primary catalysts for induction-head formation.The finding comes from qualitative inspection of samples with the highest positive influence scores.
  • Transferability of Influential Data: High-influence samples overlap strongly across induction heads but less with non-induction heads, indicating mechanism-specific rather than globally difficult data.This cross-head dissociation suggests the identified features generalize to the induction mechanism beyond individual units.
  • Temporal Dynamics: Influential samples and scores are distributed uniformly across training, while samples from early, middle, or late stages outperform random replacements.This pattern supports steady accumulation with higher signal density shortening the required accumulation period.
  • Induction Heads and ICL: Suppressing induction-head formation degrades ICL, whereas enhancing it boosts ICL performance, providing direct causal evidence of functional coupling.The paper measures ICL with the ICL Score, based on late-versus-early-token loss on WikiText-2.

6. Mechanistic Data Augmentation

The paper turns MDA findings into an automated augmentation pipeline that distills structural motifs from influential data and synthesizes mechanism-targeted examples. Across Pythia scales, synthetic augmentation accelerates induction-head formation, while insertion strategy and data regime determine effectiveness without discernible broad-capability degradation.

  • Data Augmentation Pipeline: The three-step pipeline selects high-influence samples, distills their structural motifs with an LLM, and procedurally synthesizes structurally consistent training data.Pythia-14M identifies the top N = 2000 samples; an LLM converts batches into JSON schemas and executable Python synthesis scripts.
  • Synthetic Data Generalize Across Model Scales: Synthetic augmentation consistently triggers and accelerates induction-head formation across Pythia models from 14M to 160M parameters.The experiments insert mechanistic synthetic data during localized emergence phases and compare final induction-head scores with baseline training.
  • Synthetic Data Generalize Across Model Scales: Structural motifs, rather than specific semantic content, appear to be the primary causal drivers of induction-head formation, with 14M-derived data outperforming 160M-derived data on the 160M model.This result supports cross-model consistency of mechanistic drivers and the use of lightweight proxy models.
  • Ablation Study on Insertion Strategy: MDA augmentation causes no statistically discernible degradation in language-modeling or factual-retention capabilities across the evaluated Pythia models.Figure 6a shows augmentation and baseline trajectories closely overlapping across the monitored capability axes.
  • Ablation Study on Insertion Strategy: Synthetic data outperform natural data when N ≤50, 000, but natural data outperform synthetic data at N = 100, 000.The paper attributes the crossover to finite-pattern redundancy and diversity exhaustion in synthetic data, versus broader lexical and syntactic variation in natural influential samples.
  • Ablation Study on Insertion Strategy: For natural data, dispersed insertion consistently outperforms concentrated insertion by preserving distributional alignment and reducing optimization disruption.The paper describes dispersed insertion as a localized curriculum that supports stable mechanistic integration; synthetic-data results vary by scale.

7. Discussions

The discussion qualifies MDA’s attribution scope and distinguishes individual interpretable units from broader mechanisms. Direct unit-level attribution avoids noisy path aggregation but does not yet disentangle direct from cascading upstream effects.

  • Scope of Attribution: MDA assumes that an interpretable unit’s behavior is attributable to direct training-data influence on that specific unit.Training data may also affect the unit indirectly by modulating upstream nodes in its computation graph.
  • Scope of Attribution: Path-based aggregation can fail to distinguish direct influence on an upstream previous-token head from indirect influence propagated through an induction head.The authors leave systematic decoupling of these cascading influences to future work because broader aggregation introduces substantial noise.
  • Mechanism vs Interpretable Units: High-influence samples for a specific induction head are universal across the broader induction mechanism, consistent with one abstract mechanism having multiple structurally distinct circuits.This finding motivates distinguishing mechanisms from the individual interpretable units that implement them.

8. Conclusion

The paper presents MDA as a framework linking interpretable LLM mechanisms to causal data origins. Its results identify data catalysts that generalize across model scales and connect internal mechanisms with macro-level capabilities such as ICL.

  • Conclusion: MDA traces the causal origins of interpretable LLM mechanisms back to the training corpus.The framework is presented as a basis for understanding and steering the development of internal circuits.
  • Conclusion: Identifiable data catalysts drive circuit emergence across model scales and provide a causal link between internal mechanisms and macro-level capabilities like ICL.The paper frames this link as a foundation for mechanistic alignment and precise steering or unlearning of model behaviors.

Impact Statement

MDA traces interpretable mechanisms to training data by combining component-specific influence functions with scalable curvature approximations and targeted subspace analysis. The framework supports causal investigation of circuit formation, including temporal emergence, data-pattern alignment, and extensions beyond attention heads.

  • Scalable Approximation: EK-FAC reduces inverse-curvature computation from O(d6) to O(d3) through a Kronecker-factorized Hessian approximation, with eigenvalue correction improving curvature estimation.The method retains the K-FAC eigenvector basis while correcting eigenvalues using projected per-sample gradients.
  • Mechanistic Data Attribution: MDA computes training-example influence for individual interpretable units, including attention heads, by restricting analysis to their associated parameter subspaces.The framework uses a component capability probe and an EK-FAC-approximated inverse Hessian within the selected subspace.
  • Component-Specific Formulations: MDA combines query-key, value, and output subspaces to capture the coupled attention-pattern formation of targeted heads.For previous-token heads, the active subspace concatenates the query and key matrices because attention patterns depend on their interaction.
  • Influence Interpretation: Influence scores prioritize samples whose loss gradients align with mechanism-probe gradients, with repetitive or copy-like structures often exhibiting this alignment.The influence is described as a negative inner product of natural gradients on the statistical manifold.
  • Temporal Localization: Influence analysis is restricted to the critical developmental interval in which a mechanism emerges most rapidly, rather than the full training trajectory.For induction heads, the induction score rises from an early noise floor of approximately 0.1 to a stable level of 0.4–0.5 during this window.
  • Scope: The framework is applied to induction and previous-token heads and is also explored qualitatively for sparse autoencoder features.The SAE experiments suggest that MDA can extend beyond individual attention heads to feature-level mechanisms.

E.3. Configuration for Mechanistic Data Augmentation

The section specifies controlled mechanistic data augmentation settings and examines how MDA localizes influence to induction-related mechanisms across model scales.

  • Configuration for Mechanistic Data Augmentation: Synthetic-data insertions used 100,000 samples for 14M, 20,000 for 31M and 70M, and 10,000 for 160M.Insertions occurred at steps 900, 800, 700, and 600, respectively.
  • Configuration for Mechanistic Data Augmentation: Smaller synthetic-data volumes than natural-data top-k insertions still supported the high causal density of generated mechanistic patterns.The authors provide a separate investigation of insertion quantity in Appendix H.
  • Scaling: MDA tracing scales approximately linearly with model size and remains data-parallelizable, while retraining-based causal validation is limited to models up to 160M parameters.The limitation is attributed primarily to the computational cost of repeated retraining and controlled comparisons.
  • Cross-Head Localization of Influence: The localized approximation estimates sample effects within a target unit’s parameter subspace rather than decomposing all indirect effects across the computation graph.Cross-head checks test whether target-head samples preferentially align with induction-like heads.
  • Cross-Head Localization of Influence: Samples selected from one induction head generalize to other induction heads, indicating alignment with a broader mechanism rather than exclusive influence on one head.The same data catalysts may support multiple heads implementing similar behavior.
  • Multi-Head Dynamics: In 14M, secondary induction heads retain their functional identity under interventions, whereas larger models can redistribute induction roles among heads.The 160M observations include a head whose induction score decreased by 0.26 and a possible handoff to other components.
  • Robustness: MDA rankings remain highly consistent across random seeds, supporting the stability of its identified influential examples.The cross-head projection signal is concentrated on induction heads, with most non-induction heads receiving scores roughly an order of magnitude smaller.
  • Measurement: A target head’s ablation score measures the drop in the correct token’s logit when that head is functionally removed.The score is tracked throughout training for the target head, such as Layer 3 Head 3 in the 14M model.

G.3.2. CONSISTENCY OF RESULTS

The section compares mechanistic and observational measures of induction-head formation and contextualizes their agreement within the model’s training dynamics and influence distribution.

  • Consistency of Results: Both ablation logit contribution and prefix matching score identify the same phase-transition interval and respond identically to deletion and augmentation interventions.This alignment supports induction score as a proxy for causal contribution to correct-token logits.
  • Consistency of Results: Logit difference and induction score have fully aligned ordering with the previously observed prefix matching score on the Pythia-14M synthetic task.The authors interpret this agreement as evidence that the method is reliable and robust.
  • Experimental Design: The factorial experiment varied three data types, four insertion quantities, and two scheduling strategies across 25 runs including the baseline.Data types were real high-influence, synthetic pattern-based, and random-control samples.
  • Training Dynamics: The induction mechanism undergoes a dramatic phase transition between steps 1200 and 2000, followed by minor undulations and stabilization rather than perfect monotonicity.The later behavior is attributed to competing optimization objectives involving other linguistic features.
  • Critical Window: Dense influence analysis over the full training lifecycle is intractable, motivating a critical window focused on the phase transition where the mechanism originates.The window is intended to capture the period of maximum causal density and the largest derivative of capability gain.
  • Influence Distribution: Across the critical window, cumulative positive influence substantially exceeds the absolute cumulative negative influence, and this net-positive ratio remains temporally uniform.The distribution remains uniform across varying bin sizes and continuous sliding windows.

J.2. Robustness of Temporal Uniformity

The section describes an automated pipeline that converts high-influence samples into structured mechanistic templates and executable synthetic-data generators.

  • Pipeline: The pipeline has three stages: sample mining, pattern extraction, and generator implementation.It converts raw high-influence training samples into executable data-generation scripts.
  • Sample Mining: Sample mining ranks training examples by projection score, selects the top K = 2000 positive-influence samples, and decodes them into text.The decoded samples become seeds for pattern extraction.
  • Pattern Extraction: DeepSeek-V3 extracts structural templates while ignoring specific semantic content and returns them in a JSON schema.The schema includes anchors, variable fields, generation rules, templates, and length controls.
  • Pattern Extraction: The 14M model yielded around 900 unique deduplicated mechanistic patterns.These patterns are stored in a merged JSON registry.
  • Generator Implementation: A meta-programming script prompts DeepSeek-V3 to convert each JSON pattern into a Python generator enforcing field, template, and length constraints.The generated functions implement the rules and approximate the target token length.
  • Mechanistic Patterns: The identified data drivers are long-range repetitive structures, which provide the qualitative basis for mechanistic data generation.The pipeline is presented as an operationalization of this pattern-based attribution.

L.1. Attention-Head Samples in Pythia 14M

The section inspects high-influence samples associated with attention heads and SAE features, emphasizing repetitive and increasingly abstract structures across model scales.

  • Pythia 14M Attention-Head Samples: Pythia-14M top-ranked samples from the top 0.1% span structured code, binary-like sequences, and enumeration lists while sharing long-range repetitive signatures.The samples contain recurring patterns or token sequences despite differing superficially in format.
  • Larger-Model Samples: OLMo-2-1B and OLMo-2-7B samples include repeated phrase templates, tabular sequences, query-template lists, code-like patterns, and HTML/CSS fragments.Compared with Pythia-14M, the larger-model examples appear more abstract and structured.
  • SAE Feature Samples: MDA is extended qualitatively to SAE features, where retrieved samples align with interpretable feature-level patterns.Pythia-70M examples associate Feature 9500 with LaTeX-like expressions and Feature 31939 with newline- or formatting-heavy text.
Loading 2601.21996v2…