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: The paper describes a purely attention-based GAN generator architecture where low-resolution deeper layers feature larger-scale (more distant) communication and higher-resolution layers synthesize locally using MLPs,  but still conditioned (indirectly) on the initial latent code.  State-of-the-art results are claimed for 128^2 unconditional ImageNet synthesis (but I believe this claim to be false, see below). In addition, results are shown for CelebA-HQ and FFHQ. While the architecture have some similarities to the near-concurrent work of Hudson and Zitnick (ICML 2021), lack of comparisons should not be held against this paper.  I remain unconvinced by the argumentation and how the results support them. Visually high-quality results are only shown for the simple face image datasets, for which attention mechanisms do not appear to be as necessary as for more complex ones that feature scene compositions. I feel I do not walk away feeling that I learned something.