Improved Noise Schedule for Diffusion Training

## Abstract

Abstract Diffusion models have emerged as the de facto choice for generating visual signals.
However, training a single model to predict noise across various levels poses significant challenges, necessitating numerous iterations and incurring significant computational costs.
Various approaches, such as loss weighting strategy design and architectural refinements, have been introduced to expedite convergence.
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 ⁡ SNR SNR \log\text{SNR} roman_log SNR ), theoretically equivalent to a modified noise schedule, is particularly beneficial for training efficiency when increasing the sample frequency around log ⁡ SNR = 0 SNR 0 \log\text{SNR}=0 roman_log SNR = 0 .
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.

## 1 Introduction

Diffusion models have emerged as a pivotal technique for generating visual signals across diverse domains, such as image synthesis Ramesh et al. (); Saharia et al. (); Rombach et al. () and video generation Brooks et al. ().
They are particularly adept at approximating complex distributions, where Generative Adversarial Networks (GANs) may encounter difficulties.
Despite the substantial computational resources and numerous training iterations required for convergence, improving the training efficiency of diffusion models is essential for their application in large-scale scenarios, such as high-resolution image synthesis and video generation.

Architectural enhancements offer a promising path to improve both the training speed and performance of diffusion models.
For instance, the use of Adaptive Layer Normalization Gu et al. (), when combined with zero initialization in the Transformer architecture as demonstrated by Peebles & Xie (), represents such an improvement.
Similarly, the adoption of U-shaped skip connections within Transformers, as outlined in previous works Hoogeboom et al. (); Bao et al. (); Crowson et al. (), also boosts efficiency.
In a parallel development, Karras et al. () have contributed to this endeavor by reengineering the layers of ADM UNet Dhariwal & Nichol () to preserve the magnitudes of activations, weights, and updates, ensuring a more efficient learning process.

Concurrently, various loss weighting designs have been implemented to accelerate the convergence of training. Previous works, such as eDiff-I Balaji et al. () and Min-SNR Hang et al. (), found that the training of diffusion models may encounter conflicts among various noise intensities. Choi et al. () prioritize specific noise levels during training to enhance learning of visual concepts.
Min-SNR Hang et al. () reduces weights of noisy tasks, pursuing the Pareto Optimality in different denoising tasks, validated its effectiveness on multiple datasets and architectures.
A softer version of this approach, aiming to further enhance high-resolution image synthesis within hourglass diffusion models, was introduced by Crowson et al. ().
SD3 Esser et al. () empirically found that it’s crucial to increase the weight of the intermediate noise intensities, which has demonstrated the effectiveness during training the diffusion models.

In this study, we present a novel method to enhance the training of diffusion models by strategically redefining the noise schedule, which is equivalent to importance sampling of the noise across different intensities. However, empirical evidence suggests that allocating more computation costs (FLOPs) to mid-range noise levels (around $\log\text{SNR}=0$) yields superior performance compared to increasing loss weights during the same period, particularly under constrained computational budgets. We experimentally analyze the performance of several different noise schedules, including Laplace, Cauchy, and the Cosine Shifted/Scaled, which are visualized in Figure 1. Notably, the Laplace schedule exhibits favorable performance. We recommend to choose this noise schedule in the future.

We demonstrate the effectiveness of this approach using the ImageNet benchmark, with a consistent training budget of 500K iterations. Evaluated using the FID metric, our results reveal that noise schedules with a concentrated probability density around $\log\text{SNR}=0$ consistently surpass others, as evidenced at both $256\times 256$ and $512\times 512$ resolutions with different prediction target. This research contributes to the advancement of efficient training techniques for diffusion models.

Figure: Figure 1: Illustration of the probability density functions of different noise schedules.
Refer to caption: x1.png

## 2 Method

### 2.1 Preliminaries

Diffusion models Ho et al. (); Yang et al. () learn to generate data by iteratively reversing the diffusion process. We denote the distribution of data points as $\mathbf{x}\sim p_{\text{data}}(\mathbf{x})$.
The diffusion process progressively adds noise to the data, which is defined as:

