Data Shapley in One Training Run

Authors: Jiachen T. Wang, Prateek Mittal, Dawn Song, Ruoxi Jia
Primary area: interpretability and explainable AI
Keywords: Shapley value, data valuation.

Abstract:
Data Shapley offers a principled framework for attributing the contribution of data within machine learning contexts. However, the traditional notion of Data Shapley requires re-training models on various data subsets, which becomes computationally infeasible for large-scale models. Additionally, this retraining-based definition cannot evaluate the contribution of data for a specific model training run, which may often be of interest in practice. This paper introduces a novel concept, In-Run Data Shapley, which eliminates the need for model retraining and is specifically designed for assessing data contribution for a particular model of interest. In-Run Data Shapley calculates the Shapley value for each gradient update iteration and accumulates these values throughout the training process. We present several techniques that allow the efficient scaling of In-Run Data Shapley to the size of foundation models. In its most optimized implementation, our method adds negligible runtime overhead compared to standard model training. This dramatic efficiency improvement makes it possible to perform data attribution for the foundation model pretraining stage. We present several case studies that offer fresh insights into pretraining data's contribution and discuss their implications for copyright in generative AI and pretraining data curation.

Decision: accept (oral)
Reviewer scores: 10, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730404284836

thread_type:
official_review

<thread_object>
creator: reviewer_vb7n

<content>
summary:
The paper introduces in-run Shapley, a data valuation method for SGD based models that estimates data values in one training run across different update steps. 

The authors further provide enhanced ways of computing these data values with first and second order approximations. 

The method is evaluated on the PILE dataset for GPT-2 and Pythia 4.10 against influence functions and outperforms them.

Remark: Thanks to the clarifications provided by the authors during rebuttal, I have raised Soundness from 2 to 3, Contribution from 3 to 4, and the overall score from 5 to 6.

strengths:
The paper introduces an interesting and to me novel method to estimate data values in a single training run. This makes data valuation feasible to large scale data sets. Furthermore, the method is evaluated on such a data set.

weaknesses:
The evaluation is performed against influence functions only. This way it is unclear whether the method only outperforms leave-one-out (approximated by influence functions) or is also better than the standard data Shapley value. 

Other data valuation/pruning methods are dismissed due to scalability issues (Sec 5.3.1). However, to my understanding, both 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. I'm uncertain if this generalizes to LLMS, but it should be discussed if not feasible.

It is specifically indicated that TRAK wouldn't scale to LLMs, yet according to my understanding, TRAK can function with checkpoints from a limited number of models (see Appendix "E.3 Proxies for model ensembles in compute-constrained settings" of the TRAK paper).
Given the evaluation of in-run Shapley on large datasets, 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.

Minor Remarks
--------------------
OpenDataVal is a recent benchmark for data valuation. Although it is designed for small-scale data, it offers ready-to-use implementations of various algorithms and could be employed to benchmark In-Run Shapley against other methods. This would enhance the paper by demonstrating whether In-Run Shapley performs comparably or even superior to the standard Shapley value.
Some references are missing:
* Toneva et al. (2019) demonstrated that data points learned early in training and not forgotten in later epochs are less important to model performance than points that are frequently forgotten. The former can be omitted to reduce dataset size.. This reminds of the method introduced in the paper. 
* Nguyen et al. (2023) proposed "A Bayesian Approach To Analysing Training Data Attribution In Deep Learning" to tackle the noise issue when training multiple models. This problem, referred to as "instability," is also mentioned in this paper but without citing Nguyen et al.

There are a few spelling mistakes especially in Sec. 5.3.3:
* In "data set," the word "set" is missing.
* The sentence "The results of this experiment have important implications for the current discussion on the copyright of generative AI." is repeated.

Sources
-----------
Remark: the reviewer is not the author of any of the papers below.
* Socher et. al (2022): Beyond neural scaling laws: beating power law scaling via data pruning
* Feldman and Zhang (2020): What neural networks memorize and why: Discovering the long tail via influence estimation
* Toneva et. al (2019): An Empirical Study of Example Forgetting during Deep Neural Network Learning
* OpenDataVal (2023): OpenDataVal: a Unified Benchmark for Data Valuation
* Nguyen et. al (2023): A Bayesian Approach To Analysing Training Data Attribution In Deep Learning

