Improved Noise Schedule for Diffusion Training

Thank you for your thoughtful feedback. We address your questions or concerns below.

Q1: Effectiveness of noise schedule design over loss weighting schemes

We present additional experimental comparisons between our noise scheduling approach and its derived loss weighting strategy. For this analysis, we employ the Cosine-Ply ($n=2$) schedule described in Section A.3, which builds upon the conventional Cosine schedule and demonstrates superior FID performance. While implementing Cosine-Ply as a loss weighting mechanism yields improvements over the baseline, its effectiveness remains secondary to its direct application as a noise schedule. We also include related results in revised A.6.

Under the prediction target $\mathbf{v}$, Min-SNR and Soft-Min-SNR employ loss weighting functions of $\frac{\min \{ \text{SNR}, \gamma \} }{\text{SNR} + 1}$ and $\frac{1}{\text{SNR} + 1}\frac{1}{\text{SNR}^{-1} + \gamma^{-1}}$, respectively. While these approaches emphasize intermediate noise densities through weighting (shown in Figure 7 in the revised PDF), our proposed method, which directly increases sampling frequency in these regions, demonstrates superior performance compared to such weighting-based strategies.

Q2 Sufficiency of 10k samples to calculate fid scores and more results on other datasets

The use of 10K samples for FID calculation follows established practices in recent literature, such as U-ViT [1]. It is sufficient to use FID-10k to validate the effectiveness of proposed method. Increasing sample size from 10K to 50K leads to lower FID. To further validate our results, we extended our evaluation to 50K generated samples, achieving an FID score of 5.03323. This result substantially outperforms LlamaGen [2], which reports FID scores of 7.15 and 6.09 after 100 and 300 epochs of training, respectively.

We also conduct experiments on CelebA benchmark, which is single domain containing face images. We train the model DiT-M/4 (layer=12, dim=512, num_head=8, and patch_size=4, similar to U-ViT [1]) using cosine and Laplace schedule. The training protocol consisted of 150,000 steps with a batch size of 256, employing minimal data augmentation (random flip only). The comparative FID results are as follows:

We also follow Stable Diffusion 3[1], train on a more comprehensive dataset CC12M[2] and report the FID results here. We download the dataset using webdataset. We train a DiT-base model using CLIP as text conditioner. The input images were preprocessed by cropping and resizing to $256\times 256$ resolution, then compressed into $32\times32\times4$ latent representations. The model training was conducted for 200K iterations using a batch size of 256.

The experimental results demonstrate that our proposed Laplace schedule achieves superior performance compared to Cosine schedule. Specifically, our method yields an FID score of 54.3492, representing a substantial improvement of 4.0127 points over the baseline (58.3619).

Q3: Visual results to demonstrate that this sampling method does not adversely affect the generated outputs.

FID quantifies the statistical similarity between model-generated samples and real-world reference data, where lower scores indicate more similarity to real data. It has been widely used to demonstrate the capability of visual generative models. The images sampled from our method are not adversely affected. We present some randomly generated cases in appendix (Figure 8) in modified version. In this figure, we also present results at different training iterations to demonstrate the fast convergence of our method.

About the format of Table 1 and Table 2, we have modified the format in the revised PDF.

[1] Bao, Fan, et al. "All are worth words: A vit backbone for diffusion models." CVPR. 2023.

[2] Sun P, Jiang Y, Chen S, et al. Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation[J]. arXiv preprint arXiv:2406.06525, 2024.

Thank you for your comments. We address your questions or concerns below.

Q1: Distinction between Table 1 and Table 2.

Table 1 shows the specific mathematical formulation of noise schedules. $\lambda=\log \text{SNR}$ can be used to calculate the coefficients to add noise and $p(\lambda)$ is the distribution of $\lambda$.

Table 2 show the choices of $p(\lambda)$ and $w(\lambda)$ under the unified framework VDM++ by Kingma et al [1]. For example, if we train diffusion models using cosine schedule and constant loss weight, the underlying $w(\lambda)$ is $e^{-\lambda / 2}$ and $p(\lambda)$ is $\text{sech} (\lambda / 2)$.

Q2: Details about the sampling process.

