Data Shapley in One Training Run

We thank the reviewer for the positive comments about our paper!

Q [Additional baseline comparisons] “The evaluation is performed against influence functions only ... both Memorization (Socher et al., 2022, Feldman and Zhang (2020)) and Supervised Prototypes (Socher et al., 2022) should be applicable to larger models ... It is specifically indicated that TRAK wouldn't scale to LLMs ...”

A: We thank the reviewer for the suggestion of having additional baseline comparisons, and actually this experiment was deferred to Appendix E.6 due to space constraints. In Appendix E.6, we evaluate In-Run Data Shapley on the standard tasks of mislabeled data detection and data selection, against a comprehensive set of baselines including Retraining-based Data Shapley, KNN-Shapley, Influence Functions, TRAK, Empirical Influence Functions, Datamodels, and LESS. The results demonstrate In-Run Data Shapley's competitive performance: In-Run Data Shapley's performance was among top-3 for almost all settings.

For “Memorization (Socher et al., 2022, Feldman and Zhang (2020)) and Supervised Prototypes (Socher et al., 2022) should be applicable to larger models and have been implemented on ImageNet scale by Socher et al.”, we thank the reviewer for mentioning these two works.

Memorization Scores (Feldman & Zhang, 2020): (1) Methodological Distinction: Memorization score is not a data attribution technique as it is independent of validation data, while the data attribution techniques aim to quantify the training data’s contribution to the model behavior on a validation point. The memorization score measures how much a model memorizes individual training points by computing the loss change on the training point itself when removed from training (we can think of it as replacing the validation data by the training data point itself). (2) The “influence score” proposed in (Feldman & Zhang, 2020) has already been compared with in Appendix E.6: The same paper also proposes a “influence score” for data attribution (referred to as “Empirical Influence Functions” in Appendix E.6), which evaluates contribution to validation performance and has been included in our experiments in Appendix E.6. Like other retraining-based methods, Empirical Influence Functions require multiple complete training runs (1000 in our experiments), making it impractical for foundation models where a single training run can take weeks.

Supervised Prototypes (Socher et al., 2022): (1) Methodological Distinction: Supervised Prototype is a technique for measuring data point difficulty through proximity to class distribution centers, independent of validation data. Similar to the Memorization score, this is not a data attribution technique for quantifying each training point's contribution to model performance on validation data. (2) Robustness to Noise: Although not a data attribution technique, we have additionally evaluated the performance of Supervised Prototypes on the data selection experiment in Appendix E.6. We want to stress that in our experiments, we inject noise by flipping labels for 10% of training data points to better simulate real-world scenarios with imperfect data quality. This setting differs from the original paper which assumes clean data (ImageNet is a very clean dataset for academic use). The results (shown in the last row) demonstrate that Supervised Prototypes perform poorly in such noisy settings. This is because it identifies mislabeled examples as "difficult" or "informative" cases that should be retained in the dataset. This limitation arises because Supervised Prototypes evaluate data points solely based on their relationship to class prototypes, without considering their actual impact on model performance. In other words, Supervised Prototype is unable to differentiate between genuinely difficult examples (that could benefit learning) and harmful data points (like mislabeled examples that hurt model performance). These data pruning methods are most effective on pruning pre-curated datasets where data quality issues have already been addressed, rather than raw datasets where distinguishing helpful from harmful data is crucial. In contrast, data attribution methods like ours leverage validation set performance to identify and assign negative values to harmful mislabeled data, making them more suitable for real-world scenarios where data quality cannot be guaranteed.

