Source-linked AI summary

Maximum Likelihood Training for Score-Based Diffusion ODEs by High-Order Denoising Score Matching

Cheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen, Chongxuan Li, Jun Zhu

arXiv:2206.08265v2stat.MLcs.LG

TL;DR

First-order score matching does not necessarily maximize the likelihood of score-based diffusion ODEs, leaving their relationship to the training objective unclear. The paper derives a KL-divergence bound using first-, second-, and third-order score-matching errors and proposes a high-order denoising method. Experiments report improved likelihood on synthetic data and CIFAR-10 while retaining high generation quality.

  • Problem

    First-order score matching does not necessarily maximize ScoreODE likelihood, and its relationship to ScoreODE KL divergence is not established by the known ScoreSDE bound.

    Source: §1 · p. 2, §3 · p. 3
  • Method

    The paper bounds ScoreODE KL divergence with first-, second-, and third-order score-matching errors and proposes high-order denoising score matching to train the score models.

    Source: §1 · p. 2, §4.2 · p. 5, §5 · p. 6
  • Results

    High-order score matching improves ScoreODE likelihood on synthetic data and CIFAR-10 while retaining high generation quality.

    Source: §1 · p. 2, §8 · p. 9
  • Takeaways & Limitations

    The resulting score model retains the data score as its optimal solution and can still be used with SGM sampling methods.

    Source: §8 · p. 9

    The high-order DSM analysis assumes Gaussian forward-process conditionals q(x_t|x_0).

    Source: §4 · p. 5

Abstract

from arXiv · show

Score-based generative models have excellent performance in terms of generation quality and likelihood. They model the data distribution by matching a parameterized score network with first-order data score functions. The score network can be used to define an ODE ("score-based diffusion ODE") for exact likelihood evaluation. However, the relationship between the likelihood of the ODE and the score matching objective is unclear. In this work, we prove that matching the first-order score is not sufficient to maximize the likelihood of the ODE, by showing a gap between the maximum likelihood and score matching objectives. To fill up this gap, we show that the negative likelihood of the ODE can be bounded by controlling the first, second, and third-order score matching errors; and we further present a novel high-order denoising score matching method to enable maximum likelihood training of score-based diffusion ODEs. Our algorithm guarantees that the higher-order matching error is bounded by the training error and the lower-order errors. We empirically observe that by high-order score matching, score-based diffusion ODEs achieve better likelihood on both synthetic data and CIFAR-10, while retaining the high generation quality.

1. Introduction

ScoreODEs support exact likelihood evaluation, but first-order score matching does not necessarily maximize their likelihood. The paper analyzes this gap and proposes high-order denoising score matching, reporting improved likelihood while retaining high generation quality.

1. Introduction

  • First-order score matching does not necessarily maximize ScoreODE likelihood, despite its role in training score models.The paper notes that ScoreODEs trained this way may fail to estimate likelihood even for simple distributions.Source: §1 · p. 2
  • The paper analyzes the relationship between ScoreODE maximum likelihood and score matching, deriving an equivalent KL-divergence form that exposes a gap from the first-order objective.It bounds the KL divergence using first-, second-, and third-order score-matching errors.Source: §1 · p. 2
  • The proposed high-order denoising score-matching algorithm guarantees bounded high-order matching errors and includes scale-up techniques for practice.The stated bound relates higher-order approximation errors to training error.Source: §1 · p. 2
  • High-order training improves ScoreODE likelihood on synthetic data and CIFAR-10 while retaining high generation quality.Source: §1 · p. 2

2. Score-Based Generative Models

The section introduces score-based SDE and ODE models built from learned score functions. Score matching supports maximum-likelihood training for ScoreSDEs, while the corresponding ScoreODE likelihood is not similarly guaranteed by first-order matching.

2. Score-Based Generative Models

  • Score-based generative models learn a neural score function to approximate the time-dependent data score used in reverse-time dynamics.Source: §2.1 · p. 2

2.1. ScoreSDEs and Maximum Likelihood Training

  • With the weighting function g(·)^2, minimizing the score-matching objective is equivalent to maximum-likelihood training of the reverse ScoreSDE.Source: §2.1 · p. 2