questions:
I want to confirm whether I have understood the definition of the in-run Shapley value for a data point $z$. It is the “cumulative contribution of the data point $z$ across all gradient update iterations within a single training run”. Does this mean that in each gradient update (e.g., for each batch) the Shapley data value w.r.t. to this batch is computed and averaged at the end? This would mean that for each batch $b$ the marginal contribution of $z$ would be estimated on all possible subsets of $b$. 

If so, I would be interested in an intuition how this value is related to the standard Shapley data value. This could be done conceptually but also by experimental evidence, e.g., by comparing the method to other visiting Shapley based estimates, as suggested above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730602438791

thread_type:
official_review

<thread_object>
creator: reviewer_zck2

<content>
summary:
The paper presents a novel approach called In-Run Data Shapley that eliminates the need for model retraining to evaluate data contributions in machine learning. By leveraging gradient updates within a single training run, the authors develop a computationally efficient method suitable for large-scale models, such as foundation models like GPT-2. The work introduces efficient algorithms using Taylor expansions and techniques inspired by differential privacy to minimise runtime overhead. The study also demonstrates the practical utility of In-Run Data Shapley in data curation and understanding the impact of pretraining data.

strengths:
* The introduction of In-Run Data Shapley is a significant advancement over existing retraining-based methods, offering a scalable solution to data attribution in large models.
* The case studies, particularly in data curation and copyright implications, highlight the real-world relevance of the proposed method.
* The closed-form derivations for the Shapley values using first- and second-order Taylor approximations are mathematically sound and well-explained.

weaknesses:
* The method's requirement for validation data to be available before training might limit its applicability in scenarios where such data is not readily available.
* Even though the authors address memory concerns using gradient accumulation, the practicality of this approach in extremely large models with limited resources is not fully explored.
* The first- and second-order Taylor approximations are empirically validated, but there is limited discussion on how these approximations might affect the reliability of the data attribution scores, especially in more complex models.

questions:
* The method requires validation data to be available before training. Can the authors elaborate on the potential limitations of this requirement and discuss scenarios where this may not be feasible? Are there alternative approaches that could make In-Run Data Shapley applicable when validation data is not pre-available?
* 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? Can the authors provide additional empirical results or theoretical analysis to quantify these approximation errors in diverse settings?
* While the paper demonstrates efficiency improvements, how does the proposed method scale when applied to models even larger than GPT-2 or to models with different architectures, such as transformer-based models with sparse attention mechanisms? Can the authors provide insights or preliminary results on this?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730632609991

thread_type:
official_review

<thread_object>
creator: reviewer_wuod

<content>
summary:
This paper presents a new technique called "in-run data shapley" to eleminate model retraining for assessing data contributions. This seems a new angle to assess the data attribution and importance. The results shows that this can provide efficiency improvement, and help make more insights for data contribution. The paper also identifies and discussed potential implications for copyright in genAI and pretraining data curation.

strengths:
1. The motivations, problem statement and objectives and contributions are provided clearly. 

2. The mathematical induction and proof seem solid.

3. The results and implications of the new concept/techniques are described clearly.

weaknesses:
Compared with the introduction, background, methods and mathematical induction, which are presented clearly, the results and analyses are a bit weak. It would be better to provide concrete examples with more deep analyses.

questions:
1. 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? Even using a simple example data set would be very helpful. This cannot be replaced with a loss-time curve. 

2. Is it possible to extend the method to regression and classification problems?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731162612062

thread_type:
official_review

<thread_object>
creator: reviewer_row8

<content>
summary:
This paper introduces In-Run Data Shapley, a novel principled way of attributing the data contribution for deep learning models within one training run. A classical Data Shapley technique requires retraining which is computationally prohibitive in practice for large models, and more importantly, it quantifies data contribution for a general learning algorithm, and not for a particular model at hand. Related to this, data contribution for a single model is of crucial importance for properly investigating fairness and copyright implications.