(Table copied from Appendix E.6 with other baselines' removed and Supervised Prototype added)

(continued from the previous response)

For “It is specifically indicated that TRAK wouldn't scale to LLMs, ... I would elevate my vote if the authors either compare their method to the ones mentioned above or provide reasons why those methods don't scale to the dataset.”, we thank the reviewer for mentioning TRAK paper.

Thanks for the great question!

While TRAK can indeed use checkpoints as suggested in their appendix, it actually faces significant scalability challenges when applied to foundation models. The main bottleneck is computational: computing per-sample gradients for each checkpoint effectively incurs similar costs as training with batch size 1 for an entire epoch. Since foundation models typically train for just one epoch, using multiple checkpoints becomes almost as expensive as performing multiple complete training runs. Additionally, TRAK's random projection step is memory-intensive for large models due to the gigantic projection matrix sizes.

While TRAK was already in our comparisons in Appendix E.6, during the rebuttal period we have additionally evaluated TRAK with checkpoint averaging (Trak (5 cpts), Trak (15 cpts), Trak (25 cpts) are being newly added during the rebuttal period). While increasing the number of checkpoints does improve TRAK's performance, it remains significantly below In-Run Data Shapley, which leverages information from all training iterations. We also observe diminishing returns in TRAK's performance as more checkpoints are added, which is consistent with findings from the original paper.

We sincerely thank the reviewer for the great question, which makes our experiments more comprehensive!

Q [Missing References]

A: We appreciate the reviewer highlighting relevant references.

Thank you for bringing up Toneva et al. (2018). We have cited their work in our related works section. While their approach of tracking "forgotten" examples during training shares some temporal aspects with our method, their focus is on measuring inherent example difficulty rather than data attribution with respect to validation performance. We have clarified this key distinction in our discussion and note their method's limitations to classification tasks.

We have cited Nguyen et al. (2023) in both our related works section and our discussion of Shapley value instability. While their work focuses on LOO-based data attribution methods (particularly influence functions) rather than Data Shapley's value instability, we agree it provides important context. We have included their work alongside other papers discussing influence function vulnerability (Basu et al. 2020, Søgaard et al. 2021) to provide a comprehensive view of stability challenges in data attribution methods.

Thank you for suggesting OpenDataVal. We have added citations to OpenDataVal (Jiang et al., 2024) in our related works. I have also cited (Deng et al., 2024) which is another data attribution library focusing on influence function style data attribution techniques.

Jiang, Kevin, et al. "Opendataval: a unified benchmark for data valuation." Advances in Neural Information Processing Systems 36 (2023).

Deng, Junwei, et al. "dattri: A Library for Efficient Data Attribution." Advances in Neural Information Processing Systems 36 (2024).

Toneva, Mariya, et al. "An Empirical Study of Example Forgetting during Deep Neural Network Learning." International Conference on Learning Representations. 2018.

Basu, S., P. Pope, and S. Feizi. "Influence Functions in Deep Learning Are Fragile." International Conference on Learning Representations (ICLR). 2021.

Søgaard, Anders. "Revisiting methods for finding influential examples." arXiv preprint arXiv:2111.04683 (2021). Nguyen, Elisa, Minjoon Seo, and Seong Joon Oh. "A Bayesian approach to analyzing training data attribution in deep learning." Advances in Neural Information Processing Systems 36 (2023).

Q [Typos]

A Thanks a lot for the catch! We have fixed the typo in the revised draft.

We thank the reviewer for the very positive assessment of our paper!

Q [Requirement of Validation Data]

A: Thank you for this thoughtful question about the validation data requirement. There are indeed several important scenarios where pre-available validation data could be challenging or infeasible: In online learning settings, where the data distribution may shift over time, pre-defined validation data might not effectively represent the current distribution. Similarly, in federated learning scenarios, privacy constraints might prevent access to validation data before training. Another challenging case is in few-shot or zero-shot learning tasks, where validation data from the target distribution is inherently scarce or unavailable.

For these scenarios, we envision several potential alternative approaches:

• First, as mentioned in Section 6, one could save model checkpoints during training and compute attribution scores once validation data becomes available. However, we acknowledge this approach's limitations - the choice of checkpoints can significantly impact performance, and storing multiple checkpoints for large models may be resource-intensive.
• One direction might be to develop proxy metrics that don't require explicit validation data. For instance, one could potentially use training data statistics, gradient information, or model uncertainty measures to construct surrogate utility functions. These could provide meaningful attribution scores without requiring pre-defined validation data, though careful theoretical work would be needed to establish their properties and relationship to validation-based attribution.
• Another possibility is to develop incremental updating schemes that can efficiently revise attribution scores as new validation data becomes available, rather than requiring all validation data upfront. This would be particularly valuable for online learning scenarios.

We believe exploring these alternatives represents an important direction for future research that could significantly broaden the applicability of our approach.

We agree that the requirement for validation data before training is a limitation of our current approach, and we explicitly acknowledged this in Section 6 of our paper. However, we would like to clarify that there are many practical scenarios where validation data is naturally available before training. These include working with public benchmark datasets that have established validation sets, participating in machine learning competitions where test metrics are known upfront, developing industrial applications where specific performance criteria are predetermined, and addressing regulatory requirements that specify validation metrics in advance. These examples where validation data are available upfront are not edge cases—they represent core challenges in deploying trustworthy ML systems. Our method provides an immediate and practical solution for these pressing use cases where data valuation can have a significant real-world impact.

Q [Computational efficiency for larger or different architectures?] "... the practicality of this approach in extremely large models with limited resources is not fully explored." "While the paper demonstrates efficiency improvements, ... models even larger than GPT-2 or to models with different architectures, such as transformer-based models with sparse attention mechanisms?"

A: The memory efficiency of our approach is particularly notable - we maintain constant memory overhead relative to regular training by never materializing per-sample gradient vectors. For computational efficiency, our first-order variant requires only one backpropagation pass, introducing minimal runtime overhead. During each gradient update iteration, while regular back-propagation has complexity $O(p)$, our ghost dot-product technique reduces the additional computation to $O(\sqrt{p})$ by separately computing dot products between activations and output derivatives rather than materializing full gradient vectors. This significant efficiency gain is achieved by leveraging information from regular training backpropagations, as detailed in Section 4.2.

We demonstrate this efficiency empirically in Section 5.1, where the first-order variant achieves 92.5% of regular training throughput on GPT2. For larger models (larger $p$), the relative overhead should decrease since regular training cost grows as $O(p)$ while our additional computation grows more slowly as $O(\sqrt{p})$. This favorable scaling property suggests our method becomes more efficient for larger models.

Distributed Training Considerations: However, we note that training large-scale models typically requires distributed computation across multiple GPUs. Adapting our method to such distributed training settings presents additional engineering challenges. For instance, efficiently computing gradient dot-products across model parallel shards requires careful asynchronous implementation to avoid blocking the regular training loop. While these challenges are primarily engineering in nature rather than theoretical limitations, due to our limited academic computing resources, we leave a thorough investigation of distributed implementations as important future work.

Architectural Flexibility: Although this work focuses on GPT2 and Pythia, our "ghost" techniques naturally extend to transformer-based models as they primarily consist of linear layers (query/key/value projections). For specialized variants like sparse attention, the core computations still involve linear transformations, making our method applicable with appropriate technical adaptations. For other neural network architectures involving non-linear layers such as convolutions, similar gradient decomposition methods have been developed in literature [1], suggesting our approach can also be extended to these architectures.

[1] Lee, Jaewoo, and Daniel Kifer. "Scaling up differentially private deep learning with fast per-example gradient clipping." PETS 2021

Q [Additional results for In-Run Data Shapley approximation error] “How does the use of first- and second-order Taylor approximations affect the reliability of the Shapley values in practice, especially in highly non-convex or complex loss landscapes?”

A: The accuracy of Taylor approximation is controlled by its remainder terms: $O(\eta^2)$ for first-order and $O(\eta^3)$ for second-order approximations, where $\eta$ is the learning rate. Since modern foundation model pretraining uses small learning rates (typically $\eta \leq 10^{-3}$), these remainder terms remain small, ensuring reliable approximations in non-convex landscapes.

Additional experiments on more learning rate choices. To thoroughly evaluate the impact of learning rate on approximation quality, we have conducted additional experiments during the rebuttal time (in Appendix E.2.1) varying the learning rate from our original $3 \times 10^{-4}$ to $6 \times 10^{-4}$ and $5 \times 10^{-3}$ (notably, the latter is significantly larger than typical learning rates used in LLM pretraining). While the mean squared error (MSE) increases with larger learning rates as expected, the Spearman rank correlation between approximation and (Monte Carlo-estimated) groundtruth remains remarkably high ($\approx 0.99$ for $\eta = 6 \times 10^{-4}$). Even with the extremely large learning rate of $\eta = 5 \times 10^{-3}$, we still maintain a strong Spearman rank correlation of $\approx 0.8$. This preservation of rank order is particularly important since most applications of data valuation (e.g., detecting low-quality data) rely on relative rankings rather than absolute values.

The practical effectiveness of our method is further validated through extensive experiments on downstream tasks detailed in Appendix E.6. Our method achieves comparable or better performance than most of the existing data attribution methods in both mislabeled data detection and data selection tasks. While these results do not directly compare the approximation with the groundtruth, they provide strong evidence that our approximation meaningfully captures data quality differences.

We thank the reviewer for this insightful question. The additional experiments and analyses have been included in our revised manuscript.

We thank the reviewer for the positive assessment of our paper!

Q [Example of how to use In-Run Data Shapley to calculate data contribution?] “Could you please provide an example on how your method used in-run data Shapley to calculate the data contribution? Is it possible to provide a figure?”

A: Thank you for this excellent suggestion. We have created new figures to illustrate the two key aspects of In-Run Data Shapley:

(1) Tracking data contributions throughout the training process: This figure shows the overall process of In-Run Data Shapley, which tracks data contributions throughout the training process. In each iteration, when a mini-batch of training data is selected (illustrated by the green/blue/brown arrows), we compute the contribution of each data point in that batch to the model's performance change for the single gradient update. These contributions are accumulated over time (like a “data value accountant”), as visualized by the growing bar chart at the bottom of each panel.

From the algorithmic perspective, the key idea of In-Run Data Shapley is to decompose a complete training process into individual gradient update steps. As shown in this figure, instead of evaluating data contribution across the entire training process at once (top row), we compute the Shapley value for each gradient update iteration separately (bottom row) and then aggregate them.

For each iteration $t$, we consider the "local utility function" $U^{(t)}$ that measures how much a subset of data points in batch $B_t$ contributes to reducing the validation loss in this specific update step. We calculate the Shapley value $\phi_z(U^{(t)})$ for each training point $z$ in the current batch $B_t$ with respect to this local utility function. For those $z \notin B_t$, their value will be 0 due to the null player axiom (see Section 3’s paragraph “Data Shapley for a single gradient update”).

The final contribution score for each data point $z$ is then computed as the sum of its Shapley values across all iterations where it appears: $\phi_z = \sum_{t:z\in B_t} \phi_z(U^{(t)})$. This decomposition approach makes the computation tractable while maintaining the key properties of the Shapley value through the linearity axiom.

(2) Single-update Shapley Calculation (exact calculation): Within each iteration, we compute the Shapley value for each data point in the current mini-batch using a standard Shapley value calculation. The only thing here is that the utility function is defined as the model's performance change in this single update step. For example, consider a simple binary classification scenario with a validation point and a mini-batch of two training points $\{1, 2\}$. The utility function $U(S)$ for any subset $S$ of the mini-batch would be the change in validation loss when only using points in $S$ for this gradient update. Suppose we observe:

• $U(\emptyset) = 0$ (no update)
• $U(\{1\}) = 0.3$ (gradient update on point 1 only reduces validation loss by 0.3)
• $U(\{2\}) = -0.2$ (gradient update on point 2 only increases validation loss by 0.2)
• $U(\{1,2\}) = 0.4$ (gradient update on both points reduces validation loss by 0.4)

The Shapley value for $z_1$ would be: $\phi_{z_1} = \frac{1}{2}[U(\{1\}) - U(\emptyset)] + \frac{1}{2}[U(\{1,2\}) - U(\{2\})]$ $= \frac{1}{2}(0.3 - 0) + \frac{1}{2}(0.4 - (-0.2)) = 0.35$

Similarly for $z_1$: $\phi_{z_2} = \frac{1}{2}[U(\{2\}) - U(\emptyset)] + \frac{1}{2}[U(\{1,2\}) - U(\{1\})]$ $= \frac{1}{2}(-0.2 - 0) + \frac{1}{2}(0.4 - 0.3) = -0.15$

(3) Single-update Shapley Calculation (Taylor Approximation):

However, computing $U(S)$ for all possible subsets becomes computationally intensive in practice. This motivates our development of efficient first-order and second-order approximations in Section 4, which allow us to compute these Shapley values using only gradient information without requiring actual model updates.

Specifically, in Section 4 we introduce Taylor approximations of the utility function and use them to derive the first- and second-order In-Run Data Shapley. As an example, for first-order In-Run Data Shapley, we proved that it is equal to the gradient dot product between the training and validation data. The figure illustrates the first-order In-Run Data Shapley values in a single iteration. Intuitively, if a training point's gradient aligns well with the validation gradient (like $z_1$ in the figure), it receives a positive value since updating the model in this direction would reduce the validation loss. Conversely, if the gradients point in opposite directions (like $z_2$), the point receives a negative value since it would increase validation loss.

Our paper further develops efficient "ghost" techniques to compute these gradient operations without materializing individual gradients, making the method highly practical for large-scale models.

We thank the reviewer for the very positive comments!

Q “I would like to know the authors' opinion of applying the general idea of In-Run Data Shapley to other use-cases in machine learning, such as hyperparameter importance during the HPO run”

A Thank you for this insightful question. We believe the core idea of In-Run Data Shapley—decomposing contribution analysis into per-iteration assessments—could be extended beyond data attribution to other applications in machine learning.

Hyperparameter Importance. ​​The framework could potentially be adapted to evaluate hyperparameter contributions during training. One possibility is to set a "baseline" hyperparameter value and assess how choosing a different value impacts each training iteration compared to this baseline. For instance, when evaluating learning rate choices, we could measure how using a specific learning rate value affects model updates compared to using a baseline learning rate. For differentiable hyperparameters, we could leverage Taylor expansion to approximate this difference; for non-differentiable ones, zero-order methods could potentially be used. By accumulating these contributions across training iterations, we could understand hyperparameter impact without requiring multiple complete training runs. At a high level, this view unifies our treatment of training data and hyperparameters - both can be seen as choices that influence each training iteration, where we aim to quantify their impact against baseline scenarios. While technical challenges remain in adapting our method, this direction presents an interesting opportunity for future research.

Feature attribution (e.g., context-attribution for language models). Another possible extension we envision is feature attribution for machine learning models. Feature attribution aims at understanding how each part of the input (like individual words in a sentence or pixels in an image) influences a model's final prediction. When a model makes a prediction, the input features are processed through multiple layers, with each layer transforming the information before passing it to the next layer. Feature attribution aims to track how this information flows and transforms through the network to determine each input feature's contribution to the final output. We envision a unified theoretical framework connecting training data attribution and feature attribution, drawing parallels between how information flows during training (through gradient updates) and during prediction (through layer operations). This could lead to more efficient methods for explaining model behavior, particularly for complex architectures like transformers.

Q “Do the authors envision transferring this idea to other domains, especially in the case where gradient information is not available?”

A Thank you for this interesting question!

(1) Consider scenarios where models are trained with gradient descent but the data attribution algorithm only has black-box access to intermediate checkpoints (e.g., when a third party wants to audit data influence during training). Here, zero-order methods could be used to estimate the impact of each training point by evaluating model behavior with and without that point in each iteration, without requiring access to gradients. While this approach may be computationally more expensive, it leverages the key insight of analyzing contributions iteratively rather than requiring complete retraining, and the technical adaptations can be interesting future works.

(2) Our framework could potentially extend to iterative learning algorithms that don't use gradient descent, such as k-means clustering or decision tree learning. Though these algorithms update models differently, they still proceed iteratively (e.g., k-means alternates between assignment and update steps; decision trees are built through recursive partitioning). For such algorithms, we could analyze how each data point influences these discrete update steps and aggregate these influences across iterations, similar to our approach with gradient-based training.

More broadly, we believe that the fundamental principle introduced in this work extends far beyond data attribution. As discussed in our previous response about feature attribution, this framework could offer new perspectives for understanding any iterative process where we want to quantify the impact of different components. We are excited about the potential applications of these ideas across various domains in machine learning and beyond. Thanks again for the great question!

Q [Typos]

A Thanks a lot for the catch! We have fixed the typo in the revised draft.

Thank you for your interest. You are correct that if $U(S)$ represented the utility of the entire end-to-end training process in the traditional sense, the decomposition $U(S) = \sum_t U^{(t)}(S)$ would of course not true in general. That's not what this paper is proposing. The core point is that In-Run Data Shapley defines its "global" utility function $U(S)$ differently than the traditional Retraining-based Data Shapley.

In the original Data Shapley literature, $U(S)$ typically defines utility as the final performance (e.g., accuracy, loss) of a model trained on a data subset $S$. However, things become more complicated in the context of deep learning. A specific deep learning training run is an iterative process. Defining utility $U(S)$ for a specific run makes it a sequence function rather than the pure set function required by the standard Shapley value framework. Extending Shapley axioms to sequence functions is not straightforward. The previous works often define $U(S)$ as "expected model utility across all possible training runs" instead to circumvent this conceptual issue. One contribution of this paper is initiating the study of "model-specific data attribution". This is important for applications such as training diagnosis or model decision interpretation, where you care about a specific training run instead of the expected model utility across all possible training runs.

Our paper provides a novel approach by decomposing the training process into individual iterations, where each "local utility function $U^{(t)}$" is a set function, allowing efficient Shapley value approximation. We then define the "global utility function" as $U(S) := \sum_t U^{(t)}(S)$. In other words, our $U(S)$ has a different definition compared to the original. We interpret In-Run Data Shapley as "the cumulative data contribution throughout the training run." From a game theory perspective, this approach can be regarded as a special case of "asymmetric Shapley value" (https://arxiv.org/abs/2411.00388). It also closely relates to "instrumental value" (https://arxiv.org/pdf/2412.18140) which considers context-dependent valuation.

We hope this makes sense. Please feel free to ask if there are any remaining questions. In a follow-up work (also published at ICLR'25), we did some further studies on this "model-specific data attribution" problem, in case you are interested.

PS: We appreciate your thoughtful inquiry. However, we believe the title 'Is there a theoretical mistake in this paper?' may set an unnecessarily critical tone that could influence other readers' perceptions, especially if they do not have enough time to read. We kindly suggest revising the title to a more neutral framing such as 'Question on utility decomposition' or "Question on theoretical foundations". This would better reflect the constructive nature of academic exchange.