$$ $\displaystyle\mathbf{x}_{t}=\alpha_{t}\mathbf{x}+\sigma_{t}\bm{\epsilon},\quad \text{where}\quad\bm{\epsilon}\sim\mathcal{N}(0,\mathbf{I}),$ (1) $$

where $\alpha_{t}$ and $\sigma_{t}$ are the coefficients of the adding noise process, essentially representing the noise schedule.
For the commonly used prediction target velocity: $\mathbf{v}_{t}=\alpha_{t}\bm{\epsilon}-\sigma_{t}\mathbf{x}$ Salimans & Ho (), the diffusion model $\theta$ is trained through the Mean Squared Error (MSE) loss:

$$ $\displaystyle\mathcal{L}(\theta)=\mathbb{E}_{\mathbf{x}\sim p_{\text{data}}( \mathbf{x})}\mathbb{E}_{t\sim p(t)}\left[w(t)\lVert\mathbf{v}_{\theta}(\alpha_ {t}\mathbf{x}+\sigma_{t}\bm{\epsilon},t,\mathbf{c})-\mathbf{v}_{t}\rVert_{2}^{ 2}\right],$ (2) $$

where $w(t)$ is the loss weight, $\mathbf{c}$ denotes the condition information. Common practices sample $t$ from the uniform distribution $\mathcal{U}[0,1]$. Kingma et al. () introduced the Signal-to-Noise ratio as $\text{SNR}(t)=\frac{\alpha_{t}^{2}}{\sigma_{t}^{2}}$ to measure the noise level of different states. To simplify, we denote $\lambda=\log\text{SNR}$ to indicate the noise intensities. In the Variance Preserving (VP) setting, the coefficients in Equation 1 can be calculated by $\alpha^{2}_{t}=\frac{\exp(\lambda)}{\exp(\lambda)+1}$, $\sigma^{2}_{t}=\frac{1}{\exp(\lambda)+1}$.

### 2.2 Improved Noise Schedule Design

Given that the timestep $t$ is a random variable sampled from uniform distribution, the noise schedule implicitly defines the distribution of importance sampling on various noise levels. The sampling probability of noise intensity $\lambda$ is:

$$ $\displaystyle p(\lambda)=p(t)\left|\frac{\mathrm{d}t}{\mathrm{d}\lambda}\right|.$ (3) $$

Considering that $t$ satisfies uniform distribution, and $\lambda$ is monotonically decreasing with $t$, we have:

$$ $\displaystyle p(\lambda)=-\frac{\mathrm{d}t}{\mathrm{d}\lambda}.$ (4) $$

We take cosine noise schedule Nichol & Dhariwal () as an example, where $\alpha_{t}=\cos\left(\frac{\pi t}{2}\right)$, $\sigma_{t}=\sin\left(\frac{\pi t}{2}\right)$.
Then we can deduce that $\lambda=-2\log\tan(\pi t/2)$ and $t=2/\pi\arctan e^{-\lambda/2}$.
Thus the distribution of $\lambda$ is: $p(\lambda)=-\mathrm{d}t/\mathrm{d}\lambda=\text{sech}(\lambda/2)/2\pi$.
This derivation illustrates the process of obtaining $p(\lambda)$ from a noise schedule $\lambda(t)$. On the other hand, we can derive the noise schedule from the sampling probability of different noise intensities $p(\lambda)$.
By integrating Equation 4, we have:

$$ $\displaystyle t$ $\displaystyle=1-\int_{-\infty}^{\lambda}p(\lambda)\mathrm{d}\lambda=\mathcal{P }(\lambda),$ (5) $\displaystyle\lambda$ $\displaystyle=\mathcal{P}^{-1}(t),$ (6) $$

where $\mathcal{P}(\lambda)$ represents the cumulative distribution function of $\lambda$. Thus we can obtain the noise schedule $\lambda$ by applying the inverse function $\mathcal{P}^{-1}$. In conclusion, during the training process, the importance sampling of varying noise intensities essentially equates to the modification of the noise schedules.

