Distribution-Aware Data Expansion with Diffusion Models

Summary

This paper proposed a training-free data augmentation method based on the diffusion models. To alleviate the poison phenomenon of the diffusion model, which is that the distribution of generated images will deviate from the natural distribution, this paper proposed a simple yet effective method based on the prototypes to introduce additional constraints to alleviate the deviation. The experimental results show that the proposed method achieves the SOTA performance compared to the baselines in this paper.

Strengths

1. Leveraging the prototypes as the constraint to guide the generation process is novel. Directly using the diffusion model to enrich the dataset will cause the poison phenomenon that will harm the classifier's performance since the distribution of generated images deviates from the natural image distribution. In this case, this paper leverages the cluster-based method, i.e., the prototype technology is sound and interesting. 2. This paper is easy to follow. 3. The proposed method achieves the SOTA performance compared with various baselines.

Weaknesses

1. The baselines lack the latest works, such as Brandon et al.[1] and Khawar et al [2]. For Brandon et al., although this method needs to fine-tune the pre-trained diffusion models, the time cost for fine-tuning them is very low and is close to the time cost reported in this paper. Under the closed computation cost, it should be considered the baseline. [1] Effective Data Augmentation With Diffusion Models. Brandon et al.. ICLR 2024. [2] DIFFUSEMIX: Label-Preserving Data Augmentation with Diffusion Models. Khawar et al. Arxiv:2405.14881. 2. The datasets used in this paper may not be enough. Since this paper does not only focus on small datasets, ImageNet should be considered the dataset to test the performance of the proposed method, which is similar to DIFFUSEMIX [2]. ImageNette is insufficient to replace ImageNet since it only contains 10 classes.

Questions

1. How to keep such a low time cost for the proposed method? The motivation for this question is that the author claims: "Stable Diffusion generates per sample in 12.65 seconds on average, while our DistDiff achieves the same in 13.13 seconds." The weakness of the guidance method (i.e., the energy function guidance used in this paper) is that it will dramatically increase the time cost, especially for a stable diffusion model. The reason is that the guidance method needs to calculate the gradient of the diffusion model. Concretely, please see Eq. 6. Eq. 6 needs first to calculate $\nabla_{e}D_{\theta}^{c}(z_{0|t},p_{c})$. $\nabla_{e}D_{\theta}^{c}(z_{0|t},p_{c}) = \frac{\partial D_{\theta}^{c}(z_{0|t},p_{c})}{\partial z_{0|t}}\frac{\partial z_{0|t}}{\partial e}$. By Eq. 5 and Algorithm 1, $\frac{\partial z_{0|t}}{\partial e} = \frac{\partial z_{0|t}}{\partial z_{t}}\frac{\partial z_{t}}{\partial e}$. In this condition, $\frac{\partial z_{0|t}}{\partial z_{t}}$ needs to calculate $\frac{\partial \psi}{\partial z_{t}}$ and $\psi$ is the pre-trained diffusion model. Therefore, the overall process needs to calculate the gradient of the diffusion model at least 50 times, while the time cost only increases by less than 1 second, which is hard to understand. 2. Could the author offer the ablation study for K=1 that only one group situation? The motivation for this question is that Table 6 shows that the increase of K has a limited influence on performance, which is confusing. In theory, following the storyline of this paper, more groups should increase the diversity of the generated images, which should have a positive influence. Based on the Table 6, it seems that group strategy is redundant.

Rating

5

Confidence

5

Soundness

2

Presentation

3

Contribution

3

Limitations

No additional limitations, including societal impact, should be discussed. All my concerns are listed in Weaknesses and Questions. If the author could clarify these concerns, I'm willing to increase my score.

Summary

The authors present DistDiff, a training-free data expansion framework based on a distribution-aware diffusion model. DistDiff constructs hierarchical prototypes to approximate the real data distribution, optimizing latent data points within diffusion models through hierarchical energy guidance. The framework demonstrates its capability to generate distribution-consistent samples, significantly improving data expansion tasks. DistDiff consistently enhances accuracy across a diverse range of datasets compared to models trained solely on original data.

Strengths

