Predictive data attribution

MAGIC and sketched metadifferentiation

Eric

Joint work with

Hamza Golubovic, Han Tong

Arian Maleki, Andrew Ilyas

Applications of predictive data attribution

Data curation

Which data removals would improve performance on a target task?

Failure analysis

Which training examples would change a particular prediction if removed?

Data removal

How would excluding a set of records affect held-out behavior?

Example: attributing LLM predictions

Intervention

Remove a specified subset of training text

Measurement

Next-token loss after a held-out passage

Target

How that loss would change if the model had trained without this text

Parameterizing predictive data attribution

\mathcal L(\theta;\mathbf w)=\frac{1}{n}\sum_{i=1}^{n} w_i\,\ell(x_i;\theta)

Original training set

\mathbf w=\mathbf1_n

Remove a subset S

\mathbf w=\mathbf1_n-\mathbf 1_S

Each training example receives a continuous weight

Model output function

Data weights

\mathbf w

\xrightarrow{\mathcal A}

Trained parameters

\theta(\mathbf w)

\xrightarrow{\varphi^{q}}

Final measurement

f^{q}(\mathbf w)

f^{q}(\mathbf w)=(\varphi^{q}\circ\mathcal A)(\mathbf w)

Meta-gradients

\left.\nabla_{\mathbf w}f^{q}(\mathbf w)\right|_{\mathbf w=\mathbf1_n}=\left.\nabla_{\mathbf w}\big[\varphi^{q}(\mathcal A(\mathbf w))\big]\right|_{\mathbf w=\mathbf1_n}

Gradients through the whole training procedure

MAGIC

MAGIC uses a first-order Taylor expansion

f^{q}(\mathbf w)\approx f^{q}(\mathbf1_n)+\big\langle\nabla_{\mathbf w}f^{q}(\mathbf1_n),\,\mathbf w-\mathbf1_n\big\rangle

REPLAY computes the meta-gradients

MAGIC is state of the art for predictive data attribution

ResNet-9 / CIFAR-10 · random 1% deletions · mean Spearman correlation

Challenge

Computation scales linearly with the number of queries

k\text{ queries}\quad\Longrightarrow\quad k\text{ replays}

Can we attribute the whole test dataset with fewer replays?

Attribution as matrix estimation

Y_{qi}=\left.\frac{\partial f^{q}}{\partial w_i}\right|_{\mathbf w=\mathbf1_n},\qquad Y\in\mathbb R^{k\times n}

Row

Y_q

: query

q

Column

i

: training example

i

Estimate the influence matrix with $B<k$ measurements

Two distinct objectives

Reconstruction

\min_{\widehat Y}\;\|Y-\widehat Y\|_F^2

Recover the influence scores

LDS

\max_{\widehat Y}\;\frac1k\sum_{q=1}^{k}\mathrm{LDS}_q

Predict the effects of data interventions

The chain rule factors the influence matrix

J=\left.\frac{\partial\theta(\mathbf w)}{\partial\mathbf w}\right|_{\mathbf w=\mathbf1_n}

The Jacobian of the training procedure

V=[v_1,\ldots,v_k],\qquad v_q=\nabla_{\theta}\varphi^{q}\big(\theta(\mathbf1_n)\big)

Query gradients at the trained parameters

Y=V^{\top}J

One replay can measure a combination of rows

One replay measures one query’s row

\operatorname{REPLAY}(v_q)=v_q^{\top}J=Y_q

Combining query gradients gives a measured direction $u$

u=\operatorname{REPLAY}(Vz)=(Vz)^{\top}J=z^{\top}Y

Orthogonal projection

Let $U=\operatorname{span}\{u_1,\ldots,u_B\}$ , with orthogonal $u_b$

\widehat Y_q:=Y_q\Pi_U=\sum_{b=1}^{B}\frac{\langle Y_q,u_b\rangle}{\|u_b\|_2^2}\,u_b

We know $U$ but still need $\langle Y_q,u_b\rangle$ for every query

Forward mode completes the projection

\begin{bmatrix}\langle Y_1,u_b\rangle\\[-.15em]\vdots\\[-.15em]\langle Y_k,u_b\rangle\end{bmatrix}=Yu_b^{\top}=V^{\top}(Ju_b^{\top})

Forward mode computes $Ju_b^{\top}$
Multiplying by $V^{\top}$ gives the inner products

An online viewpoint of PCA

\|Y-Y\Pi_U\|_F^2=\|Y\|_F^2-\|Y\Pi_U\|_F^2

Maximizing captured energy minimizes reconstruction error

z=v_{\max}\!\left(Y(I-\Pi_U)Y^{\top}\right)

MAGE

YY^{\top}=V^{\top}JJ^{\top}V\;\approx\;V^{\top}V

We use the query-gradient Gram matrix as a proxy

Probe its leading eigenvectors in order

From LDS to a projection objective

Start with LDS

\mathrm{LDS}_q=\text{Spearman-}\rho_S\!\left(-\langle\widehat Y_q,\mathbf1_S\rangle,\ f^{q}(\mathbf1_n-\mathbf1_S)-f^{q}(\mathbf1_n)\right)

Use a Pearson surrogate

\approx\text{Pearson-}\rho_S\!\left(-\langle\widehat Y_q,\mathbf1_S\rangle,\ f^{q}(\mathbf1_n-\mathbf1_S)-f^{q}(\mathbf1_n)\right)

Approximate with MAGIC

\approx\text{Pearson-}\rho_S\!\left(-\langle\widehat Y_q,\mathbf1_S\rangle,\ -\langle Y_q,\mathbf1_S\rangle\right)

Uniform subset sampling

=\text{Pearson-}\rho(\widehat Y_q,Y_q)

Neglect row centering

\approx\frac{\langle Y_q\Pi_U,Y_q\rangle}{\|Y_q\Pi_U\|_2\,\|Y_q\|_2}

Use orthogonality

=\frac{\|Y_q\Pi_U\|_2}{\|Y_q\|_2}

A normalized PCA objective for attribution

\sum_{q=1}^{k}\frac{\|Y_q\Pi_U\|_2^2}{\|Y_q\|_2^2}=\|\widetilde Y\Pi_U\|_F^2

\widetilde Y_q=\frac{Y_q}{\|Y_q\|_2}

The squared-correlation surrogate maximizes captured energy
of the row-normalized matrix

SPELL

Approximate PCA on normalized rows of the influence matrix

Start with MAGE measurements

Estimate row norms from the current projection

Choose the next direction from the normalized residual Gram matrix

MAGE closely tracks the reconstruction oracle

Relative Frobenius error · lower is better

$k=500,\ B=100$ · mean ± SE · rank-B SVD requires the full matrix

SPELL leads on predictive data attribution

Linear datamodeling score · higher is better

$k=500,\ B=100$ · 1% deletions · mean ± SE · oracle SVD · full MAGIC uses $B=500$

SPELL gains accuracy on queries with small meta-gradients

LDS by row-norm quartile of the influence matrix

MAGESPELLMAGIC

1% deletions · CIFAR-10: $k=1000, B=500$ · TinyStories: $k=600, B=300$

Thank you