Group Distributionally Robust Dataset Distillation with Risk Minimization

Authors: Saeed Vahidian, Mingyu Wang, Jianyang Gu, Vyacheslav Kungurtsev, Wei Jiang, Yiran Chen
Primary area: optimization
Keywords: dataset distillation, distributional robustness, generalization

Abstract:
Dataset distillation (DD) has emerged as a widely adopted technique for crafting a synthetic dataset that captures the essential information of a training dataset, facilitating the training of accurate neural models. Its applications span various domains, including transfer learning, federated learning, and neural architecture search. The most popular methods for constructing the synthetic data rely on matching the convergence properties of training the model with the synthetic dataset and the training dataset. However, using the empirical loss as the criterion must be thought of as auxiliary in the same sense that the training set is an approximate substitute for the population distribution, and the latter is the data of interest. Yet despite its popularity, an aspect that remains unexplored is the relationship of DD to its generalization, particularly across uncommon subgroups. That is, how can we ensure that a model trained on the synthetic dataset performs well when faced with samples from regions with low population density? Here, the representativeness and coverage of the dataset become salient over the guaranteed training error at inference. Drawing inspiration from distributionally robust optimization, we introduce an algorithm that combines clustering with the minimization of a risk measure on the loss to conduct DD. We provide a theoretical rationale for our approach and demonstrate its effective generalization and robustness across subgroups through numerical experiments.

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

Reviews and discussion:
<thread>
creation date:
1730547665034

thread_type:
official_review

<thread_object>
creator: reviewer_gbwm

<content>
summary:
This paper proposes a robust dataset distillation method that enhances generalization across underrepresented subgroups by using a double-layer distributionally robust optimization (DRO) approach. Traditional dataset distillation often compresses training data into synthetic sets by matching training characteristics but may struggle with subgroup generalization. To improve this, the authors cluster the training data around synthetic data points and apply a Conditional Value at Risk (CVaR) loss to minimize worst-case subgroup errors, making the synthetic dataset more representative of diverse data distributions. Experimental results show that this method significantly improves robustness under domain shifts, outperforming baseline methods on benchmarks like CIFAR-10 and ImageNet-10, particularly in challenging conditions like noise and blurring.

strengths:
- As far as I know, this paper firstly addresses a major limitation in standard dataset distillation (underrepresented or rare subgroups) by applying a double-layer distributionally robust optimization (DRO) framework. This approach ensures the synthetic dataset better represents the full diversity of the data, reducing performance drops across different data subgroups. 

- This paper provides not only empirical evidence of the proposed method's robustness against domain shifts such as noise, blur, and inversion, but also a theoretical analysis of the effectiveness of CVaR in the loss function to enhance robustness.

- The proposed method is designed to be adaptable and can integrate with various existing distillation techniques, such as gradient or distribution matching. This modularity makes it compatible with a range of distillation methods, which are still actively developing and evolving.

weaknesses:
- The proposed DRO approach, particularly with clustering and CVaR, may introduce significant computational overhead. This added complexity could be amplified for large-scale datasets due to the clustering of training (real) data. However, the paper does not discuss the computational overhead of the proposed method.

- Although the method shows promising results on benchmarks like CIFAR-10 and ImageNet-10, the experiments are limited to controlled domain shifts (e.g., noise, blur). Testing under more realistic settings, such as in transfer learning, would further validate its robustness and practical relevance. For example, one could (1) train neural networks on a synthetic dataset distilled from a coarse-grained dataset (e.g., ImageNet) and (2) fine-tune and evaluate them on a fine-grained dataset (e.g., Birds 200). This setup would better illustrate the method's effectiveness in addressing the challenges posed by rare subgroups.

questions:
Please see weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730607690287

thread_type:
official_review

<thread_object>
creator: reviewer_1jhv

<content>
summary:
This paper proposes a new data distillation approach emphasizing the subgroup coverage and generalization ability of the models trained on the synthetic dataset. This paper provides a theoretical analysis rooted in Distributionally Robust Optimization (DRO) and verifies the effectiveness of the proposed method with various experiments.

strengths:
* This paper considers the group coverage and generalization ability of the synthetic dataset in Data Distillation(DD), which is interesting and novel.

* The introduced algorithm is clear and the theoretical analysis seems solid.

* The numerical results do show the effectiveness of the proposed algorithm. Meanwhile, the figures (Fig.3, 4) seem to indicate that the proposed method improves group coverage.

weaknesses:
* While the paper emphasizes group coverage and generalization in data distillation, the experiments are mainly conducted on IID datasets. More experimental results in scenarios with distribution shift between training and testing sets (such as long-tail classification, subpopulation shift, domain adaptation, domain generalization, etc) can further validate the improvement in group coverage and generalization ability,

* According to Algorithm 1, the initialization of the synthetic dataset seems very important because it involves how the training data samples are clustered into subgroups. It may require further ablation study to verify the stability of the proposed algorithm.

* The proposed method seems like a plug-in as it only modifies the objective for data distillation and could be combined with any data distillation methods. A more sophisticated comparison between the proposed method with other objectives regarding training gradients, feature distributions, and training trajectory could help readers better understand the improvement of the proposed method.