In-Run Data Shapley is powered by the fact that model training is performed iteratively, and Shapley values can be derived analytically for each model update step via gradient dot products or gradient-Hessian-gradient products between training and validation data. The authors propose an efficient method to approximate In-Run Data Shapley values during one training run via first- and second-order Taylor expansions, and compute these values in one and two backpropagation, respectively, inspired by the ghost clipping technique. This incurs negligible computational overhead, making In-Run Data Shapley highly usable in practice. Furthermore, the authors empirically evaluate their approach on Pile dataset, highlighting several important insights re: data curation and cleanliness of the dataset, dynamic data contribution throughout the entire training process, and the degree of data contribution in copyright issues.

strengths:
The proposed method is novel, efficient and promising. The presentation of the paper is very clear and succinct; technical contributions are clearly stated and rigorously developed. The authors carefully position their work and acknowledge several limitations which they claim will be a part of future work. Empirical evaluation is thorough and the paper raises important observations re: societal implications of the large model training.

weaknesses:
I do not see any major weaknesses of the paper.
Minor (typo): line 500 - sentence "The results of this experiment..." appears twice in a row

questions:
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. More generally, do the authors envision transferring this idea to other domains, especially in the case where gradient information is not available?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732562490441

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank all reviewers for their thoughtful feedback and constructive suggestions. We are glad that our work received very positive feedback! We have carefully read the reviews and made substantial modifications to strengthen the paper. All modifications are highlighted in blue. Here's a summary of our major revisions:

- **Additional discussion on the potential of the general idea of In-Run Data Shapley (Reviewer roW8).** We have incorporated our discussion on the possible directions of generalizing the core idea of In-Run Data Shapley in **Appendix B.3**.

- **Additional figure on the core idea and algorithm overview of In-Run Data Shapley (Reviewer wuoD).** We have incorporated an additional Figure in **Appendix B** to better illustrate the core idea and algorithm of In-Run Data Shapley.

- **Additional experiments on Taylor approximation errors under various learning rates (Reviewer ZCk2).** We have conducted additional experiments during the rebuttal period in **Appendix E.2.1** varying the learning rate to systematically evaluate and demonstrate the robustness of our algorithm.

- **Additional baseline comparisons (Reviewer Vb7n).** We have conducted new experiments evaluating TRAK with different numbers of checkpoints. We have also revised the main text to highlight our comprehensive experiments comparing against a wide range of established baselines on standard benchmarks of mislabeled data detection and data selection in **Appendix E.6**.

- **Discussion on the comparison between retraining-based and In-Run Data Shapley (Reviewer Vb7n).** We have revised and enriched our discussion on the comparison between Retraining-based and In-Run Data Shapley in **Appendix B.1 and B.2**.

- **Typos & Missing References (Reviewer roW8, Vb7n).** We have fixed all of the them. Thanks so much for the catch!

As the discussion period is coming to an end, if there are any additional questions/comments about our work, we are more than happy to address them!
</content>
</thread_object>

</thread>
<thread>
creation date:
1744364401547

thread_type:
public_comment

<thread_object>
creator: ~yaoxin_ge1

<content>
title:
Question on utility decomposition.

comment:
The main idea of the In-Run Data Shapley is $U(S)=\sum_t U^{(t)}(S)$, but I found this decomposition is wrong following the construction of $U^{(t)}(S)$. Consider there are two data points $\\{A,B\\}$, and the model is updated with $A$ in the first iteration and with $B$ in the second iteration, $U(S)$ is the utility function introduced in Section 2.1. 
Then, in the first iteraion:

$U^{(1)}(\\{A\\})=\text{The performance of using \\{A\\} to train}-0=U(\\{A\\})$,

$U^{(1)}(\\{B\\})=0$,

$U^{(1)}(\\{A,B\\})=\text{The performance of using} \\{A\\}\cap \\{A,B\\} \text{ to train}-0=U(\\{A\\})$,

and, in the  second iteration:

$U^{(2)}(\\{A\\})=0$,

$U^{(2)}(\\{B\\})=\text{The performance of using \\{B\\} to update}-\text{The performance of using \\{A\\} to train}=U(\\{A,B\\})-U(\\{A\\})$,

