Improving Visual Quality of Image Synthesis by A Token-based Generator with Transformers

1 Introduction . Unconditional image synthesis generates images from latent codes by adversarial training [ 10 , 16 , 19 , 25 , 27 , 33 , 50 ] . Recent advances have been achieved by a style-based generator architecture in terms of both the visual quality and resolution of generated images [ 26 , 27 , 28 , 37 , 52 ] . In particular , the style-based generator has been widely used in many other generative tasks , including facial editing [ 9 , 40 ] , style transfer [ 1 , 38 ] , image super-resolution [ 17 , 31 ] , and image inpainting [ 2 , 52 ] . The key to the success of the style-based generator lies in the learning of the style control based on the intermediate latent spaceW [ 27 , 28 ] . Instead of feeding the input latent code z ∈ Z through the input layer only ( Figure1-a ) , the style-based generator maps the input z to an intermediate latent space w ∈ W , which then controls the “ style ” of the image at each layer via adaptive instance normalization ( AdaIN [ 21 ] ) ( Figure 1-b ) . It has been demonstrated that such a design allows a less entangled representation learning inW , leading to better generative image modeling [ 12 , 23 , 27 , 40 ] . Despite the promising results , the style-based generator can suffer from the style control via AdaIN operation [ 28 , 52 ] . Specifically , the style control is content-independent . It “ washes away ” the original information of features by normalization and assigns new styles decided by the latent codes regardless of the image/feature content . Besides , the style code w affects the entire image by scaling and biasing complete feature maps with a single value via AdaIN operation [ 35 , 44 , 54 ] . Such an imposed single style over multiple image regions can inevitably result in entangled representation of ∗This work was done while Yanhong Zeng was a research intern at Microsoft Research Asia . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . different image variations ( e.g. , hairstyle , facial expression ) [ 28 , 44 , 52 ] . These limitations for image modeling can even lead to visible artifacts in the synthesized results ( e.g. , droplet artifacts [ 27 ] ) . To get rid of the issues caused by StyleGAN ’ s style modeling , we introduce a new perspective that views image synthesis as a visual token generation problem . The visual token is a popular representation of an image patch with a predefined size and position [ 6 , 7 , 13 ] ; and has shown an impressive superiority in various tasks with the development of Transformer models [ 6 , 13 , 43 , 46 ] . Inspired by the appealing property of the token-based representation , we propose to achieve image synthesis by visual token generation . Specifically , it takes as input a sequence of latent tokens to predict the visual tokens of an image . Such a token-based representation enables a flexible local manipulation for different image regions , which makes it possible to learn content-aware and fine-grained style control for image synthesis . Under this new paradigm , we design a token-based generator , i.e. , TokenGAN , for the visual token generation problem . Specifically , we consider two different types of input tokens in the generator , i.e. , the content tokens and the style tokens . The content tokens are learned as the constant input in the generator network and the style tokens are projected from a learned intermediate latent space ( Figure1-c ) . Given a sequence of style tokens , the TokenGAN learns to control the visual token generation by rendering each content token with related style tokens according to their semantics . In particular , since the Transformer has been verified to be effective in sequence modeling in a broad range of tasks [ 6 , 43 ] , we adopt a generator network architecture from a visual Transformer to model the relations between the content tokens and the style tokens . Through such a content-dependent style modeling by the attention mechanism in Transformer , the TokenGAN is able to achieve content-aware and fine-grained style learning for image synthesis . We conduct both quantitative comparisons and qualitative analysis on several unconditional image generation benchmarks . The results show that the token-based generator has achieved comparable results to the state-of-the-art in image synthesis . We summarize our contributions as below : • We propose a new perspective of achieving image synthesis by visual token generation . Such a token-based representation enables flexible local manipulation for different image regions , leading to a better image modeling . • We propose a token-based generator ( i.e. , TokenGAN ) for the visual token generation . Specifically , the TokenGAN introduces the style tokens and the content tokens . It adopts a Transformer-based network for content-aware style modeling . • We show extensive experiments ( quantitative and qualitative comparisons , study on style editing , image inversion and image interpolation ) to verify the effectiveness of the tokenbased generator . Specifically , the token-based generator is able to synthesize high-fidelity 1024× 1024 images without any convolutions in the generator . 2 Related works . 2.1 Style-based generator . The distinguishing feature of the style-based generator is its unconventional generator architecture [ 26 , 27 , 28 , 52 ] . Typically , the style-based generator consists of a mapping network f and a synthesis network g. The mapping network is used to transform the input latent code z to an intermediate latent code w ∈ W for learning a less entangled latent space . The synthesis network follows a progressive growing design that is able to first output low-resolution images that are not affected significantly by high-resolution layers [ 25 , 27 ] . Such an architecture design can lead to an automatic , unsupervised separation of high-level attributes in different layers . For example , when trained on human faces , the low-resolution layers control coarse styles ( e.g. , pose , identity ) of images and high-resolution layers control fine styles ( e.g. , micro-structure , color scheme ) [ 25 ] . To control the attributes of the image , the style-based generator produces styles from the intermediate latent code w by affine transforms , which then control each layer of the synthesis network via adaptive instance normalization ( AdaIN ) [ 21 , 35 ] . The AdaIN operation removes styles from previous layers by normalizing each feature map to zero mean and unit deviation , and it assigns the new styles by scaling and biasing the complete normalized feature maps . However , it has been demonstrated such an AdaIN operation can destroy the magnitude information of features and thus results in the well-known droplet artifacts [ 28 ] . Besides , the scale-specific style disentanglement in StyleGAN can be limited due to the complex semantic components in each feature map [ 28 , 52 ] . In this paper , we explore a new architecture by flattening the image as a sequence of tokens , which enables fine-grained control of the image by assigning token-wise semantic-aware styles based on the attention mechanism . 2.2 Transformer in vision . The Transformer typically takes as input a sequence of vectors , called tokens [ 43 ] . The Transformer is built for sequence modeling solely on attention mechanisms over the tokens , dispensing with recurrence and convolutions entirely [ 5 , 11 , 43 , 30 ] . Due to its great success in the field of natural language processing , an increasing number of works attempt to extend Transformer for computer vision tasks [ 6 , 7 , 8 , 13 , 18 , 36 , 55 ] . For example , iGPT shows promising results in image classification by pre-training a sequence Transformer with the task of auto-regressive next pixel prediction and masked pixel prediction [ 8 ] . Dosoyvitskiy et al . propose to split an image into patches and feed these patches into a standard Transformer ( ViT ) for image classification [ 13 ] . They show that a ViT with large-scale training can trump CNNs equipped with inductive bias . DETR is a seminal work that views object detection as a direct set prediction problem , eliminating the need for hand-crafted components and achieving impressive performance by Transformer [ 6 ] . To mitigate the issue of high computation complexity associated with long sequences caused by high-resolution images , a branch of works explores lightweight Transformer ( e.g .. , Deformable DETR ) for vision tasks [ 55 ] . Transformer is also attracting increasing attention in low-level vision tasks [ 14 , 24 , 34 , 36 , 42 ] . Parmar et al . cast image generation as an autoregressive sequence generation problem and propose Image Transformer with a local self-attention mechanism [ 36 ] . However , Image Transformer can suffer from its quadratic computation cost and its long inference time due to auto-regressive prediction . Most existing works adopt a hybrid CNN-Transformer architecture , which consists of a CNN head for feature extraction , a Transformer encoder-decoder backbone , and a CNN tail for feature decoding [ 7 , 46 , 48 ] . For example , IPT fully utilizes the Transformer architecture by a large-scale pre-training and achieves promising results in several image restoration tasks [ 7 ] . In parallel work , Hudson et al . introduce bipartite structure to maintain computation of linear efficiency for long-range interactions across the image based on a convolution backbone [ 22 ] . Our token-based generator inherits the network architecture from Transformer without any convolution layers and is able to yield surprising promising results for high-resolution image synthesis . 3 Approach . In this section , we introduce the details of the proposed token-based generator ( TokenGAN ) . As depicted in Figure 2 , the TokenGAN takes as input two kinds of tokens to generate the visual tokens of images by a Transformer . We introduce the input tokens in Section 3.1 and the token-based generator architecture in Section 3.2 , following the overall optimization in Section 3.3 . 3.1 Visual tokens . It ’ s intuitive that the visual tokens of an image contain the information of both content ( i.e. , semantics ) and related styles . In our token-based generator , we choose to separate these two kinds of information by a sequence of content tokens and a sequence of style tokens , so that the interactions between them can be modeled and controlled for the image synthesis . We introduce the details of each as below . Content tokens . As depicted in Figure 2-a , the TokenGAN starts from the sequence of learned constant content tokens c ∈ Rm× ( p2×d ) through the input layer . Specifically , d is the dimension in terms of channels , p × p is the patch size of the content token , and m is the length of the content token sequence . To maintain the position information of each token , we add position encodings p ∈ Rm× ( p2×d ) to each content token following a commonly used paradigm [ 7 , 13 ] . For simplicity , we rewrite the notation for the content tokens as { c1 , · · · , cm } , where ci ∈ Rd . Style tokens . Given a latent code z from the input latent space Z , the mapping network adopts several MLPs to map the input z to a set of different style tokens { s1 , · · · , sn } ∈ S , where si ∈ Rd . As shown in Figure 2-b , the style tokens are paired with a set of learnable semantic embedding as a key-value structure in each style modulation layer of the TokenGAN . With the Transformer modeling , all the content tokens will match with the semantic embedding and then fetch the new style from the style tokens based on the matching results . The fetched new styles are used to control the values of the content tokens , which are finally decoded to images . Such a token-wise style control enables content-aware and fine-grained style learning for image synthesis .

Summary: In this paper, a transformer-based modification of the generator in the ubiquitous StyleGAN model is proposed (named _TokenGAN_). More specifically, the entire generator is replaced by a convolution-free architecture where style-content modulation is now modeled using cross-attention layers instead of AdaIN. Similar to the original StyleGAN, this should allow for decoupled control of style and content at different granularities. The paper shows that the proposed approach achieves FID scores comparable to or better than StyleGAN2 for unconditional image synthesis on FFHQ and LSUN-Churches.