**Table 1: Overview of various Noise Schedules. The table categorizes them into five distinct types: Cosine, Laplace, Cauchy, and two variations of Cosine schedules. The second column $p(\lambda)$ denotes the sampling probability at different noise intensities $\lambda$. The last column $\lambda(t)$ indicates how to sample noise intensities for training. We derived their relationship in Equation 4 and Equation 6.**
| Noise Schedule | $p(\lambda)$ | $\lambda(t)$ |
| --- | --- | --- |
| Cosine | $\text{sech}\left(\lambda/2\right)/2\pi$ | $2\log\cot\left(\frac{\pi t}{2}\right)$ |
| Laplace | $e^{-\frac{|\lambda-\mu|}{b}}/2b$ | $\mu-b\text{sgn}(0.5-t)\log(1-2|t-0.5|)$ |
| Cauchy | $\frac{1}{\pi}\frac{\gamma}{(\lambda-\mu)^{2}+\gamma^{2}}$ | $\mu+\gamma\tan\left(\frac{\pi}{2}(1-2t)\right)$ |
| Cosine Shifted | $\frac{1}{2\pi}\text{sech}\left(\frac{\lambda-\mu}{2}\right)$ | $\mu+2\log\left(\cot\left(\frac{\pi t}{2}\right)\right)$ |
| Cosine Scaled | $\frac{s}{2\pi}\text{sech}\left(\frac{s\lambda}{2}\right)$ | $\frac{2}{s}\log\left(\cot\left(\frac{\pi t}{2}\right)\right)$ |

### 2.3 Unified Formulation for Diffusion Training

VDM++ Kingma & Gao () proposes a unified formulation that encompasses recent prominent frameworks and loss weighting strategies for training diffusion models, as detailed below:

$$ $\displaystyle\mathcal{L}_{w}(\theta)=\frac{1}{2}\mathbb{E}_{\mathbf{x}\sim \mathcal{D},\bm{\epsilon}\sim\mathcal{N}(0,\mathbf{I}),\lambda\sim p(\lambda)} \left[\frac{w(\lambda)}{p(\lambda)}\left\lVert\hat{\bm{\epsilon}}_{\theta}( \mathbf{x}_{\lambda};\lambda)-\bm{\epsilon}\right\rVert_{2}^{2}\right],$ (7) $$

where $\mathcal{D}$ signifies the training dataset, noise $\bm{\epsilon}$ is drawn from a standard Gaussian distribution, and $p(\lambda)$ is the distribution of noise intensities.
Different predicting targets, such as $\mathbf{x}_{0}$ and $\mathbf{v}$, can also be re-parameterized to $\bm{\epsilon}$-prediction.
$w(\lambda)$ denotes the loss weighting strategy.
Although adjusting $w(\lambda)$ is theoretically equivalent to altering $p(\lambda)$.
In practical training, directly modifying $p(\lambda)$ to concentrate computational resources on training specific noise levels is more effective than enlarging the loss weight on specific noise levels.
Therefore, we focus on how to design $p(\lambda)$.

### 2.4 Practical Settings

Stable Diffusion 3 Esser et al. (), EDM Karras et al. (), and Min-SNR Hang et al. (); Crowson et al. () find that the denoising tasks with medium noise intensity is most critical to the overall performance of diffusion models. Therefore, we increase the probability of $p(\lambda)$ when $\lambda$ is of moderate size, and obtain a new noise schedule according to Section 2.2.

Specifically, we investigate four novel noise strategies, named Cosine Shifted, Cosine Scaled, Cauchy, and Laplace respectively. The detailed setting are listed in Table 1. Cosine Shifted use the hyperparameter $\mu$ to explore where the maximum probability should be used. Cosine Scaled explores how much the noise probability should be increased under the use of Cosine strategy to achieve better results. The Cauchy distribution, provides another form of function that can adjust both amplitude and offset simultaneously. The Laplace distribution is characterized by its mean $\mu$ and scale $b$, controls both the magnitude of the probability and the degree of concentration of the distribution. These strategies contain several hyperparameters, which we will explore in Section 3.5. Unless otherwise stated, we report the best hyperparameter results.

