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: The paper proposes a novel class of continuous-time diffusion-based generative models. Specifically, the paper proposes forward diffusions that are simple forms of stochastic Hamiltonian dynamics, unlike the previous methods, where linear stochastic differential equations (linear SDE) diffuse the data distribution. Thus, the proposed forward diffusion transport a joint distribution of data $x_0$ and auxiliary random variables $v_0$ to a prior joint distribution. Then, the generative models are their reverse-time diffusion processes on the joint space whose initial condition is the prior joint distribution.  In the Hamiltonian dynamics-type forward diffusion, the auxiliary random variable can be interpreted as velocity (as in physics), while data corresponds to position. The forward diffusion updates data by the velocities without any noise, while the velocities are updated via linear SDEs, similarly to deterministic Hamiltonian dynamics. In particular, the paper shows that the data distribution will be transported to a prior distribution and that velocities will map to a prior distribution similar to previous approaches. The prior joint distributions are commonly chosen to simple distributions, such as standard Normal distributions.  The authors emphasize two interesting properties of the proposed method. First, the paper points out that the reverse-time SDEs should only contain $\nabla_{v_t} \log p_t(v_t| x_t)$, but not logarithmic gradient wrt $x_t$, which implies that the model complexity (the size of models) is almost similar to the previous approaches. Second, the authors shows that $p_t(v_t| x_t)$ is Gaussians for all $t \in [0, T]$, including $p_0(v_0|x_0) = p_0(v_0)$. This indicates that $\nabla_{v_t} \log p_t(v_t| x_t)$ is potentially bounded unlike previous diffusion-based generative models, where the scores closer to data are possibly unbounded. As a result, esp. with common practice to learn scores at all time $t$'s by a single neural networks, training the proposed model will be more stable.  In order to learn the $\nabla_{v_t} \log p_t(v_t| x_t)$, the paper first propose to use denoising score matching (DSM) following previous diffusion-based generative models. However, observing that training the proposed model via DSM is unstable, the paper proposes a modified objective called hybrid score matching (HSM). HSM can be obtained by marginalizing out $v_0$ from score matching loss similarly to deriving DSM; the marginalization is possible as the auxiliary random variable defined to follow a known distribution, such as Normal distribution.  In addition, as the reverse-time diffusion is also a Hamiltonian SDE, the paper proposes a new integrator for generations, benefitting from the symplectic structure of the Hamiltonian dynamics. Note that it has been well appreciated that the discretization methods that utilize symplectic structure are more accurate than the Euler-Maruyama method under similar computational costs. Thus, the proposed integrator will have a better quality of samples, esp. when the number of discretization is small.  In order to show the effectiveness of the proposed method, the paper runs three main experiments: First, with toy experiments, the authors demonstrate that learning $\nabla_{v_t} \log p_t(v_t| x_t)$ is less difficult in comparison to learning the scores wrt the data directly. Then, the authors evaluate the generation qualities of the proposed models on two image modeling benchmark datasets; CIFAR-10 and CelebA-256. Finally, the paper demonstrates the proposed integrator's effectiveness compared to the Euler-Maruyama method, including additional analysis ablating the effect of hyperparameters of the proposed integrator.