Improved Noise Schedule for Diffusion Training

Authors: Tiankai Hang, Shuyang Gu, Xin Geng, Baining Guo
Primary area: generative models
Keywords: diffusion model, effcient training, noise schedule, image generation

Abstract:
Diffusion models have emerged as the de facto choice for generating high-quality visual content across multiple domains.
However, training a single model to predict noise at multiple levels presents significant challenges, requiring numerous iterations and resulting in substantial computational costs.
Various approaches, such as loss weighting strategy design and architectural refinements, have been introduced to expedite convergence and improve model performance.
In this study, we propose a novel approach to design the noise schedule for enhancing the training of diffusion models. Our key insight is that the importance sampling of the logarithm of the Signal-to-Noise ratio ($\log \text{SNR}$), theoretically equivalent to a modified noise schedule, is particularly beneficial for training efficiency when increasing the sample frequency around $\log \text{SNR}=0$. This strategic sampling allows the model to focus on the critical transition point between signal dominance and noise dominance, potentially leading to more robust and accurate predictions.
We empirically demonstrate the superiority of our noise schedule over the standard cosine schedule.
Furthermore, we highlight the advantages of our noise schedule design on the ImageNet benchmark, showing that the designed schedule consistently benefits different prediction targets.
Our findings contribute to the ongoing efforts to optimize diffusion models, potentially paving the way for more efficient and effective training paradigms in the field of generative AI.

Decision: reject
Reviewer scores: 5, 3, 5, 3, 5

Reviews and discussion:
<thread>
creation date:
1729772767228

thread_type:
official_review

<thread_object>
creator: reviewer_f1vg

<content>
summary:
This paper presents a modified noise schedule that increases the sampling frequency around $\log SNR = 0$ to accelerate convergence and generate robust, accurate predictions in diffusion models. Experimental results demonstrate the effectiveness of this approach and its advantages across various prediction targets in the ImageNet benchmark.

strengths:
S1. The problem of improving the training efficiency of diffusion models is well-motivated because diffusion models are known to be computationally intensive.

S1. The proposed method redefines the noise schedule to enhance the training of diffusion models, resulting in higher efficiency and demonstrating strong predictive performance.

S2. The effectiveness of the proposed method is confirmed through extensive experiments.

weaknesses:
W1. The theoretical support for the proposed method is lacking. While the paper establishes a connection between the sampling probability and the noise schedule, the conclusion regarding the choice of λ at a moderate strength primarily draws from prior work, lacking unique insights or innovations in its key findings.

W2. The experimental results' persuasive power is insufficient; additional datasets, such as CIFAR-10 and CelebA, are needed to validate the generation performance. Increasing the number of generated images for FID calculation would also be beneficial.

W3. Minor presentation issues. For instance, Table 1 differs in formatting from other tables.

questions:
Q1. Is there more empirical evidence to support the claim that "In practical training, directly modifying p(λ) to concentrate computational resources on specific noise levels is more effective than increasing the loss weight on those levels"?

Q2. The paper mentions calculating FID using 10K generated images. Is this sample size sufficient, and can the FID calculation be based on other widely adopted datasets such as CIFAR-10 and CelebA?

Q3. The paper uses only FID to indicate the impact on quality. Could visual examples be provided to demonstrate that this sampling method does not adversely affect the generated outputs?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730335661489

thread_type:
official_review

<thread_object>
creator: reviewer_fzj6

<content>
summary:
The paper analyzes the effect of various noise schedule choices for training diffusion models for ImageNet generation.  The authors show that noise schedule hyperparameter optimization favors putting more samples at times when the signal-to-noise ratio (SNR) is near one, and they found the model trained with Laplace schedule worked best in their experiments.  They evaluated the performance of the trained generative models by their FID score on ImageNet class-conditional image generation.

strengths:
- Overall the paper's main contribution is an empirical investigation of a set of noise schedules.  The observation that training with more samples near SNR = 1 works better in their experiments could be relevant to practice.

weaknesses:
- The results and claims are not significant contributions in my view, and some claims are unsupported.  In general, the paper lacks theoretical backing for their chosen noise schedule and the set of examined noise schedules.  The lack of theoretical grounding could be overcome by strong empirical analysis, but the empirical results show modest FID gains on ImageNet and are not conclusive enough to conclude this noise schedule should always be preferred.  If the primary argument is empirical, it would be good to see the claims generalizes across datasets.  Similarly, the claim that adjusting loss weights is worse than using a different noise schedule is too broad, as they did not optimize over different loss weighting schemes.

- The writing was also confusing.  For example, Equation 2 and Equation 7 use conflicting notation to describe the same loss optimization.  Also the title of section 2.2 includes the word "improved" but was background material on converting between time and SNR sampling; Section 2 did not clarify what was novel to the paper versus already known.  For instance, that the loss can be written in terms of sampling SNR and that loss weights are equivalent to different sampling is apparent and discussed in prior research.  The choice of figures is also surprising as figure 1 shows different distributions over log SNR but the paper contains no generated images displaying qualitative results.

