Mechanistic Training Data Attribution

This is an explainer and reimplementation of the 2026 paper by Chen et al.

Note: This post is currently incomplete!

In this post, I attempt to explain and walk through this recent paper on training data attribution.

Objective

We want to perform a kind of training data attribution. Specifically, we want to trace internal mechanisms exhibited by interpretable units of a language model (such as induction heads) back to specific training examples. We can do this through influence functions.

Background on Influence Functions

Definition. Let $z_k = (x_k, y_k)$ be an element of a dataset $\mathcal{D} = \{z_i\}_{i=1}^{N}$, and let $\epsilon$ be the weighting of $z_k$. Then the response function is the optimal solution $\theta^{*}(\epsilon)$ that minimizes the loss $\mathcal{L}$:

Definition. The influence function $\mathcal{I}_{\theta^*}(z_k)$ is a function that measures the influence of a specific training example $z_k$ on the perturbed model parameters $\theta^*$.

This influence is exactly equal to the difference between the perturbed and unperturbed optima, which can be found through a first-order Taylor expansion at 0:

We can easily obtain the expression for this partial as:

where $\mathrm{H} = \nabla_{\theta}^2 \left(\frac{1}{N} \sum_{i=1}^{N} \mathcal{L}(z_i, \theta^*) + \epsilon \mathcal{L}(z_k, \theta^*) \right)$

However, this expression alone is not very useful for our purposes. Instead, we would like to find the influence of the training example $z_k$ on the test loss $\mathcal{L}(z_{test}, \theta^*)$. Let's call this function $f$ for simplicity. Then the influence $\mathcal{I}_f$ of $z_k$ on $f$ is given by:

Substituting:

Problem. $\mathrm{H}$ is a huge matrix. Indeed, for a layer of dimensions $d_{in} \times d_{out}$, the Hessian is of size $(d_{in} \times d_{out})^2$. This is computationally infeasible to invert.

Solution. We approximate the Hessian as the Kronecker product of two smaller matrices using a clever method called EKFAC:

$\mathrm{H} \approx (U_A \otimes U_S) \Lambda (U_A \otimes U_S)^{\top}$

where $U_A$ and $U_S$ are the eigenvector matrices obtained from eigendecomposing the two K-FAC covariance factors separately: the input-activation covariance $A = \mathbb{E}[a\,a^\top] = U_A \Lambda_A U_A^\top$ (computed from that layer's input activations $a$) and the output-gradient covariance $S = \mathbb{E}[g\,g^\top] = U_S \Lambda_S U_S^\top$ (computed from the gradients $g$ flowing into that layer's outputs). The Kronecker product of the two eigenvector matrices, $U_A \otimes U_S$, then gives the eigenbasis of the full K-FAC approximation $A \otimes S$, since $(U_A \otimes U_S)$ diagonalizes $A \otimes S$ exactly whenever $U_A$ diagonalizes $A$ and $U_S$ diagonalizes $S$.

$\Lambda$ is a diagonal matrix of corrected eigenvalues. Rather than naively taking $\Lambda = \Lambda_A \otimes \Lambda_S$ as in KFAC, EK-FAC re-estimates the true eigenvalues by projecting actual per-example gradients onto the fixed eigenbasis $(U_A \otimes U_S)$ and computing their variance in that basis.

Mechanistic Data Attribution

Definition. The MDA framework is characterized by the 3-tuple $(\mu, \pi, f_{probe})$:

• The monitoring metric $\mu$ measures when the behaviour of a certain head shows variation, e.g. prefix-matching score for induction heads

• The subspace projection $\pi$ identifies a parameter subspace corresponding to an interpretable unit of the model, i.e. $\pi(\theta) = \theta_{sub} \subseteq \theta$

• $f_{probe}$ measures the performance of the interpretable unit in question on a test dataset $\mathcal{D}_{probe}$ (and thus we may have $\mu = f_{probe}$).

To measure the influence of training examples on specific interpretable units, we modify the expression for $\mathcal{I}$ as follows:

Here $\mathrm{H}^{-1}_{\theta_{sub}}$ is the EKFAC-approximated Hessian restricted to the relevant parameter subspace $\theta_{sub}$.

Implementation

I train a toy 2-layer Transformer in PyTorch on a very simple synthetic induction dataset (with each example constructed as [prefix, prefix]) with the following config:

Note that there are no MLPs in the model.

I use the standard prefix matching score as my $\mu$, $\theta_{QK}, \theta_V, \theta_O$ as $\theta_{sub}$, and the summed log-probability of the correct induction target over synthetic $[\mathrm{prefix}, \mathrm{prefix}]$ sequences as $f_{probe}$. Note that there are certain changes in my implementation:

• I don't use LayerNorm, so in my code you'll see that use hook_resid_pre, i.e. hook the pre-norm residual stream.
• $f_{probe}$ is log probs instead of the raw logit dot unembedding formulation in A.3. This is easier to get exactly right here since I know the ground-truth match position in closed form.
• Because $f_{probe}$ is maximized rather than minimized, I drop the leading negative sign in $s = \nabla_\theta \mathcal{L}(z_{train})^\top p$, $p = H^{-1}_{\theta_{sub}} \nabla_\theta f_{probe}$. Positive score still means "promotes the mechanism."

Results