PaperlayerView exact revision on arXiv

Source-linked AI summary

Score-Based Generative Modeling with Critically-Damped Langevin Diffusion

Tim Dockhorn, Arash Vahdat, Karsten Kreis

arXiv:2112.07068v4stat.MLcs.LG

TL;DR

Existing SGMs use diffusion processes that determine the difficulty of score learning, motivating a more effective design. The paper introduces critically-damped Langevin diffusion with velocity augmentation, hybrid score matching, and a tailored sampler. CLD-based models outperform prior SGMs under similar budgets, while SSCS outperforms Euler–Maruyama in CLD-based models.

  • Problem

    The forward diffusion uniquely determines the learned score function, so diffusion design is a central research target for improving SGMs.

  • Method

    CLD couples data with auxiliary velocities, injects noise only into velocities, learns a conditional velocity score with HSM, and uses a tailored SDE integrator.

  • Results

    CLD-based models outperform previous SGMs on image synthesis for similar-capacity models and sampling compute budgets, while SSCS is superior to EM.

  • Takeaways & Limitations

    CLD provides a diffusion framework with smoother data trajectories, an easier score target, fast synthesis, and improved expressivity.

  • Takeaways & Limitations

    The evaluation validates CLD on image-synthesis benchmarks, and one score-model parameterization assumes p0(x0)=δ(x).

Abstract

from arXiv · show

Score-based generative models (SGMs) have demonstrated remarkable synthesis quality. SGMs rely on a diffusion process that gradually perturbs the data towards a tractable distribution, while the generative model learns to denoise. The complexity of this denoising task is, apart from the data distribution itself, uniquely determined by the diffusion process. We argue that current SGMs employ overly simplistic diffusions, leading to unnecessarily complex denoising processes, which limit generative modeling performance. Based on connections to statistical mechanics, we propose a novel critically-damped Langevin diffusion (CLD) and show that CLD-based SGMs achieve superior performance. CLD can be interpreted as running a joint diffusion in an extended space, where the auxiliary variables can be considered "velocities" that are coupled to the data variables as in Hamiltonian dynamics. We derive a novel score matching objective for CLD and show that the model only needs to learn the score function of the conditional distribution of the velocity given data, an easier task than learning scores of the data directly. We also derive a new sampling scheme for efficient synthesis from CLD-based diffusion models. We find that CLD outperforms previous SGMs in synthesis quality for similar network architectures and sampling compute budgets. We show that our novel sampler for CLD significantly outperforms solvers such as Euler--Maruyama. Our framework provides new insights into score-based denoising diffusion models and can be readily used for high-resolution image synthesis. Project page and code: https://nv-tlabs.github.io/CLD-SGM.

1 INTRODUCTION

The diffusion process determines the score-learning problem in SGMs, motivating CLD as a smoother alternative with easier denoising targets and tailored sampling.

  • The forward diffusion uniquely determines the score function learned by an SGM, making diffusion design central to improving synthesis quality or sampling speed.
  • CLD augments data with coupled velocity variables, injects noise only into velocities, and derives a score objective for the velocity conditioned on data.
  • CLD-based models outperform previous diffusion models on CIFAR-10 for similar architectures and sampling compute budgets.
  • The authors report smoother learned networks, a CLD-specific sampler superior to Euler–Maruyama, and no difficult-to-tune hyperparameters.
  • The paper contributes CLD, a conditional score-matching objective, hybrid denoising score matching, and a tailored SDE integrator.

2 BACKGROUND

SGMs diffuse data toward a tractable prior and learn scores that reverse this process for generation.

  • An SDE defines the forward diffusion through a drift term, diffusion coefficient, and Wiener-process noise over continuous time.
  • The reverse-time diffusion uses the marginal score at the reversed time to invert the forward process.
  • Existing SGM SDEs commonly use drift and diffusion coefficients linear in data, targeting an approximately Normal terminal distribution.
  • A learned score of the diffused data enables synthesis through the reverse-time SDE.
  • For affine coefficients, the conditional perturbation distribution is analytically Normal, while λ(t) controls quality–likelihood trade-offs.

3 CRITICALLY-DAMPED LANGEVIN DIFFUSION

