AG

Mechanistic Training Data Attribution

15 minute read

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 zk=(xk,yk)z_k = (x_k, y_k) be an element of a dataset D={zi}i=1N\mathcal{D} = \{z_i\}_{i=1}^{N}, and let ϵ\epsilon be the weighting of zkz_k. Then the response function is the optimal solution θ∗(ϵ)\theta^{*}(\epsilon) that minimizes the loss L\mathcal{L}:

θ∗(ϵ)=argmin⁡θ∈RD1N∑i=1NL(zi,θ)+ϵL(zk,θ)\theta^{*}(\epsilon) = \operatorname{argmin}_{\theta \in \mathbb{R}^{\mathcal{D}}} \frac{1}{N} \sum_{i=1}^N {\mathcal{L}(z_i, \theta)} + \epsilon \mathcal{L}(z_k, \theta)

Definition. The influence function Iθ∗(zk)\mathcal{I}_{\theta^*}(z_k) is a function that measures the influence of a specific training example zkz_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:

Iθ∗(zk)=θ∗(ϵ)−θ∗(0)=∂θ∗(ϵ)∂ϵ⋅ϵ\mathcal{I}_{\theta^*}(z_k) = \theta^*(\epsilon) - \theta^*(0) = \frac{\partial \theta^*(\epsilon)}{\partial \epsilon} \cdot \epsilon

We can easily obtain the expression for this partial as:

∂θ∗(ϵ)∂ϵ=−H−1∇θL(zk,θ∗)\frac{\partial \theta^{*}(\epsilon)}{\partial \epsilon} = - \mathrm{H}^{-1} \nabla_{\theta} \mathcal{L}(z_k, \theta^*)

where H=∇θ2(1N∑i=1NL(zi,θ∗)+ϵL(zk,θ∗))\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 zkz_k on the test loss L(ztest,θ∗)\mathcal{L}(z_{test}, \theta^*). Let's call this function ff for simplicity. Then the influence If\mathcal{I}_f of zkz_k on ff is given by:

If=∂L(ztest,θ∗)∂ϵ=∇θ(L(ztest,θ∗))⊤∂θ∗(ϵ)∂ϵ\mathcal{I}_f = \frac{\partial \mathcal{L}(z_{test}, \theta^*)}{\partial \epsilon} = \nabla_{\theta}(\mathcal{L}(z_{test}, \theta^*))^{\top} \frac{\partial \theta^{*}(\epsilon)}{\partial \epsilon}

Substituting:

If(zk,ztest)=−∇θ(L(ztest,θ∗))⊤H−1∇θL(zk,θ∗)\mathcal{I}_f(z_k, z_{test}) = - \nabla_{\theta}(\mathcal{L}(z_{test}, \theta^*))^{\top} \mathrm{H}^{-1} \nabla_{\theta} \mathcal{L}(z_k, \theta^*)

Problem. H\mathrm{H} is a huge matrix. Indeed, for a layer of dimensions din×doutd_{in} \times d_{out}, the Hessian is of size (din×dout)2(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. For EKFAC specifics, see my post here.

Mechanistic Data Attribution

Definition. The MDA framework is characterized by the 3-tuple (μ,π,fprobe)(\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. π(θ)=θsub⊆θ\pi(\theta) = \theta_{sub} \subseteq \theta

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

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

I(ztrain,Dprobe)=−∇θsubL(ztrain)⊤Hθsub−1fprobe(θ,Dprobe)\mathcal{I}(z_{train}, \mathcal{D}_{probe}) = -\nabla_{\theta_{sub}} \mathcal{L}(z_{train})^{\top} \mathrm{H}_{\theta_{sub}}^{-1} f_{probe} (\theta, \mathcal{D}_{probe})

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

Implementation

In Progress!