D2C: Diffusion-Decoding Models for Few-Shot Conditional Generation

Conditional generative models of high-dimensional images have many applications , but supervision signals from conditions to images can be expensive to acquire . This paper describes Diffusion-Decoding models with Contrastive representations ( D2C ) , a paradigm for training unconditional variational autoencoders ( VAEs ) for few-shot conditional image generation . D2C uses a learned diffusion-based prior over the latent representations to improve generation and contrastive selfsupervised learning to improve representation quality . D2C can adapt to novel generation tasks conditioned on labels or manipulation constraints , by learning from as few as 100 labeled examples . On conditional generation from new labels , D2C achieves superior performance over state-of-the-art VAEs and diffusion models . On conditional image manipulation , D2C generations are two orders of magnitude faster to produce over StyleGAN2 ones and are preferred by 50 % − 60 % of the human evaluators in a double-blind study . We release our code at https : //github.com/jiamings/d2c . 1 Introduction . Generative models trained on large amounts of unlabeled data have achieved great success in various domains including images [ 8 , 50 , 75 , 42 ] , text [ 56 , 2 ] , audio [ 26 , 71 , 93 , 62 ] , and graphs [ 36 , 67 ] . However , downstream applications of generative models are often based on various conditioning signals , such as labels [ 61 ] , text descriptions [ 60 ] , reward values [ 101 ] , or similarity with existing data [ 45 ] . While it is possible to directly train conditional models , this often requires large amounts of paired data [ 57 , 74 ] that are costly to acquire . Hence , it would be desirable to learn conditional generative models using large amounts of unlabeled data and as little paired data as possible . Contrastive self-supervised learning ( SSL ) methods can greatly reduce the need for labeled data in discriminative tasks by learning effective representations from unlabeled data [ 95 , 37 , 35 ] , and have also been shown to improve few-shot learning [ 39 ] . It is therefore natural to ask if they can also be used to improve few-shot generation . Latent variable generative models ( LVGM ) are a natural candidate for this , since they already involve a low-dimensional , structured latent representation of the data they generate . However , generative adversarial networks ( GANs , [ 34 , 50 ] ) and diffusion ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . models [ 42 , 84 ] , lack explicit tractable functions to map inputs to representations , making it difficult to optimize latent variables with SSL . Variational autoencoders ( VAEs , [ 52 , 77 ] ) , on the other hand , can naturally adopt SSL through their encoder model , but they typically have worse sample quality . In this paper , we propose Diffusion-Decoding models with Contrastive representations ( D2C ) , a special VAE that is suitable for conditional few-shot generation . D2C uses contrastive self-supervised learning methods to obtain a latent space that inherits the transferrability and few-shot capabilities of self-supervised representations . Unlike other VAEs , D2C learns a diffusion model over the latent representations . This latent diffusion model ensures that D2C uses the same latent distribution for both training and generation . We provide a formal argument to explain why this approach leads to better sample quality than existing hierarchical VAEs . We further discuss how to apply D2C to few-shot conditional generation where the conditions are defined through labeled examples or manipulation constraints . Our approach combines a discriminative model providing conditioning signal and generative diffusion model over the latent space , and is computationally more efficient than methods that act directly over the image space ( Figure 1 ) . We evaluate and compare D2C with several state-of-the-art generative models over 6 datasets . On unconditional generation , D2C outperforms state-of-the-art VAEs and is competitive with diffusion models under similar computational budgets . On conditional generation with 100 labeled examples , D2C significantly outperforms state-of-the-art VAE [ 91 ] and diffusion models [ 84 ] . D2C can also learn to perform certain image manipulation tasks from as few as 100 labeled examples . Notably , for manipulating images , D2C is two orders of magnitude faster than StyleGAN2 [ 106 ] and preferred by 50 % − 60 % of human evaluations , which to our best knowledge is the first for any VAE model . 2 Background . Latent variable generative models A latent variable generative model ( LVGM ) is posed as a conditional distribution pθ : Z → P ( X ) from a latent variable z to a generated sample x , parametrized by θ . To acquire new samples , LVGMs draw random latent variables z from some distribution p ( z ) and map them to image samples through x ∼ pθ ( x|z ) . Most LVGMs are built on top of four paradigms : variational autoencoders ( VAEs , [ 52 , 77 ] ) , Normalizing Flows ( NFs , [ 28 , 29 ] ) , Generative Adversarial Networks ( GANs , [ 34 ] ) , and diffusion / score-based generative models [ 42 , 85 ] . Particularly , VAEs use an inference model from x to z for training . Denoting the inference distribution from x to z as qφ ( z|x ) , the generative distribution from z to x as pθ ( x|z ) , VAEs are trained by minimizing the following upper bound of negative log-likelihood : LVAE = Ex∼pdata [ Ez∼qφ ( z|x ) [ − log p ( x|z ) ] +DKL ( qφ ( z|x ) ‖p ( z ) ) ] ( 1 ) where pdata is the data distribution and DKL is the KL-divergence . Diffusion models Diffusion models [ 81 , 42 , 84 ] produce samples by reversing a Gaussian diffusion process . We use the index α ∈ [ 0 , 1 ] to denote the particular noise level of an noisy observation x ( α ) = √ αx + √ 1− α , where x is the clean observation and ∼ N ( 0 , I ) is a standard Gaussian distribution ; as α→ 0 , the distribution of x ( α ) converges to N ( 0 , I ) . Diffusion models are typically parametrized as reverse noise models θ ( x ( α ) , α ) that predict the noise component of x ( α ) given a noise level α , and trained to minimize ‖ θ ( x ( α ) , α ) − ‖22 , the mean squared error loss between the true noise and predicted noise . Given any non-increasing series { αi } Ti=0 between 0 and 1 , the diffusion objective for a clean sample from the data x is : ` diff ( x ; w , θ ) : = T∑ i=1 w ( αi ) E ∼N ( 0 , I ) [ ‖ − θ ( x ( αi ) , αi ) ‖ 2 2 ] , x ( αi ) : = √ αix + √ 1− αi ( 2 ) where w : { αi } Ti=1 → R+ controls the loss weights for each α . When w ( α ) = 1 for all α , we recover the denoising score matching objective for training score-based generative models [ 85 ] . Given an initial sample x0 ∼ N ( 0 , I ) , diffusion models acquires clean samples ( i.e. , samples of x1 ) through a gradual denoising process , where samples with reducing noise levels α are produced ( e.g. , x0 → x0.3 → x0.7 → x1 ) . In particular , Denoising Diffusion Implicit Models ( DDIMs , [ 84 ] ) uses an Euler discretization of some neural ODE [ 13 ] to produce samples ( Figure 2 , left ) . We provide a more detailed description for training diffusion models in Appendix A.1 and sampling from DDIM in Appendix A.2 . For conciseness , we use the notation p ( α ) ( x ( α ) ) to denote the marginal distribution of x ( α ) under the diffusion model , and p ( α1 , α2 ) ( x ( α2 ) | x ( α1 ) ) to denote the diffusion sampling process from x ( α1 ) to x ( α2 ) ( assuming α1 < α2 ) . Self-supervised learning of representations In self-supervised learning ( SSL ) , representations are learned by completing certain pretext tasks that do not require extra manual labeling [ 68 , 25 ] ; these representations can then be applied to other downstream tasks , often in few-shot or zero-shot scenarios . In particular , contrastive representation learning encourages representations to be closer between “ positive ” pairs and further between “ negative ” pairs ; contrastive predictive coding ( CPC , [ 95 ] ) , based on multi-class classification , have been commonly used in state-of-the-art methods [ 37 , 16 , 18 , 14 , 82 ] . Other non-contrastive methods exist , such as BYOL [ 35 ] and SimSiam [ 17 ] , but they usually require additional care to prevent the representation network from collapsing . 3 Problem Statement . Few-shot conditional generation Our goal is to learn an unconditional generative model pθ ( x ) such that it is suitable for conditional generation . Let C ( x , c , f ) describe an event that “ f ( x ) = c ” , where c is a property value and f ( x ) is a property function that is unknown at training . In conditional generation , our goal is to sample x such that the event C ( x , c , f ) occurs for a chosen c. If we have access to some “ ground-truth ” model that gives us p ( C|x ) : = p ( f ( x ) = c|x ) , then the conditional model can be derived from Bayes ’ rule : pθ ( x|C ) ∝ p ( C|x ) pθ ( x ) . These properties c include ( but are not limited to2 ) labels [ 61 ] , text descriptions [ 60 , 76 ] , noisy or partial observations [ 11 , 5 , 47 , 24 ] , and manipulation constraints [ 69 ] . In many cases , we do not have direct access to the true f ( x ) , so we need to learn an accurate model from labeled data [ 6 ] ( e.g. , ( c , x ) pairs ) . Desiderata Many existing methods are optimized for some known condition ( e.g. , labels in conditional GANs [ 8 ] ) or assume abundant pairs between images and conditions that can be used for pretraining ( e.g. , DALL-E [ 74 ] and CLIP [ 73 ] over image-text pairs ) . Neither is the case in this paper , as we do not expect to train over paired data . While high-quality latent representations are not essential to unconditional image generation ( e.g. , autoregressive [ 94 ] , energy-based [ 31 ] , and some diffusion models [ 42 ] ) , they can be beneficial when we wish to specify certain conditions with limited supervision signals , similar to how SSL representations can reduce labeling efforts in downstream tasks . A compelling use case is detecting and removing biases in datasets via image manipulation , where we should not only address wellknown biases a-priori but also address other hard-to-anticipate biases , adapting to societal needs [ 65 ] . 2When C refers to an event that is always true , we recover unconditioned generation . Therefore , a desirable generative model should not only have high sample quality but also contain informative latent representations . While VAEs are ideal for learning rich latent representations due to being able to incorporate SSL within the encoder , they generally do not achieve the same level of sample quality as GANs and diffusion models .

Summary: This paper proposes to learn a joint model consisting of a VAE and a latent space diffusion. The training process is augmented with a contrastive loss.  The authors show that this approach makes possible a few-shot generation of images with desired properties. To achieve this,  one needs to train a classifier in the latent space and after that sample latent codes with high classifier scores in any way.  Also, the case of semantic editing is evaluated. The manipulation is implemented using a Langevin-like approach in the latent space.