Paper: TADA: Timestep-Aware Data Augmentation for Diffusion Models (ICLR-2024)

Reviewer comment (Experiments):
As TADA uses multiple types of augmentation, including h-flip, the improvement of TADA over the h-flip baseline may come from the use of other augmentations but not the proposed weighting scheme. The author should include more experiments to show the effectiveness of the time-aware weighting scheme, for example, comparing with a baseline with a fixed augmentation strength while keeping other configurations the same as TADA.

Author response:
**The improvement may come from the use of other augmentations rather the proposed weighting scheme (comparison with fixed strength should be done)**

We agree on your point that the benefit may come from the use of various augmentation sets. However, when we apply these various augmentations with fixed strength, the problem of distribution shift arises severely, as shown in the mean face of ‘naive’ augmentation in Figure 3. That is the reason we did not put the naive augmentation (the same set of augmentations with fixed strength) into the performance comparison table. Still, to provide clear results, we will add the FID scores of naive augmentation to our paper.

Reviewer comment (Experiments):
The proposed timestep-aware weighting scheme is motivated based on the observation of two toy examples on 200 images from the FFHQ-5k dataset. Many of the parameters are selected manually: the augmentation probability, the choice of the augmentation functions, the SNR-sensitive region, and the maximum augmentation strength $\delta$. As observed from Table 9, the FID and KID increase if we use a sensitive region of [-3, 1] instead of [-3, 0]. Although the current setting is found to be effective on FFHQ and AHFA, it is unclear whether these pre-defined parameters need to be fine-tuned on any new datasets.

Author response:
**Many of the parameters are selected manually. It is unclear whether these pre-defined parameters need to be fine-tuned on any new datasets.**

For toy experiments, we reported the observations with 200 images by following the setting of P2 weighting [b]. During the design of this experiment, we confirmed that the number of chosen images does not affect the trend of this experiment. As mentioned in the paper, we mainly adopted augmentation sets from ADA [a], which is widely used in GANs literature. Therefore, we conducted our main experiments with a probability of 0.8 (that has been used as the maximum probability in ADA [a]) to demonstrate the feasibility of applying high augmentation probability with our proposed method. Since higher augmentation has been shown to be more effective in alleviating overfitting, we applied the highest probability possible. However, higher augmentation probability can lead to distribution shift, as reported in [a]. Therefore, being able to apply p=0.8 implies that our method allows usage of high probability and takes advantage of improved performance as well as alleviation of overfitting, without resulting in distribution shift. To further evaluate the effect of different probability parameters, we are currently performing probability sweeps with p=0.2, 0.4, and 0.6. We will report the results as soon as the training ends.

Our proposed method has a primary focus on data augmentation techniques that do not cause distribution shifts. As pointed out, in Tables 8 and 9 of the appendix, performance can vary depending on the sensitive region. While the region boundaries that yield the best FID scores can be fluid, we consistently set the boundary to [-3, 0] across all datasets to prioritize robustness over maximizing performance. Indeed, performance could be further improved by fine-tuning these parameters. However, we emphasize that we avoided manually tuning the parameters to maximize performance and instead conducted our experiments using safe and moderate parameters.

Reviewer comment (Experiments):
While it is acceptable that there are no solid theories supporting the timestep-aware mechanism, the method should be validated on sufficiently many scenarios, showing that the method is effective in general. I believe evaluating the method only on human faces (FFHQ) and animal faces (AHFQ) is not enough.

Author response:
**Evaluations on only human faces are not enough**

As we mainly focused on a data augmentation technique that does not lead to distribution shift, we utilized aligned datasets (e.g., subsets of FFHQ) that allow explicit evaluation of distribution shift in a qualitative manner, following ADA [a]. As reported in Tables 8 and 9 of the appendix, although the region boundaries that result in the best FID scores can be fluid, we consistently fixed the boundary to [-3, 0] across all datasets to prioritize robustness over performance maximization.

However, we agree that showing effectiveness on an unaligned dataset can imply better generalizability, we are currently training on CIFAR-10. In addition, if time allows, we will also report on the LSUN dataset.
We will make the results on these datasets available as soon as possible.

