Efficient Ensembles Improve Training Data Attribution

We thank the reviewer for the comments and address them below.

Weakness 1: Limited originality

To the best of our knowledge, this paper is the first study to propose avoiding fully independent training to achieve better efficacy-efficiency tradeoffs in TDA ensembling, while also discovering the intriguing connections among dropout, LoRA fine-tuning, and ensemble methods in the context of TDA. Going beyond the traditional (naive) ensemble is a surprising finding.

The methods we propose are well-grounded in theoretical and methodological foundations. Dropout Ensemble is motivated by the fact that Dropout was originally designed as an efficient approximation of ensemble learning, and LoRA Ensemble is directly inspired by LoRA fine-tuning as an efficient alternative to full training.

Most importantly, we would like to emphasize that in addition to our contribution on connecting Dropout/LoRA and efficient ensembling, our methods demonstrate significant performance improvements and practical applicability, making them valuable for real-world scenarios.

Weakness 2 and Question 1: Small models and datasets

We would like to highlight that our experiment settings in the original submission have covered models with tens of millions of parameters and two generative settings, i.e., the MusicTransformer (13.3M) and GPT-2 (124M) setting (added below). The experiments already involve large modern architectures and complex datasets. As shown in the table below, Dropout Ensemble could achieve comparable performance without independent retrained models as Naive Ensemble with multiple independent retrained models on GPT-2 + WikiText-2. Some additional experiments lie in Appendix K.

Quantitative evaluation of data attribution of large scale generative models is a challenging task since most of the metrics (e.g., LDS, brittleness, …) need a large number of model retraining to get the ground truth. We note that our experimental setups already match the largest setups in most of the recent papers [1][2] in the field of data attribution.

Weakness 3: Only one TDA method experimented on for LoRA

LoRA is designed for large transformer-based models, so we only perform experiments on the transformer architecture. Other than MusicTransformer, we also show the results on GPT in Figure 8 of Appendix K.

Weakness 4 and Question 2: Theory

In our theoretical analysis, we did not claim that the within-group covariance is never significantly larger than the covariance between individual predictors. In contrast, we acknowledged in line 525 that our advantage will diminish if this covariance gap is huge. However, our goal in Case 2 was to illustrate that our method can help reduce the error (negative $\Delta$) significantly, by a factor of 1/K. When we say a large K is needed for good empirical performance, we mean a relatively large number (such as 10 in our experiments). In practice, this error reduction helps to maintain a competitive performance while providing a drastic efficiency gain, as reflected in our experiments with very small K <= 5. Such results show that the covariance gap is reasonably controlled. Specifically for the MLP model on MNIST, a small K = 5 with Dropout Ensemble can outperform the naive ensemble with 25 models and save 80% of training time cost.

In reality, there is always a trade off between performance and efficiency for any ensemble method. One key advantage of our approach is to provide flexibility in balancing them, allowing users to adjust K based on their specific requirements. Smaller Ks can be used to address scenarios where efficiency is a primary concern. On the other hand, when users can afford a larger K, which is not required and may not be the most efficient trade-off depending on computational constraints, the performance improves further. Finally, our Case 1 also shows that when there is a fixed budget of how much model training can be afforded in total, our method can make a very effective trade off by doing cheap dropout and achieve performance enhancement.

References

[1] Park, S.M., Georgiev, K., Ilyas, A., Leclerc, G. & Madry, A.. (2023). TRAK: Attributing Model Behavior at Scale. <i>Proceedings of the 40th International Conference on Machine Learning</i>, in <i>Proceedings of Machine Learning Research</i> 202:27074-27113 Available from https://proceedings.mlr.press/v202/park23c.html.

[2] Bae, J., Lin, W., Lorraine, J., & Grosse, R. B. Training Data Attribution via Approximate Unrolling. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.

Dear Reviewer DqZM,

Thank you for your valuable feedback to our submission! We have carefully addressed your comments in our rebuttal and we would greatly appreciate it if you could review it and consider any further discussion. We look forward to hearing any additional thoughts you may have.

We thank the reviewer for the comments and address them below.

Weakness 1: Performance improvements in the abstract

The improvements in abstract are claimed with the performance (LDS) being controlled to be the same. They are the maximum improvements (“up to”) of each type of cost that can be achieved separately. Please see the following places in our paper for the settings where these improvements are achieved.

• Dropout Ensemble (line 405-406 in the updated paper): 80% for training and space cost.
• LoRA Ensemble (line 424-425 in the updated paper): training/serving/space by 78.4%, 60.5% and 79.6%.

Weakness 2: Comparison with using multiple checkpoints from a single training run.

Thank you for pointing this out. We have noticed that using multiple checkpoints from a single training run could also be a baseline (mentioned in TRAK [1]). We include some results in Table 2 in Appendix G.2. Dropout Ensemble could still lead to significantly better LDS.

Weakness 3: Using independently trained model and diminishing return