2.2. ScoreODEs and Exact Likelihood Evaluation

  • A ScoreODE uses the probability-flow ODE and permits exact log-density computation, unlike the reverse ScoreSDE.The ODE’s marginals match those of the associated SDE; its log-density is computed through the instantaneous change of variables.Source: §2.2 · p. 3
  • First-order matching does not theoretically guarantee a bound on ScoreODE KL divergence; prior VE ScoreODEs had about 3.45 bits/dim on CIFAR-10 versus about 3.13 for VP models.The cited comparison follows training with the score-matching objective; the VE ScoreSDE is reported to have state-of-the-art sample quality.Source: §2.2 · p. 3

3. Relationship between Score Matching and KL Divergence of ScoreODEs

The section shows that first-order score matching alone does not control ScoreODE KL divergence. It derives a bound using score-matching errors through third order, avoiding direct dependence on the model score that requires ODE computation.

3. Relationship between Score Matching and KL Divergence of ScoreODEs

  • First-order score matching can upper-bound ScoreSDE KL divergence but is insufficient to upper-bound ScoreODE KL divergence, whose relationship to score matching was previously unanalyzed.Source: §3 · p. 3

3.1. KL Divergence of ScoreODEs

  • The ScoreODE KL formulation exposes a nonzero gap from first-order score matching because its objective also includes an error involving the model score.Reducing the data-score error alone cannot control the difference between the model score and the ScoreODE score.Source: §3.1 · p. 3

3.2. Bounding the KL Divergence of ScoreODEs by High-Order Score Matching Errors

  • ScoreODE KL divergence can be bounded using first-, second-, and third-order score-matching errors, avoiding ODE-solver computation of the model score.The stated bound assumes bounded second-order score norm and error bounds for the first-, second-, and third-order estimates; the upper bound increases with each error bound.Source: §3.2 · p. 4
  • The nonparametric optimum remains the data score, so the trained score model remains usable with SGM sampling methods.Source: §3.2 · p. 4

4. Error-Bounded High-Order Denoising Score Matching (DSM)

The section extends denoising score matching to higher orders so ScoreODE likelihood bounds can be controlled through score-matching errors. Its second- and third-order objectives bound higher-order errors using training error and lower-order estimation errors.

4. Error-Bounded High-Order Denoising Score Matching (DSM)

  • The KL divergence of a ScoreODE can be upper-bounded by controlling first-, second-, and third-order score-matching errors across time.The methods in this section address a fixed time; the whole-process objective is detailed in Section 5.Source: §4 · p. 4

4.1. First-Order Denoising Score Matching

  • First-order denoising score matching trains a score model without requiring the unknown data score function directly.For a fixed time, the objective uses noisy samples xt = αt x0 + σtϵ, with Gaussian noise.Source: §4.1 · p. 5

4.2. High-Order Denoising Score Matching

  • The second-order DSM error is bounded by its training error and the first-order score-matching error.This makes the DSM objective a valid surrogate; the authors contrast its guarantee with the earlier method of Meng et al. (2021a), whose error may be unbounded.Source: §4.2 · p. 5
  • Third-order score-matching error is bounded by training error, first- and second-order errors, and the trace of the second-order error.The authors focus on second and third orders, leaving extension to higher orders for future work.Source: §4.2 · pp. 5–6

5. Training Score Models by High-Order DSM

The training method applies high-order DSM across timesteps while avoiding the costly ODE solvers needed by direct integration. Time reweighting, derivative estimators, and a positive start time address variance, computational cost, and numerical instability.

5. Training Score Models by High-Order DSM

  • The all-timestep training objective combines first-, second-, and third-order DSM, using estimated lower-order scores and their derivatives.The method uses noise prediction and the Skilling-Hutchinson trace estimator to compute high-order derivatives efficiently.Source: §5 · p. 6, §5.1 · p. 6

5.1. Variance Reduction by Time-Reweighting

  • Monte Carlo sampling of time avoids ODE solvers, while multiplying objectives by σt^2 reduces variance from differing noise scales.The authors report that the reweighted first-order objective empirically has low sample variance.Source: §5.1 · pp. 6–7

5.2. Scalability and Numerical Stability

  • Trace estimators and a small positive start time address expensive derivatives and instability near t = 0, while likelihood evaluation remains exact.Training and evaluation use t ∈ [ϵ, T], with pODEϵ used for likelihood evaluation.Source: §5.2 · p. 7

6. Related Work

The authors frame high-order score matching as a way to train ScoreODEs and neural ODEs while controlling intermediate-time distributions. They distinguish their error-bounded objectives from prior high-order DSM and conventional maximum-likelihood training.

