用于高分辨率GAN的改进型Transformer

基于注意力机制的模型，如Transformer，能够有效建模长距离依赖关系，但其自注意力操作具有二次方复杂度，导致在基于生成对抗网络（GAN）的高分辨率图像生成任务中难以应用。本文提出两种关键改进，以解决这一挑战。首先，在生成过程的低分辨率阶段，我们用所提出的多轴分块自注意力（multi-axis blocked self-attention）替代传统的全局自注意力机制，从而实现局部与全局注意力的高效融合。其次，在高分辨率阶段，我们完全摒弃自注意力机制，仅保留多层感知机（MLP），其设计灵感源自隐式神经函数（implicit neural function）。为进一步提升性能，我们引入一种基于交叉注意力的自调制（self-modulation）组件。由此构建的模型——HiT，其计算复杂度相对于图像尺寸接近线性，因而可直接扩展至高分辨率图像的合成任务。实验结果表明，所提出的HiT在无条件生成任务中，在ImageNet $128 \times 128$ 和 FFHQ $256 \times 256$ 数据集上分别取得了30.83和2.95的先进FID得分，同时保持了合理的生成吞吐量。我们认为，HiT是GAN生成器领域的一项重要里程碑，标志着完全摒弃卷积操作的生成模型的实现。相关代码已公开，地址为：https://github.com/google-research/hit-gan。