Gotta Go Fast When Generating Data with Score-Based Models

Score-based ( denoising diffusion ) generative models have recently gained a lot of success in generating realistic and diverse data . These approaches define a forward diffusion process for transforming data to noise and generate data by reversing it ( thereby going from noise to data ) . Unfortunately , current score-based models generate data very slowly due to the sheer number of score network evaluations required by numerical SDE solvers . In this work , we aim to accelerate this process by devising a more efficient SDE solver . Existing approaches rely on the Euler-Maruyama ( EM ) solver , which uses a fixed step size . We found that naively replacing it with other SDE solvers fares poorly - they either result in low-quality samples or become slower than EM . To get around this issue , we carefully devise an SDE solver with adaptive step sizes tailored to score-based generative models piece by piece . Our solver requires only two score function evaluations per step , rarely rejects samples , and leads to high-quality samples . Our approach generates data 2 to 10 times faster than EM while achieving better or equal sample quality . For high-resolution images , our method leads to significantly higher quality samples than all other methods tested . Our SDE solver has the benefit of requiring no step size tuning . 1 INTRODUCTION . Score-based generative models ( Song and Ermon , 2019 ; 2020 ; Ho et al. , 2020 ; Jolicoeur-Martineau et al. , 2020 ; Song et al. , 2020a ; Piché-Taillefer , 2021 ) have been very successful at generating data from various modalities , such as images ( Ho et al. , 2020 ; Song et al. , 2020a ) , audio ( Chen et al. , 2020 ; Kong et al. , 2020 ; Mittal et al. , 2021 ; Kameoka et al. , 2020 ) , and graphs ( Niu et al. , 2020 ) . They have further been used effectively for super-resolution ( Saharia et al. , 2021 ; Kadkhodaie and Simoncelli , 2020 ) , inpainting ( Kadkhodaie and Simoncelli , 2020 ) , source separation ( Jayaram and Thickstun , 2020 ) , and image-to-image translation ( Sasaki et al. , 2021 ) . In most of these applications , scorebased models achieved superior performances in terms of quality and diversity than the historically dominant Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) . Score-based models can be understood in two main classes : those based on a Variance Exploding ( VE ) diffusion process ( Song and Ermon , 2019 ) and those based on a Variance Preserving ( VP ) one ( Ho et al. , 2020 ) . Both diffusion processes progressively transform real data into Gaussian noise ; N ( 0 , σ2maxI ) for VE where σ2max is very large , and N ( 0 , I ) for VP . The diffusion process ( VE , VP , etc . ) is then reversed in order to generate real data from Gaussian noise . Reversing the process requires the score function , which is estimated with a neural network ( known as a score network ) . Although very powerful , score-based models generate data through an undesirably long iterative process ; meanwhile , other state-of-the-art methods such as GANs generate data from a single forward pass of a neural network . Increasing the speed of the generative process is thus an active area of research . Chen et al . ( 2020 ) and San-Roman et al . ( 2021 ) proposed faster step size schedules for VP diffusions that still yield relatively good quality/diversity metrics . Although fast , these schedules are arbitrary , require careful tuning , and the optimal schedules will vary from one model to another . Block et al . ( 2020 ) proposed generating data progressively from low to high-resolution images and show that the scheme improves speed . Similarly , Nichol and Dhariwal ( 2021 ) proposed generating low-resolution images and then upscale them since generating low-resolution images is quicker . They further suggested to accelerate VP-based models by learning dimension-specific noise rather than assuming equal noise everywhere . Note that these methods do not affect the data generation algorithm and would thus be complementary to our methods . Song et al . ( 2020a ) and Song et al . ( 2020b ) proposed removing the noise from the data generation algorithm and solve an Ordinary Differential Equation ( ODE ) rather than a Stochastic Differential Equation ( SDE ) ; they report being able to converge much faster when there is no noise . Although it improves the generation speed , Song et al . ( 2020a ) report obtaining lower-quality images when using the ODE formulation for the VE process ( Song et al. , 2020a ) . We will later show that our SDE solver generally leads to better results than ODE solvers at similar speeds . Thus , existing methods for acceleration often require considerable step size/schedule tuning ( this is also true for the baseline approach ) and do not always work for both VE and VP processes . To improve speed and remove the need for step size/schedule tuning , we propose to solve the reverse diffusion process using SDE solvers with adaptive step sizes . It turns out that off-the-shelf SDE solvers are ill-suited for generative modeling and exhibit either ( 1 ) divergence , ( 2 ) slower data generation than the baseline , or ( 3 ) significantly worse quality than the baseline ( see Appendix A ) . This can be attributed to distinct features of the SDEs that arise in score-based generative models that set them apart from the SDEs traditionally considered in the numerical SDE solver literature , namely : ( 1 ) the codomain of the unknown function is extremely high-dimensional , especially in the case of image generation ; ( 2 ) evaluating the score function is computationally expensive , requiring a forward pass of a large mini-batch through a large neural network ; ( 3 ) the required precision of the solution is smaller than usual because we are satisfied as long as the error is not perceptible ( e.g. , one RGB increment on an image ) . We devise our own SDE solver with these features in mind , resulting in an algorithm that can get around the problems encountered by off-the-shelf solvers . To address high dimensionality , we use the ` 2 norm rather than the ` ∞ norm to measure the error across different dimensions to prevent a single pixel from slowing down the solver . To address the cost of score function evaluations while still obtaining high precision , we ( 1 ) take the minimum number of score function evaluations needed for adaptive step sizes ( two evaluations ) , and ( 2 ) use extrapolation to get high precision at no extra cost . To take advantage of the reduced requirement for precision , we set the absolute tolerance for the error according to the range of RGB values . Our main contribution is a new SDE solver tailored to score-based generative models with the following benefits : • Our solver is much faster than the baseline methods , i.e . reverse-diffusion method with Langevin dynamics and Euler-Maruyama ( EM ) ; • It yields higher quality/diversity samples than EM when using the same computing budget ; • It does not require any step size or schedule tuning ; • It can be used to quickly solve any type of diffusion process ( e.g. , VE , VP ) 2 BACKGROUND . 2.1 SCORE-BASED MODELING WITH SDES . Let x ( 0 ) ∈ Rd be a sample from the data distribution pdata . The sample is gradually corrupted over time through a Forward Diffusion Process ( FDP ) , a common type of Stochastic Differential Equation ( SDE ) : dx = f ( x , t ) dt+ g ( t ) dw , ( 1 ) where f ( x , t ) : Rd × R→ Rd is the drift , g ( t ) : R→ R is the diffusion coefficient and w ( t ) is the Wiener process indexed by t ∈ [ 0 , 1 ] . Data points and their probability distribution evolve along the trajectories { x ( t ) } 1t=0 and { pt ( x ) } 1t=0 respectively , with p0 ≡ pdata . The functions f and g are chosen such that x ( 1 ) be approximately Gaussian and independent from x ( 0 ) . Inference is achieved by reversing this diffusion , drawing x ( 1 ) from its Gaussian distribution and solving the Reverse Diffusion Process ( RDP ) equal to : dx = [ f ( x , t ) − g ( t ) 2∇x log pt ( x ) ] dt+ g ( t ) dw̄ , ( 2 ) where∇x log pt ( x ) is referred to as the score of the distribution at time t ( Hyvärinen , 2005 ) and w̄ ( t ) is the Wiener process in which time flows backward ( Anderson , 1982 ) . One can observe from Equation 2 that the RDP requires knowledge of the score ( or pt ) , which we do not have access to . Fortunately , it can be estimated by a neural network ( referred to as the score network ) by optimizing the following objective : L ( θ ) = Ex ( t ) ∼p ( x ( t ) |x ( 0 ) ) , x ( 0 ) ∼pdata [ λ ( t ) 2 ∥∥sθ ( x ( t ) , t ) −∇x ( t ) log pt ( x ( t ) |x ( 0 ) ) ∥∥22 ] , ( 3 ) where λ ( t ) : R→ R is a weighting function generally chosen to be inversely proportional to : E [ ∥∥∇x ( t ) log pt ( x ( t ) |x ( 0 ) ) ∥∥22 ] . One can demonstrate that the minimizer of that objective θ∗ will be such that sθ∗ ( x , t ) = ∇x log pt ( x ) ( Vincent , 2011 ) , allowing us to approximate the reverse diffusion process . As can be seen , evaluating the objective requires the ability to generate samples from the FDP at arbitrary times t. Thankfully , as long as the drift is affine ( i.e. , f ( x , t ) = Ax + B ) , the transition kernel p ( x ( t ) |x ( 0 ) ) will always be normally distributed ( Särkkä and Solin , 2019 ) , which means that we can solve the forward diffusion in a single step . Furthermore , the score of the Gaussian transition kernel is trivial to compute , making the loss an inexpensive training objective . There are two primary choices for the FDP in the literature , which we discuss below . 2.2 VARIANCE EXPLODING ( VE ) PROCESS . The Variance Exploding ( VE ) process consists in the following FDP : dx = √ d [ σ2 ( t ) ] dt dw . Its associated transition kernel is : x ( t ) |x ( 0 ) ∼ N ( x ( 0 ) , [ σ2 ( t ) − σ2 ( 0 ) ] I ) ≈ N ( x ( 0 ) , σ2 ( t ) I ) . In practice , we let σ ( t ) = σmin ( σmax σmin ) t , where σmin = 0.01 and σmax ≈ maxi ∑N j=1 ||x ( i ) − x ( j ) || is the maximum Euclidean distance between two samples from the dataset { x ( i ) } Ni=1 ( Song and Ermon , 2020 ) . Using the maximum Euclidean distance ensures that x ( 1 ) does not depend on x ( 0 ) ; thus , x ( 1 ) is approximately distributed as N ( 0 , σ2 ( 1 ) I ) .

Summary: The paper presents a new SDE solver for the reverse process in score-based models. The algorithm is fast and offers high quality, and avoids some hyperparameter tuning. There is theoretical analysis on the stability and bias of the algorithm. The paper also has experiments comparing the proposed algorithm to several baseline methods.