Empirically, the asymptoting phenomenon and diminishing return happen to both efficient ensemble methods and naive ensembles. Experiments from our paper and TRAK [1] both demonstrate this phenomenon. Our experiments show that efficient ensembling can be used with only one independently trained model and achieve much better efficacy of TDA. A good combination of the number of independently trained models and the number of dropout masks/LORA fine-tunes could perform as the Pareto frontier.

Weakness 4 & Question 1: The overall computational cost (training+serving) of Dropout Ensemble.

We provide the table below (newly added as Table 8 in the PDF) to demonstrate that the training cost dominates the computational cost, i.e., the ratio of training cost over serving cost for one ensemble model is large. We also provide a figure (newly added as Figure 9 in the PDF) with performance (LDS) against the total computational cost (training+serving time) to show the improvement of the Dropout Ensemble over the naive independent ensemble. From the figure, it’s clear that Dropout Ensemble could bring better performance under the same computational cost in every experiment setting. The forward-only variant could further improve the performance and reduce time cost, especially for larger models.

Note that although the computational time cost may be influenced by the implementation of model training script and data attribution method, all of our comparisons across different TDA methods are under the same settings and are thus fair.

It is worth noting that the training cost and the serving cost are both for the TDA process, which are defined respectively in Section 3.2. The computational cost of TDA should consider both costs together.

Weakness 5: The competitive performance of Dropout Ensemble (forward only).

Thank you for pointing this out. We acknowledge that this phenomenon is counter-intuitive, and we do not have a thorough understanding of it.

A potential explanation to this phenomenon is that the term (Q) affects the performance “substantially” as shown in TRAK [1] paper’s ablation study. Focusing on mitigating the effect of randomness on this term and cache other terms could improve the TDA performance effectively and efficiently.

It is worth noting that this phenomenon may be closely related to some open questions existing in the TRAK algorithm. For instance, empirical evidence indicates that a particular approach to ensembling models, known as "term-wise” ensemble in the TRAK paper, plays a critical role in TRAK's superior performance. However, the underlying reasons for this remain insufficiently understood.

A thorough understanding of this particular phenomenon pointed out by the reviewer may require answers to the open questions existing in TRAK, which go beyond the scope of this paper. We will try to further investigate it in future work. We have added a brief discussion about this phenomenon in our updated paper in Appendix P.

Weakness 6: The improvement over serving cost of LoRA Ensemble.

For LoRA Ensemble, we will only use the parameters of the LoRA adapters to calculate the TDA scores instead of full parameters in the original model. This is the main reason why the serving cost has been reduced. We have added a sentence to make this clear in the main text (line 265-269).

References

[1] Park, S.M., Georgiev, K., Ilyas, A., Leclerc, G. & Madry, A.. (2023). TRAK: Attributing Model Behavior at Scale. <i>Proceedings of the 40th International Conference on Machine Learning</i>, in <i>Proceedings of Machine Learning Research</i> 202:27074-27113 Available from https://proceedings.mlr.press/v202/park23c.html.

We thank the reviewer for the comments and address them below.

Weakness 1 & Question 2: The limitation of efficient ensemble methods

Thank you for pointing out this important aspect. We agree that discussing the limitations of our methods is valuable, and we have included a detailed quantitative table in Appendix L, demonstrating the computation and space overhead. Additionally, we have provided further discussion in the main text to highlight these points in Section 3.3.

We provide the table below (newly added as Table 8 in the PDF) to demonstrate that the training cost dominates the computational cost, i.e., the ratio of training cost over serving cost for one ensemble model is large. We also provide a figure (newly added as Figure 9 in the pdf) with performance (LDS) against the total computational cost (training+serving time) to show the improvement of the Dropout Ensemble over the naive independent ensemble. From the figure, it’s clear that Dropout Ensemble could bring better performance under the same computational cost in every experiment setting. The forward-only variant could further improve the performance and reduce time cost, especially for larger models.

Note that although the computational time cost may be influenced by the implementation of model training script and data attribution method, all of our comparisons across different TDA methods are under the same settings and thus fair.

Weakness 2: Comparison between lora and dropout based methods

A comparison between Dropout and LoRA Ensemble can be found in Figure 9 (d) (newly added to the paper in Appendix M). In terms of LDS, Dropout Ensembles yield slightly better improvement compared with LoRA Ensembles and both ensembles can significantly outperform naive ensembles for transformer-based large models by achieving a better efficacy-efficiency tradeoff.

Dropout and LoRA Ensemble have different trade-offs regarding training and serving costs. The suitability of them may vary depending on the specific applications.

Question 1: training and test set in supervised case

We clarify that the notations $X_{train}$ and $X_{test}$ refer to the data SPACE instead of the datasets. In lines 126-127 in our original draft, we are only saying that the training data and test data are typically sampled from the same space. However, the sampled training and test datasets will be different.

We thank the reviewer for the comments and address them below.

Weakness 1: Comparison between Dropout and LoRA Ensemble.

A comparison between Dropout and LoRA Ensemble can be found in Figure 9 (d) (newly added to the paper in Appendix M). In terms of LDS, Dropout Ensembles yield slightly better improvement compared with LoRA Ensembles and both ensembles can significantly outperform naive ensembles for transformer-based large models by achieving a better efficacy-efficiency tradeoff.

