$\mathbb{D}^2$ Pruning: Message Passing for Balancing Diversity & Difficulty in Data Pruning

Dear Reviewer 2S5F,

We thank you for your time, effort, and insightful comments. Please see our response below:

• Time Complexity Analysis: We have revised the pdf to include a time complexity analysis in Sec. 4, to complement the GPU time estimates that we had already provided in Table 6 in the Appendix in original pdf. D2 Pruning runs in mere minutes for CIFAR10, CIFAR100 datasets and takes ~23 minutes for ImageNet-1K when we use exact $k$-NN computation for initializing the graph. However, this time estimate is significantly improved when we use approximate $k$-NN using faiss to prune the DataComp dataset consisting of 12.8 million samples, i.e., $\mathbf{D}^2$ Pruning takes ~1 hour to prune the DataComp dataset. This time overhead is negligible compared to the training time of the models on such datasets and is a good investment if it yields improvements in downstream task performances. $\mathbf{D}^2$ Pruning is a scalable method that can be used for even larger datasets with modest memory requirements ($\mathcal{O}(nk)$). In comparison, most existing coreset selection methods are not scalable to large datasets or applicable to self-supervised approaches. Many contemporary self-supervised approaches for pruning datasets also involve similarly time-intensive preprocessing steps [1, 2, 3].

• Theoretical Analysis: Our method design is motivated by the intuition that difficulty as well as diversity are needed to ensure the most representative data subset. While we do not have a theoretical analysis for this method, we provide a solid analysis of how this method works and why it is better than other methods through Figures 2,3,5, which provide insights for further research into the use of graphs for data subset selection.

• Additional reference: Thank you for the reference, we have added this to the discussion in the related work section. The CCS method that we use as one of our baselines in the paper is as good or better than the method suggested in this work.

[1] SemDeDup: Data-Efficient Learning at Web-scale through semantic deduplication

[2] T-mars: Improving Visual Representations by Circumventing Text Feature Learning

[3] Multimodal Dataset Pruning Using Image Captioning Models

Thanks to the author's comprehensive rebuttal. I've read all the responses, revised paper, and other reviewer's comments. Although I'm satisfied with the abundant content, I still have a few questions.

1. Time complexity

1.1 Why does the graph initialization take O(kn)? Could you provide a more detailed analysis of this? As far as I know, the time complexity of the KNN algorithm for a single query point is O(nd), where n is the number of training examples and d is the number of features. Thus, to find the kNN of all data examples, it would take O(d*n^2). Similarly, In Table 6, 'Graph Creation' takes longer time than 'Iterative Selection', which also takes O(kn) according to authors, for all datasets. Please convince me if I misunderstood anything.

1.2 In Table 6, it would be better to provide the training time on the full dataset, and that with D^2 pruning for further clarification.

2. Theoretical analysis

2.1 In my thought, only Figure 2 can be related to my question. D^2 pruning achieved better diversity and uncertainty than Facility Location, Graph Cut, and Moderate, while they are not a hybrid approach considering both diversity & uncertainty. Then, why D^2 pruning is better than BADGE in terms of generalization? An intuitive explanation or analysis for this would increase the novelty of this paper compared to the existing hybrid approaches like BADGE.

Thanks!

Dear Reviewer 3E93,

We thank you for your time, effort, and insightful comments. Please see our response below:

• Marginal improvements: We see relatively small gains for easy datasets like CIFAR10 and redundant NLP datasets for the finetuning task because the performance of random selection on these datasets is already quite high. However, we see larger gains for difficult datasets like CIFAR100, ImageNet-1K, and DataComp. Importantly, we think that $\mathbf{D}^2$ Pruning will be a useful framework for future research into self-supervised and unsupervised data selection approaches, which are crucial topics in contemporary research for training foundation models. due to its plug-and-play nature, where most existing work in coreset selection is no longer applicable. It will benefit from any work that investigates better importance scores or meaningful representation embeddings. Moreover, it is not only flexible but also more scalable for large datasets than many coreset selection methods that rely on sub-modular functions, as we discuss in the General Response.

• Statistical analysis: We have included a discussion of the significance of the scores in the revised pdf in Sec. 5.1. Our improvements on ImageNet-1K, DataComp, and CIFAR100 are significant (p<0.05, computed using a bootstrap of 100K samples [1, 2]). The improvement margins on CIFAR10 and NLP datasets are smaller due to random selection already doing so well on such datasets and the p-value is sometimes not significant for those improvements.

• Additional Hyperparameters: Coreset selection methods that do not operate solely on the ranking of an importance score need a hyperparameter, especially when combining the influence of two different functions, as we do in $\mathbf{D}^2$ Pruning. [3] combine facility location function for diversity with entropy score and use a weight parameter to balance them out. [4] perform a sweep over the number of bins and a cut-off ratio to remove bad examples. [5] perform bi-level optimization that occurs over multiple rounds of training and testing models. [6] tune the number of clusters used to get difficulty scores. It is definitely better to have a method without hyperparameters, but much further research is needed to get to that point.

