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

TL;DR

This work introduces Mechanistic Data Attribution (MDA), a scalable framework that links the training data origins of interpretable LLM units to their functional emergence using Influence Functions with EK-FAC. Through causal interventions on the Pythia suite, it demonstrates that removing or duplicating a small fraction of high-influence samples can significantly modulate induction and previous-token heads, while random data changes have little effect. The study reveals that repetitive structural patterns (e.g., LaTeX, XML) catalyze induction head formation, that high-influence data transfer across heads, and that induction heads causally underpin In-Context Learning (ICL). It further provides a practical mechanistic data augmentation pipeline, enabling pattern extraction and synthetic data generation that accelerates circuit convergence across model scales, offering a principled approach to steer developmental trajectories of LLMs.

Abstract

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.

Figures (9)

• Figure 1: The Mechanistic Data Attribution (MDA) framework. MDA identifies interpretable LLM units and quantifies the influence of individual training samples on their functional behavior. This enables both the discovery of mechanistic training dynamics and precise data-level interventions to steer model development.
• Figure 2: Causal validation of Mechanistic Data Attribution. Intervened retraining shows that targeted deletion and augmentation of high-influence samples (identified via MDA) significantly modulate the emergence of Induction and PT (Previous Token) heads. Head score is quantified via the metric from olsson2022context. Note the clear gap between MDA-guided interventions and the random baselines across different model scales.
• Figure 3: Distributional properties of high influence samples. a) Power-law distribution: The distribution of influence scores follows a power-law, where the top 10% of samples contribute up to 50% of the total cumulative influence. b) Cross-head consistency: High-influence samples identified by induction heads (Ihead) within the Pythia-14M model exhibit significant overlap, yet remain distinct from those identified by non-induction heads (Nhead). c) Step uniformity: The identified high-influence samples are distributed uniformly throughout the training corpus, showing no significant temporal clustering. d) Induction head scores with high influence samples replaced with those from different steps. MDA-Repl @ [$t_1$, $t_2$] represents replaced by high influence sample in step $t_1$ to $t_2$. The random replacement baseline exhibits significant deviation from the MDA replacement, exceeding three standard errors (3$\sigma$). Pruned 95% means we randomly mask the gradient of 95% samples in training, while the induction head scores still show a non-trivial increase. e) Induction scores differences of all heads from Pythia 14M. High influence samples from one head are generalizable to other induction heads (red squares).
• Figure 4: Validating the functional role of induction heads in ICL via data intervention. Under the same data augmentation and deletion settings used for induction heads, the concurrent shifts in ICL scores and induction head strength (grey dashes) provide causal evidence that these internal mechanisms are functionally coupled.
• Figure 5: Ablation of Insertion Strategy in Pythia 14M. Induction head formation is positively correlated with the quantity of mechanistic signals ($N$). The performance gap between Dispersed and Concentrated modes reveals that the temporal density of interventions significantly modulates optimization stability, particularly for high-influence real-world samples.