Table of Contents
Fetching ...

In-Run Data Shapley for Adam Optimizer

Meng Ding, Zeqing Zhang, Di Wang, Lijie Hu

TL;DR

This work demonstrates that data value is intrinsically tied to the optimizer used during training, showing that SGD-based In-Run Shapley proxies poorly reflect true contributions under Adam ($R \approx 0.058$, $\rho \approx 0.046$). It introduces Adam-aware In-Run Data Shapley with a closed-form estimator that accounts for Adam's stateful momentum and variance scaling, and a Linearized Ghost Approximation to enable scalable computation without materializing per-sample gradients. Experiments reveal near-perfect fidelity to ground-truth marginal contributions ($R > 0.99$) and maintain approximately 95% of standard training throughput, while SGD-based baselines underperform in data attribution downstream tasks like semantic source identification and SST-2 pruning. These results establish optimizer-aware data attribution as both theoretically necessary and practically feasible for large-scale adaptive optimization.

Abstract

Reliable data attribution is essential for mitigating bias and reducing computational waste in modern machine learning, with the Shapley value serving as the theoretical gold standard. While recent "In-Run" methods bypass the prohibitive cost of retraining by estimating contributions dynamically, they heavily rely on the linear structure of Stochastic Gradient Descent (SGD) and fail to capture the complex dynamics of adaptive optimizers like Adam. In this work, we demonstrate that data attribution is inherently optimizer-dependent: we show that SGD-based proxies diverge significantly from true contributions under Adam (Pearson $R \approx 0.11$), rendering them ineffective for modern training pipelines. To bridge this gap, we propose Adam-Aware In-Run Data Shapley. We derive a closed-form approximation that restores additivity by redefining utility under a fixed-state assumption and enable scalable computation via a novel Linearized Ghost Approximation. This technique linearizes the variance-dependent scaling term, allowing us to compute pairwise gradient dot-products without materializing per-sample gradients. Extensive experiments show that our method achieves near-perfect fidelity to ground-truth marginal contributions ($R > 0.99$) while retaining $\sim$95\% of standard training throughput. Furthermore, our Adam-aware attribution significantly outperforms SGD-based baselines in data attribution downstream tasks.

In-Run Data Shapley for Adam Optimizer

TL;DR

This work demonstrates that data value is intrinsically tied to the optimizer used during training, showing that SGD-based In-Run Shapley proxies poorly reflect true contributions under Adam (, ). It introduces Adam-aware In-Run Data Shapley with a closed-form estimator that accounts for Adam's stateful momentum and variance scaling, and a Linearized Ghost Approximation to enable scalable computation without materializing per-sample gradients. Experiments reveal near-perfect fidelity to ground-truth marginal contributions () and maintain approximately 95% of standard training throughput, while SGD-based baselines underperform in data attribution downstream tasks like semantic source identification and SST-2 pruning. These results establish optimizer-aware data attribution as both theoretically necessary and practically feasible for large-scale adaptive optimization.

Abstract

Reliable data attribution is essential for mitigating bias and reducing computational waste in modern machine learning, with the Shapley value serving as the theoretical gold standard. While recent "In-Run" methods bypass the prohibitive cost of retraining by estimating contributions dynamically, they heavily rely on the linear structure of Stochastic Gradient Descent (SGD) and fail to capture the complex dynamics of adaptive optimizers like Adam. In this work, we demonstrate that data attribution is inherently optimizer-dependent: we show that SGD-based proxies diverge significantly from true contributions under Adam (Pearson ), rendering them ineffective for modern training pipelines. To bridge this gap, we propose Adam-Aware In-Run Data Shapley. We derive a closed-form approximation that restores additivity by redefining utility under a fixed-state assumption and enable scalable computation via a novel Linearized Ghost Approximation. This technique linearizes the variance-dependent scaling term, allowing us to compute pairwise gradient dot-products without materializing per-sample gradients. Extensive experiments show that our method achieves near-perfect fidelity to ground-truth marginal contributions () while retaining 95\% of standard training throughput. Furthermore, our Adam-aware attribution significantly outperforms SGD-based baselines in data attribution downstream tasks.
Paper Structure (32 sections, 3 theorems, 21 equations, 5 figures, 3 tables)

This paper contains 32 sections, 3 theorems, 21 equations, 5 figures, 3 tables.

Key Result

Lemma 3.1

For any of two utility functions $U_1, U_2$ and any $\alpha_1, \alpha_2 \in \mathbb{R}$, we have $\phi_z\left(\alpha_1 U_1+\alpha_2 U_2\right)=\alpha_1 \phi_z\left(U_1\right)+\alpha_2 \phi_z\left(U_2\right)$.

Figures (5)

  • Figure 1: Comparison between SGD-based and Adam-based data Shapley values (Pearson $R = 0.0579$, Spearman $\rho = 0.0465$). Each point corresponds to a training sample. The dashed line indicates perfect agreement ($y = x$), while the dispersion reveals strong optimizer dependence.
  • Figure 2: Optimizer dependence and fidelity of data attribution. (Left) Optimizer mismatch leads to significant fidelity degradation: attribution scores computed assuming SGD fail to recover the true marginal utility when the training optimizer is Adam. (Middle) When the optimizer is correctly accounted for, Adam-aware In-Run Shapley accurately recovers ground-truth marginal utility, while SGD-based proxies exhibit substantial deviation. (Right) The linearized Ghost approximation closely matches the exact Adam-aware attribution, demonstrating that substantial efficiency gains can be achieved with minimal attribution error.
  • Figure 3: Fidelity across learning rates. Pearson correlation between predicted scores and ground-truth utility changes across different learning rates $\eta$. Adam-aware approximation maintains high fidelity ($R>0.96$) across the board, while the SGD-proxy exhibits unstable and lower correlation.
  • Figure 4: Data pruning effect: validation loss versus training steps. We compare the full dataset, random pruning (20%), and In-Run Shapley pruning (20%).
  • Figure 5: Empirical validity of the ghost approximation. (Left) Evolution of the effective perturbation magnitude $u$, showing that optimization rapidly enters and remains in a small-perturbation regime. (Right) Worst-case fidelity test conducted during the first 10 optimization steps, where $u$ is largest. Even under these early and high-magnitude conditions, the ghost approximation maintains strong agreement with the exact computation (Pearson $r=0.7502$, Spearman $\rho=0.8461$), indicating robust behavior beyond its nominal operating regime.

Theorems & Definitions (4)

  • Lemma 3.1: Linearity of the Shapley value shapley1953value
  • Theorem 4.1
  • Theorem 1.1: Restatement of Theorem \ref{['thm:adam_sv']}
  • proof