Mechanistic Data Attribution: Tracing the Training Origins of Interpretable LLM Units
Abstract
While mechanistic interpretability has identified interpretable circuits in large language models (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.
Lay Summary
(1) Problem: Large language models (LLMs) like ChatGPT are incredibly powerful, and while we can identify the internal "circuits" that make them work, we still do not know exactly which specific data from their massive training sets caused these circuits to form in the first place. (2) Solution: We developed a scalable framework called Mechanistic Data Attribution (MDA) that acts like a tracking system, tracing these internal circuits back to the precise text snippets used during training. By adding or removing a tiny fraction of these highly influential training samples (such as repetitive code or structured text), we successfully controlled the emergence rate of these circuits. We also provided evidence that these specific circuits are directly responsible for the model's ability to learn new tasks on the fly from examples. (3) Impact: Finally, we created a data pipeline that uses these insights to help AI models form these useful internal circuits much faster. This provides a structured, predictable way for researchers to steer how future AI models grow and learn.