1. The authors claim high efficiency for DistDiff. 2. The framework effectively generates samples that align with the original distribution, markedly enhancing data augmentation tasks. 3. DistDiff consistently improves accuracy across a wide variety of datasets when compared to models trained only on the original data.

Weaknesses

1. The manuscript introduces the concept of hierarchical prototypes but falls short in sufficiently explaining how these prototypes are selected or generated from the dataset. A more detailed description or examples of the prototype generation process would greatly enhance the reader’s understanding and bolster the credibility of the proposed method. 2. The authors claim high efficiency for DistDiff but do not provide supporting data. It is recommended to add a comparison table detailing DistDiff's computational time, resource usage, and scalability against other methods, to substantiate its efficiency claims. 3. The paper claims that " these two scores reinforce each other and are indispensable mutually," in Table 4, the observed impact appears minimal without P_g. Additionally, the k value for P_g seems to have little effect on the results. Further explanation and clarification are needed to substantiate these claims.

Questions

1. In Figure 6, it is unclear why the optimal number of prototypes is set at K=3, as the difference from the figure is not readily apparent and the quantitative metrics only differ by 0.01. Could this be further elaborated to justify the selection of K=3 as the optimal prototype count? 2. It is suggested that the captions be further elaborated to explain the entire pipeline process in detail. Currently, the process terms used in the caption do not correspond directly to those labeled in the figure, leading to potential confusion. A clearer alignment between the text and graphical elements would improve comprehension and the overall effectiveness of the figure. 3. It is recommended to revise Figure 1. While it depicts the complex process of DistDiff, it does not provide clear benefits. Including performance comparisons or showcasing significant differences in generated results would enhance the figure's utility and informative value.

Rating

4

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

As stated above.

Summary

This paper focuses on data augmentation or expansion by generating synthetic data from pre-trained large-scale diffusion models. To ground the samples from these large-scale diffusion models, the paper proposes an energy-based guidance approach where the energy function depends on hierarchical prototypes. In the paper, Hierarchical prototypes are essentially feature vectors that define the object classes. The hierarchical prototypes are obtained as follows - First, the features for each class are aggregated to obtain a class level representation. Further, the features of each class are clustered into K clusters to obtain sub-group representations within each class. The paper shows clear improvement for classification tasks on many standard datasets, when classifiers are trained from scratch on the augmented dataset.

Strengths

The paper addresses a very important problem - How to perform effective data augmentation using synthetic data generation models. The approach does not require any training or fine-tuning of the diffusion models to adapt the generation to the required data distribution. The paper has detailed ablation studies for each design choices. Also, the experimental results suggest considerable improvement over the prior data expansion approaches.

Weaknesses

1. There seems to be some confusion in the explanation. Section 3.3 (Transform Data Points) seems to suggest that the approach always starts with a sample from the dataset. However, the algorithm in the appendix does not include any sample from the dataset. 2. Is the approach extendable to other supervised learning tasks like segmentation, detection, etc? It looks to me like the augmentation approach is tailored for classification tasks only, with the usage of class-specific hierarchical prototypes. Whereas, traditional augmentation approaches like random cropping, rotation, etc., are generic. 3. One of the main contributions of the paper is the residual multiplicative transformation. The paper does not give a clear answer to "why should I not adjust the latent $z_t$ directly to optimize for the energy function $g_t$?" The paper shows empirical explanations in the ablation, but there is a lack of concrete reasoning as to why this approach works.

Questions

Can the authors please address the questions 2 and 3? Additionally, how is the distribution across classes chosen for the synthetic data generation process? Are all classes equally sampled? I couldn't find this information in the paper.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors have addressed the limitations.

Response to Author

Thanks for the author`s rebuttal. I have carefully read all the contents. The new experimental results show the improvement of the author's proposed method. Although the improvement in ImageNet may be fair, I think it is due to the expansion, and the time limitation of the rebuttal period does not allow us to try (5x) expansion. Most of my concerns have been addressed. Thus, I will consider increasing my score to the borderline accept. Additionally, about the group ($K$), based on the rebuttal, it seems there is a trade-off about $K$ since $K$ cannot be small and be large. Meanwhile, $K$ will be influenced by the Dataset. A natural question is raised: Is there any way to choose the $K$ in an adaptive way?