Unconditional Diffusion Guidance

1 INTRODUCTION . Diffusion models have recently emerged as an expressive and flexible family of generative models , delivering competitive sample quality and likelihood scores on image and audio synthesis tasks ( SohlDickstein et al. , 2015 ; Song & Ermon , 2019 ; Ho et al. , 2020 ; Song et al. , 2021b ; Kingma et al. , 2021 ; Song et al. , 2021a ) . These models have delivered audio synthesis performance rivaling the quality of autoregressive models with substantially fewer inference steps ( Chen et al. , 2021 ; Kong et al. , 2021 ) , and they have delivered ImageNet generation results outperforming BigGAN-deep ( Brock et al. , 2019 ) and VQ-VAE-2 ( Razavi et al. , 2019 ) in terms of FID score and classification accuracy score ( Ho et al. , 2021 ; Dhariwal & Nichol , 2021 ) . Dhariwal & Nichol ( 2021 ) proposed classifier guidance , a technique to boost the sample quality of a diffusion model using an extra trained classifier . Prior to classifier guidance , it was not known how to generate “ low temperature ” samples from a diffusion model similar to those produced by truncated BigGAN ( Brock et al. , 2019 ) or low temperature Glow ( Kingma & Dhariwal , 2018 ) : naive attempts , such as scaling the model score vectors or decreasing the amount of Gaussian noise added during diffusion sampling , are ineffective ( Dhariwal & Nichol , 2021 ) . Classifier guidance instead mixes a diffusion model ’ s score estimate with the input gradient of the log probability of a classifier . By varying the strength of the classifier gradient , Dhariwal & Nichol can trade off Inception score ( Salimans et al. , 2016 ) and FID score ( Heusel et al. , 2017 ) ( or precision and recall ) in a manner similar to varying the truncation parameter of BigGAN . We are interested in whether classifier guidance can be performed without a classifier . Because classifier guidance mixes a score estimate with a classifier gradient during sampling , classifier-guided diffusion sampling can be interpreted as attempting to confuse an image classifier with a gradientbased adversarial attack . This raises the question of whether classifier guidance is successful at boosting classifier-based metrics such as FID and Inception score ( IS ) simply because it is adversarial against such classifiers . Stepping in direction of classifier gradients also bears some resemblance to GAN training , particularly with nonparameteric generators ; this also raises the question of whether classifier-guided diffusion models perform well on classifier-based metrics because they are beginning to resemble GANs , which are already known to perform well on such metrics . To resolve these questions , we present unconditional guidance , our guidance method which avoids any classifier entirely . Rather than sampling in the direction of the gradient of an image classifier , unconditional guidance instead mixes the score estimates of a conditional diffusion model and a jointly trained unconditional diffusion model . By sweeping over the mixing weight , we attain a FID/IS tradeoff similar to that attained by classifier guidance . Our unconditional guidance results demonstrate that pure generative diffusion models are capable of synthesizing extremely high fidelity samples possible with other types of generative models . 2 BACKGROUND . We train diffusion models in continuous time ( Song et al. , 2021b ; Chen et al. , 2021 ; Kingma et al. , 2021 ) : letting x ∼ p ( x ) and z = { zλ |λ ∈ [ λmin , λmax ] } for hyperparameters λmin < λmax ∈ R , the forward process q ( z|x ) is the variance-preserving Markov process ( Sohl-Dickstein et al. , 2015 ) : q ( zλ|x ) = N ( αλx , σ2λI ) , where α2λ = 1/ ( 1 + e−λ ) , σ2λ = 1− α2λ ( 1 ) q ( zλ|zλ′ ) = N ( ( αλ/αλ′ ) zλ′ , σ2λ|λ′I ) , where λ < λ ′ , σ2λ|λ′ = ( 1− e λ−λ′ ) σ2λ ( 2 ) We will use the notation p ( z ) ( or p ( zλ ) ) to denote the marginal of z ( or zλ ) when x ∼ p ( x ) and z ∼ q ( z|x ) . Note that λ = logα2λ/σ2λ , so λ can be interpreted as the log signal-to-noise ratio of zλ , and the forward process runs in the direction of decreasing λ . Conditioned on x , the forward process can be described in reverse by the transitions q ( zλ′ |zλ , x ) = N ( µ̃λ′|λ ( zλ , x ) , σ̃2λ′|λI ) , where µ̃λ′|λ ( zλ , x ) = e λ−λ′ ( αλ′/αλ ) zλ + ( 1− eλ−λ ′ ) αλ′x , σ̃ 2 λ′|λ = ( 1− e λ−λ′ ) σ2λ′ ( 3 ) The reverse process generative model starts from pθ ( zλmin ) = N ( 0 , I ) . We specify the transitions : pθ ( zλ′ |zλ ) = N ( µ̃λ′|λ ( zλ , xθ ( zλ ) ) , ( σ̃2λ′|λ ) 1−v ( σ2λ|λ′ ) v ) ( 4 ) During sampling , we apply this transition along an increasing sequence λmin = λ1 < · · · < λT = λmax for T timesteps ; in other words , we follow the discrete time ancestral sampler of Sohl-Dickstein et al . ( 2015 ) ; Ho et al . ( 2020 ) . If the model xθ is correct , then as T →∞ , we obtain samples from an SDE whose sample paths are distributed as p ( z ) ( Song et al. , 2021b ) , and we use pθ ( z ) to denote the continuous time model distribution . The variance is a log-space interpolation of σ̃2λ′|λ and σ 2 λ|λ′ as suggested by Nichol & Dhariwal ( 2021 ) ; we found it effective to use a constant hyperparameter v rather than learned zλ-dependent v. Note that the variances simplify to σ̃2λ′|λ as λ ′ → λ , so v has an effect only when sampling with non-infinitesimal timesteps as done in practice . The reverse process mean comes from an estimate xθ ( zλ ) ≈ x plugged into q ( zλ′ |zλ , x ) ( Ho et al. , 2020 ; Kingma et al. , 2021 ) ( xθ also receives λ as input , but we suppress this to keep our notation clean ) . We parameterize xθ in terms of -prediction ( Ho et al. , 2020 ) : xθ ( zλ ) = ( zλ−σλ θ ( zλ ) ) /αλ , and we train on the objective E , λ [ ‖ θ ( zλ ) − ‖22 ] ( 5 ) where ∼ N ( 0 , I ) , zλ = αλx + σλ , and λ is drawn from a distribution p ( λ ) over [ λmin , λmax ] . This objective is denoising score matching ( Vincent , 2011 ; Hyvärinen & Dayan , 2005 ) over multiple noise scales ( Song & Ermon , 2019 ) , and when p ( λ ) is uniform , the objective is proportional to the variational lower bound on the marginal log likelihood of the latent variable model ∫ p ( x|z ) pθ ( z ) dz , ignoring the term for the unspecified decoder p ( x|z ) and for the prior at zλmin ( Kingma et al. , 2021 ) . If p ( λ ) is not uniform , the objective can be interpreted as weighted variational lower bound whose weighting can be tuned for sample quality ( Ho et al. , 2020 ; Kingma et al. , 2021 ) . We use a p ( λ ) inspired by the discrete time cosine noise schedule of Nichol & Dhariwal ( 2021 ) : we sample λ via λ = −2 log tan ( au+ b ) for uniformly distributed u ∈ [ 0 , 1 ] , where b = arctan ( e−λmax/2 ) and a = arctan ( e−λmin/2 ) − b . This represents a hyperbolic secant distribution modified to be supported on a bounded interval . For finite timestep generation , we use λ values corresponding to uniformly spaced u ∈ [ 0 , 1 ] , and the final generated sample is xθ ( zλmax ) . Because the loss for θ ( zλ ) is denoising score matching for all λ , the score θ ( zλ ) learned by our model estimates the gradient of the log-density of the distribution of our noisy data zλ , that is θ ( zλ ) ≈ σλ∇zλ log p ( zλ ) . Sampling from the learned diffusion model resembles using Langevin diffusion to sample from a sequence of distributions p ( zλ ) that converges to the conditional distribution p ( x ) of the original data x . In the case of conditional generative modeling , the data x is drawn jointly with conditioning information c , i.e . a class label for class-conditional image generation . The only modification to the model is that the reverse process function approximator receives c as input , as in θ ( zλ , c ) . 3 GUIDANCE . An interesting property of certain generative models , such as GANs and flow-based models , is the ability to perform truncated or low temperature sampling by decreasing the variance or range of noise inputs to the generative model at sampling time . The intended effect is to decrease the diversity of the samples while increasing the quality of each individual sample . Truncation in BigGAN ( Brock et al. , 2019 ) , for example , yields a tradeoff curve between FID score and Inception score for low and high amounts of truncation , respectively . Low temperature sampling in Glow ( Kingma & Dhariwal , 2018 ) has a similar effect . Unfortunately , straightforward attempts of implementing truncation or low temperature sampling in diffusion models are ineffective . For example , scaling model scores or decreasing the variance of Gaussian noise in the reverse process cause the diffusion model to generate blurry , low quality samples ( Dhariwal & Nichol , 2021 ) . 3.1 CLASSIFIER GUIDANCE . To obtain a truncation-like effect in diffusion models , Dhariwal & Nichol ( 2021 ) introduce classifier guidance , where the diffusion score θ ( zλ , c ) ≈ σλ∇zλ log p ( zλ|c ) is modified to include the gradient of the log likelihood of an auxiliary classifier model pθ ( c|zλ ) as follows : ̃θ ( zλ , c ) = θ ( zλ , c ) + wσλ∇zλ log pθ ( c|zλ ) ≈ σλ∇zλ [ log p ( zλ|c ) + w log pθ ( c|zλ ) ] , where w is a parameter that controls the strength of the classifier guidance . This modified score ̃θ ( zλ , c ) is then used in place of θ ( zλ , c ) when sampling from the diffusion model , resulting in approximate samples from the distribution p̃θ ( zλ|c ) ∝ pθ ( zλ|c ) pθ ( c|zλ ) w. Algorithm 1 Joint training a diffusion model with unconditional guidance Require : puncond : probability of unconditional training 1 : repeat 2 : ( x , c ) ∼ p ( x , c ) . Sample data with conditioning from the dataset 3 : c← ∅ with probability puncond . Randomly discard conditioning to train unconditionally 4 : λ ∼ p ( λ ) . Sample log SNR value 5 : ∼ N ( 0 , I ) 6 : zλ = αλx+ σλ . Corrupt data to the sampled log SNR value 7 : Take gradient step on ∇θ ‖ θ ( zλ , c ) − ‖2 . Optimization of denoising model 8 : until converged The effect is that of up-weighting the probability of data for which the classifier pθ ( c|zλ ) assigns high likelihood to the correct label : data that can be classified well scores high on the Inception score of perceptual quality ( Salimans et al. , 2016 ) , which rewards generative models for this by design . Dhariwal & Nichol therefore find that by setting w > 0 they can improve the Inception score of their diffusion model , at the expense of decreased diversity in their samples . Figure 2 illustrates the effect of numerically solved guidance p̃θ ( zλ|c ) ∝ pθ ( zλ|c ) pθ ( c|zλ ) w on a toy 2D example of three classes , in which the conditional distribution for each class is an isotropic Gaussian . The form of each conditional upon applying guidance is markedly non-Gaussian . As guidance strength is increased , each conditional places probability mass farther away from other classes and towards directions of high confidence given by logistic regression , and most of the mass becomes concentrated in smaller regions . This behavior can be seen as a simplistic manifestation of the Inception score boost and sample diversity decrease that occur when classifier guidance strength is increased in an ImageNet model . Applying classifier guidance with weight w + 1 to an unconditional model would theoretically lead to the same result as applying classifier guidance with weight w to a conditional model , because pθ ( zλ|c ) pθ ( c|zλ ) w ∝ pθ ( zλ ) pθ ( c|zλ ) w+1 ; or in terms of scores , θ ( zλ ) + ( w + 1 ) σλ∇zλ log pθ ( c|zλ ) ≈ σλ∇zλ [ log p ( zλ ) + ( w + 1 ) log pθ ( c|zλ ) ] = σλ∇zλ [ log p ( zλ|c ) + w log pθ ( c|zλ ) ] , but interestingly , Dhariwal & Nichol obtain their best results when applying classifier guidance to an already class-conditional model , as opposed to applying guidance to an unconditional model . For this reason , we will stay in the setup of guiding an already conditional model .

Summary: In this paper the authors propose an improvement for score-matching based generative modeling [1] resembling low temperature sampling as in GANs or flow-based models. Similarly to [2] they propose to modify the drift function used in the sampling step of the diffusion model by including the gradient of some classifier. However, contrary to [2] the classifier considered by the authors is implicit in the sense that it is purely defined by a conditional and an unconditional generative model. The authors show that using such a classifier allows to control a trade-off between IS and FID on the ImageNet dataset. The use of such implicit classifier also provides intuition on the guidance influence in score-based generative modeling: the model tries to reduce the unconditional likelihood while increasing the conditional likelihood.  [1] Song, Sohl-Dickstein, Kingma, Kumar, Ermon, Poole -- Score-based Generative Modeling through Stochastic Differential Equations [2] Dhariwal, Nichol -- Diffusion Models Beat GANs on Image Synthesis