Improved Transformer for High-Resolution GANs

Abstract

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

Accepted to NeurIPS 2021 (with updated ImageNet results). Code is available at https://github.com/google-research/hit-gan

References in corpus (13)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• Learning Transferable Visual Models From Natural Language Supervision
• GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium
• MLP-Mixer: An all-MLP Architecture for Vision
• Zero-Shot Text-to-Image Generation
• Axial Attention in Multidimensional Transformers
• TransGAN: Two Pure Transformers Can Make One Strong GAN, and That Can Scale Up
• Video Super-Resolution Transformer
• Perceiver: General Perception with Iterative Attention
• ViTGAN: Training GANs with Vision Transformers
• Explicitly disentangling image content from translation and rotation with spatial-VAE
• Fast Transformer Decoding: One Write-Head is All You Need
• Generative Models as Distributions of Functions