CLD couples data and auxiliary velocities in an extended diffusion, balancing Hamiltonian transport and friction for smooth equilibration and easier score learning.

  • 3 CRITICALLY-DAMPED LANGEVIN DIFFUSION: CLD augments each data variable with a velocity and runs a coupled diffusion in joint data–velocity space.
  • 3 CRITICALLY-DAMPED LANGEVIN DIFFUSION: The Hamiltonian component couples data and velocities, while mass, time rescaling, and friction control coupling, convergence, and velocity noise.
  • 3 CRITICALLY-DAMPED LANGEVIN DIFFUSION: Γ^2=4M yields critical damping, balancing Hamiltonian and friction effects for fastest smooth convergence without oscillations.
  • 3 CRITICALLY-DAMPED LANGEVIN DIFFUSION: Unlike prior diffusions that inject noise directly into data, CLD injects noise only into velocities and perturbs data indirectly through coupling.
  • 3 CRITICALLY-DAMPED LANGEVIN DIFFUSION: The CLD score objective requires only the velocity gradient of the joint log density, equal to the score of velocity conditioned on data.
  • 3.2 SCALABLE TRAINING: HSM samples data while analytically marginalizing the Normal initial velocity distribution, keeping the training objective tractable.
  • 3.2 SCALABLE TRAINING: HSM avoids the unbounded scores and instability associated with standard DSM near t=0 in CLD.
  • 3.3 SAMPLING FROM CLD-BASED SGMS: SSCS is a CLD-specific SDE integrator whose expensive propagator requires neural-network evaluations.

4 RELATED WORK

The paper situates CLD within statistical mechanics and distinguishes it from latent-space simplification and adaptive solver approaches.

  • Relations to Statistical Mechanics and Molecular Dynamics: CLD frames SGM diffusion as an equilibration problem inspired by thermostats and Langevin dynamics in statistical mechanics.
  • Further Related Work: CLD improves the diffusion process itself, whereas LSGM first embeds data into a smooth latent space using two training stages and extra networks.
  • Further Related Work: GGF uses adaptive standard Euler-based methods, while CLD's SSCS is derived from first principles and performs well without adaptive step sizes.
  • Further Related Work: The paper notes conceptual similarities between velocity augmentation and augmented normalizing flows or Hamiltonian Monte Carlo.

5 EXPERIMENTS

Experiments evaluate CLD-SGMs on CIFAR-10 and CelebA-HQ-256, comparing solvers, sampling budgets, and CLD hyperparameters. CLD achieves strong FID with similar model capacity and compute, while its sampler and training choices improve performance.

  • 5.1 IMAGE GENERATION: CLD-SGM achieves FID 2.25 with probability-flow ODE sampling and 2.23 with generative SDE sampling on CIFAR-10.These results are reported for the main CIFAR-10 model.
  • 5.1 IMAGE GENERATION: CLD-SGM obtains stronger FID than LSGM-100M while using approximately 100M parameters instead of LSGM’s approximately 475M-parameter model.The larger LSGM model marginally outperforms CLD, but the smaller fair-comparison model does not.
  • 5.1 IMAGE GENERATION: Under network capacities and NFE budgets ≤500, CLD-SGM outperforms published results in CIFAR-10 FID.The comparison uses similar network capacity and sampling compute budgets.
  • 5.2 SAMPLING SPEED AND SYNTHESIS QUALITY TRADE-OFFS: With one exception, CLD-SGM outperforms adaptive and fixed-step baselines across sampling-speed and synthesis-quality comparisons.The exception is VESDE with 2,000 NFE.
  • 5.2 SAMPLING SPEED AND SYNTHESIS QUALITY TRADE-OFFS: SSCS significantly outperforms Euler–Maruyama under limited NFE budgets, while both samplers perform similarly with fine SDE discretization.The convergence behavior is attributed to both methods’ errors becoming negligible at high NFE.
  • 5.3 ABLATION STUDIES: Mixed-score training improves FID from 3.56 to 3.14, whereas regular DSM often becomes unstable and HSM is treated as crucial for CLD-SGM training.The instability is associated with unbounded scores near t=0 under standard DSM.

6 CONCLUSIONS