Reviewer comment (Evaluation):
The proposed TADA can only apply a timestep-aware global strength for all augmentations. The method cannot assign strength and probability values for individual augmentations, which may result in limited augmentation flexibility. In Table 8, applying only the color transformation leads to the best result. However, TADA cannot discover such an augmentation policy.

Author response:
**The proposed TADA can only apply a timestep-aware global strength for all augmentations. The method cannot assign strength and probability values for individual augmentations, which may result in limited augmentation flexibility.**

We can assign different strengths and probability to individual augmentation. This can be checked on the provided source code. Specifically, a user can set the values of probability and adjust the strength as a hyperparameter. Furthermore, we will make our code publicly available and user-friendly.

**In Table 8, applying only the color transformation leads to the best result. However, TADA cannot discover such an augmentation policy.**

Our primary focus is to show the effectiveness of our method with a standard augmentation set, by adopting an augmentation set that is widely used in GANs literature. As mentioned in the limitations of our paper, the augmentation set we chose may not be optimal for diffusion models. Although not every augmentation has a positive effect, as shown in table 8, the augmentation overall gave benefits on the performance compared to the baseline. The optimal set can lead to additional improvement in performance. Thus, we believe that further research on finding the best set for diffusion models can be an interesting avenue.

Reviewer comment (Theory):
While the observations by Toy Examples 1 and 2 provided by the paper are interesting, the paper fails to explain why the distribution shifts arise in the diffusion model from these results (-> **Q1** and **Q2**).

Author response:
**Q1. Why do the results in Fig. 1 indicate a distribution shift at sensitive timesteps?**

For each model, we measured the perceptual similarity between the predicted $\hat{x}_0$ and the original image $x_0$ at each timestep of the reverse process. We then analyzed the two models' difference in the perceptual similarity.

The baseline model, which was trained on the original dataset without data augmentation, is highly likely to produce in-distribution samples. Therefore, the perceptual similarity between the predicted $\hat{x}_0$ and the original image $x_0$ from this model can be thought of as the lower limit at each timestep. We then compare this with the augmented model. Since this model was trained with an augmented dataset, the generated samples will encompass both the in-distribution and out-of-distribution samples. This can be confirmed when the perceptual similarity between $\hat{x}_0$ and $x_0$ shows a noticeably higher perceptual difference if the predicted sample is an out-of-distribution sample.

In summary, the difference in perceptual similarity between the two models is an indication of the degree to which the augmented model is producing out-of-distribution samples. The larger the difference, the more likely it is that the augmented model is producing samples that are not representative of the original dataset. The results in Fig. 1 align with this explanation and suggest that data augmentation can cause diffusion models to produce out-of-distribution samples, especially at sensitive timesteps.

**Q2. What does Toy Example 2 imply?**

Toy experiment 2 illustrates that exchanging the reverse process between the baseline and augmented models during the intermediate timesteps can modify the sampling trajectory, ultimately determining whether the generated samples end up within the augmented data distribution, even when they initially originated from the original data distribution, and vice versa. This highlights the influence of sensitive timestep on the sample distribution.

The observation that exchanging the trajectory at other timesteps had no impact on the generated data distribution further strengthens the result of toy experiment 1 that the middle timesteps are the most susceptible to data augmentation (Figures 8 and 9). This implies that the model's ultimate generation to adhere to the original or augmented data distribution is predominantly made during these middle timesteps.

Reviewer comment (Theory):
The theoretical foundation for the proposed method is weak. It is not clear why the proposed method solves the problems of distribution shift and overfitting (-> **Q3**).

Author response:
**Q3. What is the theoretical background of the proposed method and what is the expected semantic behavior as a data augmentation? Why does the strength of the transformation at each time step help to suppress distribution shifts and overfitting?**

We acknowledge that our data augmentation technique does not have a solid theoretical basis, and we view this as a limitation of our paper. Meanwhile, our paper is built upon our empirical findings that distribution shift occurs due to data augmentation and it appears in specific timestep intervals. To address this, we devised a method that adjusts the augmentation strength based on the timesteps. The reason that we view that the weak strength helps to suppress distribution shifts is that, we keep x_t of each timestep to not deviate too far and let the model learn how to construct the actual contents of an image, similar to findings from P2 weighting [a]. Effect on overfitting is not exclusively attributed to our proposed method (suppressing the augmentation strength) but can also result from the use of data augmentation itself [b]. The point we wanted to convey through our effectiveness on overfitting is to show that although we weaken the augmentation strength along specific timesteps, we still maintain the effectiveness on alleviating overfitting.

