Source-linked AI summary

Maximum Likelihood Training of Score-Based Diffusion Models

Yang Song, Conor Durkan, Iain Murray, Stefano Ermon

arXiv:2101.09258v4stat.MLcs.LG

TL;DR

Score-based diffusion models have tractable likelihoods through continuous normalizing flows, but weighted score matching does not directly optimize likelihood and CNF maximum likelihood is expensive. The paper introduces likelihood weighting, which upper-bounds negative log-likelihood and enables approximate maximum likelihood training; this consistently improves likelihood and reaches 2.83 and 3.76 bits/dim on CIFAR-10 and ImageNet 32ˆ32 without augmentation.

  • Problem

    Weighted score matching trains efficiently but does not directly promote high likelihood, while direct CNF maximum likelihood requires an expensive ODE solve at every optimization step.

  • Method

    The paper modifies weighted score matching with likelihood weighting so the objective upper-bounds negative log-likelihood for score-based diffusion models.

  • Results

    Likelihood improves consistently across datasets, architectures, and SDEs; the best models reach 2.83 bits/dim on CIFAR-10 and 3.76 bits/dim on ImageNet 32ˆ32 without data augmentation.

  • Takeaways & Limitations

    The bound provides an efficient route to approximate maximum likelihood training and makes score-based diffusion models competitive alternatives to continuous normalizing flows.

  • Takeaways & Limitations

    There is no theoretical guarantee that improving SDE likelihood improves ODE likelihood, and sampling and likelihood computation remain slow.

Abstract

from arXiv · show

Score-based diffusion models synthesize samples by reversing a stochastic process that diffuses data to noise, and are trained by minimizing a weighted combination of score matching losses. The log-likelihood of score-based diffusion models can be tractably computed through a connection to continuous normalizing flows, but log-likelihood is not directly optimized by the weighted combination of score matching losses. We show that for a specific weighting scheme, the objective upper bounds the negative log-likelihood, thus enabling approximate maximum likelihood training of score-based diffusion models. We empirically observe that maximum likelihood training consistently improves the likelihood of score-based diffusion models across multiple datasets, stochastic processes, and model architectures. Our best models achieve negative log-likelihoods of 2.83 and 3.76 bits/dim on CIFAR-10 and ImageNet 32x32 without any data augmentation, on a par with state-of-the-art autoregressive models on these tasks.

1 Introduction

Score-based diffusion models unify score-based and diffusion probabilistic models through stochastic processes that perturb data and learn the reverse path. The paper seeks an efficient objective that improves likelihood while retaining score matching’s training efficiency.

  • Background: Score-based diffusion models unify score-based generative models and diffusion probabilistic models within one stochastic-process framework.Both model families perturb data with noise distributions and generate samples by learning the reverse path from noise to data.
  • Background: The reverse-time process depends on a time-dependent score that can be estimated by a neural network trained with weighted score matching losses.The reverse-time SDE has a closed-form expression determined by the score of the perturbed data distribution.
  • Motivation: Score matching is more efficient than continuous normalizing flow maximum likelihood, but it does not necessarily improve likelihood values.CNF maximum likelihood requires an expensive ODE solve at every optimization step, whereas weighted score matching does not.
  • Contribution: Likelihood weighting makes the weighted score matching objective an upper bound on negative log-likelihood for approximate maximum likelihood training.The bound becomes tight when the score model equals the true time-dependent gradient field of a certain reverse-time SDE.
  • Results: 2.83 bits/dim on CIFAR-10 and 3.76 bits/dim on ImageNet 32ˆ32 are achieved without data augmentation using variational dequantization.The method consistently improves likelihood across multiple datasets, architectures, and SDEs, with only slight degradation of Fréchet Inception distances.

2 Score-based diffusion models

Score-based diffusion models diffuse data toward a tractable noise prior with an SDE, then learn its time reversal through the perturbed distribution’s score. Neural score models are trained using weighted score matching objectives that can be estimated efficiently.

  • 2.1 Diffusing data to noise with an SDE: An SDE gradually diffuses the data distribution toward an analytically tractable prior noise distribution.The drift, diffusion coefficient, and terminal time are chosen so pT(x) approaches π(x).
  • 2.1 Diffusing data to noise with an SDE: The forward SDE has drift f(x,t), diffusion coefficient g(t), and a Wiener process representing infinitesimal Gaussian noise.Its solution is a diffusion process with marginal distribution pt and transition distribution p0t.
  • 2.2 Generating samples with the reverse SDE: The reverse-time SDE generates samples by starting from π at t=T and solving backward toward t=0 using the time-dependent score.With the correct score and initialization from pT, the reverse process reproduces the forward diffusion in reverse.
  • 2.2 Generating samples with the reverse SDE: A neural network sθ(x,t) estimates the score ∇x log pt(x) across times and data points.The score-based model need not satisfy density integral constraints, making it easier to parameterize than many likelihood-based models.
  • 2.2 Generating samples with the reverse SDE: Weighted score matching minimizes expected squared error between the true score and the neural score model.When the transition density is tractable Gaussian, denoising score matching enables Monte Carlo estimates using dataset samples and perturbed samples.

