MAGIC: Near-Optimal Data Attribution for Deep Learning (2504.16430v1)

Abstract: The goal of predictive data attribution is to estimate how adding or removing a given set of training datapoints will affect model predictions. In convex settings, this goal is straightforward (i.e., via the infinitesimal jackknife). In large-scale (non-convex) settings, however, existing methods are far less successful -- current methods' estimates often only weakly correlate with ground truth. In this work, we present a new data attribution method (MAGIC) that combines classical methods and recent advances in metadifferentiation to (nearly) optimally estimate the effect of adding or removing training data on model predictions.

• The paper introduces MAGIC, a metagradient-based method that computes the exact influence function for precise single-model data attribution.
• It leverages the Replay algorithm to differentiate through iterative and smooth training processes, achieving near-perfect alignment with true model behavior for small data drops.
• The approach enables detailed data debugging, model unlearning, and data selection while highlighting its computational cost and the necessity for smooth training dynamics.

Predictive data attribution aims to understand how changes in the training data affect the predictions of a trained machine learning model. While straightforward for convex models using techniques like the influence function, this task becomes significantly more challenging for large-scale, non-convex deep learning models. Existing methods in deep learning settings often provide estimates that only weakly correlate with the true impact of data changes.

The paper introduces a new setting called "single-model predictive data attribution" to address a key challenge in deep learning: the non-determinism of training. Standard training procedures often involve randomness (e.g., initialization, data shuffling), meaning training on the same dataset can yield different models. Previous data attribution methods typically predict the average behavior of a learning algorithm when exposed to different data, not the behavior of a specific trained model. The single-model setting removes this training randomness, allowing for the possibility of perfectly predicting how a given model's outputs would change if its training data had been different. This is motivated by the practical need to understand the impact of data on a particular deployed model.

The core contribution of the paper is "Magic" (Metagradient-based Attribution via Ground-truth Influence Computation), a novel data attribution method designed for this single-model setting in deep learning. Magic achieves near-optimal performance by calculating the exact influence function, which is the gradient of the model's output with respect to the training data weights, evaluated at the standard training data weights. The influence function provides a first-order Taylor approximation of how the model output changes when data weights are perturbed: $$\hat{f}(w) = f(\mathbf{1}_n) + \left(\frac{\partial f(w)}{\partial w}\bigg|_{w=\mathbf{1}_n}\right)^\top (w - \mathbf{1}_n)$$, where $f(w)$ is the model output function for dataset weighted by $w$, and $\mathbf{1}_n$ represents the standard training dataset where all examples have a weight of 1.

Calculating this exact influence function in non-convex deep learning is difficult because there is no simple closed-form solution. Magic overcomes this by leveraging recent advances in metagradient calculation, specifically the Replay algorithm [(2504.16430), engstrom2025optimizing]. Replay is an algorithm capable of computing the exact gradient of a model's final output with respect to hyperparameters or data weights, provided the learning algorithm is iterative and smooth.

For Magic to be applicable, the learning algorithm $A$ must meet two conditions:

1. Iterative: The algorithm proceeds in discrete steps, updating an optimizer state (which includes model parameters) based on gradients computed on minibatches (as shown in Eq 2 in the paper). This describes standard deep learning training with optimizers like SGD or Adam.
2. Smooth: Small changes in data weights $w$ result in only small changes in the gradient $\nicefrac{\partial f(w)}{\partial w}$. Standard deep learning algorithms are not always smooth due to phenomena like phase transitions or sudden changes in optimization dynamics. However, the paper notes that "smooth counterparts" often exist or can be constructed, for example, by performing initial pretraining steps or carefully selecting hyperparameters. This smoothness is crucial because the influence function is a gradient, and a non-smooth landscape makes a first-order approximation unreliable.

Magic implements the influence function calculation using the Replay algorithm, which essentially differentiates through the entire training process computation graph. This involves a backward pass through the sequence of optimizer states and gradients computed during training, propagating gradients from the final model output back to the initial data weights. The paper outlines the Replay process as a series of backward steps from the final state $s_T$ to the initial state $s_0$, accumulating the contribution of data weights $w$ at each step $t$.

The paper evaluates Magic using the Linear Datamodeling Score (LDS) [(2504.16430), ilyas2022datamodels, park2023trak], which measures the Spearman correlation between predicted and true model outputs across different random training data subsets (simulated by dropping samples). Experiments on ResNet-9 trained on CIFAR-10, GPT-2 fine-tuned on Wikitext, and Gemma-2B fine-tuned with LoRA on instruction tuning data demonstrate that Magic achieves near-perfect LDS (close to 1.0) for small drop fractions (e.g., 1%). In contrast, baselines like TRAK [(2504.16430), park2023trak] and EK-FAC [(2504.16430), bae2022if, grosse2023studying] show significantly lower correlations. Magic's prediction accuracy decreases as the drop fraction increases, as expected from a first-order Taylor approximation further away from the expansion point ($w=\mathbf{1}_n$).

From a practical perspective:

• Performance: Magic provides highly accurate data attribution for small changes to the training dataset of a specific model. This is particularly useful for tasks requiring precise understanding of individual data point influence.
• Computational Cost: Magic's cost scales linearly with the number of test samples ($n$) and the cost of training ($T \cdot \text{cost_per_step}$, where $T$ is the number of steps and cost_per_step is dominated by minibatch gradient computation). Specifically, calculating the influence for $n$ test samples costs roughly $n \times (3 \text{ to } 5) \times \text{training_cost}$. This is efficient when attributing the impact on a few test examples but can become prohibitively expensive for attributing the impact on many test examples compared to baselines like TRAK/EK-FAC, whose cost for computing influence scores for all training points scales differently (e.g., $O(N \cdot k)$ for TRAK) and is relatively constant regardless of $n$.
• Requirements: Implementing Magic requires access to the full training history (optimizer states and minibatches) and the learning algorithm must be smooth or adaptable to be smooth. The paper suggests pretraining for a few hundred steps can induce sufficient smoothness in some cases.

Potential applications unlocked by Magic's high fidelity include:

• Data Debugging: Identifying training data points that negatively impact specific model predictions.
• Model Unlearning: More precisely estimating the effect of removing data to verify unlearning procedures.
• Data Summarization/Selection: Finding representative or influential data points.

Limitations primarily involve the computational cost for attributing effects on a large number of test samples and the requirement for the learning algorithm to be smooth, which may not hold for standard training out-of-the-box, although workarounds exist. Future research directions include developing more efficient metagradient computation methods and techniques to scale Magic to many test samples simultaneously.