The paper presents CLD as a diffusion process that makes SGMs smoother, easier to denoise, and more effective for image synthesis. It also identifies tailored objectives and integration methods as promising foundations for future work.

  • CLD diffuses data more smoothly and easily than previous SGMs, producing smoother score functions, fast synthesis, and improved expressivity.
  • CLD outperforms previous SGMs on image synthesis for similar-capacity models and sampling compute budgets.
  • The paper derives HSM, a CLD-specific denoising objective, and a tailored SDE integrator for sampling.
  • Future work includes extending CLD beyond images, accelerating sampling, adapting it toward maximum likelihood, and exploring other thermostating methods.

7 ETHICS AND REPRODUCIBILITY

The paper frames its contribution as a fundamental algorithmic advance and notes that ethical effects depend on how image generation is applied. It also provides public code and detailed training information for reproducibility.

  • Ethics: The proposed CLD does not imply immediate ethical concerns, but image-generation applications can be used for both creative and nefarious purposes.
  • Ethics: The ethical impact depends on the application domain and task.
  • Reproducibility: The authors publicly release source code, reproduction instructions, and detailed training information and hyperparameters.
  • Langevin dynamics: Underdamped dynamics produce oscillations, overdamped dynamics slow convergence through strong noise injection, and critical damping balances these effects.
  • Langevin dynamics: Critical damping mixes and converges quickly, and is described as optimal in terms of convergence.

A.2 VERY HIGH FRICTION LIMIT AND CONNECTIONS TO PREVIOUS SDES IN SGMS

The appendix connects common SGM diffusions to high-friction Langevin dynamics and explains why CLD instead uses critically damped joint data-velocity diffusion. It also derives CLD’s equilibrium and velocity-only score objective.

  • Langevin formulation: A Langevin diffusion can be written with time-dependent β(t), potential, and friction terms governing the joint data-velocity dynamics.
  • Very high friction limit: At maximum friction, the current-step velocity is completely cancelled, yielding high-friction overdamped Langevin dynamics.
  • Connections to previous SDEs: The VPSDE is exactly the transition kernel obtained in the high-friction limit of the more general Langevin diffusion.
  • Connections to previous SDEs: Previous high-friction diffusions require more noise to reach the prior, whereas critically damped CLD can reach equilibrium with shorter diffusion or less injected noise.
  • CLD construction: CLD augments data with auxiliary velocities, injects noise only into velocity, and couples data and velocity in joint space.
  • CLD construction: Critical damping sets M = Γ^2/4, while affine drift and diffusion make the joint state analytically tractable under Gaussian initialization.
  • Equilibrium: CLD converges to an equilibrium with independent standard-normal data and velocity distributed as N(0, M I_d).
  • CLD objective: Because noise enters through velocity, the derived objective requires gradients only with respect to the velocity component.

B.4 CLD-SPECIFIC IMPLEMENTATION DETAILS

The appendix details numerical stabilization, sampling formulations, Hamiltonian structure, likelihood estimation, and the relationship between HSM and DSM. It also explains why Hamiltonian mixing can reduce score-learning difficulty.

  • Numerical implementation: HSM can diverge as t approaches zero and may be numerically unstable, so the authors add ε_num I_2d before Cholesky decomposition.
  • Numerical implementation: With double precision and ε_num = 10^-9, HSM works in practice for the stated experimental setup.
  • Sampling: CLD samples can be generated by simulating the reverse-time SDE with SSCS, Euler–Maruyama, or other solvers.
  • Sampling: The probability-flow ODE offers an alternative sampler, but imperfect score models make its generated distribution not exactly equivalent to reverse-time SDE sampling.
  • Probability-flow ODE: The CLD probability-flow ODE combines Hamiltonian dynamics with a score term, without the Ornstein–Uhlenbeck term.
  • Likelihood estimation: The probability-flow ODE supports log-likelihood estimation, but the resulting lower bound can be potentially quite loose because the input includes an auxiliary velocity.
  • Hamiltonian motivation: Hamiltonian dynamics can mix distributions while remaining analytically invertible, reducing the need to learn a complex score to reverse that component.
  • HSM and DSM: HSM analytically averages over the initial velocity distribution, reducing training-gradient variance relative to sample-based DSM estimates.