6. Related Work

  • Unlike conventional neural ODE maximum-likelihood training, high-order score matching matches distributions between the forward process and the neural ODE at each time.The authors report that the resulting dynamics are empirically smooth and that Monte Carlo objective estimates avoid black-box ODE solvers.Source: §6 · p. 7
  • The proposed high-order DSM bounds matching error by training error and lower-order estimation errors, unlike the potentially unbounded error of Meng et al. (2021a).The method can also train separate high-order score models for applications such as uncertainty quantification and Ozaki sampling.Source: §6 · p. 7

7. Experiments

Across synthetic and image data, high-order score matching improves ScoreODE density or likelihood while retaining high sample quality. The experiments compare training orders on a Gaussian mixture and checkerboard, then evaluate image datasets.

7. Experiments

  • The experiments use VE-type diffusion models, with first- and second-order matching defined through the corresponding settings of λ1 and λ2.The reported experimental comparison focuses on likelihood improvement and sample quality.Source: §7 · p. 7

7.1. Example: 1-D Mixture of Gaussians

  • On a 1-D Gaussian mixture, third-order score matching produces a model density closer to the data than first-order matching.First-order matching keeps ℓSM(t) small even as ℓFisher(t) remains large for most t; higher orders reduce ℓFisher(t).Source: §7.1 · pp. 7–8

7.2. Density Modeling on 2-D Checkerboard Data

  • On 2-D checkerboard data, second-order matching improves model density, while third-order matching reduces Fisher divergence and achieves excellent density.The models are trained from scratch without pretrained lower-order score models.Source: §7.2 · p. 8, Figure 4 · p. 8

7.3. Density Modeling on Image Datasets

  • High-order score matching improves ScoreODE likelihood on CIFAR-10 and ImageNet while retaining the corresponding ScoreSDEs' high sample quality.The authors report acceptable computational costs using trace estimators and Jacobian-vector products.Source: §7.3 · p. 8

8. Conclusion

The paper proposes high-order denoising score matching for ScoreODE likelihood training and bounds the KL divergence using first-, second-, and third-order score errors. It reports improved ScoreODE likelihood while retaining sample quality, and leaves extensions to other SGM types for future work.

8. Conclusion

  • The method improves ScoreODE likelihood on several density-modeling benchmarks while retaining high sample quality, with extensions to other SGM types left for future work.The conclusion identifies VP, latent-space, and critically-damped Langevin diffusion models as possible extensions.Source: §8 · p. 9
  • The KL divergence upper bound is controlled by first-, second-, and third-order score-matching errors, and the high-order DSM error is bounded by training error and lower-order errors.The optimal score model remains the same as under the original SGM training objective.Source: §8 · p. 9

A. Assumptions

  • The analysis adopts regularity and differentiability assumptions to ensure reverse-time SDE and probability-flow ODE existence, valid integration by parts, and existence of the third score function.The appendix lists these assumptions explicitly.Source: A. Assumptions · p. 11

10. Novikov’s condition: E

  • The listed Novikov condition imposes exponential-tail bounds on the time-dependent distributions.The passage gives a bound of order O(e^-∥x∥^k) for q_t and p^SDE_t.Source: §10 · p. 11

B. Distribution gap between score-based diffusion SDEs and ODEs

The analysis distinguishes the probability-flow ODE associated with a ScoreSDE from the ScoreODE and relates KL divergence bounds to score-matching errors. It also develops score evolution and KL derivations for the ODE setting.

B. Distribution gap between score-based diffusion SDEs and ODEs

C. KL divergence and variational gap of ScoreSDEs

D. Instantaneous change of score function of ODEs

E.1. Proof of Theorem 3.1

  • The proof of Theorem 3.1 rewrites the KL divergence from the data distribution to the ScoreODE distribution in integral form.Source: §E.1 · p. 14

E.2. Proof of Theorem 3.2

  • The KL-divergence upper bound depends on the first-, second-, and third-order score errors, which can be minimized to reduce that bound.The bound also depends on constants associated with the forward process.Source: §E.2 · p. 17

E.3. Maximum likelihood training of ScoreODE by ODE solvers

Direct ODE-solver likelihood training is difficult to scale and does not ensure the learned score follows the data score. The paper develops error-bounded high-order DSM formulations and connects their training objectives to errors in the second-order score and its trace.

