DRoP: Distributionally Robust Data Pruning

Authors: Artem M Vysogorets, Kartik Ahuja, Julia Kempe
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: data pruning, classification bias, robustness

Abstract:
In the era of exceptionally data-hungry models, careful selection of the training data is essential to mitigate the extensive costs of deep learning. Data pruning offers a solution by removing redundant or uninformative samples from the dataset, which yields faster convergence and improved neural scaling laws. However, little is known about its impact on classification bias of the trained models. We conduct the first systematic study of this effect and reveal that existing data pruning algorithms can produce highly biased classifiers. We present theoretical analysis of the classification risk in a mixture of Gaussians to argue that choosing appropriate class pruning ratios, coupled with random pruning within classes has potential to improve worst-class performance. We thus propose DRoP, a distributionally robust approach to pruning and empirically demonstrate its performance on standard computer vision benchmarks. In sharp contrast to existing algorithms, our proposed method continues improving distributional robustness at a tolerable drop of average performance as we prune more from the datasets.

Decision: accept (spotlight)
Reviewer scores: 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730676745990

thread_type:
official_review

<thread_object>
creator: reviewer_kas6

<content>
summary:
The paper studies classification bias in data pruning. It first provides a systematic study of the existing methods and shows that they are all biased to some extent. To this end, the authors propose DRoP, a distributionally robust approach for pruning that preserves good robustness against classification bias. Their method achieves state-of-the-art robustness performance when combined with random subsampling across different settings.

strengths:
- The proposed method DRoP is quite concise, which is highly appreciated since it can potentially increase the accessibility and applicability of the method beyond the current scope.
- Despite its simplicity, DRoP + random sampling achieves superior performance across various settings.
- The paper conducts a systematic evaluation of the existing data pruning methods and reveals a valuable conclusion that none of the existing methods actually has good robustness in terms of classification bias.
- The evaluation covers a wide range of scenarios beyond the classic data pruning, which enhances the importance of the method as it could lead to potential influence that extends beyond traditional use cases.
- The discussion of the limitations and future work is sound and makes sense.

weaknesses:
- The flow of the paper can be improved. At the moment it feels that the related work is scattered throughout the paper, which disrupts readability and makes it challenging to grasp. It would be clearer to bring all of the related work into one section and then refer back to it when needed in the experiment section. For example, the authors may group related work into categories like "Data Pruning Methods", "Robustness and Fairness", "Long-Tailed Recognition Techniques", and "Robust Data Pruning Methods", and place this consolidated section after the introduction.

- Since there is no standalone related work section, it is difficult to assess the novelty of the proposed method and the relation with existing data pruning techniques. It would be great to also include a brief discussion of the most recent advancements in data pruning and how DRoP is similar to/different from them.

- The robustness evaluation on imbalanced datasets and group distributional robustness is not compared with the methods from the corresponding fields. It would be great if the performance of some task-specific methods is also compared in Figures 8 and 9, which helps to better assess the method's relative performance in context. For example, [1,2] for "imbalanced datasets" and [3,4] for "group distributional robustness"

[1] Kang, Bingyi, et al. "Decoupling representation and classifier for long-tailed recognition." ICLR 2020.

[2] Cao, Kaidi, et al. "Learning imbalanced datasets with label-distribution-aware margin loss." NeurIPs 2019.

[3] Liu, Evan Z., et al. "Just train twice: Improving group robustness without training group information." ICML 2021.

[4] Sagawa, Shiori, et al. "Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization." ICLR 2020.

questions:
- It is quite impressive that Random + DRoP outperforms other methods when the overall dataset density is low (e.g. in Figure 5). But do the authors have an idea why in many cases, it gives inferior worst-class performance when the overall dataset density is high? For example, for ResNet18 + TinyImageNet at 90% density, Random + DRoP is the second worst performance (only better than Forgetting) in terms of the worst-class performance in Figure 5.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693363191

thread_type:
official_review

<thread_object>
creator: reviewer_xskc

<content>
summary:
The paper studies the effect of dataset pruning on classification bias for several existing methods using the metrics of worst-class accuracy, difference between best and worst class accuracy, and standard deviation. Due to the purported worsening of worst-class accuracy at the cost of better average performance, the paper proposes random pruning based on error rates on a holdout set which is motivated by a theoretical analysis of a simple 1-D mixture of two Gaussians. Experimental analysis is conducted on CIFAR10, CIFAR100, TinyImageNet and ImageNet datasets for the various settings.

strengths:
1. The presented theoretical analysis on the GMM for minimizing worst-case statistical risk and average risk is good and aligns with the proposed pruning method.
2. The experimental analysis is comprehensive and convincing.

weaknesses:
1. While the Random+DRoP method does reduce classification bias compared to existing dataset pruning strategies, it generally seems to do slightly worse on average accuracy (Figure 5). The combined Strategy+DRoP seems to mitigate this a bit but the trend is not clear.
2. The experiments do not always indicate such a clear and striking trend towards robustness as the authors suggest. Some aggregate scalar numbers might be beneficial to understand which strategy does better across all data densities.

questions:
1. in Figure 6, right, why does a higher negative correlation of class accuracy to class density indicate more robustness? Ideally class accuracy should have zero correlation with class density for a model invariant to class size.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730725038586

thread_type:
official_review

<thread_object>
creator: reviewer_jvus

<content>
summary:
This paper presents a dataset pruning method that helps the models trained on the pruned datasets be more fair in their class accuracies. The proposed method is called "DRoP", and it selects class-specific pruning ratios based on the error rates, improving performance on underrepresented classes.
The emprical evaluations are performed using VGGs, ResNets on CIFAR-10, CIFAR-100, TinyImageNet, and WaterBirds datasets.
Most of the evaluations are very favourable for the proposed method.