For the detailed sampling process, the sampled signal-to-noise ratios (SNRs) are aligned across different noise schedules. Due to varying noise schedules, the same timestep $t$ corresponds to different SNRs under different noise schedules, as shown in Table 1. Specifically, for the cosine noise schedule, we utilize a uniformly sampled set of timesteps during the generation process: $\{t_0, t_1, \ldots, t_s\}$, which correspond to the SNRs: $\{\text{SNR}_0, \text{SNR}_1, \ldots, \text{SNR}_s\}$. For other noise schedules, such as the Laplace schedule, we compute the corresponding timesteps $\{t'_0, t'_1, \ldots, t'_s\}$ based on the SNRs: $\{\text{SNR}_0, \text{SNR}_1, \ldots, \text{SNR}_s\}$. This computation is feasible because the relationship is defined by a specific mathematical form and the function is monotonic. In practice, when using DDIM for sampling with the Laplace schedule, the timesteps employed are $\{t'_0, t'_1, \ldots, t'_s\}$.

Q3: Why is the Laplace schedule recommended over the Cosine-Scaled schedule?

We adopt such schedule for several reasons. 1. The optimal performance is slightly better than cosine-scaled. 2. $p(\lambda)$ of Laplace is similar to that of cosine schedule with importance sampling as shown in Figure 4. 3. As we summarize in Sec 3.5, the specific mathematical formula is not that important. Even though the performances of these two schedules are superior to standard cosine schedule. , we want to offer a new choice to the community.

Q4: Why not include other noise schedules in Figure 2?

In Figure 2, our primary objective is to demonstrate the superior effectiveness of our noise schedule design compared to existing loss weight approaches. Both Min-SNR and Soft-Min-SNR strategies assign higher weights to intermediate noise intensities and can serve as robust benchmarks. Including additional schedules would potentially obscure the key comparison and complicate the visual interpretation of the results.

Q5: Limited performance improvement and Generalization to other datasets

In the context of class-conditional generation on ImageNet, improvement from 10.85 to 7.96 cannot be seemed as limited improvement. We extend our sample images to 50K to calculate the FID score. We obtain a score of 5.03. Notably, our base model, trained for merely 100 epochs, outperforms the advanced LLamaGen [2] model (which achieves FID scores of 6.09 and 7.15 at 300 and 100 epochs, respectively).

To further validate the effectiveness of our proposed approach, we conduct additional experiments on the CelebA facial dataset, with results presented below:

Additionally, we extend our evaluation to the larger-scale CC12M dataset:

The experimental results from both unconditional image generation on CelebA and text-to-image generation on CC12M demonstrate the consistent superiority of our approach across different domains and tasks.

Thanks for your feedback on our manuscript. We address your questions or concerns below.

Q1: Why does Laplace schedule work and why should we increase the sampling frequency.

We explored the importance of different training time intervals in a diffusion model by dividing the time range $(0,1)$ into four equal segments $\text{bin}_i = \left(\frac{i}{4}, \frac{i + 1}{4}\right), i=0,1,2,3$. A base model $\mathbf{M}$ was trained over the full range, followed by specialized fine-tuning on each segment to create four checkpoints $\mathbf{m}_0$ through $\mathbf{m}_3$. The evaluation process involved using these specialized checkpoints within their respective time intervals while defaulting to the base model elsewhere. Results from FID measurements across these configurations revealed that optimizing the middle level noise (segments 1 and 2) led to better performance, highlighting that under the same training budget, putting more compute to middle level noise densities contributes to better performance of the diffusion model. The related results can be found in revised A.5.

Furthermore, we present an additional derivation for Flow Matching with logit-normal sampling. Research on scaling laws for diffusion transformers [4] reveals that the cosine schedule outperforms vanilla flow matching, while flow matching with logit-normal sampling demonstrates superior performance compared to the cosine schedule, which aligns with our discovering. From the logSNR distribution perspective, flow matching with logit-normal sampling exhibits greater concentration around logSNR=0 compared to both the cosine schedule and vanilla flow matching. We include the derivation and visualization of $p(\lambda)$ in revised Appendix A.4.

Q2: Experiments on additional datasets.

ImageNet, comprising over one million natural images, has been widely used to validate improvements in diffusion models [2, 3]. To validate the generality of our method, we conduct additional experiments on two more datasets, including CelebA (single domain of face images) and CC12M (text to image dataset).

Using the CelebA dataset at 64×64 resolution, we trained a DiT architecture in pixel space. The architecture comprises 12 layers with an embedding dimension of 512, 8 attention heads, and a patch size of 4, implemented across different noise schedules. We evaluated this unconditional single-domain generation model using the FID metric, with results shown below:

Following stable diffusion 3[1], we train diffusion models on a more complicated dataset CC12M[2] dataset , which comprises over 12 million image-text pairs. We download the dataset using webdataset. We train a DiT-base model using CLIP as text conditioner. The images are cropped and resized to $256\times 256$ resolution, compressed to $32\times 32\times 4$ latents and trained for 200k iterations at batch size 256. The experimental results are presented below:

