Improved Transformer for High-Resolution GANs (2106.07631v3)

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

Analysis on "Improved Transformer for High-Resolution GANs"

This paper introduces HiT (Hierarchical Transformer), an innovative Transformer-based architecture designed to mitigate the challenges faced by attention-based models, specifically when applied to high-resolution Generative Adversarial Networks (GANs). The advent of HiT marks a significant stride in leveraging Transformer architectures for generative tasks, especially in scaling to high-definition image synthesis without relying on convolutions. This summary will explore the key methodologies, empirical findings, and potential implications of adopting HiT in the domain of high-resolution image generation.

Key Contributions and Methodology

HiT addresses the inherent computational inefficiencies of traditional Transformers by introducing two primary modifications:

1. Multi-Axis Blocked Self-Attention: In the low-resolution stages, HiT substitutes global self-attention with a more efficient multi-axis blocked self-attention mechanism. This approach parallelizes local and global attention in a sparse manner, thereby significantly reducing computational overhead while maintaining representational richness.
2. The Use of MLPs in High-Resolution Stages: Departing from the conventional self-attention mechanism, HiT adopts multi-layer perceptrons (MLPs) in the high-resolution stages. This simplification rests on the assumption that spatial features have been adequately captured in the preceding stages. Here, MLPs serve to generate the detailed pixel-level information necessary for high-resolution images.

Additionally, to further enhance performance, HiT incorporates a cross-attention module for self-modulation, allowing intermediate features to be dynamically re-weighted based on initial latent inputs. This interaction provides crucial global contextual information, particularly beneficial when self-attention is absent.

Empirical Evaluation

The empirical evaluation affirms HiT’s proficiency in image synthesis, demonstrating state-of-the-art Fréchet Inception Distance (FID) scores particularly on the ImageNet $128\times128$ and FFHQ $256\times256$ datasets with FIDs of 30.83 and 2.95, respectively. These scores underscore HiT’s efficacy, reinforcing its competitive standing against traditional convolutional GANs like StyleGAN2.

Moreover, HiT exhibits promising throughput, enabling practical applications without sacrificing computational feasibility. This is particularly impressive given HiT’s avoidance of convolution operations, marking a distinctive achievement among Transformer-based GAN implementations.

Implications and Future Directions

HiT indeed sets a precedent for convolution-free generative models, reflecting an evolving step towards integrating Transformers in GANs. Practically, HiT’s computational efficiency and enhanced synthesis capabilities could accelerate transformative applications in image generation, including creative industries and virtual reality environments that require high-fidelity visual content.

Theoretically, HiT raises intriguing possibilities about alternative architectural designs for GANs. An area of future research could explore optimizing Transformer-based discriminators in tandem with HiT, potentially culminating in fully-attention-based GAN frameworks. Furthermore, the exploration of advanced regularization strategies tailored for Transformers could yield additional performance gains, particularly for ultra-high-resolution tasks.

In conclusion, while HiT demonstrates considerable advancements in Transformer-based generative modeling, its success also underscores the broader perspective that Transformers can transcend traditional applications, paving the way for robust, high-performing architectures across computationally demanding domains. The insights provided by HiT will undoubtedly inform future endeavors aiming to marry efficiency with complexity in generative models.