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: In this paper, the authors introduce a novel approach for training a score-based generative model.  With prior score based networks, generation is performed by solving a stochastic differential equation (SDE) based on Langevin Dynamics, using an estimate of gradient of the log likelihood of the underlying signal distribution.  In the present paper, the authors present a novel forward process, where diffusion is run in a joint data-velocity space.  The noise term is only applied to the velocity component.  This reduces the learning problem to only needing to learn the score of the conditional distribution of velocity given data, which is easier than learning the score of the data distribution directly.  The paper shows that the novel scheme (CLD) yields higher quality for image generation when compared to prior models of similar capacity and number of neural network evaluations.