UFC-BERT: Unifying Multi-Modal Controls for Conditional Image Synthesis

1 Introduction . Conditional image synthesis aims to create an image according to the given control signals . With the increasing demand for flexible conditional image synthesis , various kinds of control signals have been introduced into this field , which can be divided into three main modalities : ( i ) textual controls ( TC ) , including the class labels [ 1 ] and natural language descriptions [ 62 , 54 ] ; ( ii ) visual controls ( VC ) , such as a spatially-aligned sketch map for reference [ 17 , 60 ] or another image for style transfer [ 15 , 27 ] ; ( iii ) preservation controls ( PC ) , which require the synthesized image to preserve some given image blocks , e.g. , image outpainting and inpainting [ 63 , 69 ] . However , control signals of various modalities possess different characteristics . Existing works [ 62 , 26 , 61 ] hence typically design separate methods customized for each control modality . Moreover , most of these approaches only utilize one type of control signal and can not simultaneously combine multiple types of controls in a concise and versatile model . This begs the question : can we integrate any number of multi-modal control signals into a unified framework for flexible conditional image synthesis ? There are two inevitable challenges in this setting : ( i ) how to unify the multi-modal controls and represent them in a unified form , especial when employing multiple control signals from different modalities concurrently ; ( ii ) how to guarantee the fulfillment of the multi-modal controls while ensuring the fidelity of the synthesized image . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Recently , two-stage image synthesis [ 42 , 48 , 3 , 13 , 47 ] has made great progress . The first stage learns a convolutional autoencoder with quantized latent representations for converting an image into a sequence of discrete tokens , e.g. , for compressing a 256×256 image into a sequence of 32 × 32 tokens where each token correlates mainly with an 8×8 block of the image . Converting a sequence of tokens back into an image is also supported . The second stage then typically adopts an autoregressive model , e.g. , PixelCNN [ 41 ] or a unidirectional Transformer decoder [ 55 ] , to capture the distribution over sequences of tokens . Particularly , the Transformer-based methods [ 3 , 13 , 47 ] exploit the global expressivity of Transformer to capture long-range relationships between local constituents . In this paper , we make two key observations about the two-stage framework . First , the two-stage framework has the advantage that it can potentially unify the multi-modal control signals and the generated image into a single sequence of discrete tokens . However , existing works [ 42 , 48 , 13 , 47 ] largely neglect this advantage of the two-stage framework over the traditional one-stage approaches such as those based mainly on the generative adversarial networks ( GAN ) [ 18 ] . Second , the autoregressive ( AR ) approach to sequence generation , adopted by the existing two-stage methods such as DALL-E [ 47 ] and VQGAN [ 13 ] , brings undesirable shortcomings : ( i ) the token-by-token synthesis procedure leads to slow generation speed , especially for the heavyweight Transformer [ 3 , 13 , 47 ] ; ( ii ) each generated token can only catch sight of the previously generated tokens and can not incorporate bidirectional contexts , which may affect the holistic consistency of image synthesis ; ( iii ) the fixed left-to-right order of autoregressive decoding can not respond to the preservation control signals unless the image blocks to be preserved are at the beginning of the sequence . Notably , different from AR generation , non-autoregressive ( NAR ) sequence generation with bidirectional Transformer , i.e. , BERT [ 9 ] , can naturally avoid the three shortcomings . Based on the aforementioned observations , we propose UFC-BERT , a novel BERT-based two-stage framework to UniFy any number of multi-modal Controls for conditional image synthesis . Concretely , the textual , visual , and preservation control signals , as well as the generated image , are uniformly represented as a sequence of discrete tokens , as shown in Figure 2 . The textual control consists of word tokens for class labels or natural language descriptions . The visual control ( s ) and the generated image are both represented as discrete tokens due to the first stage , where each token corresponds to a block within the reference image ( s ) or the generated image . Zero , one , or more reference images are supported . To preserve a given image block within the generated image , we encode the given image block into discrete tokens and fix corresponding parts of the generated sequence to the tokens . We train UFC-BERT via the masked sequence modeling task , which predicts a masked subset of the target image ’ s tokens conditioned on both the multi-modal control signals and the generation target ’ s unmasked tokens . During inference , we adopt Mask-Predict , a NAR generation algorithm [ 16 , 21 , 7 ] , which predicts all target tokens at the first iteration and then iteratively re-mask and re-predict a subset of tokens with low confidence scores . To further improve upon the NAR generation algorithm , we exploit the discriminative capability of the BERT architecture [ 11 , 70 ] and add two estimators ( see Figure 2 ) , where one estimator estimates the relevance between the generated image and the control signals , and the other one estimates the image ’ s fidelity . The two estimators help improve the quality of the synthesized image , because at each iteration we can generate multiple samples and keep only the highly-scored one before starting the next iteration . The two estimators also help save the number of iterations needed , since the algorithm can dynamically terminate if running for more iterations no longer improves the scores . The extensive experiments on M2C-Fashion , a newly collected clothing dataset with tens of millions of image-text pairs , as well as on Multi-Modal CelebA-HQ [ 28 , 61 ] , a public facial dataset , demonstrate UFC-BERT can synthesize high-quality images that comply with various multi-modal controls . 2 Related Works . We have discussed the connection between our work and Two-Stage Image Synthesis in the Introduction . In this section , we further discuss related works from other fields . Conditional Image Synthesis . A variety of control signals have been introduced into conditional image synthesis . The class-conditional generation task [ 1 , 39 ] adopts class labels as control signals . The text-to-image synthesis task [ 64 , 65 , 30 , 62 , 72 , 54 ] further employs natural language descriptions as controls . The image-to-image translation task generates photo-realism images from visual controls , such as a sketch map [ 17 , 60 ] , semantic label map [ 26 , 4 , 58 ] , human pose [ 37 ] or another image for style transfer [ 15 , 27 ] . Moreover , image outpainting and inpainting [ 25 , 63 ] can be regarded as image synthesis conditioned on preservation control signals , where some image blocks of the desired image are already specified and need to be preserved in the generated image . However , these works only utilize one kind of control signal and design their methods customized for each kind of control . Text-guided image manipulation [ 10 , 40 , 69 , 31 , 61 ] semantically edits an image , where the text description and the original image serve as control signals . But they still fail to unify multiple modalities in a universal form and can not easily extend to more control modalities . To promote versatility and extensibility , we propose UFC-BERT to unify any number of multi-modal controls . Visual-Language Transformer . With great progress in language tasks [ 55 , 44 , 45 , 2 ] , the transformer architecture is being rapidly transferred to other fields such as vision [ 3 , 68 , 12 ] and audio [ 6 ] . Recently , pretraining visual-language transformer [ 43 , 24 , 70 , 7 , 35 , 53 , 67 ] ( e.g . multi-modal BERT ) has achieved significant improvements on a variety of downstream tasks , e.g . visual question answering , image captioning [ 70 ] , and text-to-image generation [ 7 ] . Among them , the single-stream architecture [ 52 , 32 , 46 , 5 , 43 , 24 ] uses a single transformer to jointly model a pair of text and image , while the two-stream architecture [ 35 , 36 , 53 ] applies two transformers to separately learn the representations of the text and the image , respectively . Our UFC-BERT is also a variant of the single-stream visual-language transformer , but focuses on flexible multi-modal image synthesis instead of multi-modal pretraining . Non-Autoregressive Sequence Generation . Though it is natural to autoregressively predict tokens from left to right when generating a sequence , autoregressive decoding suffers from the slow speed and sequential error accumulation . Thus , the non-autoregressive generation ( NAR ) paradigm is proposed to avoid these drawbacks in neural machine translation [ 19 , 20 , 29 , 16 ] , image captioning [ 14 , 22 , 70 ] , and speech synthesis [ 50 , 49 ] . These approaches often employ the bidirectional Transformer ( i.e . BERT ) as it is not trained with a specific generation order . Our progressive NAR generation algorithm improves upon the Mask-Predict non-autoregressive algorithm [ 57 , 16 , 38 , 33 ] , by introducing the relevance estimator and the fidelity estimator to facilitate sample selection and dynamic termination . 3 UFC-BERT For Multi-Modal Image Synthesis . 3.1 Background : Two-Stage Image Synthesis . In this section , we review the two-stage architecture [ 42 , 48 , 13 , 47 ] for image synthesis . At the first stage , a codebook Z = { zk } Kk=1 for vector quantization is learned , where zk ∈ Rnz is the k-th code-word in the codebook and K is the number of code-words . An image X ∈ RH×W×3 can be transformed into ( or from ) a collection of code-words Z ∈ Rh×w×nz . Concretely , a convolutional encoder E first encodes the original image X as Ẑ = E ( X ) ∈ Rh×w×nz . Then an element-wise quantization step q ( · ) is applied to each element Ẑij to obtain the element ’ s closest code-word zk , i.e. , q ( Ẑij ) = argminzk∈Z‖Ẑij − zk‖ . For reconstruction , a convolutional decoder D is also learned for recovering image X̂ ∈ RH×W×3 from Z such that X̂ is close to X . The first stage can be denoted by Z = q ( E ( X ) ) , X̂ = D ( Z ) . ( 1 ) Due to the convolutional layers , each of the h× w elements of Ẑ mainly correlates with an Hh × W w block of the image , though its receptive field may be larger if multiple convolutions are stacked . At the second stage , image X ’ s quantized representation Z can be rewritten as a sequence of codes I ∈ { 0 , . . . , |Z|−1 } NI , composed ofNI ( = h×w ) indices from the codebook Z . Thus , image synthesis can be formulated as autoregressive sequence generation , i.e . predicting the distribution Pr ( Ii|I < i , C ) of the next token Ii conditioned on the preceding tokens I < i and the control signals C. The distribution is typically modeled using a unidirectional Transformer . The likelihood is then Pr ( I|C ) =∏ i Pr ( Ii|I < i , C ) . Parameters are learned by minimizing LAR = EI∼data [ − log Pr ( I|C ) ] . We focus on improving the second stage . Specifically , the autoregressive paradigm adopted by the existing two-stage works [ 42 , 48 , 13 , 47 ] suffers from slow generation speed , fails to capture bidirectional contexts , and can not fully support preservation control signals . We thus propose UFCBERT , a novel NAR approach for stage two , to unify any number of multi-modal controls and tackle the shortcomings of AR . As for stage one , we directly follow VQGAN ’ s design [ 47 ] , which improves upon VQVAE [ 42 ] by incorporating a perceptual loss [ 27 ] and patch-based adversarial training [ 26 ] .

Summary: This paper proposes a two-stage non-autoregressive (NAR) transformer-based approach for multi-modal conditional image generation in the discrete latent space of a pretrained VQGAN, called UFC-BERT. More specifically, the proposed model draws inspiration from BERT [1] and Mask-Predict [2] and uses masked language (token) modeling to (i) incorporate bidirectional context and (ii) significantly speed up generation over autoregressive baselines. Sampling from the masked model is performed via an extension of the progressive mask-predict algorithm, where two newly introduced discriminative heads of the transformer model guide the sampling process to generate samples that are consistent with the given conditioning and have high fidelity.