By re-allocating the computation resources at different noise intensities, we can train the complete denoising process. During sampling process, we standardize the sampled SNR to align with the cosine schedule, thereby focusing our exploration solely on the impact of different strategies during training.
It is important to note that, from the perspective of the noise schedule, how to allocate the computation resource during inference is also worth reconsideration. We will not explore it in this paper and leave this as future work.

## 3 Experiments

### 3.1 implementation Details

Dataset. We conduct experiments on ImageNet Deng et al. () with $256\times 256$ and $512\times 512$ resolution. For each image, we follow the preprocessing in Rombach et al. () to center crop and encode images to latents. The shape of compressed latent feature is $32\times 32\times 4$ for $256^{2}$ images and $64\times 64\times 4$ for $512^{2}$ images.

Network Architecture.
We adopt DiT-B from Peebles & Xie () as our backbone.
We replace the last AdaLN Linear layer with vanilla linear.
Others are kept the same as the original implementation.

Training Settings.
We adopt the Adam optimizer with learning rate $1\times 10^{-4}$.
We set the batch size to 256 as Peebles & Xie (); Hang et al. ().
Each model is trained for 500K iterations if not specified. Our implementation is mainly built on OpenDiT Zhao et al. () and experiments are mainly conducted on 8$\times$16G V100 GPUs.

Baselines and Metrics.
We compare our proposed noise schedule with several baseline settings in Table 2. For each setting, we sample images using DDIM Song et al. () with 50 steps. Despite the noise strategy for different settings may be different, we ensure they are the same at each sampling step. This approach is adopted to exclusively investigate the impact of the noise strategy during the training phase. Moreover, we report results with different classifier-free guidance scales, and the FID is calculated using 10K generated images.

**Table 2: Comparison of different methods and related loss weighting strategies. The $w(\lambda)$ is introduced in Equation 7.**
| Method | $w(\lambda)$ | $p(\lambda)$ |
| --- | --- | --- |
| Cosine | $e^{-\lambda/2}$ | $\text{sech}(\lambda/2)$ |
| Min-SNR | $e^{-\lambda/2}\cdot\min\{1,\gamma e^{-\lambda}\}$ | $\text{sech}(\lambda/2)$ |
| Soft-Min-SNR | $e^{-\lambda/2}\cdot\gamma/(e^{\lambda}+\gamma)$ | $\text{sech}(\lambda/2)$ |
| FM-OT | $(1+e^{-\lambda})\text{sech}^{2}(\lambda/4)$ | $\text{sech}^{2}(\lambda/4)/8$ |
| EDM | $(1+e^{-\lambda})(0.5^{2}+e^{-\lambda})\mathcal{N}(\lambda;2.4,2.4^{2})$ | $(0.5^{2}+e^{-\lambda})\mathcal{N}(\lambda;2.4,2.4^{2})$ |

### 3.2 Comparison with baselines and loss weight designs

This section details the principal findings from our experiments on the ImageNet-256 dataset, focusing on the comparative effectiveness of various noise schedules and loss weightings in the context of CFG values. Table 3 illustrates these comparisons, showcasing the performance of each method in terms of the FID-10K score.

The experiments reveal that our proposed noise schedules, particularly Laplace, achieve the most notable improvements over the traditional cosine schedule, as indicated by the bolded best scores and the blue numbers representing the reductions compared to baseline’s best score of 10.85.

We also provide a comparison with methods that adjust the loss weight, including Min-SNR and Soft-Min-SNR. We find that although these methods can achieve better results than the baseline, they are still not as effective as our method of modifying the noise schedule. This indicates that deciding where to allocate more computational resources is more efficient than adjusting the loss weight. Compared with other noise schedules like EDM Karras et al. () and Flow Lipman et al. (), we found that no matter which CFG value, our results significantly surpass theirs under the same training iterations.

