Paper: $\mathbb{D}^2$ Pruning: Message Passing for Balancing Diversity & Difficulty in Data Pruning (ICLR-2024)

Reviewer comment (Evaluation):
The performance improvement seems marginal. In most cases, the improvement is less than $1\%$.

Author response:
* **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.

Reviewer comment (Evaluation):
With such performance differences, repeated evaluation can be suggested to mitigate the variance in the model training.

Author 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.

Reviewer comment (Experiments):
$D^2$ introduces some additional hyper-parameter, which may increase the coreset selection cost.

Author response:
* **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.

Reviewer comment (Reproducibility):
How many iterations will the forward message passing phase have?

Author response:
* **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.

Reviewer comment (Evaluation):
What is the importance score in fig 2 (what metrics used)?

Author response:
* **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.

Reviewer comment (Experiments):
It seems that $D^2$ can be combined with other metrics. Do you run an ablation study to see how $D^2$ performs with other importance scores?

Author response:
* **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.

Reviewer comment (Evaluation):
No time complexity analysis. The proposed algorithm based on message-passing seems to take quite a lot of time. The author should provide the time-complexity analysis with the exact GPU time taken because a data pruning method that takes too long time is less practical.

Author response:
* **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].

**1) Time complexity** 

Thank you for identifying the mistake, the time complexity for graph initialization is erroneously mentioned as $\mathcal{O}(kn)$ in Sec. 4. It is indeed $\mathcal{O}(vn^2)$ where v is the vector dimension when the kNN is computed using the full dataset. We have revised the paper to reflect this change. Correspondingly, the time estimates in Table 6 for Graph initialization are higher than iterative selection because graph initialization is a more time-intensive step than iterative selection. The time estimates for graph creation of different datasets are not strictly comparable to each other because we run jobs of different batch sizes according to the size of the datasets to prevent OOM issues while computing the similarity matrix. We have also edited Table 6 to include training time for each dataset after data selection.

Reviewer comment (Theory):
No theoretical analysis. How this message-passing algorithm can guarantee better generalization than other baselines? Although the author provides some intuition (data pruning should consider both diversity and difficulty), why it should be achieved by the message-passing and how it can reduce the generalization error in Eq.(1) is missing.

Author response:
* **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.

**2) Novelty of $\mathbf{D}^2$ Pruning**

Importance scores distributions tend to be skewed towards low scores, an example of the distribution of forgetting scores can be seen in Fig. 5(a). 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 [2] 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 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 heard-to-learn areas in the embedding space via forward message passing and then imposes diversity in data selection via reverse message passing. 