C.1 GRADIENT VARIANCE REDUCTION VIA HSM

The section compares hybrid score matching (HSM) with denoising score matching (DSM), finding that HSM reduces Monte Carlo gradient variance, especially at larger γ.

  • Motivation: HSM avoids unbounded scores as t →0 and yields lower-variance Monte Carlo gradients than DSM.The section motivates HSM both theoretically and practically.
  • Analysis: The variance analysis compares gradients of HSM and DSM using samples from the same distribution, isolating differences in their squared norm terms.The covariance study evaluates K_HSM and K_DSM on CIFAR-10 training data.
  • Analysis: The analysis measures the trace of covariance matrices as the total variation of the corresponding random variables.This scalar summarizes the estimated covariance matrices used in Fig. 7.
  • Results: For DSM, Tr(Cov) is very large at small t, while HSM shows barely any visual difference between γ = 0.04 and γ = 1.DSM is also reported to become unstable for small γ.
  • Results: HSM demonstrates clear variance reduction, particularly for large γ, because it integrates out the initial velocity distribution instead of using noisy initial-velocity samples.This provides the stated explanation for HSM’s empirical advantage.

D.2 DERIVATION AND ANALYSIS

The section derives SSCS by splitting the CLD generative dynamics into analytically tractable and numerically approximated components. Although SSCS is first-order, splitting can make it more accurate than Euler–Maruyama when the score term is not dominant.

  • Derivation: SSCS uses symmetric splitting and Fokker–Planck operators to construct an integrator for the otherwise intractable CLD propagator.The propagator is discretized into N steps of size δt and decomposed into separately tractable terms.
  • Splitting terms: The A component evolves the joint state stochastically, while the S component updates only the velocity component through an ODE.The S component is approximated with an Euler step.
  • Error analysis: SSCS has local error ∼δt^2 and global error ∼δt, making it a first-order integrator like Euler’s method.The additional Euler step in the S component determines these convergence orders.
  • Error analysis: SSCS is expected to be more accurate than Euler–Maruyama when the dynamics are not dominated by the S component, because the analytically tractable part is split off.The learned conditional velocity score is close to Normal for much of the diffusion, suggesting that S is small in practice.
  • Sampling behavior: SSCS and Euler–Maruyama perform similarly at very small step sizes, while experiments show SSCS is suited to relatively few synthesis steps.The section attributes the few-step advantage to smaller practical integration errors.
  • Limitations: SSCS is expected to have weak and strong convergence order 1 under additive noise and appropriate smoothness conditions, with detailed analysis left for future work.Without additive noise, the stated strong-convergence order would be 0.5.

E.1 SCORE AND JACOBIAN EXPERIMENTS

The experiments test whether CLD simplifies learned scores and neural networks, then evaluate image synthesis and sampling choices. CLD produces simpler scores and lower network complexity, while SSCS improves limited-budget sampling.

  • Experimental setup: CLD and VPSDE are compared on analytically tractable mixtures of Normals, enabling exact access to diffused marginals and ground-truth scores.This setup supports controlled score and sampler analyses unavailable for image data.
  • Score experiment: ξ_CLD(t) is smaller than ξ_VPSDE(t) for every t ∈[0,T], with the largest difference at small t.The result indicates that CLD scores remain closer to Normal scores than VPSDE scores.
  • Complexity experiment: CLD yields significantly lower neural-network complexity than VPSDE, and mixed score parameterization further reduces complexity for both diffusions.Complexity is measured using the squared Frobenius norm of network Jacobians.
  • Ablations: Denoising only slightly influences FID for CLD with SSCS, unlike its reported importance for direct data-space SDEs.The probability-flow ODE experiments also found denoising important for preventing solver blow-up near t →0.
  • Ablations: Mixed score parameterization substantially improves competitive FIDs, whereas γ and M have only small effects on CIFAR-10 FID scores.This observation comes from the ablation experiments.
  • Baselines: LSGM-100M achieves an FID of 4.60, an NLL bound of 2.96 bpd, and requires 131 NFE on average for sampling.These values are reported for the recommended Runge–Kutta 4(5) probability-flow ODE setting.
  • Analytical sampling: For each n ∈{20,50,100,200}, CLD with SSCS outperforms both VPSDE and CLD with EM in negative log-likelihood on the mixture-of-Normals benchmark.The benchmark uses an analytically known ground-truth score.