E.3. Maximum likelihood training of ScoreODE by ODE solvers

  • Direct ScoreODE likelihood training requires ODE solvers at each optimization step, is hard to scale to large SGM networks, and does not ensure score or dynamics match the data process.The cited example reports 2–3 minutes to evaluate log likelihood for one batch; direct maximum likelihood may also produce complex dynamics.Source: §E.3 · pp. 17–18

F. Error-bounded High-Order Denoising Score Matching

  • Expectation formulations for first-, second-, and third-order score functions support the design of high-order denoising score matching objectives.The text specifically identifies the first-order formulation, second-order formulation, and third-order formulation used to design the method.Source: F. Error-bounded High-Order Denoising… · pp. 18–19

F.1. Proof of Theorem 4.1

  • The second-order score matrix’s distance from its optimal predictor can be viewed as the training error.Source: §F.1 · p. 21

F.2. Proof of Corollary 4.2

  • The analogous deviation for the trace of the second-order score is also treated as training error.Source: §F.2 · p. 22

F.3. Proof of Theorem 4.3

The appendix shows why small first-order score error does not ensure small higher-order error, then develops bounded, estimable high-order DSM objectives and their training procedure.

F.3. Proof of Theorem 4.3

  • The DSM objective analyzed by Meng et al. (2021a) can have arbitrarily large second-order score error despite zero training error and arbitrarily small first-order error.The construction considers an unbounded first-order score function.Source: §F.3 · pp. 23–24
  • The proposed high-order DSM has a bounded-error property: its second-order matching error is bounded by the squared first-order error, avoiding the unbounded-error problem.The third-order objective’s training error can be represented by the distance from its optimal solution.Source: §F.3 · pp. 23–24

G. Estimated objectives of high-order DSM for high-dimensional data

  • Hutchinson’s trace estimator and automatic differentiation reduce trace and Frobenius-norm computations from O(d^2) to O(d) for high-dimensional data.The estimator is unbiased when the random vector has zero mean and identity covariance; Gaussian and Rademacher distributions are given as choices.Source: G. Estimated objectives of high-order… · p. 24

G.1. Relationship between the estimated objectives and the original objectives for high-order DSM

  • The estimated second-order DSM objective is equivalent to the original, while the estimated trace-second-order and third-order objectives upper-bound their originals.The appendix also establishes equivalence between two formulations of the third-order objective with respect to optimizing the model parameters.Source: §G.1 · pp. 25, 27

H. Training algorithm

  • Training samples data, noise, and time, forms noisy inputs, and uses Jacobian-vector and vector-Jacobian products to calculate the score-model objectives.The algorithm specifies Gaussian or Rademacher vectors and stops gradients through the first-order score estimate and its Jacobian-vector product.Source: H. Training algorithm · p. 28

I.1. Choosing of λ1, λ2

  • The authors train from default initialization without lower-order pretraining and use λ1=0.5, λ2=0.1 for synthetic data and λ1=λ2=1 for CIFAR-10.The appendix also gives objective-specific λ choices for second- and third-order DSM.Source: §I.1 · p. 28, §I.2 · p. 29

I.4. CIFAR-10 experiments

On CIFAR-10, high-order DSM improves ScoreODE likelihood with training costs near or below twice first-order DSM. The ImageNet 32x32 appendix describes baseline and fine-tuning settings but gives no outcome in the supplied passages.

I.4. CIFAR-10 experiments

  • Second-order DSM costs are close to first-order training, while third-order DSM costs less than twice as much; the method improves ScoreODE likelihood on CIFAR-10.The reported timing is for the shallow VE model; Table 3 benchmarks computation time, iteration counts, and memory.Source: §I.4 · p. 30, Table 3 · p. 30

I.5. ImageNet 32x32 experiments

  • For ImageNet 32x32, the authors train VE and VE (deep) first-order baselines from scratch, then further train both high-order variants for 100k iterations.The reported test-set metric is average bpd over five repeats.Source: §I.5 · p. 30

J. Additional results for VE, VP and subVP types

On CIFAR-10, third-order DSM improves VP and subVP ScoreODE likelihood for start times above a threshold, though not at very small start times. The added score matching also improves VP model density on 1-D mixture-of-Gaussians data and retains visually similar VE samples.

J. Additional results for VE, VP and subVP types

Loading 2206.08265v2…