Learning to Efficiently Sample from Diffusion Probabilistic Models

1 INTRODUCTION . Denoising Diffusion Probabilistic Models ( DDPMs ) have emerged as a powerful class of generative models ( Sohl-Dickstein et al. , 2015 ; Ho et al. , 2020 ) . DDPMs model the data distribution through an iterative denoising process , and have been applied successfully to a variety of applications , including unconditional image generation ( Song & Ermon , 2019 ; Ho et al. , 2020 ; Song et al. , 2021 ; Nichol & Dhariwal , 2021 ) , shape generation ( Cai et al. , 2020 ) , text-to-speech ( Chen et al. , 2021 ; Kong et al. , 2020 ) and single image super-resolution ( Saharia et al. , 2021 ; Li et al. , 2021 ) . DDPMs are easy to train , featuring a simple denoising objective ( Ho et al. , 2020 ) with noise schedules that successfully transfer across different models and datasets . This contrasts to Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) , which require an inner-outer loop optimization procedure that often entails instability and requires careful hyperparameter tuning . DDPMs also admit a simple non-autoregressive inference process ; this contrasts to autoregressive models with often prohibitive computational costs on high dimensional data . The DDPM inference process starts with samples from the corresponding prior noise distribution ( e.g. , standard Gaussian ) , and iteratively denoises the samples under the fixed noise schedule . However , DDPMs often need hundreds-tothousands of denoising steps ( each involving a feedforward pass of a large neural network ) to achieve strong results . While this process is still much faster than autoregressive models , this is still often computationally prohibitive , especially when modeling high dimensional data . There has been much recent work focused on improving the sampling speed of DDPMs . WaveGrad ( Chen et al. , 2021 ) introduced a manually crafted schedule requiring only 6 refinement steps ; however , this schedule seems to be only applicable to the vocoding task where there is a very strong conditioning signal . Denoising Diffusion Implicit Models ( DDIMs ) ( Song et al. , 2020 ) accelerate sampling from pre-trained DDPMs by relying on a family of non-Markovian processes . They accelerate the generative process through taking multiple steps in the diffusion process . However , DDIMs sacrifice the ability to compute log-likelihoods . Nichol & Dhariwal ( 2021 ) also explored the use of ancestral sampling with a subsequence of the original denoising steps , trying both a uniform stride and other hand-crafted strides . San-Roman et al . ( 2021 ) improve few-step sampling further by training a separate model after training a DDPM to estimate the level of noise , and modifying inference to dynamically adjust the noise schedule at every step to match the predicted noise level . All these fast-sampling techniques rely on a key property of DDPMs – there is a decoupling between the training and inference schedule . The training schedule need not be the same as the inference schedule , e.g. , a diffusion model trained to use 1000 steps may actually use only 10 steps during inference . This decoupling characteristic is typically not found in other generative models . In past work , the choice of inference schedule was often considered a hyperpameter selection problem , and often selected via intuition or extensive hyperparmeter exploration ( Chen et al. , 2021 ) . In this work , we view the choice of the timesteps of the inference schedule ( which we just call an inference path ) as an independent optimization problem , wherein we attempt to learn the best schedule . Our approach relies on the observation that we can solve this optimization problem with dynamic programming . Given a fixed budget of K refinement steps and a pre-trained DDPM , we find the set of timesteps that maximizes the corresponding evidence lower bound ( ELBO ) . As an optimization objective , the ELBO has a key decomposability property : the total ELBO is the sum of individual KL terms , and for any two inference paths , if the timesteps ( s , t ) contiguously occur in both , they share a common KL term , therefore admitting memoization ( see Section 4.1 for a precise definition ) . Our main contributions are the following : • We introduce a method that that finds the likelihood-optimal inference paths with a simple dynamic programming algorithm for all possible computation budgets of K refinement steps . The algorithm searches over T > K timesteps , only requiring O ( T ) neural network forward passes . It only needs to be applied once to a pre-trained DDPM , does not require training or retraining a DDPM , and is applicable to both time-discrete and time-continuous DDPMs . • We experiment with DDPM models from prior work . On both Lsimple CIFAR10 and Lhybrid ImageNet 64x64 , we discover schedules which require only 32 refinement steps , yet sacrifice only 0.1 bits per dimension compared to their original counterparts with 1,000 and 4,000 steps , respectively . • We show that our method can be applied to any decomposable set of objectives . In particular , optimizing a reweighted ELBO can favourably bias our algorithm towards solutions with better FID scores , as we find that optimizing the exact variational lower bound may lead to worse FID scores , which is consistent with prior work on unconditional image generation . 2 BACKGROUND ON DENOISING DIFFUSION PROBABILISTIC MODELS . Denoising Diffusion Probabilistic Models ( DDPMs ) ( Ho et al. , 2020 ; Sohl-Dickstein et al. , 2015 ) are defined in terms of a forward Markovian diffusion process q and a learned reverse process pθ . The forward diffusion process gradually adds Gaussian noise to a data point x0 through T iterations , q ( x1 : T | x0 ) = ∏T t=1 q ( xt | xt−1 ) , ( 1 ) q ( xt | xt−1 ) = N ( xt | √ αt xt−1 , ( 1− αt ) I ) , ( 2 ) where the scalar parameters α1 : T determine the variance of the noise added at each diffusion step , subject to 0 < αt < 1 . The learned reverse process aims to model q ( x0 ) by inverting the forward process , gradually removing noise from signal starting from pure Gaussian noise xT , p ( xT ) = N ( xT | 0 , I ) ( 3 ) pθ ( x0 : T ) = p ( xT ) ∏T t=1 pθ ( xt−1 | xt ) ( 4 ) pθ ( xt−1 | xt ) = N ( xt−1 | µθ ( xt , t ) , σ2t I ) . ( 5 ) The parameters of the reverse process can be optimized by maximizing the following variational lower bound on the training set : Eq log p ( x0 ) ≥ Eq [ log pθ ( x0|x1 ) − T∑ t=2 DKL ( q ( xt−1|xt , x0 ) ‖pθ ( xt−1|xt ) ) − LT ( x0 ) ] ( 6 ) where LT ( x0 ) = DKL ( q ( xT |x0 ) ‖ p ( xT ) ) . Nichol & Dhariwal ( 2021 ) have demonstrated that training DDPMs by maximizing the ELBO yields competitive log-likelihood scores on both CIFAR10 and ImageNet 64×64 achieving 2.94 and 3.53 bits per dimension respectively . Two notable properties of Gaussian diffusion process that help formulate DDPMs tractably and efficiently include : q ( xt | x0 ) = N ( xt | √ γt x0 , ( 1− γt ) I ) , where γt = ∏t i=1 αi , ( 7 ) q ( xt−1 | x0 , xt ) = N ( xt−1 ∣∣∣ √γt−1 ( 1− αt ) x0 +√αt ( 1− γt−1 ) xt 1− γt , ( 1− γt−1 ) ( 1− αt ) 1− γt I ) . ( 8 ) Given the marginal distribution of xt given x0 in ( 7 ) , one can sample from the q ( xt | x0 ) independently for different t and perform SGD on a randomly chosen KL term in ( 6 ) . Furthermore , given that the posterior distribution of xt−1 given xt and x0 is Gaussian , one can compute each KL term in ( 6 ) between two Gaussians in closed form and avoid high variance Monte Carlo estimation . 3 LINKING DDPMS TO CONTINUOUS TIME AFFINE DIFFUSION PROCESSES . Before describing our approach to efficiently sampling from DDPMs , it is helpful to link DDPMs to continuous time affine diffusion processes , as it shows the compatibility of our approach to both time-discrete and time-continuous DDPMs ( Song et al. , 2021 ; Kingma et al. , 2021 ) . Let x0 ∼ q ( x0 ) denote a data point drawn from the empirical distribution of interest and let q ( xt|x0 ) denote a stochastic process for t ∈ [ 0 , 1 ] defined through an affine diffusion process through the following stochastic differential equation ( SDE ) : dXt = fsde ( t ) Xtdt+ gsde ( t ) dBt , ( 9 ) where fsde , gsde : [ 0 , 1 ] → [ 0 , 1 ] are integrable functions satisfying fsde ( 0 ) = 1 and gsde ( 0 ) = 0 . Following Särkkä & Solin ( 2019 ) ( section 6.1 ) , we can compute the exact marginals q ( xt|xs ) for any 0 ≤ s < t ≤ 1 . We get : q ( xt | xs ) = N ( xt ∣∣∣ψ ( t , s ) xs , ( ∫ t s ψ ( t , u ) 2g ( u ) 2du ) I ) ( 10 ) where ψ ( t , s ) = exp ∫ t s f ( u ) du . Since these integrals are difficult to work with , we instead propose ( in parallel to Kingma et al . ( 2021 ) ) to define the marginals directly : q ( xt | x0 ) = N ( xt | f ( t ) x0 , g ( t ) 2I ) ( 11 ) where f , g : [ 0 , 1 ] → [ 0 , 1 ] are differentiable , monotonic functions satisfying f ( 0 ) = 1 , f ( 1 ) = 0 , g ( 0 ) = 0 , g ( 1 ) = 1 . Then , by implicit differentiation it follows that the corresponding diffusion is dXt = f ′ ( t ) f ( t ) Xtdt+ √ 2g ( t ) ( g′ ( t ) − f ′ ( t ) g ( t ) f ( t ) ) dBt . ( 12 ) We provide a proof for Equation 12 in the appendix ( A.1 ) . To complete our formulation , let fts = f ( t ) f ( s ) and gts = √ g ( t ) 2 − f2tsg ( s ) 2 . Then , it follows that for any 0 < s < t ≤ 1 we have that q ( xt | xs ) = N ( xt | ftsxs , g2tsI ) , ( 13 ) q ( xs | xt , x0 ) = N ( xs ∣∣∣ 1 g2t0 ( fs0g 2 tsx0 + ftsg 2 s0xt ) , g2s0g 2 ts g2t0 I ) , ( 14 ) We include proofs for ( 13 ) and ( 14 ) in the appendix ( A.2 ) . These equations show that we can perform inference with any ancestral sampling path ( i.e. , the timesteps can attain continuous values ) by formulating the reverse process in terms of the posterior distribution as pθ ( xs | xt ) = q ( xs | xt , x̂0 = 1ft0 ( xt − gt0 θ ( xt , t ) ) ) , ( 15 ) justifying the compatibility of our main approach with time-continuous DDPMs . We note that this reverse process is also mathematically equivalent to a reverse process based on a time-discrete DDPM derived from a subsequence of the original timesteps as done by Song et al . ( 2020 ) ; Nichol & Dhariwal ( 2021 ) . For the case of s = 0 in the reverse process , we follow the parametrization of Ho et al . ( 2020 ) to obtain discretized log likelihoods and compare our log likelihoods fairly with prior work . Algorithm 1 : Given a matrixL ∼ ( T+1 ) × ( T+1 ) of precomputed L ( · , · ) terms , find the likelihoodoptimal schedules for all step budgets . def vectorized_dp_all_budgets ( L ) : T = len ( L ) - 1 D = np.full ( L.shape , -1 ) C = np.full ( L.shape , np.inf ) C [ 0 , 0 ] = 0 for k in range ( 1 , T + 1 ) : bpds = C [ k - 1 , None ] + L C [ k ] = np.amin ( bpds , axis=-1 ) D [ k ] = np.argmin ( bpds , axis=-1 ) return D Algorithm 2 : Fetch the shortest path of K steps from the dynamic programming results implicitly returned by Algorithm 1. def fetch_shortest_path ( D , K ) : optpath = [ ] t = K for k in reversed ( range ( K ) ) : optpath.append ( t ) t = D [ k , t ] return optpath

Summary: Samples are generated from DDPMs by solving an SDE (often in "discrete time", which is used to refer to specifically the Euler--Maruyama discretisation). This necessitates a choice for where to make numerical steps. Each choice of step locations has a corresponding ELBO. This paper demonstrates that (on a pretrained model) the optimal ELBO may be obtained via a dynamic programming algorithm for the location of the steps.