When SGD Proxies Fail: Optimizer Aware Data Attribution for Modern Deep Learning

Introduction

Data attribution has become a critical bottleneck in modern machine learning pipelines. As training datasets grow to billions of samples and models scale to trillions of parameters, the ability to identify which specific training examples drive model behavior is no longer a luxury but a necessity. Without precise attribution, we cannot detect harmful biases, defend against data poisoning, or prune redundant samples to reduce computational waste. The Shapley value, originating from cooperative game theory, provides the theoretical gold standard for fair attribution, distributing model performance among data points based on their marginal contributions. Yet the exact computation of Shapley values requires retraining models on exponentially many subsets, rendering it intractable for deep learning.

Recent work has introduced In-Run Data Shapley methods, which estimate contributions dynamically during a single training pass rather than through retrospective retraining. While this approach promises scalability, it carries a hidden assumption that has gone largely unchallenged until now. The existing framework is built exclusively for Stochastic Gradient Descent (SGD), relying on the linear structure of SGD updates to enable efficient computation. However, modern deep learning practice has largely abandoned vanilla SGD in favor of adaptive optimizers, particularly Adam. This raises a fundamental question that Ding et al. address in their paper, "In-Run Data Shapley for Adam Optimizer": Is data attribution invariant to the choice of optimizer, or do we need optimizer specific attribution mechanisms?

The Optimizer Dependence Problem

The core contribution of Ding et al. is demonstrating that data attribution is fundamentally optimizer dependent. This finding contradicts the implicit assumption that one can approximate Shapley values using SGD based proxies even when the target model is trained with Adam. The authors reveal that SGD based In-Run Shapley correlates with true marginal contributions under Adam with a Pearson coefficient of merely 0.11. In practical terms, this means attribution scores computed via SGD proxies are essentially uncorrelated with actual data influence in Adam trained models.

This divergence stems from the structural differences between SGD and Adam updates. Standard In-Run methods rely on a first order Taylor expansion of the utility function that assumes parameter updates are linear combinations of data gradients. SGD satisfies this property; the update is simply the gradient scaled by a learning rate. Adam, however, introduces history dependent moments and adaptive variance scaling. The update rule divides the momentum term by the square root of the second moment estimate, creating a non linear dependence on the gradient history. This non linearity breaks the linear structure required by existing In-Run approximations.

The implications extend beyond theoretical inconsistency. In downstream tasks such as data pruning or noisy sample detection, using SGD proxies for Adam trained models leads to systematically incorrect rankings of data importance. The paper demonstrates that practitioners relying on SGD based attribution for Adam training pipelines are not merely introducing approximation error; they are working with noise. This finding suggests that much of the prior work on data attribution may have limited applicability to contemporary training regimes, where Adam and its variants dominate.

Deriving Adam Aware Attribution

To address this gap, Ding et al. propose Adam Aware In-Run Data Shapley, deriving a closed form estimator that explicitly accounts for Adam's adaptive dynamics. The technical innovation rests on two pillars: a fixed state assumption that redefines per iteration utility, and the Linearized Ghost Approximation that