$U^{(2)}(\\{A,B\\})=U(\\{A,B\\})-U(\\{A\\})$,

so $U^{(2)}(B)+U^{(1)}(B)\neq U(B)$. That means the linear combination does not match the utility function, and in this example, the In-Run Data Shapley is just calculating the marginal contribution for each data point. Did I miss or misunderstand something?
</content>
</thread_object>

</thread>
<thread>
creation date:
1749610258761

thread_type:
public_comment

<thread_object>
creator: ~allen_schmaltz1

<content>
title:
"Instance/prediction-specific" data attribution methods

comment:
An older approach for data attribution with the models with non-identifiable parameters (e.g., LLMs) is to make a semi-supervised connection to the observed data, conditional on the output prediction, by adding a bottleneck ("exemplar") layer to the model, and re-casting the prediction as a function over the training set's labels and representation-space via a metric-learner approximation. How do we know that the matched exemplars are actually relevant, or equivalently, that the approximation is faithful to the original model? One simple (but meaningful) metric is whether the prediction of the metric-learner approximation matches the class of the prediction of the original model, and if they do not, the discrepancies should be concentrated in low probability regions. Remarkably, relatively simple functions over the representation space and labels achieve that behavior. (In this way, such metric-learner approximations are a fundamentally different concept than Shapley-based metric-learner approaches, and provide an alternative approach for addressing the questions in the final sentence of Appendix B.1.1. Since In-Run Data Shapley requires validation sets, such alternative approaches are applicable in the same settings.) This line of work was introduced in the following paper, which appeared in the journal Computational Linguistics and was presented at EMNLP 2021: "Detecting Local Insights from Global Labels: Supervised & Zero-Shot Sequence Labeling via a Convolutional Decomposition" https://doi.org/10.1162/coli_a_00416.

Note that these older methods also address feature attribution, a proposed future extension of In-Run Data Shapley mentioned in Appendix B.3. When projection to a lower resolution than the available labels is desired for data analysis (e.g., from the document level to the word level for error detection), the aforementioned method provides a straightforward means of defining the applicable inductive bias via the loss (e.g., encouraging one feature detection per document, or multi-label detections, as shown in subsequent work) over the exemplar layer decomposition. In other words, we can decompose an LLM's document-level prediction to word-level feature detections, and then map to the labeled data in the support set. This works with models of arbitrary scale; the dense matching via the distilled representations from the filter applications of the exemplar layer requires computation on the order of commonly used dense retrieval mechanisms, and the other transforms require minimal additional overhead.

The punchline is that the key advantage relative to alternative data-attribution methods is this then leads to methods with which we can close the loop on the connection between the data, the representation space, [the features,] the predictions, and the predictive uncertainty at the instance level. See my works from 2022 and more recently for further details (e.g., https://arxiv.org/abs/2502.20167), including incorporating uncertainty-aware verification and interpretability-by-exemplar as intrinsic properties of the LLMs, themselves (i.e., for both train- and test-time search/compute/generation), rather than as a post-hoc approximation.

I would argue that data attribution tied to estimators of the predictive uncertainty is more principled than achieving the Shapley properties 1-3 described on page 1 (even if they were exactly achieved, rather than via an approximation), since it directly translates to the quantities of interest needed when using LLMs in practice. It also then becomes possible to estimate class-conditional contributions and to avoid marginalizing over distinctions across high/low-uncertainty partitions and high/low-variance partitions of the data, with concomitant benefits for constraining the output, or e.g., making targeted expansions of the support for a given instance. That is, in practice, we typically seek data attribution that is instance/prediction-specific, rather than algorithm-specific or model-specific (using the terminology of this paper).
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents an approach aimed to evaluate the contributions of the data samples to the model. 
This goal (attributing a relevance score to data samples w.r.t. the eventual model) is at the core of scientific, technical and legal concerns (privacy, copyright). 
The method is based on the accumulation of Shapley value at each gradient iteration, with affordable memory and computational extra requirements.  
One limitation is the availability at training time of validation samples.