• Iterations in message passing: We use a single iteration of message passing in our main results. We conducted ablation experiments where ran multiple rounds of message passing and found that it does not benefit data selection due to aggressive smoothing of importance scores over the spatial dimension. Please see results in Table 8 and Figure 5 in Appendix of the revised pdf, and the General Response for a discussion.

• Importance Score in Fig. 2: We use the forgetting score in Fig. 2, we have fixed the caption in the revised pdf to answer this question.

• Other importance scores: Yes, we ran experiments with scores other than forgetting scores and we report those results in Table 8 in the Appendix. Please see General Response for the discussion.

[1] Computer-intensive methods for testing hypotheses

[2] An introduction to the bootstrap

[3] Accelerating batch active learning using continual learning techniques.

[4] Coverage-centric coreset selection for high pruning rates.

[5] Glister: Generalization based data subset selection for efficient and robust learning.

[6] Beyond neural scaling laws: beating power law scaling via data pruning.

Dear Reviewer niHR,

We thank you for your time, effort, and insightful comments. Please see our response below:

• Distance in embedding space as an indicator of diversity: Embeddings from pre-trained models have been widely used as representations of the semantic content of a sample within a high-dimensional sample for vision and text modalities. Distance or similarity computations between such embeddings have been successfully used for performing semantic similarity tasks like k-nearest neighbor selection [1], sentence similarity [2], etc. In our work, we refer to diversity in a data subset as the diversity of semantic content of the data samples in the subset. Hence, we use cosine similarity between embeddings as a mark of semantic similarity of two data samples and aim to select samples with maximum diversity based on this definition. We perform ablation experiments using different sources of embeddings and find that the embedding from the last layer of the ResNet18 (polling layer before classifier) works better than features from the deeper layer for vision models, and the [CLS] token embedding in RoBERTa works better than non-[CLS] token embeddings for NLP datasets. See general response and Sec. D.2, Table 8 in Appendix in the revised pdf.

• Importance Score: Similarly, we also perform ablation experiments for different importance scores, and find that the entropy score benefits selection from CIFAR10 while the EL2N score improves performance on the Adversarial NLI dataset. See general response.

[1] Beating power law scaling via data pruning

[2] Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks

Dear Reviewer 4iU9,

We thank you for your time, effort, and insightful comments. Please see our response below:

• Baselines: We would like to thank the reviewer for putting together a detailed list of the works relevant to ours and outlining the methods for our convenience. We agree that these methods are relevant to our work. So, we picked four additional baselines for experiments (see general response).

• Discussion on suggested baselines: We have added discussion on new baselines and other suggested references in Sec. 5.1, Sec. 5.3 and Sec. 6, as much as possible in the limited space. Additionally, we have updated Figure 2 to show results from data selection using diversity-only approaches using facility location (FL) and graph-cut (GC) sub-modular functions [1]. As we discuss in Sec. 5.1, the main drawback of using FL or GC submodular functions in any form is the need for access to the complete similarity matrix that causes a memory requirement of $\mathcal{O}(n^2)$ [2], which is not scalable for larger datasets.

• Moving Averages: Yes, we consider the moving average of El2N, AUM scores, etc.

• Statistical analysis: We have included a discussion of the significance of the scores in the revised pdf in Sec. 5.1. Our improvements on ImageNet-1K, DataComp, and CIFAR100 are significant (p<0.05, computed using bootstrap of 100K samples [3, 4]). The improvement margins on CIFAR10 and NLP datasets are smaller due to random selection already doing so well on such datasets and the p-value is sometimes not significant for those improvements.

• BADGE: We use the true labels in our implementation of BADGE to make a fair comparison.

• $x_k$, $s_k$: Thank you for catching this error, we have fixed this in the revised version.

• Message passing for more than one round: Please see general response.

• Eq 1. Theta: Thank you for pointing this out, we have fixed this in the revised pdf.

[1] Submodular combinatorial information measures with applications in machine learning

[2] https://apricot-select.readthedocs.io/en/latest/index.html

[3] Computer-intensive methods for testing hypotheses

[4] An introduction to the bootstrap

I have a few questions:

1. For the Facility location function it is possible to use sparse matrices (in fact as long as entries are non-negative function is submodular). Therefore, did the experiments use the sparse kernels (the sparse similarity matrix that was used in the graph for D^2)? Therefore, I think it is not okay to write FL always needs $\mathcal{O}(n^2)$ space.

2. CAL-SDS2 can be used with any hardness metric, and in general - $F(A) = FL(A) + \lambda \operatorname{log}(1 + m(A))$ where $m(A)$ is a function for hardness -- entropy, margin, EL2N, etc. So I still feel that additional exploration on that front can be done.

Thank you for the follow-up comments!

