Source-linked AI summary
Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation
Diederik P. Kingma, Ruiqi Gao
TL;DR
Diffusion objectives appear different from maximum likelihood and ELBO objectives, despite their use in high-quality image generation. The paper shows that weighted diffusion losses are integrals of ELBOs across noise levels and, under monotonic weighting, equal the ELBO with additive-noise data augmentation.
Problem
State-of-the-art diffusion models use objectives that look different from maximum likelihood and ELBO objectives, motivating an analysis of their relationship.
Method
The paper analyzes diffusion objectives as weighted losses over noise levels and proves an equivalence between monotonic weighting and the ELBO with additive-noise data augmentation.
Results
Weighted diffusion losses have an interpretation as weighted integrals of ELBOs, while monotonic weightings correspond to the ELBO with additive-noise data augmentation.
Takeaways & Limitations
The equivalence enables apples-to-apples comparisons between diffusion models and other likelihood-based models using held-out objectives rather than FID scores.
Takeaways & Limitations
Comparisons with other likelihood-based models using held-out objectives are left for future work.
Abstract
from arXiv · showhide
To achieve the highest perceptual quality, state-of-the-art diffusion models are optimized with objectives that typically look very different from the maximum likelihood and the Evidence Lower Bound (ELBO) objectives. In this work, we reveal that diffusion model objectives are actually closely related to the ELBO. Specifically, we show that all commonly used diffusion model objectives equate to a weighted integral of ELBOs over different noise levels, where the weighting depends on the specific objective used. Under the condition of monotonic weighting, the connection is even closer: the diffusion objective then equals the ELBO, combined with simple data augmentation, namely Gaussian noise perturbation. We show that this condition holds for a number of state-of-the-art diffusion models. In experiments, we explore new monotonic weightings and demonstrate their effectiveness, achieving state-of-the-art FID scores on the high-resolution ImageNet benchmark.
1 Introduction
Diffusion models achieved strong image-generation quality with objectives that appear different from maximum likelihood and ELBO training. This paper unifies these objectives through weighted losses and connects monotonic weighting to ELBO training with Gaussian-noise augmentation.
- Diffusion models evolved from ELBO-optimized generative models into state-of-the-art systems for image and other generation tasks.
- Sample-quality metrics such as FID favored denoising score matching and noise-prediction objectives over traditionally used maximum-likelihood and ELBO objectives.
- The paper shows that commonly used diffusion objectives are closely related to the ELBO through weighted losses over noise levels.
- A monotonic weighting makes the diffusion objective equivalent to the ELBO with Gaussian noise perturbation as data augmentation.
- The paper proposes comparing diffusion models with likelihood-based models using equivalent augmented-data objectives, while leaving such experiments for future work.
2 Model
The model combines a Gaussian forward diffusion process with a learned reverse generative model. Training learns a score network across noise levels, while sampling approximately reverses the process from a simple terminal distribution.
- The diffusion model contains latent variables z_t, a forward conditional process q(z_0,...,z_1|x), and a reverse joint model p(z_0,...,z_1).
- The Gaussian forward process maps each timestep t to a noise level through a strictly decreasing log-SNR schedule λ=f_λ(t).
- Uniform timestep sampling induces a noise-level density p(λ)=−dt/dλ, and the schedule determines how sampling time is distributed across noise levels.
- Sampling starts from z_1∼p(z_1) and approximately solves the reverse SDE using the estimated score.
- A score network s_θ(z;λ) approximates ∇_z log q_t(z); an exact score would permit exact reversal of the forward process.
3 Diffusion Model Objectives
Diffusion training objectives can be expressed as weighted losses over noise levels, with the weighting determined by the objective and training schedule. The resulting integral is schedule-invariant, but schedule choice still affects estimator variance and optimization efficiency.
- Denoising score matching and noise prediction train a score network over noise scales, with noise prediction parameterizing the score as −ε̂_θ(z;λ)/σ_λ.
- Existing diffusion objectives are special cases of a weighted loss with objective-specific weighting functions, summarized in Table 1.
- The ELBO corresponds to uniform weighting, w(λ_t)=1, while noise prediction uses the implied noise-level density p(λ_t).
- The weighted-loss integral depends on the noise-schedule endpoints but not on the schedule shape between them; weighting is the substantive distinction among objectives.
- Although the integral is schedule-invariant, the Monte Carlo estimator and its gradients depend on schedule-induced variance, affecting optimization efficiency.
4 The weighted loss as the ELBO with data augmentation
The paper proves that monotonic weighting transforms the weighted diffusion objective into an ELBO with additive Gaussian-noise data augmentation. This gives a unified interpretation for several diffusion objectives and models.
- Theorem 1 states that a monotonic weighting makes the weighted diffusion objective equivalent to the ELBO with additive-noise data augmentation.
- Monotonic weighting means w(λ_t) increases with t and therefore decreases with λ.
- The proof rewrites the weighted loss using a probability distribution p_w(t) determined by the derivative of the weighting function.
- The intermediate KL term equals the expected negative ELBO of noise-perturbed data up to a constant.
- The result applies to objectives with monotonic implied weighting and represents additive noise as a DistAug-style training augmentation.
5 Experiments
Experiments test monotonic weighting functions and adaptive noise schedules across ImageNet resolutions, finding improved or competitive scores and state-of-the-art unguided high-resolution results.
- Experimental setup: VDM++ combines proposed monotonic weightings and adaptive noise schedules with diffusion models, extending the VDM objective with data augmentation.The models are named VDM++ because their objective is equivalent to the VDM objective combined with data augmentation.
- ImageNet 64x64: Replacing the non-monotonic sech(λ/2) weighting with sigmoid(−λ + k) improved FID and Inception scores for all tested k except k = 5, with k = 2 best.The adaptive training noise schedule produced similar scores with faster training, while the EDM sampler improved scores further.
- High-resolution ImageNet: On ImageNet 128x128, the proposed weightings achieved better or comparable FID and IS scores than the Simple Diffusion baseline.The comparison changed only the training noise schedule and weighting function relative to the baseline setup.
- ImageNet 64x64: Sigmoidal and ‘EDM-monotonic’ weightings slightly outperformed the original weighting for v-prediction models.This matches the experiments with ϵ-prediction and EDM parameterizations.
- High-resolution ImageNet: With shifted ‘EDM-monotonic’ weighting, the method achieved state-of-the-art FID and IS scores at 128, 256, and 512 resolutions without guidance.With classifier-free guidance, it outperformed all diffusion-based approaches at 128 and 512 resolutions and substantially improved over Simple Diffusion across settings.
6 Conclusion and Discussion
The paper establishes that weighted diffusion objectives are weighted integrals of ELBOs, and that monotonic weighting gives an ELBO interpretation with noise-perturbation data augmentation. This equivalence enables direct comparison with other likelihood-based models, while those comparisons remain future work.
- Conclusion: Weighted diffusion loss is a weighted integral of ELBO objectives, with one ELBO corresponding to each noise level.The weighting function determines how the ELBO contributions across noise levels are combined.
- Conclusion: Monotonic weighting makes the diffusion objective equivalent to the ELBO with Gaussian noise perturbation as data augmentation.The augmentation uses a distribution of noise levels.
- Discussion: This equivalence permits apples-to-apples comparisons between diffusion models and other likelihood-based models using held-out objectives rather than FID scores.The paper specifically identifies autoregressive transformers as a possible comparison target, but leaves those experiments for future work.
A Main proof
The proof rewrites weighted diffusion losses through KL divergences and integration by parts, showing that each noise level contributes an ELBO for noise-perturbed data and that monotonic weights yield an expected ELBO.
- Weighted-loss decomposition: The weighted loss is decomposed by integration by parts into boundary terms and an integral of ELBO terms weighted by changes in w(λt).The geometric interpretation treats both sides as equal areas formed from infinitesimal rectangles.
- Boundary terms: The proof treats endpoint contributions as typically small because the relevant KL divergence or endpoint weighting is small by design.These terms are parameter-independent or suppressed by small boundary weights.
- KL–ELBO relationship: At each noise level, L(t; x) equals the expected negative ELBO of noise-perturbed data plus a constant.The joint KL divergence is identified with this expected negative ELBO and an entropy term independent of the model parameters.
- Conclusion of proof: The weighted loss therefore equals the negative expected ELBO of noise-perturbed data up to a constant, so minimizing it is equivalent to maximizing that expected ELBO.The expectation is over both the weighting distribution and the perturbed data distribution.
D Derivation of weighting functions for previous works
The section rewrites diffusion objectives from prior work as weighted losses indexed by noise level, showing how parameterizations and schedules determine their weighting functions. It also establishes equivalences among several objectives, including Flow Matching and v-prediction with a cosine schedule.
- General framework: Prior diffusion objectives can be expressed as special cases of a weighted loss with a weighting function w(λ) determined by the training objective and noise schedule.The weighting specifies the contribution assigned to each noise level.
- EDM: The EDM objective is a special case of the weighted loss after rewriting its noise scale as λ = −2 log ˜σ.Its weighting is given by ˜w(λ) = (e^−λ + 0.52)/(e^−λ · 0.52), up to constant scaling.
- Parameterizations: The v-prediction loss corresponds to w(λ) = (e^−λ + 1)p(λ), while its VP-SDE form follows from the relationship between v-, ϵ-, x-, and score prediction.The section derives these conversions through the parameterization identities.
- Flow Matching and InDI: Flow Matching with the Optimal Transport flow path is equivalent to v-prediction with a cosine schedule, while InDI induces a related weighting that emphasizes lower noise levels.The InDI weighting is e^−λsech^2(λ/4)/8, giving slightly more weight to lower noise than the FM-OT weighting.
- Optimization: The noise schedule affects the variance of the Monte Carlo loss estimator and its gradients, so it influences optimization efficiency even when the integrated loss is schedule-invariant.This distinction separates the exact weighted loss from its training-time estimator.
G.1 Comparison with Theorem 1 by Lyu [2012]
The comparison distinguishes the paper’s result from Lyu’s related theorem by emphasizing weaker assumptions and relevance during optimization. In particular, the paper does not assume that the learned and forward conditional processes are equal.
- Comparison with Lyu [2012]: Lyu’s result assumes a variance-exploding diffusion process and identical forward conditionals p(z_t|x) = q(z_t|x).The cited setup uses t = σ^2_λ and λ = −log(t).
- Comparison with Lyu [2012]: The paper’s corresponding theorem requires fewer assumptions and does not impose p(z_t|x) = q(z_t|x).That equality may fail before or during optimization and cannot simply be assumed after optimization either.
- Optimization relevance: Because the loss is used while parameters are being optimized, the paper considers its own theorem more relevant to optimization than Lyu’s equality-based comparison.The distinction concerns properties of the loss during training, not only the optimized model.
- Implementation: The implementation uses low-discrepancy time sampling to reduce diffusion-loss estimator variance and accelerate optimization.Training also uses Adam with default hyperparameters and global gradient-norm clipping at 1.
I.1 The shape of d dλL(λ; x) for low-bit data
The section analyzes how diffusion-loss contributions vary with bit precision using a univariate approximation, then identifies a weighting that makes 8-bit training resemble unweighted 5-bit training. The resulting weighting concentrates more strongly on low noise levels than the paper’s other weightings.
- Approximation: The analysis approximates data as uniformly distributed over 2^n pixel values and models the noisy marginal as a univariate Gaussian mixture.The data are normalized to [−1, 1], and L(λ; x) is defined through a KL divergence.
- Bit contributions: The study plots expected loss and per-bit contributions across λ for bit precisions n = 1, ..., 8.Subtracting the n−1-bit curve from the n-bit curve isolates the contribution of each additional bit.
- Bit contributions: Each bit-contribution curve has area exactly 1 bit.This provides the stated normalization for interpreting the plotted contributions.
- 5-bit equivalence: 5-bit training resembles 8-bit training with a sigmoidal weighting that drops sharply between λ = 7.5 and λ = 10.The 5-bit unweighted loss curve is reported as very similar to the weighted 8-bit loss curve.
- Weighting comparison: The 5-bit-like weighting assigns substantially more weight to low noise levels than the weighting functions used elsewhere in the paper.The comparison includes ELBO, v-prediction, FM-OT, EDM-monotonic, and sigmoid weightings.
J Limitations
The authors note that their empirical findings depend on hyperparameter choices and that generative-model advances may introduce misuse and fairness risks.
- Changing the dataset or spatial resolution generally requires re-tuning optimization hyperparameters, architectural choices, and weighting functions.Such re-tuning can be time consuming and costly.
- More efficient and effective generative models could be misused, including through realistic synthetic media used for disinformation, identity theft, or blackmail.
- Generative models trained on large datasets could inherit and reproduce dataset biases, potentially producing unfair outcomes or harmful stereotypes.
- Suggested mitigations include gated model release, AI-generated-media detection, and monitoring mechanisms.
L Samples from our model trained on 512 × 512 ImageNet
The paper presents random samples from its highest-resolution 512x512 ImageNet model, distinguishing samples generated without guidance from those using guidance strength 4.
- The samples come from the authors’ highest-resolution 512x512 model trained on ImageNet and were not cherry-picked, except for removing depictions of humans.
- Figures 5 and 6 show random samples generated without guidance.
- Figures 7 and 8 show random samples generated with guidance strength 4.