[论文解读] Efficient Image Generation with Variadic Attention Heads

Steven Walton, Ali Hassani+3 more|arXiv (Cornell University)|Nov 10, 2022

Generative Adversarial Networks and Image Synthesis被引用 8

一句话总结

本文提出 StyleNAT，一种基于 Transformer 的图像生成框架，采用 Hydra-NA——一种通过邻域注意力（NA）实现不同感受野的注意力头分区新方法——以高效捕捉局部与全局特征。其在 FFHQ-256 上取得 2.046 的新 SOTA FID 分数，在 FFHQ-1024 上取得 4.174 的新 SOTA FID 分数，优于先前的 GAN 与 Transformer 模型，参数量减少 28%，采样吞吐量提高 56%。

ABSTRACT

While the integration of transformers in vision models have yielded significant improvements on vision tasks they still require significant amounts of computation for both training and inference. Restricted attention mechanisms significantly reduce these computational burdens but come at the cost of losing either global or local coherence. We propose a simple, yet powerful method to reduce these trade-offs: allow the attention heads of a single transformer to attend to multiple receptive fields. We demonstrate our method utilizing Neighborhood Attention (NA) and integrate it into a StyleGAN based architecture for image generation. With this work, dubbed StyleNAT, we are able to achieve a FID of 2.05 on FFHQ, a 6% improvement over StyleGAN-XL, while utilizing 28% fewer parameters and with 4$ imes$ the throughput capacity. StyleNAT achieves the Pareto Frontier on FFHQ-256 and demonstrates powerful and efficient image generation on other datasets. Our code and model checkpoints are publicly available at: https://github.com/SHI-Labs/StyleNAT

研究动机与目标

• 解决标准 Transformer 在高分辨率图像生成中效率低下且适应性有限的问题。
• 克服如 Swin 和 NA 等局部注意力机制中感受场固定的局限性。
• 通过扩张和滑动窗口 NA 将注意力头划分为不同感受场，实现灵活、数据自适应的注意力机制。
• 在减少模型规模和推理时间的同时，实现 FID 指标的 SOTA 表现。
• 开发一种可视化方法，用于解释基于 Swin 和 NA 的模型中的局部注意力图。

提出的方法

• 提出 Hydra-NA，一种通过邻域注意力（NA）将多头注意力划分为具有不同感受场的独立头的方法。
• 采用 NA 的滑动窗口与扩张变体，使不同头能够关注局部邻域或扩展的全局上下文。
• 设计一种受 StyleGAN 启发的风格化生成器，与 Hydra-NA 模块结合，实现分层特征学习。
• 使用映射网络将随机潜在向量嵌入为风格码，用于条件化跨层的特征图。
• 应用可学习的位置编码，以在潜在空间中保留空间结构。
• 开发一种可视化技术，用于在不同分辨率层级上解释注意力图，识别如水印等记忆化的数据集特征。

实验结果

研究问题

• RQ1 将注意力头划分为不同感受场是否能提升基于 Transformer 的图像生成器的表达能力与效率？
• RQ2 使用可变卷积核大小与扩张率的局部注意力机制，如何影响高分辨率图像生成中的长距离依赖建模？
• RQ3 单一架构在无需大量超参数调优的情况下，能在多大程度上泛化至多样化的数据集？
• RQ4 注意力图能否揭示生成图像中对数据集特异性特征（如水印或 ID）的记忆化现象？
• RQ5 与先前的 GAN 和 Transformer 相比，所提方法是否在计算成本更低的前提下实现了 SOTA 性能？

主要发现

• StyleNAT 在 FFHQ-256 上实现 2.046 的新 SOTA FID 分数，优于包括 StyleGAN-XL 在内的卷积模型与基于 Transformer 的模型。
• 在 FFHQ-1024 上，StyleNAT 在基于 Transformer 的 GAN 中实现 4.174 的新 SOTA FID 分数，展现出卓越的全局特征建模能力。
• 尽管性能更优，StyleNAT 相较于 StyleGAN-XL 参数量减少 28%，采样吞吐量提升 56%。
• 注意力可视化显示，模型会记忆数据集特异性特征（如水印与图像 ID），尤其在低分辨率（如 16x16）时更为明显。
• 注意力图中水印与 ID 的存在与视觉上实际存在的水印强相关，表明模型可能对训练数据存在过拟合。
• 尽管生成图像质量高，FID 分数在 Church 数据集上仍表现不佳，表明模型更倾向于记忆分布性特征而非实现泛化。