1. Thanks for the suggestion, we are running experiments using the sparse similarity matrix where only k-nearest neighbor distances are non-zero and everything else is 0. We have also revised our paper to remove the lines that say that FL always needs $\mathcal{O(n^2)}$ space. We have also expanded our discussion of the Top-$k$ Voting method in the Related Work section. Top-$k$ voting is similar to the reverse message passing step in our method and only ensures diversity in selection. $\mathbf{D}^2$ Pruning first combines the influence of importance scores and local neighborhood in the forward message passing step and then imposes diversity in the reverse message passing step.

2. We also agree that CAL-SDS2 warrants a proper search for hyper-parameters as well as experiments with various importance scores to get the best result. In our current experiments, we are performing a grid search over $\lambda$ = $\[1, 2, ... 10]$ and $\sigma$ = $[0.1, 0.2, 0.3, ... 1.0]$ for entropy, forgetting and EL2N scores. The best results so far for CIFAR10 are as follows:

We see improved performance from CAL-SDS2 with Forgetting score at all pruning rates and CAL-SDS2 performs quite well as low pruning rates, but it still struggles at high pruning rates, even with Forgetting score as the difficulty metric. We will continue tuning our results with CAL-SDS2 and definitely add the improved results to our paper when the experiments end.

We would like to reiterate the intuition behind $\mathbf{D}^2$ Pruning:

Importance score distributions tend to be skewed towards low scores, an example of the distribution of forgetting scores can be seen in Fig. 5(a) in the Appendix. However, an easy sample surrounded by many other easy samples in the embedding space is less important than an easy sample that is surrounded by many difficult samples in the embedding space. This is because the area containing many easy samples is easy to learn for the model whereas the area containing difficult samples is hard to learn for the model. Intuitively, we want to include many samples from the hard-to-learn areas (even the easy samples) while sufficiently representing the easy-to-learn areas to prevent drops in accuracy for samples lying in that area [1]. Existing methods like BADGE and CAL-SDS2 select the most important sample that is also representative of the complete set of samples; they do not make a distinction between easy-to-learn and hard-to-learn areas in the data representation space. Whereas, the forward message passing step in $\mathbf{D}^2$ Pruning increases the importance of a sample by an amount that is proportional to the importance scores of the samples surrounding it, thus ranking an easy sample in a hard-to-learn area higher than that in an easy-to-learn area. Effectively, the forward message passing step recalibrates the importance scores by taking the difficulty of the sample's local neighborhood into account and results in a less skewed distribution as we demonstrate in Fig. 5(b). The samples are ranked by this recalibrated importance score (represented by the updated node features in the graph) for the reverse message passing step.

Additionally, in the reverse message passing step, after the selection of each data sample, only its nearest neighbor samples are updated. This allows us to maintain a constant selection time for the same data subset size, with increasing size of the original dataset, and makes $\mathbf{D}^2$ Pruning relatively more scalable for large-scale datasets, as we demonstrate with the DataComp dataset.

To summarize, $\mathbf{D}^2$ Pruning is a scalable data selection algorithm that effectively re-ranks samples based on easy-to-learn and hard-to-learn areas in the embedding space via forward message passing and then imposes diversity in data selection via reverse message passing.

Dear Reviewer DbMJ,

We thank you for your time, effort, and insightful comments. Please our response below:

• Additional baselines: Please see general response.
• K-shot setting: Please see general response, under Ablation Experiments.
• Very small selection ratios: Thank you for the suggestion, we ran experiments along these lines and report our results in Table 7, Sec. D.1 in the Appendix. We see large improvements using $\mathbf{D}^2$ Pruning in some cases, such as 3% over CCS at 99.5% pruning of CIFAR10. However, the improvements are not consistent, either from CCS or from our method, suggesting that diversity isn't the important factor at extremely low data budgets.
• Inconsistent numberings: Thank you for catching this error, we have fixed it in the revised version.
• References to sections in Appendix: Thank you for the suggestion, we have fixed this in the revised version.
• Section 5.1 Parameter: The parameter $\gamma$ refers to $\gamma_{r}$. We have fixed this in the revised version.
• Algorithm: Thank you for the suggestion, we have included a pseudo-code of $\mathbf{D}^2$ Pruning for data selection in Algorithm 1 in the Appendix. Please let us know if it doesn't answer your question correctly, we will be happy to provide a different answer according to your suggestion.
• Incremental gains: We see relatively small gains for easy datasets like CIFAR10 and redundant NLP datasets for the finetuning task because the performance of random selection on these datasets is already quite high. However, we see larger gains for difficult datasets like CIFAR100, ImageNet-1K, and DataComp. Importantly, we think that D2 Pruning will be a useful framework for future research into self-supervised and unsupervised data selection approaches, which are crucial topics in contemporary research for training foundation models. due to its plug-and-play nature, where most existing work in coreset selection is no longer applicable. It will benefit from any work that investigates better importance scores or meaningful representation embeddings. Moreover, it is not only flexible but also more scalable for large datasets than many coreset selection methods that rely on sub-modular functions, as we discuss in the General Response.