PaperlayerView exact revision on arXiv

Source-linked AI summary

Common Diffusion Noise Schedules and Sample Steps are Flawed

Shanchuan Lin, Bingchen Liu, Jiashi Li, Xiao Yang

arXiv:2305.08891v4cs.CV

TL;DR

Common diffusion schedules and samplers can leave training inconsistent with inference, where the model receives pure Gaussian noise. The paper introduces schedule, sampling, prediction, and guidance fixes, and reports that they resolve Stable Diffusion’s brightness restriction while aligning training and inference. A supported limitation is that offset noise can produce brightness values that do not fit the true data distribution.

  • Problem

    Common schedules may retain signal at the final training timestep, while inference supplies pure Gaussian noise, creating a mismatch that restricts Stable Diffusion’s brightness range.

  • Method

    The paper rescales schedules for zero terminal SNR, starts sampling at the last timestep, trains with v prediction, and rescales classifier-free guidance.

  • Results

    The corrected design completely resolves the reported issue and enables images with a diverse brightness range, including responses to explicit white and black-background prompts.

  • Takeaways & Limitations

    Aligning terminal noise and sampling with inference allows diffusion models to generate samples across the brightness range supported by the prompts.

  • Takeaways & Limitations

    Offset noise can produce brightness that does not fit the true data distribution, making it a non-fundamental alternative to correcting the diffusion process.

Abstract

from arXiv · show

We discover that common diffusion noise schedules do not enforce the last timestep to have zero signal-to-noise ratio (SNR), and some implementations of diffusion samplers do not start from the last timestep. Such designs are flawed and do not reflect the fact that the model is given pure Gaussian noise at inference, creating a discrepancy between training and inference. We show that the flawed design causes real problems in existing implementations. In Stable Diffusion, it severely limits the model to only generate images with medium brightness and prevents it from generating very bright and dark samples. We propose a few simple fixes: (1) rescale the noise schedule to enforce zero terminal SNR; (2) train the model with v prediction; (3) change the sampler to always start from the last timestep; (4) rescale classifier-free guidance to prevent over-exposure. These simple changes ensure the diffusion process is congruent between training and inference and allow the model to generate samples more faithful to the original data distribution.

1. Introduction

Common diffusion schedules and sampling implementations create a training–inference mismatch that restricts Stable Diffusion to medium-brightness images. The paper proposes simple corrections that restore the full brightness range and resolve the issue.

  • Stable Diffusion typically produces medium-brightness images, even for prompts requesting solid black or white backgrounds.Generated images have mean brightness around 0 on a -1 to 1 scale.
  • Common noise schedules fail to enforce zero SNR at the final timestep, leaving signal information during training.The leaked signal includes low-frequency information such as each channel’s mean.
  • At inference, pure Gaussian noise replaces the partially informative final training state, creating a training–inference discrepancy that restricts generated brightness.The discrepancy causes the model to continue generating images according to the mean associated with the final training timestep.
  • The paper proposes rescaling schedules for zero terminal SNR, starting samplers at the last timestep, and rescaling classifier-free guidance to address over-exposure.These changes are intended to align training and inference while controlling excessive guidance scale.
  • Experiments report that the corrected design resolves the issue and enables Stable Diffusion to generate much darker, brighter, and more cinematic images.The paper presents the changes as broadly relevant to diffusion-model designs, not only Stable Diffusion.

2. Background

Diffusion models gradually add Gaussian noise through a forward process and learn a Gaussian reverse process to restore data step by step. The formulation uses a variance schedule, closed-form noisy samples, and an SNR defined from the cumulative signal coefficient.

  • The forward process gradually destroys information by adding Gaussian noise according to a manually defined variance schedule β1, ..., βT.The discrete variance-preserving formulation uses αt := 1 − βt and cumulative signal coefficients.
  • A noisy sample at arbitrary timestep t is drawn from a Gaussian conditional distribution with signal coefficient √¯αt and noise coefficient √(1 − ¯αt).Equivalently, xt = √¯αt x0 + √(1 − ¯αt) ε, where ε is standard Gaussian noise.
  • The signal-to-noise ratio at timestep t is defined as SNR(t) := ¯αt / (1 − ¯αt).
  • The reverse process restores information step by step using Gaussian transitions whose means are predicted by neural models.Models commonly predict noise ε instead of directly predicting the reverse-step mean.
  • The reverse trajectory factorizes into the prior p(xT) followed by learned transitions pθ(xt−1|xt) from t = T down to 1.The transition variance can be calculated from forward-process posteriors.

3. Methods