questions:
- Can the authors provide more details on the distinction between the contents of Table 1 and Table 2?

- The description of the inference strategy in Section 2.4 was brief but is important in understanding the results.  Can you expand on the precise procedure used at sampling time?

- The advantage of Laplace over Cosine Scaled in Table 3 appears limited.  Cosine Scaled is more stable with respect to CFG than Laplace, while almost meeting the optimized performance of Laplace.  Why suggest Laplace over Cosine Scaled given that Cosine Scaled seems more robust? 

- In Figure 2, why are the other noise schedules not included?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707999486

thread_type:
official_review

<thread_object>
creator: reviewer_rshk

<content>
summary:
This work aims to improve the efficiency of the training of diffusion models. The authors show that noise schedule is equivalent to importance sampling over noise intensities. Based on this, the authors derive some novel noise schedules to effectively assign more importance to mid-range noise intensities. The proposed method is experimentally demonstrated to be effective.

strengths:
1.	The proposed method enjoys clean formulation.
2.	Judging from the experimental results provided, the proposed method shows consistent improvement over competitors in various scenarios.
3.	The formulation to relate noise schedule to noise importance sampling is rather universal, implying potential future extension.

weaknesses:
Generally speaking, while the proposed method has the potential to make great contributions, the presentation of the manuscript makes it hard for the community to learn valuable new knowledge from this work. More concretely:

For Writing:

1. What is the most important key insight or takeaway for readers? While method details are extensive, there lacks a more general summary of the key insight, especially in the introduction section.
2. The authors should provide more background for noise schedule in the introduction section, letting readers know the actual meaning and context of noise schedule as early as possible.
3. What is the relation between the proposed method and the two previous approaches (loss weighting and architectural changes)? What inspires the authors to devise the new method given the existence of previous approaches?

For Evaluation:

4. The comparison of loss weighting and noise schedule might not be entirely fair. I would recommend converting the same distribution into both loss weighting and noise schedule for a more insightful comparison.
5. There can be some qualitative visualization evaluation. For example, showcasing image samples across different training steps with different methods.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730719870984

thread_type:
official_review

<thread_object>
creator: reviewer_em87

<content>
summary:
This paper presents a novel method for enhancing the training of diffusion models by strategically redefining the noise schedule. The key insight is the importance sampling of the logarithm of the Signal-to-Noise ratio (log SNR) theoretically equivalent to a modified noise schedule. Empirical results show that the proposed Laplace noise schedule, which focuses computational resources on mid-range noise levels, yields superior performance compared to adjusting loss weights under constrained computational budgets.

strengths:
1. The paper presents a novel approach to design the noise schedule for enhancing the training of diffusion models.
2. The proposed method is effective and easy to apply.
3. The findings contribute to the ongoing efforts to optimize diffusion models, potentially paving the way for more efficient and effective training paradigms.
4. The paper is overall well-written and easy to follow.

weaknesses:
1. The paper lacks a formal theoretical guarantee for the effectiveness of the proposed method. While the experimental results are promising, providing a theoretical foundation would strengthen the validity of the approach. It would be better to add more analysis on why Laplace schedule is working and why we should increase the sample frequency around $log SNR = 0$.
2. The experiments are only conducted on the ImageNet dataset with different resolutions. However, the generalizability of the proposed method to other datasets or tasks remains untested.
3. The effectiveness of this method seems not stable enough. According to Table 3, the proposed Laplace is ineffective when CFG=1.5.

questions:
Please see the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730813908823

thread_type:
official_review

<thread_object>
creator: reviewer_i2dm

<content>
summary:
This paper presents an approach to improving the training efficiency of diffusion models, focusing on the design of the noise schedule. The authors propose a strategy based on importance sampling the logarithm of the Signal-to-Noise Ratio (log SNR), specifically concentrating on noise levels around log SNR = 0. This targeted focus on the transition between signal dominance and noise dominance helps improve the model’s ability to make more robust predictions. 

The paper compares various noise schedules (Laplace, Cauchy, Cosine Shifted, and Cosine Scaled) and empirically shows that the Laplace noise schedule outperforms the traditional cosine schedule, achieving better results in terms of training efficiency. The proposed method demonstrates its effectiveness on benchmarks like ImageNet and provides insights into optimizing diffusion model training under constrained computational resources.

strengths:
1. The paper proposes a noise schedule design, focusing computational resources on medium noise levels (log SNR = 0). It provides a perspective on optimizing diffusion model training, potentially paving the way for more efficient generative AI methods.

2. The authors conduct a thorough set of experiments, comparing various noise schedules and loss weighting strategies. The experimental setup is robust, testing the methods on multiple resolutions and prediction targets, and demonstrating Laplace schedule’s performance.

3. By improving training efficiency and robustness under constrained computational budgets, this work addresses a challenge for applying diffusion models to large-scale applications like high-resolution image synthesis and long video generation.