questions:
I have stated many of my suggestions and concerns in the weaknesses section. Below, I have further questions with some minor issues.

* While theoretical analyses have been provided, I wonder whether the proposed method would affect the convergence rate and make it more difficult to find an optimal solution for the proposed objective. I would appreciate an empirical time complexity analysis regarding the proposed method.

* In Eq.10, what does the *N* represent? I did not see any introduction of the *N*.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699006702

thread_type:
official_review

<thread_object>
creator: reviewer_p8rd

<content>
summary:
The work proposes a robust dataset distillation approach that incorporates distributional robust optimization (DRO) to enhance generalization and performance across subgroups. This method combines clustering with risk-minimized loss to conduct dataset distillation. By prioritizing representativeness and coverage over training error guarantees, the approach enhances the models trained on synthetic datasets for real-world scenarios. Both theoretical analysis and empirical validation on multiple standard benchmarks are provided, demonstrating the effectiveness of the proposed approach.

strengths:
1. The paper proposes applying distributional robust optimization to dataset distillation, providing a reasonable approach to enhance generalization.

2. Drawing on distributional robust optimization theory, this work establishes a theoretical foundation to support the proposed approach to dataset distillation.

3. The paper is well-structured, with clear algorithm block and effective visualizations that enhance the presentation of the work.

weaknesses:
1. The empirical experiments focus primarily on CIFAR-10, ImageNet-10. Extending the evaluation to larger, real-world datasets with a greater number of classes would better demonstrate the effectiveness of proposed approach in generalization and robustness under real-world conditions.

2. Comparing the proposed approach with additional baseline method addressing generalization and robustness would provide a more comprehensive assessment of the proposed approach, including comparisons with baseline methods such as [1, 2, 3].

3. It would be valuable to visualize the robust inference tasks using real-world data rather than illustrative visuals, which could provide more insight and observations in real-world scenarios.

Reference:

[1] Domain generalization with adversarial feature learning. CVPR 2018.

[2] Self-challenging Improves Cross-Domain Generalization. ECCV 2020.

[3] HYPO: Hyperspherical Out-of-Distribution Generalization. ICLR 2024.

questions:
Please refer to the detailed suggestions provided in the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721637520

thread_type:
official_review

<thread_object>
creator: reviewer_daax

<content>
summary:
This paper proposes an algorithm for dataset distillation by incorporating distributionally robust optimization into it. There is theoretical justification and empirical validation of the proposed algorithm.

strengths:
There are both theoretical and experimental demonstrations of the effectiveness of the algorithm.

weaknesses:
1. There seems to be a mismatch between the motivation and the experiments. The motivation emphasizes regions with low population density, which usually correspond to the worst-group performance in subpopulation shift [1]. However, the main experiments related to distribution shift are conducted on test sets with perturbations or the worst group induced by an additional clustering process. It would be better to conduct more experiments on subpopulation datasets included in [1]. 
2. The introduction has too many paragraphs, which makes the logic of the introduction tedious with poor readability. 

Some minor issues:

- In Line 049, "some technique". 
- In Line 129 and 131, "Algorithm" "Numerical Results" their first letters do not need to be capitalized. 
- In the last paragraph of introduction, Section 3 is not mentioned.

[1] Yang, Yuzhe, et al. "Change is Hard: A Closer Look at Subpopulation Shift." *International Conference on Machine Learning*. PMLR, 2023.

questions:
In Line 017, what does "targeting the training dataset" mean?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732165787183

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We want to express sincere gratitude to the reviewers for their detailed and constructive comments and advice on the manuscript. 
We are grateful that all reviewers acknowledge our theoretical analysis of the connection between distributionally robust optimization and dataset distillation robustness, as well as the experimental results to support the effectiveness of the proposed method. 
We have carefully revised the manuscript and marked the modified parts as **blue**, with the mark "**NEW**" on the right of the modified contents. 

Specifically, we have made the following major revisions:

1. We have included more realistic evaluation scenarios in **Section D.1** of the revised appendix. Now the section contains:

- Domain transfer from pre-training to downstream fine-tuning on fine-grained classification.

- Domain generalization from one ImageNet subset to another subset.

- Subpopulation shift benchmark involving Spurious Correlations.

2. We have included the comparison on more domain generalization training methods in **Section D.7**. 

3. We have added the ablation study on the initialization of synthetic samples in **Section D.6**, where the proposed RDD method is robust enough to handle different initialization. 

4. We have rewritten some parts of the manuscript to enhance the clarity of our idea and fixed some vague expressions. 

We hope that these modifications have addressed the concerns raised by the reviewers. 
We want to thank the reviewers again for their insightful opinions that help further refine this paper. We welcome any feedback and discussion from you. If you fell our response cannot address your concerns, please kindly let us know.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission develops a novel formulation of dataset distillation that makes use of ideas from distributionally robust optimisation. A method based on segmenting the input space into different subgroups and ensuring that a distilled dataset leads to a model that has good worst-case performance on these subgroups is provided. The reviewers agree that this direction is interesting, novel, and well-executed. However, all reviewers had some concerns about the particular datasets distribution shifts used in the experimental validation of the new method.

I agree with the consensus the reviewers have come to; despite some issues about realism of the experimental evaluation, the paper makes a good contribution and should be accepted.