As our method is built upon empirical evidence, we apologize that at this point, we cannot exactly provide answers with profound theoretical proofs. We will  endeavor to perform a deeper theoretical analysis on our future research. Still, we would like to emphasize the importance of this empirical study since we believe that empirical research can provide valuable insights and practical benefits, regardless of formal theoretical understanding. Indeed, a lot of theoretical studies are inspired and informed by empirical observations. One of the widely-used studies is CutMix [c], despite lacking rigorous theoretical investigations, has been utilized by many practitioners and following studies. CutMix has provided motivation for many theoretical studies such as [d, e]. We hope that empirical study is regarded as a valuable and essential form of research.

[a] Choi, Jooyoung, et al. "Perception prioritized training of diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[b] Karras, Tero, et al. "Training generative adversarial networks with limited data." Advances in Neural Information Processing Systems 33 (2020): 12104-12114.

[c] Yun, Sangdoo, et al. "CutMix: Regularization strategy to train strong classifiers with localizable features." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2019.

[d] Park, Chanwoo, Sangdoo Yun, and Sanghyuk Chun. "A unified analysis of mixed sample data augmentation: A loss function perspective." Advances in Neural Information Processing Systems 35 (2022): 35504-35518.

[e] Oh, Junsoo, and Chulhee Yun. "Provable benefit of mixup for finding optimal decision boundaries." International Conference on Machine Learning. PMLR, 2023.

Reviewer comment (Experiments):
The experimental setup is unfair. The proposed method introduces color transformations that are not used in the baseline (AR), making it impossible to evaluate the proposed method with the experimental results. Also, in some of the experimental evaluations (Table 4 and Fig. 4), there is no comparison with the AR and the presentation is selective (-> **Q6** and **Q7**).

Author response:
**Q6. Why does the proposed method add transformations that are not used in AR?**

We agree on your point that we employed color transformations that were not used in AR. This was because we tried to perform a fair comparison by presenting the best performance of AR methods, adhering to the specific augmentations that were found to be effective for EDM (color transformations were not used as authors of EDM mentioned that they found out color corruptions were rather harmful (Section F.2 of EDM [a])). 

Still, we agree on your point that the exact augmentation sets used for AR and TADA differ. Therefore, we are currently training TADA on the EDM augmentation set and training EDM on the augmentations that we used for TADA. We will report the results as soon as we obtain the corresponding results.

[a] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." Advances in Neural Information Processing Systems 35 (2022): 26565-26577.

**Q7. Why are the AR results not listed in Table 4 and Fig. 4?**

The purpose of Table 4 was to demonstrate the generalizability of TADA to various settings. Therefore, we performed a simple comparison with the most frequently used data augmentation for diffusion models (h-flip). Similarly, Figure 4 aimed to illustrate the effectiveness of TADA in addressing overfitting (albeit potentially trivial due to the well-known ability of data augmentations to mitigate overfitting), rather than directly comparing it to other data augmentation methods.

Throughout our experiments, our primary focus was not to assert that TADA outperforms AR in all scenarios, but rather to emphasize its adaptability and ease of integration into various diffusion model settings, while maintaining comparable performance to AR. We have tried to emphasize this point in our paper, but this point was not explicitly emphasized in our paper. We will revise our paper to make this clear.

Once again, we deeply appreciate the time you dedicated to reviewing our paper. We will incorporate your constructive feedback into revised versions of our paper. We respectfully ask if you consider raising the score. If you have any additional feedback that could further enhance our work, please share it with us at your convenience.

Reviewer comment (Presentation):
The quality of the presentation with figures and tables is poor. Fig. 1(a), 4, and 5 cannot be accurately understood by a colorblind person like me.

Author response:
Furthermore, we apologize for inconvenience due to some figures that were not carefully made with appropriate colors. We have revised the corresponding figures (Fig.1(a), 4 and 5).