3 Likelihood of score-based diffusion models

Score-based diffusion models define likelihoods through both a reverse SDE and an associated probability-flow ODE. The ODE yields tractable CNF likelihoods but makes maximum likelihood training expensive, motivating an efficient lower-bound approach for the SDE likelihood.

  • SDE likelihood: The reverse SDE defines a probabilistic model pSDEθ from the score model, prior, and forward SDE.Samples are obtained by numerically solving the approximate reverse-time SDE from noise at t=T.
  • ODE likelihood: The associated probability-flow ODE has the same marginal distributions as the SDE at every time.Replacing the true score with sθ produces deterministic dynamics defining the ODE-based model.
  • ODE likelihood: The ODE-based model is a continuous normalizing flow with tractable log-likelihood computed through volume-change dynamics.Samples can be generated by numerically solving the ODE backward from a prior sample.
  • Training cost: Maximum likelihood for the ODE model requires an ODE solve at every optimization step and can be prohibitively expensive for large-scale models.Weighted score matching is much more efficient but does not generally promote high likelihood for either model.
  • Efficient likelihood training: A lower bound on log pSDEθ enables efficient likelihood evaluation and optimization.The bound is used as the basis for the paper’s training objective rather than directly optimizing the ODE likelihood.

4 Bounding the likelihood of score-based diffusion models

The paper connects likelihood training to score matching by showing that likelihood weighting yields an efficient upper bound on the data-to-model KL divergence. Under specific score conditions, the bound becomes tight, while experiments show improved likelihood across settings, with limitations for ODE likelihood and individual-point evaluation.

  • 4.1 Bounding the KL divergence with likelihood weighting: Likelihood weighting λ(t)=g(t)^2 makes the weighted score-matching objective upper-bound D_KL(p || p^SDE_θ), providing an efficient maximum-likelihood proxy.With a fixed prior, optimizing this objective is equivalent to minimizing an upper bound on the KL divergence and therefore supports likelihood optimization.
  • 4.1 Bounding the KL divergence with likelihood weighting: The upper bound becomes tight when the learned score equals the true time-dependent score of a reverse-time diffusion initialized from the prior.Theorem 2 assumes matching continuous second-order distributions and s_θ(x,t)=∇_x log q_t(x) for all t.
  • 4.1 Bounding the KL divergence with likelihood weighting: Likelihood weighting consistently improves p^ODE_θ likelihood when the score model well-approximates the true score across x and t.The paper reports this empirical improvement in experiments, while the approximation condition describes when the theoretical connection applies.
  • 4.2 Bounding the log-likelihood on individual datapoints: Theorem 3 supplies bounds for individual datapoints because the distribution-level bound is suitable for training but is not directly computable pointwise.The score-matching objective contains the unknown true score and is evaluable only up to an additive constant through denoising score matching.
  • 4.2 Bounding the log-likelihood on individual datapoints: The framework leaves joint training of the drift, diffusion, and score model for future work.The stated bounds can be inserted into objectives involving individual-data likelihood, including variational dequantization.

5 Improving the likelihood of score-based diffusion models

The section evaluates likelihood weighting, importance sampling, and variational dequantization across datasets, SDEs, and architectures, finding improved likelihood with small FID degradation.

  • 5.1 Variance reduction via importance sampling: Importance sampling significantly reduces the variance introduced by likelihood weighting in the Monte Carlo training objective.Loss variance after convergence decreases from 98.48 to 0.068 on CIFAR-10 and from 0.51 to 0.043 on ImageNet.
  • 5.3 Experiments: The experiments compare likelihood weighting, importance sampling, and variational dequantization across DDPM++ architectures, VP and subVP SDEs, and CIFAR-10 and ImageNet 32ˆ32.VE SDE experiments are omitted because likelihood weighting matches the original weighting and performs worse in the reported CIFAR-10 experiments.
  • 5.3 Experiments: Likelihood weighting with importance sampling uniformly improves likelihood across the tested model architectures, SDEs, and datasets, with only slight FID degradation.The experiments report that better likelihood tends to accompany worse FID, while visual quality shows no obvious difference.
  • 5.2 Variational dequantization: Variational dequantization uniformly improves the pSDEθ likelihood bound in all tested settings, regardless of likelihood weighting.It narrows the gap between continuous and discrete probability models.
  • 5.3 Experiments: 2.83 bits/dim on CIFAR-10 and 3.76 bits/dim on ImageNet 32ˆ32 are achieved by ScoreFlows using likelihood weighting, importance sampling, and variational dequantization.The CIFAR-10 result uses no horizontal flipping; the results are on a par with state-of-the-art autoregressive models and outperform existing normalizing flow models.

