TransGAN: Two Pure Transformers Can Make One Strong GAN, and That Can Scale Up

1 Introduction . Generative adversarial networks ( GANs ) have gained considerable success on numerous tasks [ 1 , 2 , 3 , 4 , 5 , 6 , 7 ] . Unfortunately , GANs suffer from the notorious training instability , and numerous efforts have been devoted to stabilizing GAN training , introducing various regularization terms [ 8 , 9 , 10 , 11 ] , better losses [ 1 , 12 , 13 , 14 ] , and training recipes [ 15 , 16 ] . Among them , one important route to improving GANs examines their neural architectures . [ 17 , 8 ] reported a large-scale study of GANs and observed that when serving as ( generator ) backbones , popular neural architectures perform comparably well across the considered datasets . Their ablation study suggested that most of the variations applied in the ResNet family resulted in very marginal improvements . Nevertheless , neural architecture search ( NAS ) was later introduced to GANs and suggests enhanced backbone designs are also important for improving GANs , just like for other computer vision tasks . Those works are consistently able to discover stronger GAN architectures beyond the standard ResNet topology 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . [ 18 , 19 , 20 ] . Other efforts include customized modules such as self-attention [ 21 ] , style-based generator [ 22 ] , and autoregressive transformer-based part composition [ 23 ] . However , one last “ commonsense '' seems to have seldomly been challenged : using convolutional neural networks ( CNNs ) as GAN backbones . The original GAN [ 24 , 25 ] used fully-connected networks and can only generate small images . DCGAN [ 26 ] was the first to scale up GANs using CNN architectures , which allowed for stable training for higher resolution and deeper generative models . Since then , in the computer vision domain , every successful GAN relies on CNN-based generators and discriminators . Convolutions , with the strong inductive bias for natural images , crucially contribute to the appealing visual results and rich diversity achieved by modern GANs . Can we build a strong GAN completely free of convolutions ? This is a question not only arising from intellectual curiosity , but also of practical relevance . Fundamentally , a convolution operator has a local receptive field , and hence CNNs can not process long-range dependencies unless passing through a sufficient number of layers . However , that is inefficient , and could cause the loss of feature resolution and fine details , in addition to the difficulty of optimization . Vanilla CNN-based models are therefore inherently not well suited for capturing an input image ’ s “ global '' statistics , as demonstrated by the benefits from adopting self-attention [ 21 ] and non-local [ 27 ] operations in computer vision . Moreover , the spatial invariance possessed by convolution poses a bottleneck on its ability of adapting to spatially varying/heterogeneous visual patterns , which also motivates the success of relational network [ 28 ] , dynamic filters [ 29 , 30 ] and kernel prediction [ 31 ] methods . 1.1 Our Contributions . This paper aims to be the first pilot study to build a GAN completely free of convolutions , using only pure transformer-based architectures . We are inspired by the recent success of transformer architectures in computer vision [ 32 , 33 , 34 ] . Compared to parallel generative modeling works [ 21 , 23 , 35 ] that applied self-attention or transformer encoder in conjunction with CNN-based backbones , our goal is more ambitious and faces several daunting gaps ahead . First and foremost , although a pure transformer architecture applied directly to sequences of image patches can perform very well on image classification tasks [ 34 ] , it is unclear whether the same way remains effective in generating images , which crucially demands the spatial coherency in structure , color , and texture , as well as the richness of fine details . The handful of existing transformers that output images have unanimously leveraged convolutional part encoders [ 23 ] or feature extractors [ 36 , 37 ] . Moreover , even given well-designed CNN-based architectures , training GANs is notoriously unstable and prone to mode collapse [ 15 ] . Training vision transformers are also known to be tedious , heavy , and data-hungry [ 34 ] . Combining the two will undoubtedly amplify the challenges of training . In view of those challenges , this paper presents a coherent set of efforts and innovations towards building the pure transformer-based GAN architectures , dubbed TransGAN . A naive option may directly stack multiple transformer blocks from raw pixel inputs , but that would scale poorly due to memory explosion . Instead , we start with a memory-friendly transformer-based generator by gradually increasing the feature map resolution in each stage . Correspondingly , we also improve the discriminator with a multi-scale structure that takes patches of varied size as inputs , which balances between capturing global contexts and local details , in addition to enhancing memory efficiency more . Based on the above generator-discriminator design , we introduce a new module called grid self-attention , that alleviates the memory bottleneck further when scaling up TransGAN to high-resolution generation ( e.g . 256× 256 ) . To address the aforementioned instability issue brought by both GAN and Transformer , we also develop a unique training recipe in association with our innovative TransGAN architecture , that effectively stabilizes its optimization and generalization . That includes showings the necessity of data augmentation , modifying layer normalization , and replacing absolute token locations with relative position encoding . Our contributions are outlined below : • Novel Architecture Design : We build the first GAN using purely transformers and no convolution . TransGAN has customized a memory-friendly generator and a multi-scale discriminator , and is further equipped with a new grid self-attention mechanism . Those architectural components are thoughtfully designed to balance memory efficiency , global feature statistics , and local fine details with spatial variances . • New Training Recipe : We study a number of techniques to train TransGAN better , including leveraging data augmentation , modifying layer normalization , and adopting relative position encoding , for both generator and discriminator . Extensive ablation studies , discussions , and insights are presented . • Performance and Scalability : TransGAN achieves highly competitive performance compared to current state-of-the-art GANs . Specifically , it sets the new state-of-the-art inception score of 10.43 and FID score of 18.28 on STL-10 . It also reaches competitive 9.02 inception score and 9.26 FID on CIFAR-10 , and 5.28 FID score on CelebA 128× 128 , respectively . Meanwhile , we also evaluate TransGAN on higher-resolution ( e.g. , 256× 256 ) generation tasks , where TransGAN continues to yield diverse and impressive visual examples . 2 Related Works . Generative Adversarial Networks . After its origin , GANs quickly embraced fully convolutional backbones [ 26 ] , and inherited most successful designs from CNNs such as batch normalization , pooling , ( Leaky ) ReLU and more [ 38 , 39 , 40 , 18 ] . GANs are widely adopted in image translation [ 3 , 4 , 41 ] , image enhancement [ 7 , 42 , 43 ] , and image editing [ 44 , 45 ] . To alleviate its unstable training , a number of techniques have been studied , including the Wasserstein loss [ 46 ] , the stylebased generator [ 22 ] , progressive training [ 16 ] , lottery ticket [ 47 ] , and spectral normalization [ 48 ] . Transformers in Computer Vision . The original transformer was built for NLP [ 49 ] , where the multi-head self-attention and feed-forward MLP layer are stacked to capture the long-term correlation between words . A recent work [ 34 ] implements highly competitive ImageNet classification using pure transformers , by treating an image as a sequence of 16× 16 visual words . It has strong representation capability and is free of human-defined inductive bias . In comparison , CNNs exhibit a strong bias towards feature locality , as well as spatial invariance due to sharing filter weights across all locations . However , the success of original vision transformer relies on pretraining on large-scale external data . [ 50 , 51 ] improve the data efficiency and address the difficulty of optimizing deeper models . Other works introduce the pyramid/hierarchical structure to transformer [ 52 , 53 , 54 ] or combine it with convolutional layers [ 55 , 56 ] . Besides image classification task , transformer and its variants are also explored on image processing [ 37 ] , point cloud [ 57 ] , semantic segmentation [ 58 ] , object detection [ 32 , 59 ] and so on . A comprehensive review is referred to [ 60 ] . Transformer Modules for Image Generation . There exist several related works combining the transformer modules into image generation models , by replacing certain components of CNNs . [ 61 ] firstly formulated image generation as autoregressive sequence generation , for which they adopted a transformer architecture . [ 62 ] propose sparse factorization of the attention matrix to reduce its complexity . While those two works did not tackle the GANs , one recent ( concurrent ) work [ 23 ] used a convolutional GAN to learn a codebook of context-rich visual parts , whose composition is subsequently modeled with an autoregressive transformer architecture.The authors demonstrated success in synthesizing high-resolution images . However , the overall CNN architecture remains in place ( including CNN encoder/decoder for the generators , and a fully CNN-based discriminator ) , and the customized designs ( e.g , codebook and quantization ) also limit their model ’ s versatility . Another concurrent work [ 35 ] employs a bipartite self-attention on StyleGAN and thus it can propagate latent variables to the evolving visual features , yet its main structure is still convolutional , including both the generator and discriminator . To our best knowledge , no other existing work has tried to completely remove convolutions from their generative modeling frameworks . 3 Technical Approach : A Journey Towards GAN with Pure Transformers . In this section , we start by introducing the memory-friendly generator and multi-scale discriminator , equipped with a novel grid self-attention . We then introduce a series of training techniques to stabilize its training procedure , including data augmentation , the modified normalization , and injecting relative position encoding to self-attention . To start with , we choose the transformer encoder [ 49 ] as our basic block and try to make minimal changes . An encoder is a composition of two parts . The first part is constructed by a multi-head self-attention module and the second part is a feed-forward MLP with GELU non-linearity . The normalization layer is applied before both of the two parts . Both parts employ residual connection . 3.1 Memory-friendly Generator . The task of generation poses a high standard for spatial coherency in structure , color , and texture , both globally and locally . The transformer encoders take embedding token words as inputs and calculate the interaction between each token recursively . [ 63 , 34 ] . The main dilemma here is : what is the right “ word '' for image generation tasks ? If we similarly generate an image in a pixel-by-pixel manner through stacking transformer encoders , even a low-resolution image ( e.g . 32× 32 ) can result in an excessively long sequence ( 1024 ) , causing the explosive cost of self-attention ( quadratic w.r.t . the sequence length ) and prohibiting the scalability to higher resolutions . To avoid this daunting cost , we are inspired by a common design philosophy in CNN-based GANs , to iteratively upscale the resolution at multiple stages [ 25 , 16 ] . Our strategy is hence to increase the input sequence and reduce the embedding dimension gradually . Figure 2 ( left ) illustrates a memory-friendly transformer-based generator that consists of multiple stages . Each stage stacks several transformer blocks . By stages , we gradually increase the feature map resolution until it meets the target resolution H ×W . Specifically , the generator takes the random noise as its input , and passes it through a multiple-layer perceptron ( MLP ) to a vector of length H0 ×W0 × C. The vector is reshaped into a H0 ×W0 resolution feature map ( by default we use H0 =W0 = 8 ) , each point a C-dimensional embedding . This “ feature map '' is next treated as a length-64 sequence of C-dimensional tokens , combined with the learnable positional encoding . To scale up to higher-resolution images , we insert an upsampling module after each stage , consisting of a reshaping and resolution-upscaling layer . For lower-resolution stages ( resolution lower than 64 × 64 ) , the upsampling module firstly reshapes the 1D sequence of token embedding back to a 2D feature map Xi ∈ RHi×Wi×C and then adopts the bicubic layer to upsample its resolution while the embedded dimension is kept unchanged , resulting in the output X ′ i ∈ R2Hi×2Wi×C . After that , the 2D feature map X ′ i is again reshaped into the 1D sequence of embedding tokens . For higher-resolution stages , we replace the bicubic upscaling layer with the pixelshuffle module , which upsamples the resolution of feature map by 2× ratio and also reduces the embedding dimension to a quarter of the input . This pyramid-structure with modified upscaling layers mitigates the memory and computation explosion . We repeat multiple stages until it reaches the target resolution ( H , W ) , and then we will project the embedding dimension to 3 and obtain the RGB image Y ∈ RH×W×3 .

Summary: This work successfully takes an important step in bringing transformer architecture to the family of GAN models, so far dominated by convolutional architectures. Another contribution is the localised attention ('Grid Self-Attention') that addresses the problem of increased memory requirements of attention at higher-resolution feature map. The empirical evaluation shows competitive results on several benchmarks, including state-of-the-art scores on STL-10; however, it is missing an important benchmark - ImageNet.   Despite such limitation and missing references and comparisons to some relevant prior work (which hopefully could be addressed in rebuttal), this work seems to meet the quality and novelty requirements for NeurIPS. I am willing to improve the rating further if my concerns are addressed.