**Table 3: Comparison of various noise schedules and loss weightings on ImageNet-256, showing the performance (in terms of FID-10K) of different methods under different CFG values. The best results highlighted in bold and the blue numbers represent the improvement when compared with the baseline FID 10.85. The line in gray is our suggested noise schedule.**
| Method | CFG=1.5 | CFG=2.0 | CFG=3.0 |
| --- | --- | --- | --- |
| Cosine Nichol & Dhariwal () | 17.79 | 10.85 | 11.06 |
| EDM Karras et al. () | 26.11 | 15.09 | 11.56 |
| FM-OT Lipman et al. () | 24.49 | 14.66 | 11.98 |
| Min-SNR Hang et al. () | 16.06 | 9.70 | 10.43 |
| Soft-Min-SNR Crowson et al. () | 14.89 | 9.07 | 10.66 |
| Cosine Shifted Hoogeboom et al. () | 19.34 | 11.67 | 11.13 |
| Cosine Scaled | 12.74 | 8.04 | 11.02 |
| Cauchy | 12.91 | 8.14 | 11.02 |
| Laplace | 16.69 | 9.04 | 7.96 (-2.89) |

Furthermore, we investigate the convergence speed of these method, and the results are shown in Figure 2. It can be seen that adjusting the noise schedule converges faster than adjusting the loss weight. Additionally, we also notice that the optimal training method may vary when using different CFG values for inference, but adjusting the noise schedule generally yields better results.

Figure: Figure 2: Comparison between adjusting the noise schedule, adjusting the loss weights and baseline setting. The Laplace noise schedule yields the best results and the fastest convergence speed.
Refer to caption: x2.png

### 3.3 Robustness on different predicting targets

We evaluate the effectiveness of our designed noise schedule across three commonly adopted prediction targets: $\bm{\epsilon}$, $\mathbf{x}_{0}$ and $\mathbf{v}$.
The results are shown in Table 4.

We observed that regardless of the prediction target, our proposed Laplace strategy significantly outperforms the Cosine strategy. It’s noteworthy that as the Laplace strategy focuses the computation on medium noise levels during training, the extensive noise levels are less trained, which could potentially affect the overall performance. Therefore, we have slightly modified the inference strategy of DDIM to start sampling from $t_{\max}=0.99$.

**Table 4: Effectiveness evaluated using FID-10K score on different predicting targets. The proposed Laplace schedule performs better than the baseline Cosine schedule along with training iterations.**
| Predict Target | Noise Schedule | 100K | 200k | 300k | 400k | 500k |
| --- | --- | --- | --- | --- | --- | --- |
| $\mathbf{x}_{0}$ | Cosine | 35.20 | 17.60 | 13.37 | 11.84 | 11.16 |
|  | Laplace (Ours) | 21.78 | 10.86 | 9.44 | 8.73 | 8.48 |
| $\mathbf{v}$ | Cosine | 25.70 | 14.01 | 11.78 | 11.26 | 11.06 |
|  | Laplace (Ours) | 18.03 | 9.37 | 8.31 | 8.07 | 7.96 |
| $\bm{\epsilon}$ | Cosine | 28.63 | 15.80 | 12.49 | 11.14 | 10.46 |
|  | Laplace (Ours) | 27.98 | 13.92 | 11.01 | 10.00 | 9.53 |

### 3.4 Robustness on high resolution images

To explore the robustness of the adjusted noise schedule to different resolutions, we also designed experiments on Imagenet-512. As pointed out by Chen (), the adding noise strategy will cause more severe signal leakage as the resolution increases. Therefore, we need to adjust the hyperparameters of the noise schedule according to the resolution.

Specifically, the baseline Cosine schedule achieves the best performance when the CFG value equals to 3. So we choose this CFG value for inference.
Through systematic experimentation, we explored the appropriate values for the Laplace schedule’s parameter $b$, testing within the range {0.5, 0.75, 1.0}, and determined that $b=0.75$ was the most effective, resulting in an FID score of 9.09. This indicates that despite the need for hyperparameter tuning, adjusting the noise schedule can still stably bring performance improvements.