Reviewer comment (Experiments):
This paper seems to evaluate the proposed method on only a limited data setting. Since data augmentation for discriminative models tends to improve performance when using full datasets besides limited data size settings, I would like to see the performance of the proposed method when a dataset has a sufficiently large size. If the proposed method does not improve the performance in this situation, it is the limitation of the proposed method.

Author response:
Thank you for your response. We have tested our model using the full FFHQ-70k dataset and acquired the following results:

| h-flip | tada |
|--------|-------|
| 15.22   | 16.67 |

Consistent with the findings of ADA (Karras et al., 2020), our experiments indicate that data augmentation does not yield significant performance gains, implying that it may not be a crucial factor when training on a large and diverse dataset.

Reviewer comment (Experiments):
Does the proposed method work well if a dataset has a sufficiently large size?

Author response:
As data augmentation is primarily employed in scenarios with limited data, our focus has been on evaluating its effectiveness in this context. Consistent with prior studies that demonstrate the benefits of data augmentation for small datasets, we presented results for datasets up to a size of 30k, following DiffAug [a] and ADA [b].  
However, we acknowledge the value of evaluating performance on the full FFHQ dataset to assess effectiveness and generalizability even under a large dataset. Accordingly, we are currently training TADA on the full FFHQ dataset and will provide the results as soon as the training is complete.

Reviewer comment (Experiments):
The paper would benefit from a more detailed comparison with existing data augmentation methods, particularly in how the results translate to qualitative improvements in the generated samples.

Author response:
While extensive research has been conducted on data augmentation for generative adversarial networks (GANs), limited research has explored this for auto-regressive and diffusion models. Specifically, no prior studies have investigated data augmentation for diffusion models, with only a rudimentary method introduced in EDM, leaving no other model for comparison. We believe this gap in research presents a significant opportunity for further exploration and advancement of data augmentation techniques for diffusion models.

The reason that our method improves the robustness and generalization of diffusion models is not solely attributed to our proposed method but can also result from the use of data augmentation itself. Our method's primary impact lies in alleviating distribution shifts by adaptively adjusting the augmentation strength during specific timesteps. This strategy involves suppressing augmentation steps in regions where the diffusion model is susceptible to distribution shift, effectively mitigating this issue. Simultaneously, it allows for the application of diverse augmentation techniques with high augmentation probability, leveraging the benefits of data augmentation to promote better generalizability.

Reviewer comment (Theory):
Could you provide more detail on how the augmentation strength is adjusted during different timesteps? Is there a theoretical framework or empirical evidence that supports the chosen method?

Author response:
As illustrated in Figure 2, augmentation strength (w_t) can be expressed as a quadratic curve across all timesteps. Specifically, since SNR can be determined as a function of timestep, we put the SNR value into Equation 3 to compute the augmentation strength w_t. Consequently, we adjust the augmentation strength by multiplying w_t to the maximum augmentation strength applied to training data. We newly provided an additional pseudo code describing the whole procedure in Appendix D.3.

While the strength of the data augmentation can be kept constant across all timesteps, toy experiments suggest that diffusion models trained with augmented training data become more susceptible to producing unintended transformed samples during the intermediate timesteps. Based on this empirical evidence, we designed an augmentation strength schedule that suppresses the augmentation during the intermediate timesteps.
The rationale behind our choice of a quadratic function is twofold:

* Simplicity: The quadratic function provides a straightforward and intuitive approach to controlling the augmentation strength across timesteps. This simplicity makes it easy for users to adjust the augmentation intensity of a sensitive region without the need for complex hyperparameter tuning.

* Smooth Transition: The quadratic function facilitates a seamless transition between strong and weak augmentation depending on the timestep axis. Toy experiments have demonstrated that data augmentation applied to sensitive timestep can significantly impact the generated samples. However, precisely identifying these sensitive regions is challenging. Therefore, we opted for a smooth function that ensures weak augmentation around sensitive timesteps while still enabling data augmentation to be applied to a substantial portion of the timesteps during training.

This allows us to attain a balance between simplicity and effectiveness, making it a reasonable choice for the proposed method. 

