Improved Transformer for High-Resolution GANs

By Long Zhao, Zizhao Zhang, Ting-Li Chen, Dimitris N. Metaxas, Han Zhang. 2021

114 citations recorded by the source corpus.

Attention-based models, exemplified by the Transformer, can effectively model long range dependency, but suffer from the quadratic complexity of self-attention operation, making them difficult to be adopted for high-resolution image generation based on Generative Adversarial Networks (GANs). In this paper, we introduce two key ingredients to Transformer to address this challenge. First, in low-resolution stages of the generative process, standard global self-attention is replaced with the proposed multi-axis blocked self-attention which allows efficient mixing of local and global attention. Second, in high-resolution stages, we drop self-attention while only keeping multi-layer perceptrons reminiscent of the implicit neural function. To further improve the performance, we introduce an additional self-modulation component based on cross-attention. The resulting model, denoted as HiT, has a nearly linear computational complexity with respect to the image size and thus directly scales to synthesizing high definition images. We show in the experiments that the proposed HiT achieves state-of-the-art FID scores of 30.83 and 2.95 on unconditional ImageNet $128 \times 128$ and FFHQ $256 \times 256$, respectively, with a reasonable throughput. We believe the proposed HiT is an important milestone for generators in GANs which are completely free of convolutions. Our code is made publicly available at https://github.com/google-research/hit-gan

Subfields: Computer Vision and Pattern Recognition, Generative Models, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Vision and Pattern Recognition

In Generative Models

In Model Architectures

Citations in the corpus

TransGAN: Two Pure Transformers Can Make One Strong GAN, and That Can Scale Up (2021, methodology): “Specifically, we apply the self-modulation [85, 22, 35] to our generator and use cross-attention [53, 86] to map the latent space to the global region.”
Muse: Text-To-Image Generation via Masked Generative Transformers (2023, methodology): “To address this issue, we experimented with two architectures, the Swin transformer (Liu et al., 2021) and Multi-axis transformer (Zhao et al., 2021), to incorporate local-global attention.”
MAXIM: Multi-Axis MLP for Image Processing (2022, methodology): “Our work is inspired by the multi-axis blocked selfattention proposed in [120], which performs attention on more than a single axis.”
MaxViT: Multi-Axis Vision Transformer (2022, methodology): “Since Batch Normalization [37,98] achieves better results on image generation, we replaced all Layer Norm with Batch Norm under this setting.”