**Table 5: FID-10K results on ImageNet-512. All models are trained for 500K iterations.**
| Noise Schedule | Cosine | Laplace |
| --- | --- | --- |
| FID-10K | 11.91 | 9.09 (-2.82) |

### 3.5 Ablation Study

We conduct an ablation study to analyze the impact of hyperparameters on various distributions of $p(\lambda)$, which are enumerated below.

Laplace distribution is easy to implement and we adjust the scale to make the peak at the middle timestep.
We conduct experiments with different Laplace distribution scales $b\in\{0.25,0.5,1.0,2.0,3.0\}$. The results are shown in Figure 3. The baseline with standard cosine schedule achieves FID score of 17.79 with CFG=1.5, 10.85 with CFG=2.0, and 11.06 with CFG=3.0 after 500K iterations.
We can see that the model with Laplace distribution scale $b=0.5$ achieves the best performance 7.96 with CFG=3.0, which is relatively 26.6% better than the baseline.

Figure: Figure 3: FID-10K results on ImageNet-256 with different Laplace distribution scales $b$ in $\{0.25,0.5,1.0,2.0,3.0\}$. The location parameter $\mu$ is fixed to 0. Baseline denotes standard cosine schedule.
Refer to caption: x3.png

Cauchy distribution is another heavy-tailed distribution that can be used for noise schedule design. The distribution is not symmetric when the location parameter is not 0.
We conduct experiments with different Cauchy distribution parameters and the results are shown in Table 6.
Cauchy(0, 0.5) means $\frac{1}{\pi}\frac{\gamma}{(\lambda-\mu)^{2}+\gamma^{2}}$ with $\mu=0,\gamma=0.5$.
We can see that the model with $\mu=0$ achieve better performance than the other two settings when fixing $\gamma$ to 1.
It means that the model with more probability mass around $\lambda=0$ performs better than others biased to negative or positive directions.

**Table 6: FID-10k results on ImageNet-256 with different Cauchy distribution parameters.**
|  | Cauchy(0, 0.5) | Cauchy(0, 1) | Cauchy(-1, 1) | Cauchy(1, 1) |
| --- | --- | --- | --- | --- |
| CFG=1.5 | 12.91 | 14.32 | 18.12 | 16.60 |
| CFG=2.0 | 8.14 | 8.93 | 10.38 | 10.19 |
| CFG=3.0 | 11.02 | 11.26 | 10.81 | 10.94 |

Cosine Shifted Hoogeboom et al. () is the shifted version of the standard cosine schedule. We evaluate the schedules with both positive and negative $\mu$ values.
Shifted with $\mu=1$ achieves FID-10k score $\{19.34,11.67,11.13\}$ with CFG $\{1.5,2.0,3.0\}$.
Results with shifted value $\mu=-1$ are $\{19.30,11.48,11.28\}$.
Comparatively, both scenarios demonstrate inferior performance relative to the baseline cosine schedule ($\mu=0$). Additionally, by examining the data presented in Table 6, we find concentrated on $\lambda=0$ can best improve the results.

Cosine Scaled is also a modification of Cosine schedule. When $s$ is equal to 1, it becomes the standard Cosine version. $s>1$ means sampling more heavily around $\lambda=0$ while $s<1$ means sampling more uniformly of all $\lambda$. We report related results in Table 7.
Larger values of $s(s>1)$ outperform the baseline; however, $s$ should not be excessively large and must remain within a valid range. A model trained with $s=2$ attains a score of 8.04, representing a $\mathbf{25.9\%}$ improvement over the baseline.

**Table 7: FID-10k results on ImageNet-256 with different scales of Cosine Scaled distribution.**
| $1/s$ | 1.3 | 1.1 | 0.5 | 0.25 |
| --- | --- | --- | --- | --- |
| CFG=1.5 | 39.74 | 22.60 | 12.74 | 15.83 |
| CFG=2.0 | 23.38 | 12.98 | 8.04 | 8.64 |
| CFG=3.0 | 13.94 | 11.16 | 11.02 | 8.26 |

