Score-Based Generative Modeling with Critically-Damped Langevin Diffusion

1 INTRODUCTION . Score-based generative models ( SGMs ) and denoising diffusion probabilistic models have emerged as a promising class of generative models ( Sohl-Dickstein et al. , 2015 ; Song et al. , 2021c ; b ; Vahdat et al. , 2021 ; Kingma et al. , 2021 ) . SGMs offer high quality synthesis and sample diversity , do not require adversarial objectives , and have found applications in image ( Ho et al. , 2020 ; Nichol & Dhariwal , 2021 ; Dhariwal & Nichol , 2021 ; Ho et al. , 2021 ) , speech ( Chen et al. , 2021 ; Kong et al. , 2021 ; Jeong et al. , 2021 ) , and music synthesis ( Mittal et al. , 2021 ) , image editing ( Meng et al. , 2021 ; Sinha et al. , 2021 ; Furusawa et al. , 2021 ) , super-resolution ( Saharia et al. , 2021 ; Li et al. , 2021 ) , image-to-image translation ( Sasaki et al. , 2021 ) , and 3D shape generation ( Luo & Hu , 2021 ; Zhou et al. , 2021 ) . SGMs use a diffusion process to gradually add noise to the data , transforming a complex data distribution to an analytically tractable prior distribution . A neural network is then utilized to learn the score function—the gradient of the log probability density—of the perturbed data . The learnt scores can be used to solve a stochastic differential equation ( SDE ) to synthesize new samples . This corresponds to an iterative denoising process , inverting the forward diffusion . In the seminal work by Song et al . ( 2021c ) , it has been shown that the score function that needs to be learnt by the neural network is uniquely determined by the forward diffusion process . Consequently , the complexity of the learning problem depends , other than on the data itself , only on the diffusion . Hence , the diffusion process is the key component of SGMs that needs to be revisited to further improve SGMs , for example , in terms of synthesis quality or sampling speed . Inspired by statistical mechanics ( Tuckerman , 2010 ) , we propose a novel forward diffusion process , the critically-damped Langevin diffusion ( CLD ) . In CLD , the data variable , xt ( time t along the diffusion ) , is augmented with an additional “ velocity ” variable vt and a diffusion process is run in the joint data-velocity space . Data and velocity are coupled to each other as in Hamiltonian dynamics , and noise is injected only into the velocity variable . Similarly as in Hamiltonian Monte Carlo ( Neal , 2011 ) , the Hamiltonian component helps to efficiently traverse the joint data-velocity space and to transform the data distribution into the prior distribution more smoothly . We derive the corresponding score matching objective and show that for CLD the neural network is tasked with learning only the score of the conditional distribution of velocity given data∇vt log pt ( vt|xt ) , which is arguably easier than learning the score of diffused data distribution directly . Using techniques from molecular dynamics ( Bussi & Parrinello , 2007 ; Tuckerman , 2010 ; Leimkuhler & Matthews , 2013 ) , we also derive a novel SDE integrator tailored to CLD ’ s reverse-time synthesis SDE . We extensively validate CLD and the new SDE solver : ( i ) We show that the neural networks learnt in CLD-based SGMs are smoother than those of previous SGMs . ( ii ) On the CIFAR-10 image modeling benchmark , we demonstrate that CLD-based models outperform previous diffusion models in synthesis quality for similar network architectures and sampling compute budgets . We attribute these positive results to the Hamiltonian component in the diffusion and to CLD ’ s easier score function target , the velocity-data conditional distribution . ( iii ) We show that our novel sampling scheme for CLD significantly outperforms the popular Euler–Maruyama method . ( iv ) We perform ablations on various aspects of CLD and find that CLD does not have difficult-to-tune hyperparameters . In summary , we make the following technical contributions : ( i ) We propose CLD , a novel diffusion process for SGMs . ( ii ) We derive a score matching objective for CLD , which requires only the conditional distribution of velocity given data . ( iii ) We propose a new type of denoising score matching ideally suited for scalable training of CLD-based SGMs . ( iv ) We derive a tailored SDE integrator that enables efficient sampling from CLD-based models . ( v ) Overall , we provide novel insights into SGMs and point out important new connections to statistical mechanics . 2 BACKGROUND . Consider a diffusion process ut ∈ Rd defined by the Itô SDE dut = f ( ut , t ) dt+G ( ut , t ) dwt , t ∈ [ 0 , T ] , ( 1 ) with continuous time variable t ∈ [ 0 , T ] , standard Wiener process wt , drift coefficient f : Rd × [ 0 , T ] → Rd and diffusion coefficientG : Rd× [ 0 , T ] → Rd×d . Defining ūt : = uT−t , a corresponding reverse-time diffusion process that inverts the above forward diffusion can be derived ( Anderson , 1982 ; Haussmann & Pardoux , 1986 ; Song et al. , 2021c ) ( with positive dt and t ∈ [ 0 , T ] ) : dūt = [ −f ( ūt , T − t ) +G ( ūt , T − t ) G ( ūt , T − t ) > ∇ūt log pT−t ( ūt ) ] dt+G ( ūt , T − t ) dwt , ( 2 ) where∇ūt log pT−t ( ūt ) is the score function of the marginal distribution over ūt at time T − t. The reverse-time process can be used as a generative model . In particular , Song et al . ( 2021c ) model data x , setting p ( u0 ) =pdata ( x ) . Currently used SDEs ( Song et al. , 2021c ; Kim et al. , 2021 ) have drift and diffusion coefficients of the simple form f ( xt , t ) =f ( t ) xt andG ( xt , t ) =g ( t ) Id . Generally , f and G are chosen such that the SDE ’ s marginal , equilibrium density is approximately Normal at time T , i.e. , p ( uT ) ≈N ( 0 , Id ) . We can then initialize x0 based on a sample drawn from a complex data distribution , corresponding to a far-from-equilibrium state . While we let the state x0 relax towards equilibrium via the forward diffusion , we can learn a model sθ ( xt , t ) for the score ∇xt log pt ( xt ) , which can be used for synthesis via the reverse-time SDE in Eq . ( 2 ) . If f and G take the simple form from above , the denoising score matching ( Vincent , 2011 ) objective for this task is : min θ Et∼ [ 0 , T ] Ex0∼p ( x0 ) Ext∼pt ( xt|x0 ) [ λ ( t ) ‖sθ ( xt , t ) −∇xt log pt ( xt|x0 ) ‖22 ] ( 3 ) If f and G are affine , the conditional distribution pt ( xt|x0 ) is Normal and available analytically ( Särkkä & Solin , 2019 ) . Different λ ( t ) result in different trade-offs between synthesis quality and likelihood in the generative model defined by sθ ( xt , t ) ( Song et al. , 2021b ; Vahdat et al. , 2021 ) . 3 CRITICALLY-DAMPED LANGEVIN DIFFUSION . We propose to augment the data xt ∈ Rd with auxiliary velocity1 variables vt ∈ Rd and utilize a diffusion process that is run in the joint xt-vt-space . With ut = ( xt , vt ) > ∈ R2d , we set f ( ut , t ) : = ( 0 βM−1 −β −ΓβM−1 ) ⊗ Id ut , G ( ut , t ) : = ( 0 0 0 √ 2Γβ ) ⊗ Id , ( 4 ) where ⊗ denotes the Kronecker product . The coupled SDE that describes the diffusion process is ( dxt dvt ) = ( M−1vt −xt ) βdt︸ ︷︷ ︸ Hamiltonian component= : H + ( 0 −ΓM−1vt ) βdt+ ( 0√ 2Γβ ) dwt︸ ︷︷ ︸ Ornstein-Uhlenbeck process= : O , ( 5 ) which corresponds to Langevin dynamics in each dimension . That is , each xi is independently coupled to a velocity vi , which explains the blockwise structure of f and G. The mass M ∈ R+ is a hyperparameter that determines the coupling between the xt and vt variables ; β ∈ R+ is a constant time rescaling chosen such that the diffusion converges to its equilibrium distribution within t ∈ [ 0 , T ] ( in practice , we set T=1 ) when initialized from a data-defined non-equilibrium state and is analogous to β ( t ) in previous diffusions ( we could also use time-dependent β ( t ) , but found constant β ’ s to work well , and therefore opted for simplicity ) ; Γ ∈ R+ is a friction coefficient that determines the strength of the noise injection into the velocities . Notice that the SDE in Eq . ( 5 ) consists of two components . The H term represents a Hamiltonian component . Hamiltonian dynamics are frequently used in Markov chain Monte Carlo methods to accelerate sampling and efficiently explore complex probability distributions ( Neal , 2011 ) . The Hamiltonian component in our diffusion process plays a similar role and helps to quickly and smoothly converge the initial joint data-velocity distribution to the equilibrium , or prior ( see Fig . 1 ) . Furthermore , Hamiltonian dynamics on their own are trivially invertible ( Tuckerman , 2010 ) , which intuitively is also beneficial in our situation when using this diffusion for training SGMs . TheO term corresponds to an OrnsteinUhlenbeck process ( Särkkä & Solin , 2019 ) in the velocity component , which injects noise such that the diffusion dynamics properly converge to equilibrium for any Γ > 0 . It can be shown that the equilibrium distribution of this diffusion is pEQ ( u ) = N ( x ; 0d , Id ) N ( v ; 0d , MId ) ( see App . B.2 ) . There is a crucial balance between M and Γ ( McCall , 2010 ) : For Γ2 < 4M ( underdamped Langevin dynamics ) the Hamiltonian component dominates , which implies oscillatory dynamics of xt and vt that slow down convergence to equilibrium . For Γ2 > 4M ( overdamped Langevin dynamics ) the Oterm dominates which also slows down convergence , since the accelerating effect by the Hamiltonian component is suppressed due to the strong noise injection . For Γ2=4M ( critical damping ) , an ideal balance is achieved and convergence to pEQ ( u ) occurs as fast as possible in a smooth manner without oscillations ( also see discussion in App . A.1 ) ( McCall , 2010 ) . Hence , we propose to set Γ2=4M and call the resulting diffusion critically-damped Langevin diffusion ( CLD ) ( see Fig . 1 ) . Diffusions such as the VPSDE ( Song et al. , 2021c ) correspond to overdamped Langevin dynamics with high friction coefficients Γ ( see App . A.2 ) . Furthermore , in previous works noise is injected directly into the data variables ( pixels , for images ) . In CLD , only the velocity variables are subject to direct noise and the data is perturbed only indirectly due to the coupling between xt and vt .

Summary: This paper proposes a critically-damped version of score-based diffusion models, by extending the inference process to an augmented state-space and diffusing the data coupled with an auxiliary velocity variable. The authors further proposed a hybrid score matching loss for training the reverse generative process, which provides an empirical advantage of learning a conditional score function that does not blow up (e.g. to infinity) and therefore stabilizes training. For sampling, the authors derived a new numerical integration method based on first principles of statistical mechanics and MD, which they found has a better convergence behavior in practice if the computing budget is limited.