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 paper proposes to build a GAN model using only transformer blocks, and completely free of convolutions. Several architecture designs are proposed and a unique training recipe is also introduced for stably training the proposed TransGAN. TransGAN  can scale up to high resolution (e.g. 256 × 256), and delivers state-of-the-art results on multiple benchmarks.