## 4 Related Work

### Efficient Diffusion Training

Generally speaking, the diffusion model uses a network with shared parameters to denoise different noise intensities. However, the different noise levels may introduce conflicts during training, which makes the convergence slow. Min-SNR Hang et al. () seeks the Pareto optimal direction for different tasks, achieves better convergence on different predicting targets.
HDiT Crowson et al. () propose a soft version of Min-SNR to further improve the efficiency on high resolution image synthesis. Stable Diffusion 3 Esser et al. () puts more weight on the middle timesteps by multiplying the distribution of logit normal distribution.
On the other hand, architecture modification is also explored to improve diffusion training. DiT Peebles & Xie () proposes adaptive Layer Normalization with zero initialization to improve the training of Transformer architectures. A more robust ADM UNet with better training dynamics is proposed in EDM2 Karras et al. () by preserving activation, weight, and update magnitudes.

### Noise Schedule Design for Diffusion Models

The design of the noise schedule plays a critical role in training diffusion models. In DDPM, Ho et al. () propose linear schedule for the noise level, which is later used in Stable Diffusion Rombach et al. () version 1.5 and 2.0. iDDPM Nichol & Dhariwal () introduces a cosine schedule aimed at bringing the sample with the highest noise level closer to pure Gaussian noise. EDM Karras et al. () proposes a new continuous framework and make the logarithm of noise intensity sampled from a Gaussian distribution. Flow matching with optimal transport Lipman et al. (); Liu et al. () linearly interpolates the noise and data point as the input of flow-based models.  Chen () underscored the need for adapting the noise schedule according to the token length, and several other works Lin et al. (); Tang et al. () emphasize that it’s important to prevent signal leakage in the final step.

## 5 Conclusion

In this technical report, we present a novel method for enhancing diffusion model training by redefining the noise schedule. We theoretically analyzed that this approach equates to performing importance sampling on the noise.
Empirical results show that our proposed Laplace noise schedule, focusing computational resources on mid-range steps, yields superior performance compared to the adjustment of loss weights under constrained budgets.
This study not only contributes significantly to developing efficient training techniques for diffusion models but also offers potential for future large-scale applications.

## Appendix A: Detailed Implementation for Noise Schedule

We provide a simple PyTorch implementation for the Laplace noise schedule and its application in training. This example can be adapted to other noise schedules, such as the Cauchy distribution, by replacing the laplace_noise_schedule function. The model accepts noisy samples $\mathbf{x}_{t}$, timestep $t$, and an optional condition tensor $\mathbf{c}$ as inputs. This implementation supports prediction of $\{\mathbf{x}_{0},\mathbf{v},\bm{\epsilon}\}$.

## Appendix B: Details for Sampling Process

As we mentioned before, choosing which noise schedule for sampling worth exploration. In this paper, we focus on exploring what kind of noise schedule is needed for training. Therefore, we adopted the same inference strategy as the cosine schedule to ensure a fair comparison. Specifically, first we sample $\{t_{0},t_{1},\ldots,t_{s}\}$ from uniform distribution $\mathcal{U}[0,1]$, then get the corresponding SNRs from Cosine schedule: $\{\frac{\alpha^{2}_{t_{0}}}{\sigma^{2}_{t_{0}}},\frac{\alpha^{2}_{t_{1}}}{
\sigma^{2}_{t_{1}}},\ldots,\frac{\alpha^{2}_{t_{s}}}{\sigma^{2}_{t_{s}}}\}$. According to Equation 6, we get the corresponding $\{t^{\prime}_{0},t^{\prime}_{1},\ldots,t^{\prime}_{s}\}$ by inverting these SNR values through the respective noise schedules. Finally, we use DDIM Song et al. () to sample with these new calculated $\{t^{\prime}\}$.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
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.
soundness: 2
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.
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.
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?

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see the weakness section.
soundness: 2
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: See weaknesses.
soundness: 2
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
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?
soundness: 2
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.
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.
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.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
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?
soundness: 3
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.
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.
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.

Meta-review: 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.

Decision: Reject