weaknesses:
1. While the experiments focus on ImageNet and high-resolution image tasks, the paper does not explore how the proposed noise schedules would perform in other domains or more complex real-world scenarios which would be useful for understanding the broader applicability of the approach.

2. Figure/tables might not be fully convincing, and it would be great to explain if they are statistically significant. In Figure 2, we could see that as number of training iterations increase to 500k, the gap between the proposed approach and benchmark shrinks significantly, so a nature question is if this still remains superior as number of training iterations increase more, beyond 500k. In Table 3, it shows that Laplace has the best performance when CFG=3.0, but that is not the case for other CFG values. Should we expect them to outperform for any given CFG value or should we anchor only on its performance at CFG=3.0. If only at CFG=3., how does it impact its adaptivability to different settings or datasets?

questions:
While the results on ImageNet are promising, further experiments on other types of data could help validate the generalizability of the approach. also, what would the impact of the proposed noise schedule be on diffusion models trained in more dynamic or non-stationary environments? It would be useful to explore how the noise schedule adapts in settings where the distribution of data might shift.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732638578085

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response (1/2)

comment:
Dear AC and reviewers,

We sincerely thank the reviewers for their valuable and constructive feedback to help improve our submission! We are encouraged by the recognition of our work's contributions: Multiple reviewers acknowledged the importance of addressing training efficiency in diffusion models and appreciated our thorough experimental validation. Reviewer **eM87** highlighted the ease of implementation, while Reviewer **rSHk** noted our clean formulation. Reviewer **i2Dm** commended our approach's effectiveness under computational constraints. Reviewers **rSHk** and **fzj6** recognized the value of our unified framework and empirical insights. Reviewers **i2Dm**, **eM87**, and **f1Vg** appreciated our method's consistent performance improvements across various experiments.

We have addressed each reviewer's questions individually. Below, we present a concise summary of our responses to some common concerns:

1. **Additional Experiments on Other Datasets**. We conducted additional experiments on CelebA (unconditional generation) and CC12M (text-to-image generation), demonstrating consistent improvements across different domains and tasks. These experiments further demonstrate the *generalization of our method*.
2. **Visual Results and Evaluation**. We showcase our method's *fast convergence* by displaying results at various training iterations using randomly generated cases in the appendix. We also extended our quantitative evaluation to 50K samples (achieving FID 5.03), outperforming recent methods like LlamaGen even with fewer training epochs.
3. **Stability and CFG Analysis**. We clarified that different models naturally favor different CFG values, as acknowledged in recent literature. We experimented with different CFG scales to better illustrate the model's performance.
4. **Comparison with Loss Weighting**. We conducted additional experiments comparing noise schedules with their equivalent loss weight implementations. Results show that directly modifying the noise schedule (FID 7.98) is more effective than implementing it as a loss weight (FID 8.88).
5. **Writing Improvements**. We revised the introduction to better explain the role of the noise schedule in diffusion models and to strengthen the claim of our contribution. We also improved the method section by adjusting titles for better suitability and ensuring smoother transitions between sections.

We have carefully considered each reviewer's feedback and made corresponding revisions. Here is a summary of updates (marked blue in revised PDF):

In the main paper,

1. In the introduction, we revised the second and third paragraphs to simplify the discussion on structural design and loss weight design. This revision ensures a more concise presentation without delving too deeply into related work, resulting in a smoother narrative flow. 
2. In the introduction, we provided a *clearer summary of our contributions* compared to the previous version. This revision makes it easier for readers to follow and understand our key points. 
3. In the method section 2.1, we added an explanation for why we study noise schedules in the PRELIMINARIES section. 
4. In the method section 2.2, we changed the subtitle from `Improved Noise Schedule Design` to `Noise Schedule Design from a Probability Perspective.` This makes the section more specific and adds an explanation of the generality of our approach, providing a more intuitive understanding of existing noise schedules. 
5. In the method section 2.3, we explained the purpose of introducing a unified framework: within this framework, we can now apply these theoretical insights to practical settings. 
6. In the method section 2.4, we added a sentence to make our sampling process description more specific. 
7. In section 3.1, we revised the text to explain why we tested multiple CFG values and their corresponding FID scores. 
8. In Section 3.5, we further summarized the experimental results: While these different schedules take various mathematical forms, they all achieve similar optimal performance when given equivalent training budgets. The specific mathematical formulation is less crucial than the underlying design philosophy: increasing the sampling probability of intermediate noise levels. This principle provides a simple yet effective guideline for designing noise schedules. 
9. For Table 1, we changed the format to align with Table 2.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies noise schedule for training diffusion models that improves efficiency by increasing the sampling frequency around a log SNR of 0, leading to faster convergence and enhanced performance. After the rebuttal phase, the overall comments from our reviewers are negative. Some concerns still need to be addressed, and two major ones are the need for improved writing and the weakness of technical contributions. I have to recommend rejecting this paper, hoping reviewers' feedback will be helpful for future submissions.