[1] Beyond neural scaling laws: beating power law scaling via data pruning
[2] Accelerating Batch Active Learning Using Continual Learning Techniques (TMLR/DMLR@ICML'23)

Reviewer comment (Experiments):
I think this work is also related but missed in discussion/comparison. [a] Active learning is a strong baseline for data subset selection. NeurIPS workshop, 2022

Author response:
* **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.

Reviewer comment (Novelty):
I feel that the prime weakness of this work is from the angle of related works and baselines. Highlighting the importance of balancing diversity and difficulty is not new, and indeed if one is using not well-tuned diversity selection methods, then they won't be able to give a full representation of the dataset (minor modes which are not outliers but are difficult).

Author response:
* **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.

Reviewer comment (Experiments):
I need an explanation, and if possible, results on the following baselines for general supervised cases: CRAIG [1], GLISTER [2], GradMatch [3], Top-k method [4], and CREST [5].

Author response:
* **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.

Reviewer comment (Novelty):
A combination of submodularity with difficulty has been explored in prior work, including FASS [6], CAL-SDS2 [7], MCL [8], and DIHCL [9], and these should be considered.

Author response:
1. Thank you for clarifying the possibility of using sparse matrices in facility location functions for us. We used the full similarity matrix for Facility location, but we are re-running our experiments using the sparse similarity matrix. We implemented our version of CAL-SDS2 using pre-defined functions and optimizers in [1]; if the reviewer is aware of a codebase that directly implements CAL-SDS2 using sparse matrices, we request them to kindly point us to it.

2. We have launched jobs for tuning the hyperparameters of CAL-SDS2 [$\sigma$, $\lambda$] under different importance scores and nearest-neighbor similarity matrices as used for $\mathbf{D}^2$ Pruning, and so far we haven't seen better results than what we have reported using entropy + full similarity matrix. We will report the results again as soon as our experiments are completed. We would also like to mention that we did not come across CAL-SDS2 in our literature review for coreset selection methods that balance difficulty as well as diversity because the paper focuses on active/continual learning. We wish we had sufficient time to explore CAL-SDS2, however, the rebuttal period is not sufficient to conduct *fair* comparative experiments. We agree with the reviewer that submodular functions can be used to balance difficulty and diversity, and we will amply cover this in our experiments and discussion section. However, the goal of our work is to show that representing datasets as graphs *also* works as an effective framework for balancing difficulty + diversity in data selection, not just in the scenario of one-shot coreset selection but in self-supervised data pruning and unsupervised data filtering as well. It is an interesting line of research because it opens up possibilities for using graph neural networks (GNNs) to not just prune but also understand datasets, and our work is a proof of concept in this direction. Therefore, we request the reviewer to consider our work independently of the work on sub-modular functions.

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

Reviewer comment (Experiments):
More recent works on diversity-based selection for NLP, such as MILO [10] and INGENIOUS [11], should be discussed or used as baselines.

Author response:
* **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.

Reviewer comment (Novelty):
On SSL, see SAS [12], which should be considered for inclusion or discussion.

Author response:
* **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.

Reviewer comment (Novelty):
On multimodality, see T-MARS [13], which should be considered for inclusion or discussion.

Author response:
* **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.

Reviewer comment (Experiments):
I believe the paper should discuss these works and should include some of them as baseline.

Author response:
* **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.

Reviewer comment (Experiments):
For the difficulty-based methods in supervised settings, methods such as forgetting-event clearly use the learning dynamics. However, methods such as entropy, EL2N, and area under margin score can be computed during any step of training. Does the paper consider the moving average of the dynamics (of some pre-trained models, for which they’ve assumed access), or is it taken at the end of training? If it is taken at the end of the training, then training set methods like entropy/margin/EL2N can be very wrong in judging the hardness.

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

Reviewer comment (Evaluation):
Can the authors provide standard deviation or statistical analysis in cases where the second-best technique is very close?

Author 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 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.

Reviewer comment (Experiments):
For BADGE, does the author use true labels instead of pseudo labels? Since BADGE was proposed in Active Learning, it is important to ensure it does not have a disadvantage in supervised setting comparisons.

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

Reviewer comment (Writing):
k and s_k are overloaded as expressions in the paragraph above equation 6.

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

Reviewer comment (Writing):
In Eq. 1, theta should rather be written as $\theta^*(S')$ to show that it is a solution to the optimization.

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

Reviewer comment (Novelty):
The paper demonstrates very incremental gains in performance over State-of-the-Art method (like Ash et al., 2019).

Author response:
* **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.

Reviewer comment (Presentation):
The paper uses inconsistent numberings (A and then 2, 3) in section 1 which should be rectified.

Author response:
* **Inconsistent numberings**: Thank you for catching this error, we have fixed it in the revised version.

Reviewer comment (Experiments):
An important investigation aspect for this paper would be to demonstrate performance on very small values of selection ratios $k$ as discussed in (Guo et al, 2021) and perform more than 1 message passing (K-shot setting).

Author response:
* **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.

Reviewer comment (Writing):
It is unclear if the parameter $\gamma$ mentioned in section 5.1 refers to $\gamma_f$ or $\gamma_r$.

Author response:
* **Section 5.1 Parameter**: The parameter $\gamma$ refers to $\gamma_{r}$. We have fixed this in the revised version.

Reviewer comment (Reproducibility):
Although optional, including an algorithmic view of the proposed approach would be interesting to clarify how $\mathbb{D}^2$-PRUNING fits into the training and evaluation process of deep-learning models.

Author response:
* **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.