CogView: Mastering Text-to-Image Generation via Transformers

1 Introduction . “ There are two things for a painter , the eye and the mind ... eyes , through which we view the nature ; brain , in which we organize sensations by logic for meaningful expression. ” ( Paul Cézanne [ 17 ] ) 1Codes and models are at https : //github.com/THUDM/CogView . We also have a demo website of our latest model at https : //wudao.aminer.cn/CogView/index.html ( without post-selection ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . As contrastive self-supervised pretraining has revolutionized computer vision ( CV ) [ 24 , 21 , 8 , 32 ] , visual-language pretraining , which brings high-level semantics to images , is becoming the next frontier of visual understanding [ 38 , 30 , 39 ] . Among various pretext tasks , text-to-image generation expects the model to ( 1 ) disentangle shape , color , gesture and other features from pixels , ( 2 ) understand the input text , ( 2 ) align objects and features with corresponding words and their synonyms and ( 4 ) learn complex distributions to generate the overlapping and composite of different objects and features , which , like painting , is beyond basic visual functions ( related to eyes and the V1–V4 in brain [ 22 ] ) , requiring a higher-level cognitive ability ( more related to the angular gyrus in brain [ 3 ] ) . The attempts to teach machines text-to-image generation can be traced to the early times of deep generative models , when Mansimov et al . [ 35 ] added text information to DRAW [ 20 ] . Then Generative Adversarial Nets [ 19 ] ( GANs ) began to dominate this task . Reed et al . [ 42 ] fed the text embeddings to both generator and discriminator as extra inputs . StackGAN [ 54 ] decomposed the generation into a sketch-refinement process . AttnGAN [ 51 ] used attention on words to focus on the corresponding subregion . ObjectGAN [ 29 ] generated images following a text→boxes→layouts→image process . DM-GAN [ 55 ] and DF-GAN [ 45 ] introduced new architectures , e.g . dyanmic memory or deep fusion block , for better image refinement . Although these GAN-based models can perform reasonable synthesis in simple and domain-specific dataset , e.g . Caltech-UCSD Birds 200 ( CUB ) , the results on complex and domain-general scenes , e.g . MS COCO [ 31 ] , are far from satisfactory . Recent years have seen a rise of the auto-regressive generative models . Generative Pre-Training ( GPT ) models [ 37 , 4 ] leveraged Transformers [ 48 ] to learn language models in large-scale corpus , greatly promoting the performance of natural language generation and few-shot language understanding [ 33 ] . Auto-regressive model is not nascent in CV . PixelCNN , PixelRNN [ 47 ] and Image Transformer [ 36 ] factorized the probability density function on an image over its sub-pixels ( color channels in a pixel ) with different network backbones , showing promising results . However , a real image usually comprises millions of sub-pixels , indicating an unaffordable amount of computation for large models . Even the biggest pixel-level auto-regressive model , ImageGPT [ 7 ] , was pretrained on ImageNet at a max resolution of only 96× 96 . The framework of Vector Quantized Variational AutoEncoders ( VQ-VAE ) [ 46 ] alleviates this problem . VQ-VAE trains an encoder to compress the image into a low-dimensional discrete latent space , and a decoder to recover the image from the hidden variable in the stage 1 . Then in the stage 2 , an auto-regressive model ( such as PixelCNN [ 47 ] ) learns to fit the prior of hidden variables . This discrete compression loses less fidelity than direct downsampling , meanwhile maintains the spatial relevance of pixels . Therefore , VQ-VAE revitalized the auto-regressive models in CV [ 41 ] . Following this framework , Esser et al . [ 15 ] used Transformer to fit the prior and further switches from L2 loss to GAN loss for the decoder training , greatly improving the performance of domain-specific unconditional generation . The idea of CogView comes naturally : large-scale generative joint pretraining for both text and image ( from VQ-VAE ) tokens . We collect 30 million high-quality ( Chinese ) text-image pairs and pretrain a Transformer with 4 billion parameters . However , large-scale text-to-image generative pretraining could be very unstable due to the heterogeneity of data . We systematically analyze the reasons and solved this problem by the proposed Precision Bottleneck Relaxation and Sandwich Layernorm . As a result , CogView greatly advances the quality of text-to-image generation . A recent work DALL-E [ 39 ] independently proposed the same idea , and was released earlier than CogView . Compared with DALL-E , CogView steps forward on the following four aspects : • CogView outperforms DALL-E and previous GAN-based methods at a large margin according to the Fréchet Inception Distance ( FID ) [ 25 ] on blurred MS COCO , and is the first open-source large text-to-image transformer . • Beyond zero-shot generation , we further investigate the potential of finetuning the pretrained CogView . CogView can be adapted for diverse downstream tasks , such as style learning ( domain-specific text-to-image ) , super-resolution ( image-to-image ) , image captioning ( image-to-text ) , and even text-image reranking . • The finetuned CogView enables self-reranking for post-selection , and gets rid of an additional CLIP model [ 38 ] in DALL-E . It also provides a new metric Caption Loss to measure the quality and accuracy for text-image generation at a finer granularity than FID and Inception Score ( IS ) [ 43 ] . • We proposed PB-relaxation and Sandwich-LN to stabilize the training of large Transformers on complex datasets . These techniques are very simple and can eliminate overflow in forwarding ( characterized as NaN losses ) , and make CogView able to be trained with almost FP16 ( O22 ) . They can also be generalized to the training of other transformers . 2 Method . 2.1 Theory . In this section , we will derive the theory of CogView from VAE3 [ 26 ] : CogView optimizes the Evidence Lower BOund ( ELBO ) of joint likelihood of image and text . The following derivation will turn into a clear re-interpretation of VQ-VAE if without text t. Suppose the dataset ( X , T ) = { xi , ti } Ni=1 consists of N i.i.d . samples of image variable x and its description text variable t. We assume the image x can be generated by a random process involving a latent variable z : ( 1 ) ti is first generated from a prior p ( t ; θ ) . ( 2 ) zi is then generated from the conditional distribution p ( z|t = ti ; θ ) . ( 3 ) xi is finally generated from p ( x|z = zi ; ψ ) . We will use a shorthand form like p ( xi ) to refer to p ( x = xi ) in the following part . Let q ( z|xi ; φ ) be the variational distribution , which is the output of the encoder φ of VAE . The log-likelihood and the evidence lower bound ( ELBO ) can be written as : log p ( X , T ; θ , ψ ) = N∑ i=1 log p ( ti ; θ ) + N∑ i=1 log p ( xi|ti ; θ , ψ ) ( 1 ) ≥ − N∑ i=1 ( − log p ( ti ; θ ) ︸ ︷︷ ︸ NLL loss for text + E zi∼q ( z|xi ; φ ) [ − log p ( xi|zi ; ψ ) ] ︸ ︷︷ ︸ reconstruction loss + KL ( q ( z|xi ; φ ) ‖p ( z|ti ; θ ) ) ︸ ︷︷ ︸ KL between q and ( text conditional ) prior ) . ( 2 ) The framework of VQ-VAE differs with traditional VAE mainly in the KL term . Traditional VAE fixes the prior p ( z|ti ; θ ) , usually as N ( 0 , I ) , and learns the encoder φ . However , it leads to posterior collapse [ 23 ] , meaning that q ( z|xi ; φ ) sometimes collapses towards the prior . VQ-VAE turns to fix φ and fit the prior p ( z|ti ; θ ) with another model parameterized by θ . This technique eliminates posterior collapse , because the encoder φ is now only updated for the optimization of the reconstruction loss . In exchange , the approximated posterior q ( z|xi ; φ ) could be very different for different xi , so we need a very powerful model for p ( z|ti ; θ ) to minimize the KL term . Currently , the most powerful generative model , Transformer ( GPT ) , copes with sequences of tokens over a discrete codebook . To use it , we make z ∈ { 0 , ... , |V | − 1 } h×w , where |V | is the size of codebook and h× w is the number of dimensions of z . The sequences zi can be either sampled from q ( z|xi ; φ ) , or directly zi = argmaxz q ( z|xi ; φ ) . We choose the latter for simplicity , so that q ( z|xi ; φ ) becomes a one-point distribution on zi . The Equation ( 2 ) can be rewritten as : − N∑ i=1 ( E zi∼q ( z|xi ; φ ) [ − log p ( xi|zi ; ψ ) ] ︸ ︷︷ ︸ reconstruction loss − log p ( ti ; θ ) ︸ ︷︷ ︸ NLL loss for text − log p ( zi|ti ; θ ) ︸ ︷︷ ︸ NLL loss for z ) . ( 3 ) The learning process is then divided into two stages : ( 1 ) The encoder φ and decoder ψ learn to minimize the reconstruction loss . ( 2 ) A single GPT optimizes the two negative log-likelihood ( NLL ) losses by concatenating text ti and zi as an input sequence . As a result , the first stage degenerates into a pure discrete Auto-Encoder , serving as an image tokenizer to transform an image to a sequence of tokens ; the GPT in the second stage undertakes most of the modeling task . Figure 3 illustrates the framework of CogView . 2meaning that all computation , including forwarding and backwarding are in FP16 without any conversion , but the optimizer states and the master weights are FP32 . 3In this paper , bold font denotes a random variable , and regular font denotes a concrete value . See this comprehensive tutorial [ 12 ] for the basics of VAE . 2.2 Tokenization . In this section , we will introduce the details about the tokenizers in CogView and a comparison about different training strategies about the image tokenizer ( VQVAE stage 1 ) . Tokenization for text is already well-studied , e.g . BPE [ 16 ] and SentencePiece [ 28 ] . In CogView , we ran SentencePiece on a large Chinese corpus to extract 50,000 text tokens . The image tokenizer is a discrete Auto-Encoder , which is similar to the stage 1 of VQ-VAE [ 46 ] or d-VAE [ 39 ] . More specifically , the Encoder φ maps an image x of shape H ×W × 3 into Encφ ( x ) of shape h× w × d , and then each d−dimensional vector is quantized to a nearby embedding in a learnable codebook { v0 , ... , v|V |−1 } , ∀vk ∈ Rd . The quantized result can be represented by h× w indices of embeddings , and then we get the latent variable z ∈ { 0 , ... , |V | − 1 } h×w . The Decoder ψ maps the quantized vectors back to a ( blurred ) image to reconstruct the input . In our 4B-parameter CogView , |V | = 8192 , d = 256 , H =W = 256 , h = w = 32 . The training of the image tokenizer is non-trivial due to the existence of discrete selection . Here we introduce four methods to train an image tokenizer . • The nearest-neighbor mapping , straight-through estimator [ 2 ] , which is proposed by the original VQVAE . A common concern of this method [ 39 ] is that , when the codebook is large and not initialized carefully , only a few of embeddings will be used due to the curse of dimensionality . We did not observe this phenomenon in the experiments . • Gumbel sampling , straight-through estimator . If we follow the original VAE to reparameterize a categorical distribution of latent variable z based on distance between vectors , i.e . p ( zi×w+j = vk|x ) = e −‖vk−Encφ ( x ) ij‖2/τ∑|V |−1 k=0 e −‖vk−Encφ ( x ) ij‖2/τ , an unbiased sampling strategy is zi×w+j = argmaxk gk−‖vk−Encφ ( x ) ij‖2/τ , gk ∼ Gumbel ( 0 , 1 ) , where the temperature τ is gradually decreased to 0 . We can further use the differentiable softmax to approximate the one-hot distribution from argmax . DALL-E adopts this method with many other tricks to stabilize the training . • The nearest-neighbor mapping , moving average , where each embedding in the codebook is updated periodically during training as the mean of the vectors recently mapped to it [ 46 ] . • The nearest-neighbor mapping , fixed codebook , where the codebook is fixed after initialized . Comparison . To compare the methods , we train four image tokenizers with the same architecture on the same dataset and random seed , and demonstrate the loss curves in Figure 2 . We find that all the methods are basically evenly matched , meaning that the learning of the embeddings in the codebook is not very important , if initialized properly . In pretraining , we use the tokenizer of moving average method . The introduction of data and more details about tokenization are in Appendix A .

Summary: The paper proposes CogView - a cross-modal Transformer with VQ-VAE as the tokenizer. The model is primarily designed for text-to-image generation but also shows promising results on various downstream tasks.  The paper also provides many insights on the effective training tricks (including PB-Relax and Sandwich-LN) to stabilize the training of text-to-image transformers. The experimental results show that CogView outperforms the recent state-of-the-art approach DALL-E in the zero-shot settings.