Dropout and LoRA Ensemble have different trade-offs regarding training and serving costs. The suitability of them may vary depending on the specific applications.

Weakness 2: Comparison between Dropout Ensemble and TracIN.

Thank you for pointing out the connection. Dropout Ensemble is also available for the TracIN. The result in Figure 4 for Grad-Dot could be seen as a special and efficient case to TracIN-CP. Dropout Ensemble could also improve its performance. The result in Table 2 (Appendix G.2) also shows that Dropout Ensemble could further improve the result by using intermediate checkpoints. Additionally, Dropout Ensemble does not require additional space cost for the checkpoints needed by TracIN.

Weakness 3: LORA ENSEMBLE vs GEX

Thank you for pointing out this relevant work that we missed. GEX leverages geometric ensembles to collect different ensemble models. The advantages of LoRA Ensemble over GEX lie in all the three types of costs as outlined in our paper.

• Training cost: geometric ensemble needs to change the learning rate during the ensemble model collection process, which means that it can not use pre-trained checkpoints of large models.
• Serving cost: LoRA Ensemble only uses the parameters of the LoRA adapter to calculate the attribution score, while GEX still uses the full parameters.
• Space cost: All the ensemble models collected by geometric ensembles needs to be stored, which costs more space compared to LoRA Ensemble (that only saves the LoRA adapter).

Weakness 4: More experiments on downstream data-centric tasks

Thank you for suggesting this experiment—It is indeed an insightful addition to evaluate the methods on downstream data-centric tasks such as noisy label detection. As shown in the PDF, we have included new quantitative results in Appendix O, demonstrating how Dropout Ensemble improves the noisy label detection performance evaluated by AUC metric.

Weakness 5: Typos

Thank you for pointing this out. We have fixed all of them in our updated paper.

Dear Reviewer 8sXq,

Thank you for your valuable feedback to our submission! We have carefully addressed your comments in our rebuttal and we would greatly appreciate it if you could review it and consider any further discussion. We look forward to hearing any additional thoughts you may have.

We thank the reviewer for the comments and address them below.

Weakness 1: The storage cost of TRAK and test of our method on larger settings

The storage cost of TRAK comes from two sources. The first source is the checkpoints of independent retrained models. Another source comes from the caching of projected gradients.

Efficient ensemble methods could reduce the cost from the first source significantly (up to 80%). Some experiments are stated in Figure 3 and Figure 6(d).

The second source is small when the number of data points is not too large. Some tricks such as block-wise matrix multiplication and data filtering can be applied to reduce the memory cost on cached projected gradients and are orthogonal to the ensemble. While the development of these tricks is out of the scope of this paper, it is worth noting that efficient ensembles are orthogonal to these tricks, and can be straightforwardly combined with these tricks.

Additionally, we have applied Dropout Ensemble on GPT-2 (124M) + WikiText-2 dataset, which is a much larger setting. As shown in the table below, Dropout Ensemble could achieve comparable performance without independent retrained models as Naive Ensemble with multiple independent retrained models.

Weakness 2: Comparison between Dropout and LoRA Ensemble

A comparison between Dropout and LoRA Ensemble can be found in Figure 9 (d) (newly added to the paper in Appendix M). In terms of LDS, Dropout Ensembles yield slightly better improvement compared with LoRA Ensembles and both ensembles can significantly outperform naive ensembles for transformer-based large models by achieving a better efficacy-efficiency tradeoff.

Dropout and LoRA Ensemble have different trade-offs regarding training and serving costs. The suitability of them may vary depending on the specific applications.

LoRA is designed for large transformer-based models, so we only perform experiments on the transformer architecture. Other than MusicTransformer, we also show the results on GPT in Figure 8 of Appendix K.

Weakness 3: Challenges and experiment on large scale generative models

We have added a new experiment on GPT-2 (124M) + WikiText-2 dataset as shown in the reply to Weakness 1.

Quantitative evaluation of data attribution of large scale generative models is a challenging task since most of the metrics (e.g., LDS, brittleness, …) need a large number of model retraining to get the ground truth.

Weakness 4: Limitation of the approaches

Thank you for pointing out this important aspect. We agree that discussing the limitations of our ensemble methods is valuable, and we have included a detailed qualitative table in Appendix L, demonstrating the computation and space overhead. Additionally, we have provided further discussion in the main text to highlight these points in Section 3.3. More empirical results are demonstrated in Appendix M for computation overhead and Figure 6(d) for space overhead.

Weakness 5: qualitative examples

Thank you for suggesting this experiment—it is indeed an insightful addition to visually demonstrate the effectiveness of our ensemble method. As shown in the PDF, we have included new qualitative examples in Appendix N, demonstrating how Dropout Ensemble improves the TDA results without requiring additional training costs.

Dear Reviewer SJ8M,

Thank you for your valuable feedback to our submission! We have carefully addressed your comments in our rebuttal and we would greatly appreciate it if you could review it and consider any further discussion. We look forward to hearing any additional thoughts you may have.