Improved Transformer for High-Resolution GANs

1 Introduction . Attention-based models demonstrate notable learning capabilities for both encoder-based and decoderbased architectures [ 66 , 69 ] due to their self-attention operations which can capture long-range dependencies in data . Recently , Vision Transformer [ 12 ] , one of the most powerful attention-based models , has achieved a great success on encoder-based vision tasks , specifically image classification [ 12 , 60 ] , segmentation [ 37 , 64 ] , and vision-language modeling [ 46 ] . However , applying the Transformer to image generation based on Generative Adversarial Networks ( GANs ) is still an open problem . The main challenge of adopting the Transformer as a decoder/generator lies in two aspects . On one hand , the quadratic scaling problem brought by the self-attention operation becomes even worse when generating pixel-level details for high-resolution images . For example , synthesizing a high definition image with the resolution of 1024× 1024 leads to a sequence containing around one million pixels in the final stage , which is unaffordable for the standard self-attention mechanism . On the other hand , generating images from noise inputs poses a higher demand for spatial coherency in structure , color , and texture than discriminative tasks , and hence a more powerful yet efficient self-attention mechanism is desired for decoding feature representations from inputs . In view of these two key challenges , we propose a novel Transformer-based decoder/generator in GANs for high-resolution image generation , denoted as HiT . HiT employs a hierarchical structure of Transformers and divides the generative process into low-resolution and high-resolution stages , focusing on feature decoding and pixel-level generating , respectively . Specifically , its low-resolution ∗This work was done while Long Zhao was a student researcher at the Google Brain team . Correspondence to : Long Zhao ( lz311 @ rutgers.edu ) and Han Zhang ( zhanghan @ google.com ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . stages follow the design of Nested Transformers [ 73 ] but enhanced by the proposed multi-axis blocked self-attention to better capture global information . Assuming that spatial features are well decoded after low-resolution stages , in high-resolution stages , we drop all self-attention operations in order to handle extremely long sequences for high definition image generation . The resulting high-resolution stages of HiT are built by multi-layer perceptrons ( MLPs ) which have a linear complexity with respect to the sequence length . Note that this design aligns with the recent findings [ 58 , 59 ] that pure MLPs manage to learn favorable features for images , but it simply reduces to an implicit neural function [ 41 , 42 , 43 ] in the case of generative modeling . To further improve the performance , we present an additional cross-attention module that acts as a form of self-modulation [ 8 ] . In summary , this paper makes the following contributions : • We propose HiT , a Transformer-based generator for high-fidelity image generation . Standard self-attention operations are removed in the high-resolution stages of HiT , reducing them to an implicit neural function . The resulting architecture easily scales to high definition image synthesis ( with the resolution of 1024× 1024 ) and has a comparable throughput to StyleGAN2 [ 28 ] . • To tame the quadratic complexity and enhance the representation capability of self-attention operation in low-resolution stages , we present a new form of sparse self-attention operation , namely multi-axis blocked self-attention . It captures local and global dependencies within nonoverlapping image blocks in parallel by attending to a single axis of the input tensor at a time , each of which uses a half of attention heads . The proposed multi-axis blocked self-attention is efficient , simple to implement , and yields better performance than other popular self-attention operations [ 37 , 62 , 73 ] working on image blocks for generative tasks . • In addition , we introduce a cross-attention module performing attention between the input and intermediate features . This module re-weights intermediate features of the model as a function of the input , playing a role as self-modulation [ 8 ] , and provides important global information to high-resolution stages where self-attention operations are absent . • The proposed HiT obtains competitive FID [ 16 ] scores of 30.83 and 2.95 on unconditional ImageNet [ 49 ] 128 × 128 and FFHQ [ 28 ] 256 × 256 , respectively , highly reducing the gap between ConvNet-based GANs and Transformer-based ones . We also show that HiT not only works for GANs but also can serve as a general decoder for other models such as VQ-VAE [ 61 ] . Moreover , we observe that the proposed HiT can obtain more performance improvement from regularization than its ConvNet-based counterparts . To the best of our knowledge , these are the best reported scores for an image generator that is completely free of convolutions , which is an important milestone towards adopting Transformers for high-resolution generative modeling . 2 Related work . Transformers for image generation . There are two main streams of image generation models built on Transformers [ 63 ] in the literature . One stream of them [ 14 , 44 ] is inspired by auto-regressive models that learn the joint probability of the image pixels . The other stream focuses on designing Transformer-based architecture for generative adversarial networks ( GANs ) [ 15 ] . This work follows the spirit of the second stream . GANs have made great progress in various image generation tasks , such as image-to-image translation [ 57 , 74 , 75 , 79 ] and text-to-image synthesis [ 70 , 71 ] , but most of them depend on ConvNet-based backbones . Recent attempts [ 23 , 32 ] build a pure Transformer-based GAN by a careful design of attention hyper-parameters as well as upsampling layers . However , such a model is only validated on small scale datasets ( e.g. , CIFAR-10 [ 31 ] and STL-10 [ 10 ] consisting of 32× 32 images ) and does not scale to complex real-world data . To our best knowledge , no existing work has successfully applied a Transformer-based architecture completely free of convolutions for high-resolution image generation in GANs . GANformer [ 21 ] leverages the Transformer as a plugin component to build the bipartite structure to allow long-range interactions during the generative process , but its main backbone is still a ConvNet based on StyleGAN [ 28 ] . GANformer and HiT are different in the goal of using attention modules . GANformer utilizes the attention mechanism to model the dependences of objects in a generated scene/image . Instead , HiT explores building efficient attention modules for synthesizing general objects . Our experiments demonstrate that even for simple face image datasets , incorporating the attention mechanism can still lead to performance improvement . We believe our work reconfirms the necessity of using attention modules for general image generation tasks , and more importantly , we present an efficient architecture for attention-based high-resolution generators which might benefit the design of future attention-based GANs . Attention models . Many efforts [ 2 , 17 , 65 ] have been made to reduce the quadratic complexity of self-attention operation . When handling vision tasks , some works [ 37 , 62 , 73 ] propose to compute local self-attention in blocked images which takes advantage of the grid structure of images . [ 6 ] also presents local self-attention within image blocks but it does not perform self-attention across blocks as in HiT . The most related works to the proposed multi-axis blocked self-attention are [ 17 , 65 ] where they also compute attention along axes . But our work differs notably in that we compute different attentions within heads on blocked images . Some other works [ 7 , 21 , 22 ] avoid directly applying standard self-attention to the input pixels , and perform attention between the input and a small set of latent units . Our cross-attention module differs from them in that we apply cross-attention to generative modeling designed for pure Transformer-based architectures . Implicit neural representations . The largest popularity of implicit neural representations/functions ( INRs ) is studied in 3D deep learning to represent a 3D shape in a cheap and continuous way [ 41 , 42 , 43 ] . Recent studies [ 1 , 13 , 34 , 54 ] explore the idea of using INRs for image generation , where they learn a hyper MLP network to predict an RGB pixel value given its coordinates on the image grid . Among them , [ 1 , 54 ] are closely related to our high-resolution stages of the generative process . One remarkable difference is that our model is driven by the cross-attention module and features generated in previous stages instead of the hyper-network presented in [ 1 , 54 ] . 3 Approach . 3.1 Main Architecture . In the case of unconditional image generation , HiT takes a latent code z ∼ N ( 0 , I ) as input and generates an image of the target resolution through a hierarchical structure . The latent code is first projected into an initial feature with the spatial dimension of H0 × W0 and channel dimension of C0 . During the generation process , we gradually increase the spatial dimension of the feature map while reducing the channel dimension in multiple stages . We divide the generation stages into low-resolution stages and high-resolution stages to balance feature dependency range in decoding and computation efficiency . The overview of the proposed method is illustrated in Figure 1 . In low-resolution stages , we allow spatial mixing of information by efficient attention mechanism . We follow the decoder form of Nested Transformer [ 73 ] where in each stage , the input feature is first divided into non-overlapping blocks where each block can be considered as a local patch . After being combined with learnable positional encoding , each block is processed independently via a shared attention module . We enhance the local self-attention of Nested Transformer by the proposed multiaxis blocked self-attention that can produce a richer feature representation by explicitly considering local ( within blocks ) as well as global ( across blocks ) relations . We denote the overall architecture of these stages as multi-axis Nested Transformer . Assuming that spatial dependency is well modeled in the low-resolution stages , the high-resolution stages can focus on synthesizing pixel-level image details purely based on the local features . Thus in the high-resolution stage , we remove all self-attention modules and only maintain MLPs which can further reduce computation complexity . The resulting architecture in this stage can be viewed as an implicit neural function conditioned on the given latent feature as well as positional information . We further enhance HiT by adding a cross-attention module at the beginning of each stage which allows the network to directly condition the intermediate features on the initial input latent code . This kind of self-modulation layer leads to improved generative performance , especially when selfattention modules are absent in high-resolution stages . In the following sections , we provided detailed descriptions for the two main architectural components of HiT : ( i ) multi-axis blocked self-attention module and ( ii ) cross-attention module , respectively . 3.2 Multi-Axis Blocked Self-Attention . Different from the blocked self-attention in Nested Transformer [ 73 ] , the proposed multi-axis blocked self-attention performs attention on more than a single axis . The attentions performed in two axes correspond to two forms of sparse self-attention , namely regional attention and dilated attention . Regional attention follows the spirit of blocked local self-attention [ 37 , 62 ] where tokens attend to their neighbours within non-overlapped blocks . To remedy the loss of global attention , dilated attention captures long-range dependencies between tokens across blocks : it subsamples attention regions in a manner similar to dilated convolutions with a fixed stride equal to the block size . Figure 2 illustrates an example of these two attentions . To be specific , given an input of image with size ( H , W , C ) , it is blocked into a tensor X of the shape ( b× b , Hb × W b , C ) representing ( b , b ) non-overlapping blocks each with the size of ( H b , W b ) . Dilated attention mixes information along the first axis of X while keeping information along other axes independent ; regional attention works in an analogous manner over the second axis of X . They are straightforward to implement : attention over the i-th axis of X can be implemented by einsum operation which is available in most deep learning libraries . We mix regional and dilated attention in a single layer by computing them in parallel , each of which uses a half of attention heads . Our method can be easily extended to model more than two axes by performing attention on each of the axis in parallel . Axial attention [ 17 , 20 , 65 ] can be viewed as a special case of our method , where the blocking operation before attention is removed . However , we find blocking is the key to achieve significant performance improvement in the experiment . This is because the blocked structure reveals a better inductive bias for images . Compared with [ 17 , 37 , 65 ] where different attention modules are interleaved in consecutive layers , our approach aggregates local and global information in a single round , which is not only more flexible for architecture design but also shown to yield better performance than interleaving in our experiment . Balancing attention between axes . In each multi-axis blocked self-attention module , the input feature is blocked in a balanced way such that we have b× b ≈ Hb × W b . This ensures that regional and dilated attention is computed on an input sequence of a similar length , avoiding half of the attention heads are attended to a too sparse region . In general , performing dot-product attention between two input sequences of length N = H × W requires a total of O ( N2 ) computation . When computing the balanced multi-axis blocked attention on an image with the block size of S , i.e. , S = √ N , we perform attention on S sequences of length S , which is a total ofO ( S ·S2 ) = O ( N √ N ) computation . This leads to an O ( √ N ) saving in computation over standard self-attention .

Summary: This paper proposes a new transformer-based generator for high-resolution image generation. It addresses the quadratic scaling problem of the attention operator with multi-axis blocked self-attention which considering the within blocks and across blocks attention. The proposed cross-attention for input and intermediate features is also novel. The evaluation is done in ImageNet and FFHQ datasets.  Promising results are provided.