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: This work presents a method to efficiently sample from a pre-trained DDPM by solving a dynamic programming problem that can maximize the log likelihood of the data samples given a fixed computational budget. This is done by defining a least-cost path problem to select a reduced set of time steps among a full grid of potential time steps across different possible step budget sizes, where the ELBO is used as the cost function. The authors show that their method can identify DDPM schedules that can achieve significantly higher log likelihood (i.e. lower bits/dim) than prior DDPM schedules in the regime where about a hundred steps or fewer are used.