F.1.2 MAXIMUM LIKELIHOOD TRAINING

The maximum-likelihood experiments find that CLD produces simpler score networks and captures challenging multi-scale structure better than VPSDE in toy settings. Extended image experiments further report efficient sampling and strong sample quality across solvers and tolerances.

  • Maximum-likelihood comparison: CLD significantly outperforms VPSDE on a challenging toy dataset under maximum-likelihood training, while high-dimensional maximum-likelihood image training remains future work.The authors describe CLD as promising in this setting but explicitly limit the demonstrated scope to toy experiments.
  • Maximum-likelihood complexity: CLD models with larger γ generally learn less complex networks, particularly at smaller times t.This conclusion comes from the maximum-likelihood complexity experiment.
  • Challenging toy dataset: CLD captures the multi-scale toy dataset’s five rolls and fine structure, whereas VPSDE captures only the main modes and connects them undesirably.Both models use the same simple ResNet architecture with fewer than 100k parameters.
  • CIFAR-10 sampling: For any n ≥150, CLD with SSCS and quadratic striding outperforms DDIM, while DDIM performs better at n = 50.The comparison concerns fixed-step sampling compute budgets and includes the qualification that DDIM was designed specifically for few-step sampling.
  • Additional samples: FID 2.23 is obtained with EM-QS using 2000 function evaluations, while SSCS-QS obtains FID 3.07 using only 150 evaluations.These are additional CLD-SGM samples reported for the two solver configurations.
  • CelebA-HQ-256 sampling: At solver error tolerance 10^-3, the adaptive Runge–Kutta 4(5) solver still produces very good samples using an average of 129 neural-network evaluations.The accompanying visual comparison reports little difference between tolerances 10^-5 and 10^-4, low-frequency artifacts at 10^-3, and deterioration beginning at 10^-2.
  • Generation paths: CLD generation paths show velocity variables encoding data at intermediate times and converging to Normal distributions as t approaches either endpoint.At t = 1, both variables follow the equilibrium distribution; as t approaches 0, data approaches the data distribution while velocity approaches another Normal distribution.
  • Generation paths: Hamiltonian coupling makes negative velocities resemble noisy, sign-flipped image data that pull pixel values toward the equilibrium mean.The authors attribute this behavior to the negative sign on x_t in the Hamiltonian term.

G PROOFS OF PERTURBATION KERNELS

This section proves the stated perturbation-kernel solutions for the forward diffusion and the analytical splitting term used in SSCS. The derivations verify the corresponding mean and covariance evolutions, including the constant-β implementation used in practice.

  • Forward diffusion: The appendix establishes ODEs governing the evolution of the forward diffusion’s mean and covariance matrix.The claimed mean and covariance solutions are checked by substitution into the corresponding ODEs.
  • Forward diffusion: Substitution of the claimed mean solution into the forward-diffusion ODE completes the proof of mean correctness.The derivation decomposes the resulting expression into two equations before concluding the verification.
  • Forward diffusion: Substitution of the claimed covariance solution into the forward-diffusion ODE completes the proof of covariance correctness.The appendix verifies the covariance expressions through successive simplified identities.
  • Analytical splitting term: The analytical splitting term has ODEs similar to the forward diffusion, differing by flipped signs in the off-diagonal drift terms.The derivation presents corresponding mean and covariance solutions for this term.
  • Analytical splitting term: For the analytical splitting term, the initial covariance is zero because each update starts from an updated sample without uncertainty.The derivation allows a general initial time while retaining this zero-covariance initialization in practice.
  • Analytical splitting term: The claimed mean and covariance solutions for the analytical splitting term are verified by substitution into their respective ODEs.The appendix concludes separate correctness proofs for the mean and covariance.
  • Implementation: With constant β, T = 1, and small step size δt, the general derivation reduces to the expressions used in SSCS.The appendix relates B(t) to β and δt for the practical implementation.
Loading 2112.07068v4…