strengths:
The idea of DRoP is quite interesting and to the best of my understanding very novel. 

The method seems to be simple and straightforward.

The paper provides detailed theoritical insights and this helped understand the method.

weaknesses:
The framing of the paper is currently unclear to me. 
(W1) Firstly, It would significantly help in include "Dataset Pruning" rather than just "Pruning" in the title of the paper, because given the dominence of "model pruning" methods, it is very understandable that one might get confused by the title. One suggestion would be DRoPD (.... Pruning Datasets).

(W2) Secondly, the problem of "classification bias" is never really defined in the paper, that is, it is never explained what is meant by "classification bias" and why is it a problem and how datasets affect this so-called "classification bias".

(W3) Next, the use of the word "Robustness" seems incorrect in this context. It is not really "robustness" being evaluated but class-wise fairness. Might be prudent to frame the paper differently around this.

(W4) Lastly, the evaluations are limited to non-robust small models like VGG-16, VGG-19, ResNets up to ResNet50. However, it has been seen that large models like ResNet101, ConvNeXt-B onwards, and ViT based large models are more robust to distribution changes. They might be more robust to underrepresentation of classes as well, or might not show gains when using the proposed "DRoP" method. Thus having experimental evaluations with these models would help better understand the applicability of the proposed method.

questions:
Q1- In Figure 9 (leftmost plot), when dataset sparisty is more than 50% (density less than 0.5) why is the proposed DRoP significantly less stable than all the other methods? Some explaination for this would be very helpful.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731712929340

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for taking time to read our paper and for their valuable comments. We are happy to see the uniform acknowledgement of the importance of our work. We revised the manuscript as per your suggestions, with all changes highlighted in red in the updated file. We list these changes below. 
1. The title of the paper is changed to "DRoP: Distrbutionally Robust **Data** Pruning" to avoid any confusion with model pruning, as was suggested by Reviewer JvuS.
2. As proposed by multiple reviewers, Section 2 is split into two parts to accomodate a separate "Related Work" section in order to improve the flow of the paper and make all relevant background information more accessible. 
3. We added a sentence describing the issue of classification bias to Section 1 where it is first referenced.
4. As requested by Reviewer kAs6, we extend the comparison of our method to prior art in distributional robustness (page 10). In Table 2 of their paper, Pezeshki et al. (2024) report the performance of ResNet-50 on Waterbirds for a variety of existing algorithms, including JTT (Zheran Liu et al., 2021); we copy these numbers verbatim and plot them in Figure 9 to compare with Random+DRoP and other related pruning methods studied in our work. Since we are not using group attributions for validation and tuning hyperparameters for Random+DRoP, we copy the results reported for this exact setup (Pezeshki et al., 2024).
5. We provide a comprehensive explanation for the observed sudden degradation of the average accuracy of pruning with DRoP at low dataset densities in Figure 9 (ResNet-50 on Waterbirds) in the figure caption. Further, we put additional discussion and visuals about this in appendix F.
6. To provide an alternative way to examine the effectiveness of DRoP and the emerging trends, we create Table 2 (Appendix D) where we report the majority of our experimental results. 

It has been pointed out by one reviewer that our evaluation lacks large-scale models. We cannot agree with this statement as we consider a variety of vision architectures including ResNet-50 trained on ImageNet from scratch. There is an abundance of highly impactful related works that were recently published in  reputable venues with empirical validation on similar or weaker models (Paul et al., 2021; Good et al., 2022; Ma et al., 2022; Ayed & Hayou, 2023; Zhao & Bilen, 2023; Feng et al., 2024), to name a few. Still, in an attempt to address this concern and reconfirm our conclusions with even larger models, we currently run additional experiments with ResNet-101 on CIFAR-100. These results are still pending and are currently not reflected in the revised manuscript. If these experiments finish before the end of the rebuttal period, we will submit another revision of the paper and post an update about it here. Note that results with larger models on ImageNet are out of reach for us: already the results we present on ImageNet with ResNet50 took over a month to run with the resources at our disposal.

In addition to submitting a revised paper, we reply to each reviewer individually to address their questions and concerns. We hope that these responses together with the revised manuscript clear up any confusion and resolve all issues that the reviewers had, and that you will consider increasing your rating for our paper.

**Referenes** \
[1] Sagawa, S., et al. "Distributionally Robust Neural Networks". ICLR 2020. \
[2] Paul, M., et al. "Deep Learning on a Data Diet: Finding Important Examples Early in Training". NeurIPS 2021. \
[3] Good, A., et al. "Recall Distortion in Neural Network Pruning and the Undecayed Pruning Algorithm". NeurIPS 2022. \
[4] Ma, X., et al. "On the Tradeoff Between Robustness and Fairness". NeurIPS 2022. \
[5] Ayed, F., Hayou, S. "Data Pruning and Neural Scaling Laws: Fundamental Limitations of Score-based Algorithms". TMLR 2023. \
[6] Zhao, B., Bilen, H. "Dataset Condensation with Distribution Matching". ICCV 2023. \
[7] Feng, Y., et al. "Embarrassingly Simple Dataset Distillation". ICLR 2024. \
[8] Pezeshki, M., et al. "Discovering Environments with XRM". ICML 2024. \
[9] Zheran Liu, E., et al.,"Just Train Twice: Improving Group Robustness without Training Group Information". ICML 2021.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a method, DRoP, for data pruning. The proposed approach is based on the solid theoretical and empirical analysis of previous data pruning methods that are shown to potentially result in high classification biases. DRoP is therefore phrased as distributionally robust pruning. The proposed approach is quite simple and effective and addresses an important problem of data pruning. The submitted revision is also improved in terms of presentation and clarity.  All reviewers agree that the paper is a valuable contribution to ICLR.