The paper repairs diffusion training–inference mismatches by enforcing zero terminal SNR, training with v prediction, starting sampling at the last timestep, and rescaling classifier-free guidance.

  • 3.1. Enforce Zero Terminal SNR: Common schedules, including Stable Diffusion’s, do not reach zero terminal SNR, leaving signal at the final training timestep and creating train/inference inconsistency.The leaked signal includes low-frequency information such as channel means, while inference begins with pure Gaussian noise.
  • 3.1. Enforce Zero Terminal SNR: Existing schedules are rescaled in √ᾱ_t space by preserving √ᾱ_1, setting √ᾱ_T to zero, and linearly adjusting intermediate timesteps.Cosine schedules can instead remove β_t clipping; future schedules should set β_T = 1.
  • 3.2. Train with V Prediction and V Loss: With zero terminal SNR, ε prediction becomes trivial, so the model is trained with v prediction and v loss instead.At t = T, v_T = x_0, allowing pure-noise input to predict the data-distribution mean conditioned on the prompt; v-loss fine-tuning has similar visual quality to ε loss.
  • 3.3. Sample from the Last Timestep: Samplers should include the last timestep because implementations that omit it expose the model to inputs inconsistent with inference.Trailing selects the last timestep first and is reported to use sample steps more efficiently when S is small.
  • 3.4. Rescale Classifier-Free Guidance: The proposed changes completely resolve the reported Stable Diffusion brightness issue in the authors’ experiments.The guidance-rescale implementation is provided as Algorithm 2.

4. Evaluation

The evaluation compares a reference Stable Diffusion model with a model trained and sampled using the proposed fixes. The fixed model generates a broader brightness range and shows improved FID/IS relative to the reference.

  • Experimental Setup: The fixed model was finetuned on filtered Laion data using the same training configurations as an unchanged reference model.Both models were trained for 50k iterations with batch size 2048, learning rate 1e-4, and EMA decay 0.9999.
  • Qualitative Results: The fixed model generates diverse brightness levels, whereas the flawed reference remains limited to medium brightness.The fixed model follows prompts such as “white background” and “Solid black background” more accurately.
  • Quantitative Results: The quantitative evaluation samples 10k COCO 2014 validation images with corresponding captions and reports FID and IS.The table compares Stable Diffusion with the authors’ model under DDIM sampling.
  • Qualitative Results: The qualitative comparison uses matched seeds while applying DDIM with 50 steps, trailing timestep selection, guidance weight w = 7.5, and rescale factor ϕ = 0.7.Different negative prompts are used for the paired images.
  • Quantitative Results: The authors report improved FID/IS for their model, suggesting better fit to the image distribution and greater visual appeal.

5. Ablation

The ablations examine timestep selection, model behavior at zero terminal SNR, classifier-free guidance rescaling, and offset noise. Trailing steps help most at very small sample counts, guidance rescaling controls overexposure, and offset noise reaches extreme brightness but may misfit the true distribution.

  • 5.1. Comparison of Sample Steps: Trailing timestep selection is noticeably better than linspace at S = 5, while their difference is subtle at S = 25.The comparison uses DDIM and the same seed; the result is reported for a model trained with a zero terminal SNR schedule.
  • 5.2. Analyzing Model Behavior with Zero SNR: At t = T, the zero-terminal-SNR model predicts almost the same result across different noise inputs, with variation beginning at the next sample step.This matches the ideal-model analysis, where the first prediction is invariant to the noise input and later steps introduce variation.
  • 5.4. Comparison to Offset Noise: Offset noise enables very bright and dark samples but may produce brightness values that do not fit the true data distribution.The authors characterize it as addressing the symptom rather than the fundamental diffusion-process issue.

6. Implementation

The section identifies sampler implementation pitfalls under zero terminal SNR and specifies direct formulations that avoid singularities. It also emphasizes correct timestep selection and the model’s behavior at the terminal step.

  • Zero terminal SNR is valid mathematically, but sampler implementations must avoid common timestep and formulation pitfalls.The section frames these issues as implementation concerns affecting diffusion sampling.
  • At t = T, the model generates the prompt-conditioned mean of the data distribution, while variation begins at the next sample step.Figure 5 visualizes this behavior across timestep t for an astronaut-on-horse prompt.
  • Some DDPM implementations convert v prediction to epsilon before sampling, causing lost signal information and division by zero when α_t = 0.The problematic conversion is especially severe at the terminal step under zero SNR.
  • The correct DDPM procedure converts v prediction to x_0 and samples directly with the x_0 formulation.This avoids the singularity introduced by the epsilon formulation.
  • For DDIM, v prediction is converted to epsilon and x_0 before applying the sampler’s x_0-based sampling equation.The section presents this as the corresponding implementation path for DDIM.

7. Conclusion

The paper concludes that diffusion models should use zero-terminal-SNR schedules and begin sampling at the last timestep to align training with inference. It also proposes schedule rescaling and guidance rescaling as practical fixes.

  • Diffusion models should use zero terminal SNR and start sampling from the last timestep to align training behavior with inference.
  • The paper proposes rescaling existing schedules and classifier-free guidance to enforce zero terminal SNR and counter image over-exposure.
Loading 2305.08891v4…