Pseudo Numerical Methods for Diffusion Models on Manifolds

1 INTRODUCTION . Denoising Diffusion Probabilistic Models ( DDPMs ) ( Sohl-Dickstein et al. , 2015 ; Ho et al. , 2020 ) is a class of generative models which model the data distribution through an iterative denoising process reversing a multi-step noising process . DDPMs have been applied successfully to a variety of applications , including image generation ( Ho et al. , 2020 ; Song et al. , 2020b ) , text generation ( Hoogeboom et al. , 2021 ; Austin et al. , 2021 ) , 3D point cloud generation ( Luo & Hu , 2021 ) , textto-speech ( Kong et al. , 2021 ; Chen et al. , 2020 ) and image super-resolution ( Saharia et al. , 2021 ) . Unlike Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) , which require careful hyperparameter tuning according to different model structures and datasets , DDPMs can use similar model structures and be trained by a simple denoising objective which makes the models fit the noise in the data . To generate samples , the iterative denoising process starts from white noise and progressively denoises it into the target domain according to the noise predicted by the model at every step . However , a critical drawback of DDPMs is that DDPMs require hundreds to thousands of iterations to produce high-quality samples and need to pass through a network at least once at every step , which makes the generation of a large number of samples extremely slow and infeasible . In contrast , GANs only need one pass through a network . There have been many recent works focusing on improving the speed of the denoising process . Some works search for better variance schedules , including Nichol & Dhariwal ( 2021 ) and Watson et al . ( 2021 ) . Some works focus on changing the inference equation , including Song et al . ( 2020a ) ∗Corresponding author 1Our implementation is available at https : //github.com/luping-liu/PNDM . and Song et al . ( 2020b ) . Denoising Diffusion Implicit Models ( DDIMs ) ( Song et al. , 2020a ) relying on a non-Markovian process accelerate the denoising process by taking multiple steps every iteration . Probability Flows ( PFs ) ( Song et al. , 2020b ) build a connection between the denoising process and solving ordinary differential equations and use numerical methods of differential equations to accelerate the denoising process . Additionally , we introduce more related works in Appendix A.1 . However , this direct connection between DDPMs and numerical methods ( e.g. , forward Euler method , linear multi-step method and RungeKutta method ( Timothy , 2017 ) ) has weaknesses in both speed and effect ( see Section 3.1 ) . Some numerical methods are straightforward , like the forward Euler method , but they can only trade quality for speed . Some numerical methods can accelerate the reverse process without loss of quality , like the Runge-Kutta method , but they need to propagate forward more times along a neural network at every step . Furthermore , we also notice that numerical methods can introduce noticeable noise at a high speedup rate , which makes high-order numerical methods ( e.g. , Runge-Kutta method ) even less effective than DDIMs . This phenomenon is also mentioned in Salimans & Ho ( 2022 ) . To figure out the reason for the performance degradation in classical numerical methods , we conduct some analyses and find that classical numerical methods may sample data far away from the main distribution area of the data , and the inference equations of DDPMs do not satisfy a necessary condition of numerical methods at the last several steps ( see Section 3.2 ) . To tackle these problems , we design new numerical methods called pseudo numerical methods for diffusion models ( PNDMs ) to generate samples along a specific manifold in Rn , which is the highdensity region of the data . We first compute the corresponding differential equations of diffusion models directly and self-consistently , which builds a theoretical connection between DDPMs and numerical methods . Considering that classical numerical methods can not guarantee to generate samples on certain manifolds , we provide brand-new numerical methods called pseudo numerical methods based on our theoretical analyses . We also find that DDIMs are simple cases of pseudo numerical methods , which means that we also provide a new way to understand DDIMs better . Furthermore , we find that the pseudo linear multi-step method is the fastest method for diffusion models under similar generated quality . Besides , we provide a detailed theoretical analysis of our new theory and give visualization results to support our theory intuitively . According to our experiments , our methods have several advantages : • Our methods combine the benefits of DDIMs and high-order numerical methods successfully . We theoretically prove that our new methods PNDMs are second-order convergent while DDIMs are first-order convergent , which makes PNDMs 20x faster without loss of quality on Cifar10 and CelebA . • Our methods can reduce the best FID of pre-trained models with even shorter sampling time . With only 250 steps , our new denoising process can reduce the best FID by around 0.4 points Cifar10 and CelebA . We achieve a new SOTA FID score of 2.71 on CelebA . • Our methods work well with different variance schedules , which means that our methods have a good generalization and can be used together with those works introducing better variance schedules to accelerate the denoising process further . 2 BACKGROUND . In this section , we introduce some backgrounds . Firstly , we present the classical understanding of DDPMs . Then we provide another understanding based on Song et al . ( 2020b ) , which inspires us to use numerical methods to accelerate the denoising process of diffusion models . After that , we introduce some background on numerical methods used later in this paper . 2.1 DENOISING DIFFUSION PROBABILISTIC MODELS . DDPMs model the data distribution from Gaussian distribution to image distribution through an iterative denoising process . Let x0 be an image , then the diffusion process is a Markov process and the reverse process has a similar form to the diffusion process , which satisfies : xt+1 ∼N ( √ 1− βtxt , βtI ) , t = 0 , 1 , · · · , N − 1. xt−1 ∼N ( µθ ( xt , t ) , βθ ( xt , t ) I ) , t = N , N − 1 , · · · , 1 . ( 1 ) Here , βt controls the speed of adding noise to the data , calling them the variance schedule . N is the total number of steps of the denoising process . µθ and βθ are two neural networks , and θ are their parameters . Ho et al . ( 2020 ) get some statistics estimations of µθ and βθ . According to the properties of the conditional Gaussian distribution , we have : q ( xt|x0 ) = N ( √ ᾱtx0 , ( 1− ᾱt ) I ) , q ( xt−1|xt , x0 ) = N ( µ̄t ( xt , x0 ) , β̄tI ) . ( 2 ) Here , αt = 1 − βt , ᾱt = ∏t i=1 αi , µ̄t = √ ᾱt−1βt 1−ᾱt x0 + √ αt ( 1−ᾱt−1 ) 1−ᾱt xt and β̄t = 1−ᾱt−1 1−ᾱt βt . Then this paper sets βθ = β̄t and designs a objective function to help neural networks to represent µθ . Objective Function The objective function is defined by : Lt−1 = Eq [ ||µ̄t ( xt , x0 ) − µθ ( xt , t ) ||2 ] = Ex0 , ϵ [ || 1√ αt ( xt ( x0 , ϵ ) − βt√ 1− ᾱt ϵ ) − µθ ( xt ( x0 , ϵ ) , t ) ||2 ] = Ex0 , ϵ [ β2t αt ( 1− ᾱt ) ||ϵ− ϵθ ( √ ᾱtx0 + √ 1− ᾱtϵ , t ) ||2 ] . ( 3 ) Here , xt ( x0 , ϵ ) = √ ᾱtx0+ √ 1− ᾱtϵ , ϵ ∼ N ( 0 , 1 ) , ϵθ is an estimate of the noise ϵ . The relationship between µθ and ϵθ is µθ = 1√αt ( xt − βt√ 1−ᾱt ϵθ ) . Because ϵ ∼ N ( 0 , 1 ) , we assume that the mean and variance of ϵθ are 0 and 1 . 2.2 STOCHASTIC DIFFERENTIAL EQUATION . According to Song et al . ( 2020b ) , there is another understanding of DDPMs . The diffusion process can be treated as solving a certain stochastic differential equation dx = ( √ 1− β ( t ) − 1 ) x ( t ) dt +√ β ( t ) dw . According to Anderson ( 1982 ) , the denoising process also satisfies a similar stochastic differential equation : dx = ( ( √ 1− β ( t ) − 1 ) x ( t ) − β ( t ) ϵθ ( x ( t ) , t ) ) dt+ √ β ( t ) dw̄ . ( 4 ) This is Variance Preserving stochastic differential equations ( VP-SDEs ) . Here , we change the domain of t from [ 1 , N ] to [ 0 , 1 ] . When N tends to infinity , { βi } Ni=1 , { xi } Ni=1 become continuous functions β ( t ) and x ( t ) on [ 0 , 1 ] . Song et al . ( 2020b ) also show that this equation has an ordinary differential equation ( ODE ) version with the same marginal probability density as Equation ( 4 ) : dx = ( ( √ 1− β ( t ) − 1 ) x ( t ) − 1 2 β ( t ) ϵθ ( x ( t ) , t ) ) dt . ( 5 ) This different denoising equation with no random item and the same diffusion equation together is Probability Flows ( PFs ) . These two denoising equations show us a new possibility that we can use numerical methods to accelerate the reverse process . As far as we know , DDIMs first try to remove this random item , so PFs can also be treated as an acceleration of DDIMs , while VP-SDEs are an acceleration of DDPMs . 2.3 NUMERICAL METHOD . Many classical numerical methods can be used to solve ODEs , including the forward Euler method , Runge-Kutta method and linear multi-step method ( Timothy , 2017 ) . Forward Euler Method For a certain differential equation satisfying dxdt = f ( x , t ) . The trivial numerical method is forward Euler method satisfying xt+δ = xt + δf ( xt , t ) . Runge-Kutta Method Runge-Kutta method uses more information at every step , so it can achieve higher accuracy 2 . Runge-Kutta method satisfies :  k1 = f ( xt , t ) , k2 = f ( xt + δ 2k1 , t+ δ 2 ) k3 = f ( xt + δ 2k2 , t+ δ 2 ) , k4 = f ( xt + δk3 , t+ δ ) xt+δ = xt + δ 6 ( k1 + 2k2 + 2k3 + k4 ) . ( 6 ) Linear Multi-Step Method Linear multi-step method is another numerical method and satisfies : xt+δ = xt + δ 24 ( 55ft − 59ft−δ + 37ft−2δ − 9ft−3δ ) , ft = f ( xt , t ) . ( 7 ) 3 PSEUDO NUMERICAL METHOD FOR DDPM . In this section , we first compute the corresponding differential equations of diffusion models to build a direct connection between DDPMs and numerical methods . As a byproduct , we can directly use pre-trained models from DDPMs . After establishing this connection , we provide detailed analyses on the weakness of classical numerical methods . To solve the problems in classical numerical methods , we dive into the structure of numerical methods by dividing their equations into a gradient part and a transfer part and define pseudo numerical methods by introducing nonlinear transfer parts . We find that DDIMs can be regarded as simple pseudo numerical methods . Then , We explore the pros and cons of different numerical methods and choose the linear multi-step method to make numerical methods faster . Finally , we summarize our findings and analyses and safely propose our novel pseudo numerical methods for diffusion models ( PNDMs ) , which combine our proposed transfer part and the gradient part of the linear multi-step method . Furthermore , we analyze the convergence order of pseudo numerical methods to demonstrate the effectiveness of our methods theoretically .

Summary: Highlighting the high computational complexity for sampling from Denoising Diffusion Probabilistic Models (DDPMs) (e.g. wrt GANs), authors build on the connection between diffusion processes and ODEs to propose efficient (pseudo-)numerical methods so as to sample data from the data manifold. The main idea is to combine the discrete update proposed in DDIMs, with a fourth-order gradient estimation given by the Runge-Kutta or linear multi-step methods. The motivation being that such gradient estimator should yield trajectories that stay closer to the data manifold. They empirically assess their method(s) on CIFAR10 and CelebA in terms of sample quality (measured by FID), and show that they get a ~x20 speed-up wrt DDIMs or a significant improvement in FID with the same number of steps.