ImageBART: Bidirectional Context with Multinomial Diffusion for Autoregressive Image Synthesis

1 Introduction . Spurred by the increasingly popular attention mechanism , a remarkably simple principle has driven progress in deep generative modeling over the past few years : Factorizing the likelihood of the data in an autoregressive ( AR ) fashion p ( x ) = ∏ i pθ ( xi|x < i ) ( 1 ) and subsequently learning the conditional transition probabilities with an expressive neural network such as a transformer [ 75 ] . The success of this approach is evident in domains as diverse as language modeling [ 7 ] , music generation [ 16 ] , neural machine translation [ 46 , 76 ] , and ( conditional ) image synthesis [ 54 , 8 ] . However , especially for the latter task of image synthesis , which is also the focus of this work , the high dimensionality and redundancy present in the data challenges the direct applicability of this approach . Missing Bidirectional Context Autoregressive models which represent images as a sequence from the top-left to the bottom-right have demonstrated impressive performance in sampling novel images and completing the lower half of a given image [ 8 , 21 ] . However , the unidirectional , fixed ordering ∗The first three authors contributed equally to this work . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of sequence elements not only imposes a perceptually unnatural bias to attention in images by only considering context information from left or above . It also limits practical applicability to image modification : Imagine that you only have the lower half of an image and are looking for a completion of the upper half then these models fail at this minor variation of the completion task . The importance of contextual information from both directions [ 36 ] has also been recognized in the context of language modeling [ 14 , 45 ] . However , simply allowing bidirectional context as in [ 14 ] does not provide a valid factorization of the density function for a generative model . Furthermore , the sequential sampling strategy introduces a gap between training and inference , as training relies on so-called teacher-forcing [ 3 ] ( where ground truth is provided for each step ) and inference is performed on previously sampled tokens . This exposure bias can introduce significant accumulations of errors during the generation process , affecting sample quality and coherence [ 57 ] . Global Context & Control via Multinomial Diffusion We propose a coarse-to-fine approach that addresses the unidirectional bias of generative autoregressive models and their exposure bias as well as the lacking global context . We formulate learning the data density as a hierarchical problem . A coarser stage provides compressed contextual side information about the entire image for the autoregressive process on the next finer stage . We utilize a diffusion process to gradually eliminate information and compress the data , yielding a hierarchy of increasingly abstract and compact representations . The first scale of this approach is a discrete representation learning task ( cf . [ 74 , 58 , 16 , 21 , 78 , 56 ] ) . Subsequently , we further compress this learned representation via a fixed , multinomial diffusion process [ 65 , 30 ] . We then invert this process by training a Markov chain to recover the data from this hierarchy . Each Markovian transition is modeled autoregressively but it simultaneously attends to the preceding state in the hierarchy , which provides crucial global context to each individual autoregressive step . As each of this steps can also be interpreted as learning a denoising cloze task [ 45 ] , where missing tokens at the next finer stage are “ refilled ” with a bidirectional encoder and an autoregressive decoder , we dub our approach ImageBART . Contributions of our work Our approach tackles high-fidelity image synthesis with autoregressive models by learning to invert a fixed multinomial diffusion process in a discrete space of compact image representations to successively introduce context . This reduces both the often encountered exposure bias of AR models and also enables locally controlled , user-interactive image editing . Additionally , our model effectively handles a variety of conditional synthesis tasks and our introduced hierarchy corresponds to a successively compressed image representation . We observe that our model sample visually plausible images while still enabling a trade-off between reconstruction capability and compression rate . 2 Related Work . Latent Variable Models Among likelihood-based approaches , latent variable models represent a data distribution with the help of unobserved latent variables . For example , Variational Autoencoders ( VAEs ) [ 38 , 59 ] encode data points into a lower dimensional latent variable with a factorized distribution . This makes them easy to sample , interpolate [ 44 , 37 ] and modify [ 77 ] . In a conditional setting [ 39 ] , latent variables which are independent from the conditioning lead to disentangled representations [ 31 , 69 , 48 , 60 , 5 ] . A hierarchy of latent variables [ 66 ] gives mutli-scale representations of the data . Unfortunately , even the deepest instantiations of these models [ 47 , 71 , 10 ] lack in sample quality compared to other generative models and are oftentimes restricted to highly regular datasets . Autoregressive Models AR models represent a distribution as a product of conditional , learnable factors via the chain rule of probability densities . While this makes them powerful models for density estimation [ 70 , 24 ] , their samples often lack global consistency . Especially on image data modeled with convolutional architectures [ 73 , 62 ] , this has been attributed to a locality bias of convolutional neural networks ( CNNs ) which biases the model towards strong local correlations between neighboring pixels at the expense of a proper modeling of coherence [ 40 , 22 ] . This leads to samples resembling texture patterns without discernible global structure . Attempts to fix this properties by including explicit latent variables [ 27 , 9 , 22 ] have not been overly successful , mainly due the expressiveness of AR models , providing little incentive for learning additional latent variables . Generative Models on Improved Representations Another successful line of work first learn an improved image representation and subsequently learn a generative model for this representation [ 74 , 12 ] . Most works [ 58 , 21 , 56 ] learn a discrete representation which is subsequently modeled autoregressively but approaches using continuous representations in combination with VAEs [ 12 ] , or normalizing flows [ 1 , 60 , 20 , 4 , 18 ] , exist too . Learning a compact representation enables the use of transformers for autoregressive modeling [ 8 ] , which avoids the locality bias of CNNs , can be used for the synthesis of complex scenes conditioned on text as in DALL-E [ 56 ] , and , when combined with adversarial learning [ 25 ] , enables sampling of coherent high-resolution images [ 21 ] . However , AR modeling of a learned representation still limits applications compared to latent variable models . Their samples can still exert artifacts resulting from a sequential modeling of components , and , since these models are always trained by “ teacher-forcing ” , they are susceptible to an exposure bias [ 3 , 57 , 26 , 63 , 43 ] . Diffusion Probabilistic Models Diffusion probabilistic models revert a fixed , diffusion process with a learned Markov Chain [ 65 ] . Being directly applied in pixel space , however , downstream analysis reveals that these models tend to optimize subtle details of the modeled data , which have little contribution to the sample quality [ 29 , 15 ] , particularly hindering applications on high-resolution and -complexity datasets . By using a multinomial diffusion process [ 30 ] ( recently generalized by [ 2 ] ) on a compressed , discrete representation of images , we circumvent these issues . Diffusion probabilistic models require a very large number of diffusion steps in order to model the reverse process with a model distribution that factorizes over components . Because our approach uses autoregressively factorized models for the reverse process , we can reduce the required number of steps and obtain significant improvements in sampling speed and the ability to model complex datasets . 3 Method . 3.1 Hierarchical Generative Models . To tackle the difficult problem of modeling a highly complex distribution p ( x ) of high-dimensional images x , we ( i ) introduce bidirectional context into an otherwise unidirectional autoregressive factorization of p ( x ) as in Eq . ( 1 ) and ( ii ) reduce the difficulty of the learning problem with a hierarchical approach . To do so , we learn a sequence of distributions ( ptθ ) T t=0 , such that each distribution pt−1θ models a slightly more complex distribution with the help of a slightly simpler distribution ptθ one level above . This introduces a coarse-to-fine hierarchy of image representations x0 : T : = ( xt ) T t=0 , such that an xt−1 is modeled conditioned on xt , i.e . xt−1 ∼ pt−1θ ( xt−1|xt ) and defines a reverse Markov Chain for x = : x0 as pθ ( x0 ) = pTθ ( xT ) ∏T t=1 p t−1 θ ( xt−1|xt ) . Since our goal is to approximate the original distribution p ( x ) with pθ ( x0 ) , we introduce a forward Markov Chain , qθ ( x1 : T |x0 ) = ∏T t=1 q t θ ( xt|xt−1 ) , to obtain a tractable upper bound on the Kullback-Leibler ( KL ) divergence between p and pθ , KL ( p ( x0 ) ‖pθ ( x0 ) ) = : KL , using the evidence lower bound ( ELBO ) . With qTθ ( xT |xT−1 ) : = pTθ ( xT ) , we obtain KL ≤ Ex0 , x1 log p ( x0 ) p0θ ( x0|x1 ) ︸ ︷︷ ︸ = : L1→ discrete repr . learning + T∑ t=2 Ex0 , xtKL ( q t−1 θ ( xt−1|xt , x0 ) ‖p t−1 θ ( xt−1|xt ) ) ︸ ︷︷ ︸ = : Lt→ decoupled with diffusion process ( 2 ) We use L1 to learn a compressed and discrete representation of images , such that subsequent stages of the hierarchy do not need to model redundant information ( Sec . 3.2 ) . With Lt , t > 1 we learn a model that can rely on global context from a coarser representation xt to model the representation xt−1 ( Sec . 3.3 ) . See Fig . 1 for an overview of the proposed model . 3.2 Learning a compact , discrete representation for images . Since the first stage of the hierarchical process is the one that operates directly on the data , we assign it a separate role . To avoid that the optimization of Lt ( t = 1 , . . . , T ) in Eq . ( 2 ) unnecessarily wastes capacity on redundant details in the input images—which is an often encountered property of pixel-based likelihood models [ 74 , 21 , 50 ] —we take L1 = Ep ( x0 ) q1θ ( x1|x0 ) log p ( x0 ) p0θ ( x0|x1 ) to be the reconstruction term for a discrete autoencoder model . This has the advantage that we can directly build on work in neural discrete representation learning , which has impressively demonstrated that discrete representations can be used for high-quality synthesis of diverse images while achieving strong compression . In particular , [ 49 ] and [ 21 ] have shown that adding an adversarial realism prior to the usual autoencoder objective helps to produce more realistic images at higher compression rates by locally trading reconstruction fidelity for realism . More specifically , we follow [ 21 ] to encode images into a low-dimensional representation which is then vector-quantized with a learned codebook of size K to obtain { 0 , . . . , K − 1 } h×w 3 x1 ∼ q1θ ( x1|x0 ) deterministically as the index of the closest codebook entry . The encoder is a convolutional neural network ( CNN ) with four downsampling steps , such that h = H/16 and w = W/16 for any input image x0 ∈ RH×W×3 . For downstream autoregressive learning , this representation is then unrolled into a discrete sequence of length N = h · w. To recover an image from x1 , we utilize a CNN decoder G , such that the reverse model is specified as − log p0θ ( x0|x1 ) ∝ frec ( x0 , Gθ ( x1 ) ) + logDφ ( Gθ ( x1 ) ) = : Lrec ( x0 , x1 ; θ ) + Ladv ( x1 ; θ , φ ) ( 3 ) Here , frec denotes the perceptual similarity metric [ 23 , 32 , 19 , 80 ] ( known as LPIPS ) and Dφ denotes a patch-based adversarial discriminator [ 25 ] . Note that , due to the deterministic training , the likelihood in Eq . ( 3 ) is likely to be degenerate . Dφ is optimized to differentiate original images x0 from their reconstruction Gθ ( x1 ) using simultaneous gradient ascent , such that the objective for learning the optimal parameters { θ∗ , φ∗ } reads : { θ∗ , φ∗ } = arg min θ max φ ( Lrec ( x0 , x1 ; θ ) − Ladv ( x1 ; θ , φ ) + logDφ ( x0 ) + Lcb ( θ ) ) ( 4 ) The optimization of θ via this objective includes the parameters of the encoder and decoder in addition to the parameters of the learned codebook , trained via the codebook loss Lcb as in [ 74 , 21 ] .

Summary: This paper presents ImageBART, a generative model of images that uses multinomial diffusion to model the distribution of discrete latent codes extracted from a VQ-GAN encoder. Unlike prior diffusion based models, the reverse diffusion processes are modelled autoregressive, which enables a large reduction in the number of diffusion steps required. The model can condition on contexts such as a caption or semantic segmentation map, and unlike typical autoregressive models, is capable of image completion for arbitrary masked regions.