6 Conclusion

The paper proposes efficient approximate maximum-likelihood training for score-based diffusion models and reports competitive likelihoods. It also identifies limitations involving ODE likelihoods, sampling speed, and data support.

  • 6 Conclusion: The likelihood-weighted score-matching objective upper bounds negative log-likelihood, enabling approximate maximum-likelihood training.The upper bound is analogous to the evidence lower bound used for variational autoencoders.
  • 6 Conclusion: The method consistently improves likelihood across multiple model architectures, SDEs, and datasets.The reported improvement comes with competitive likelihoods when combined with variational dequantization.
  • 6 Conclusion: Score-based diffusion models are presented as competitive alternatives to continuous normalizing flows, with tractable likelihood computation and more efficient maximum-likelihood training.
  • 6 Conclusion: The work has no theoretical guarantee that improving SDE likelihood improves ODE likelihood, and score-based diffusion models can have slow sampling.The paper also states that the current formulation supports continuous data and requires dequantization for discrete data.

Checklist

The checklist records disclosures about theoretical assumptions, proofs, experimental reproducibility, compute, assets, and broader-impact considerations. It also notes that experiments were not repeated because resources were insufficient.

  • Checklist: The paper states that its theoretical results include full assumptions and complete proofs in the appendices.
  • Checklist: The authors state that code, data, instructions, and training details needed to reproduce experiments are provided or released.Code is released at the score_flow repository, while experimental details are in Appendix C.
  • Checklist: The authors report compute and resource types in Appendix C but did not report error bars because experiments were not repeated multiple times.
  • Checklist: The paper cites existing datasets and states that all datasets used are publicly available, while standard licenses are available online.
  • Checklist: The paper states that limitations and potential negative societal impacts are discussed, including privacy issues associated with ImageNet.

A Proofs

The appendix states regularity assumptions for the SDEs and score models, then derives relationships among forward and reverse SDEs, KL divergences, score matching, entropy, and likelihood estimation.

  • Assumptions: The proofs assume smoothness, growth, Lipschitz, nondegeneracy, and integrability conditions for drift, diffusion, scores, and distributions.The assumptions include Novikov’s condition and bounds on score and drift functions.
  • KL and score matching: The KL divergence between initial distributions is expressed through a time integral of score-related terms plus a terminal KL divergence.The derivation uses the chain rule for KL divergences, the Fokker–Planck equation, and integration by parts.
  • Reverse-time SDE: The reverse-time SDE replaces the true score ∇x log pt(x) with the learned score sθ(x,t) while retaining the diffusion coefficient.
  • Theorem 2: Under qT = π and sθ(x,t) = ∇x log qt(x), the learned reverse process has the same stochastic process as the corresponding forward-time SDE.
  • Entropy estimation: Theorem 4 expresses initial differential entropy using terminal entropy and a time-dependent score function, yielding estimators from i.i.d. samples.The resulting entropy estimators use a trained score-based model and do not require a density model.

B Numerical stability

The numerical-stability appendix starts diffusion at a small positive time to avoid instability and derives a corrected likelihood bound for that truncated time horizon.

  • Numerical stability: Starting the SDE at a small nonzero time ε avoids numerical instability associated with t = 0.The practical time horizon is changed from [0,T] to [ε,T].
  • Likelihood bound: The corrected bound is used for test-time likelihood bounds and for training the flow model used in variational dequantization.It is not used to train the time-dependent score-based model.
  • Implementation choices: The authors choose different ε values for VP and subVP SDEs because ε = 10^-5 can cause stiffness for subVP ODE solvers.For subVP SDEs, ε = 10^-2 avoids numerical issues without hurting performance according to the passage.
  • Likelihood computation: The correction is applied when reporting SDE likelihood bounds but not when computing ODE likelihoods.The passage states that ODE likelihood remains valid on the truncated time horizon.

C Experimental details

Experiments use CIFAR-10 and down-sampled ImageNet at 32×32 resolution, with specified architectures, training settings, likelihood confidence intervals, and FID-based sample-quality evaluation.

  • Datasets: CIFAR-10 contains 50,000 training and 10,000 test images, while down-sampled ImageNet contains 1,281,149 training and 49,999 test images.
  • Model architectures: The variational dequantization model follows Flow++ architecture, and ImageNet score-based models use no dropout.
  • Confidence intervals: All likelihood values average results over approximately 50,000 test datapoints sampled with replacement.
  • Confidence intervals: The typical 95% confidence-interval radius is 0.006 bits/dim on CIFAR-10 and 0.008 bits/dim on ImageNet.
  • Sample quality: FID values are computed from 50,000 samples generated with numerical ODE solvers and compared with training and test data.
  • Sample quality: Likelihood weighting with importance sampling slightly increases FID scores, while producing samples with comparable visual quality.
Loading 2101.09258v4…