We apologize for any ambiguity in the paper. We will revise the manuscript to improve the clarity of the methodology.

Reviewer comment (Experiments):
In your experiments, did you explore the impact of the timestep-aware augmentation on the convergence speed and stability of the training process for diffusion models?

Author response:
Our proposed framework exhibits improved convergence speed and stability compared to the baseline (h-flip) and alleviates overfitting. As shown in Figure 4, FID scores computed every 10k training iterations indicate that our framework converges slightly faster than the baseline models. Additionally, our framework effectively alleviates overfitting, as evidenced by the lower FID scores at later stages of training.

Reviewer comment (Experiments):
How does TADA compare with other state-of-the-art data augmentation methods in terms of qualitative and quantitative results?

Author response:
While extensive research has been conducted on data augmentation for generative adversarial networks (GANs), limited research has explored this for auto-regressive and diffusion models. Specifically, no prior studies have investigated data augmentation for diffusion models, with only a rudimentary method introduced in EDM, leaving no other model for comparison. We believe this gap in research presents a significant opportunity for further exploration and advancement of data augmentation techniques for diffusion models.

Reviewer comment (Evaluation):
Could you provide insights or metrics on how the proposed method improves the robustness and generalization of diffusion models in scenarios with limited data?

Author response:
The reason that our method improves the robustness and generalization of diffusion models is not solely attributed to our proposed method but can also result from the use of data augmentation itself. Our method's primary impact lies in alleviating distribution shifts by adaptively adjusting the augmentation strength during specific timesteps. This strategy involves suppressing augmentation steps in regions where the diffusion model is susceptible to distribution shift, effectively mitigating this issue. Simultaneously, it allows for the application of diverse augmentation techniques with high augmentation probability, leveraging the benefits of data augmentation to promote better generalizability.

Reviewer comment (Evaluation):
What are the potential limitations or drawbacks of your proposed method?

Author response:
As we mentioned in the discussion section of the paper, limitation lies in the lack of a profound theoretical basis. The majority of diffusion models that are currently used and studied are based on the same objective and architecture with variations on hyperparameters [a, b, c]. Therefore, we believe TADA can be seamlessly integrated into various diffusion models. To confirm this, we are currently training under VP [a]  and will provide the results under this setting.

Meanwhile, a recent study [d] found that latent diffusion has different timestep-wise task affinity compared to ADM (that we mainly used). Therefore, to apply our TADA to latent diffusion models, the boundary regions should be adjusted.   

Furthermore, we believe that our method is applicable across wide datasets with limited data settings. However, as we mentioned in the discussion of our paper, it can be less effective when an abundant amount of data is available.

Reviewer comment (Evaluation):
Are there specific types of diffusion models or data domains where TADA might be less effective?

Author response:
The majority of diffusion models that are currently used and studied are based on the same objective and architecture with variations on hyperparameters [a, b, c]. Therefore, we believe TADA can be seamlessly integrated into various diffusion models. To confirm this, we are currently training under VP [a]  and will provide the results under this setting.

Meanwhile, a recent study [d] found that latent diffusion has different timestep-wise task affinity compared to ADM (that we mainly used). Therefore, to apply our TADA to latent diffusion models, the boundary regions should be adjusted.   

Furthermore, we believe that our method is applicable across wide datasets with limited data settings. However, as we mentioned in the discussion of our paper, it can be less effective when an abundant amount of data is available.

Reviewer comment (Experiments):
How scalable is the TADA method with respect to model size and dataset complexity?

Author response:
We varied the model size (i.e., the number of parameters), and found that our method showed noticeable performance improvement in smaller model sizes (Table 4). This implies that TADA offers huge benefits when large computational resources and abundant training data are not available. 
In addition to the results on various sizes of the FFHQ dataset, we demonstrated our scalability on the AFHQ dataset. To further show the scalability of our method, we are currently training on CIFAR10. We will report the results as soon as the training is completed.

Reviewer comment (Experiments):
Have you tested the method on out-of-distribution samples or on datasets with significant class imbalances?

Author response:
We have not yet tested our method on out-of-distribution samples or class imbalance datasets. We appreciate your feedback on this point and hope that we can incorporate this aspect into our further research.