Our method's generalizability is demonstrated through its performance on both unconditional image generation (CelebA) and text-to-image generation (CC12M) tasks.

Q3: Stability of results across different classifier-free guidance (CFG) scales.

In the context of classifier-free guidance, a strong model does not necessarily achieve the lowest FID across different CFG scales. Recent prominent diffusion models search the classifier-free guidance scale to demonstrate their best performance. For example, DiT [2] chooses CFG = 1.5 to report its best performance, while EDM2 [3] selects CFG = 1.4. We conducted sampling at three CFG values and selected the optimal one to better evaluate the actual performance of different models. The model trained with cosine schedule favors CFG=2.0 while the model trained with Laplace schedule favors CFG=3.0.

[1] Bao, Fan, et al. "All are worth words: A vit backbone for diffusion models." CVPR. 2023.

[2] Peebles W, Xie S. Scalable diffusion models with transformers[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 4195-4205.

[3] Karras T, Aittala M, Lehtinen J, et al. Analyzing and improving the training dynamics of diffusion models[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024: 24174-24184.

[4] Liang Z, He H, Yang C, et al. Scaling Laws For Diffusion Transformers[J]. arXiv preprint arXiv:2410.08184, 2024.

Thanks for your reply.

We have conducted additional experiments and substantially revised the manuscript to address your feedback. While we note some discrepancies between your initial and current assessment (with lower rating), we would like to reference your earlier evaluation that highlighted several strengths of our work:

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

We have carefully addressed your current concerns in the revised manuscript and look forward to your evaluation of these improvements. We remain committed to resolving any remaining issues within the revision timeline.

Best,
Authors of submission 3059

Thank you for your feedback. We address your questions or concerns below.

Q1: Key insight / takeaway for readers.

Takeaway: This framework simplifies noise schedule design in diffusion models - instead of complex calculations, researchers can now achieve better results by simply identifying and emphasizing intermediate noise levels in the $p(\lambda)$ distribution.

Our insights align with recent methods that employ flow matching with logit normal sampling. Specifically, we demonstrate how flow matching with logit normal sampling assigns higher intermediate step probabilities than both cosine scheduling and standard flow matching approaches. This aligns with empirical results from our analysis (detailed in Section A.4 and Figure 5). The effectiveness of this approach is evident in recent leading models like Stable Diffusion 3 [1] and Movie Gen [2], where it has improved training efficiency.

Q2: Add more background about noise schedule.

Noise schedule in diffusion models is a function that determines how much noise is added to the input data at each timestep $t$ during the training process, controlling the distribution of noise levels that the neural network learns to remove.

We have added it in introduction in revised manuscript, highlighted in blue for easy reference.

Q3: Relationship between architecture design and loss weight schemes and motivation to design new method.

Our method is orthogonal to architecture design. Our method shares similar insights like loss weight design but achieve superior performance. Specifically, using $\mathbf{v}$ as predict target, the Min-SNR / Soft-Min-SNR loss weight is $\frac{\min\{\text{SNR}, \gamma\}}{\text{SNR} + 1}$ / $\frac{1}{\text{SNR} + 1}\frac{1}{\text{SNR}^{-1} + \gamma^{-1}}$ respectively, which puts more weight on middle level noise densities (as shown in Figure 7).

The loss weight determines the relative importance of each noise level in the training dynamics of diffusion models. Our approach directly allocates computational resources by assigning more training FLOPs to specific noise densities.

We investigated the significance of training intervals in diffusion models by partitioning the time domain $(0,1)$ into four equal segments $\text{bin}_i = \left(\frac{i}{4}, \frac{i + 1}{4}\right), i=0,1,2,3$. After training a base model $\mathbf{M}$ across the entire range, we fine-tuned it on individual segments to create specialized checkpoints $\mathbf{m}_0$ through $\mathbf{m}_3$. The evaluation combined these specialized checkpoints within their corresponding time intervals while utilizing the base model elsewhere. FID measurements demonstrated that optimizing intermediate noise levels (segments 1 and 2) yielded superior performance, indicating that allocating more computational resources to mid-level noise densities enhances the diffusion model's effectiveness under fixed training constraints. The related results can be found in revised A.5.

Q4: Comparison between noise schedule and loss weight from $p(\lambda)$

We have added experiments here to compare with noise schedule and it performed as loss weight. We choose Cosine-Ply as in Equation 22 because it is based on standard cosine schedule and can achieve relative low FID.

We can see that even though such loss weight performs better than standard cosine schedule, but it still lags behind than that is trained as noise schedule.

Q5: Visualization Results

We have provided several visual examples in revised manuscript. We compare the samples generated from DiT-B trained with Cosine schedule and Laplace schedule at different training iterations. Our results demonstrate faster convergence and better visual qualities which aligns with the FID results.

[1] Esser P, Kulal S, Blattmann A, et al. Scaling rectified flow transformers for high-resolution image synthesis[C]//Forty-first International Conference on Machine Learning. 2024.

[2] Polyak A, Zohar A, Brown A, et al. Movie gen: A cast of media foundation models[J]. arXiv preprint arXiv:2410.13720, 2024.

Thank you for your thoughtful feedback. We address your questions or concerns below.

Q1: Performance in other domains or more complex real-world scenarios

ImageNet, comprising over one million natural images, has been widely adopted as a benchmark dataset for validating improvements in diffusion models [4, 5, 6]. In addition to ImageNet, we evaluate our approach on the CelebA dataset (64×64 resolution in pixel space), which consists of face images. We employ a DiT architecture (12 layers, embedding dimension of 512, 8 attention heads, and patch size of 4) using different noise schedules. This is an unconditional generation setting within a single domain. We present FID results as follows:

We also follow Stable Diffusion 3 [1], train on a more complicated dataset CC12M[2] dataset (over 12M image-text pairs) and report the FID results here. We download the dataset using webdataset. We train a DiT-base model using CLIP as text conditioner. The images are cropped and resized to $256\times 256$ resolution, compressed to $32\times 32\times 4$ latents and trained for 200k iterations at batch size 256.

Our method demonstrated strong generalization capabilities across both unconditional image generation using the CelebA dataset and text-to-image generation using the CC12M dataset.

Q2: Gap after longer training.

As we presented in the paper, our goal is to obtain a strong diffusion model under limited and fixed training budget. With longer training, the gap is going to be smaller. That is also because of our model already achieves relative good results. We generate more samples and extend FID-10k to FID-50k, the score comes to 5.03323. Even with only 100 epoch training, our model outperforms LlamaGen [3] with longer training (FID 7.151 at 100 epoch, FID 6.092 at 300 epoch, in table 9 [3]).

Q3: Effects of CFG selection and performance variability across different CFG values

A strong model does not mean the best FID scores over different CFG values. Recent famous diffusion models search the classifier-free guidance scale to demonstrate their performance, e.g., DiT [6] choose 1.5 to report their performance and EDM2 [7] chooses 1.4. We conduct sampling on three cfg values to and choose the best one to better demonstrate the performance of different models. The model trained with cosine schedule favors CFG=2.0 while the model trained with Laplace schedule favors CFG=3.0. For different models, different training datasets, the suitable classifier free guidance differs and we should try different values instead of directly using a pre-defined value.

[1] Esser P, Kulal S, Blattmann A, et al. Scaling rectified flow transformers for high-resolution image synthesis[C]//Forty-first International Conference on Machine Learning. 2024.

[2] Changpinyo S, Sharma P, Ding N, et al. Conceptual 12m: Pushing web-scale image-text pre-training to recognize long-tail visual concepts[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021: 3558-3568.

[3] Sun P, Jiang Y, Chen S, et al. Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation[J]. arXiv preprint arXiv:2406.06525, 2024.

[4] Karras T, Aittala M, Aila T, et al. Elucidating the design space of diffusion-based generative models[J]. Advances in neural information processing systems, 2022, 35: 26565-26577.

[5] Dhariwal P, Nichol A. Diffusion models beat gans on image synthesis[J]. Advances in neural information processing systems, 2021, 34: 8780-8794.

[6] Peebles W, Xie S. Scalable diffusion models with transformers[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 4195-4205.

[7] Karras T, Aittala M, Lehtinen J, et al. Analyzing and improving the training dynamics of diffusion models[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024: 24174-24184.

In the appendix,

1. In A.4, we added a derivation of the distribution $p(\lambda)$ for Flow Matching with Logit Normal Sampling, a recently popular training sampling method, and compared it with the cosine schedule and vanilla flow matching. We explained its superiority from the perspective of $p(\lambda)$.
2. In A.5, we conducted experiments to verify which intervals are more important for diffusion learning. We confirmed that learning intermediate noise strengths is more crucial than learning at the extremes.
3. In A.6, we added experiments comparing noise schedules with implementing them as loss weights, demonstrating the effectiveness of noise schedules.
4. In A.7, we added experiments on more datasets, including unconditional generation on the CelebA dataset, which contains only facial data, and conditional generation on the more complex text-image dataset CC12M. Our approach was effective in both cases, demonstrating its generalizability.
5. In A.8, we added more visual comparisons, showing that our approach not only performs well in FID metrics but also presents superior results and convergence speed in generated samples.

Best,
Authors of submission 3059