[論文レビュー] 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

ひとこと要約

この論文は、局所的およびグローバルな特徴を効率的に捉えるために、近接注意（NA）を用いた注意ヘッドの新しい分割法であるHydra-NAを用いる、Transformerベースの画像生成フレームワークであるStyleNATを提案する。FFHQ-256では2.046の新SOTA FIDスコア、FFHQ-1024では4.174の新SOTA FIDスコアを達成し、パラメータ数を28%削減し、推論スループットを56%向上させながら、先行するGANおよびTransformerモデルを上回る性能を発揮した。

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 のような局所的注意機構における固定受容 field の制限を克服すること。
• 拡張およびスライディングウィンドウ型 NA を用いて、注意ヘッドを多様な受容 field に分割することで、データに適応する柔軟な注意メカニズムを実現すること。
• FIDスコアでSOTA性能を達成するとともに、モデルサイズと推論時間を削減すること。
• Swin や NA ベースのモデルにおける局所的注意マップを解釈可能な可視化手法を開発すること。

提案手法

• 近接注意（NA）を用いて、異なる受容 field を持つ複数の注意ヘッドに多ヘッド注意を分割する手法であるHydra-NAを導入する。
• スライディングウィンドウ型および拡張型 NA を採用し、ヘッドが局所的近傍または拡張されたグローバルな文脈に注目できるようにする。
• StyleGAN をインspired したスタイリングベースのジェネレータを設計し、階層的特徴学習を可能にするHydra-NAブロックと統合する。
• ランダムな潜在ベクトルをスタイルコードに埋め込むためのマッピングネットワークを用い、層間の特徴マップを条件づける。
• 潜在空間における空間的構造を保持するために、学習可能な位置エンコーディングを適用する。
• 異なる解像度レベルでの注意マップを解釈可能な可視化技術を開発し、画像に埋め込まれたデータセット固有のアーティファクト（例：透かし）を同定する。

実験結果

リサーチクエスチョン

• RQ1 異なる受容 field を持つ注意ヘッドに分割することで、Transformerベースの画像生成器の表現力と効率性が向上するか？
• RQ2 可変なカーネルサイズおよび拡張を用いた局所的注意は、高解像度画像生成における長距離依存関係のモデリングにどのように影響するか？
• RQ3 豊富なハイパーパrameterチューニングなしで、1つのアーキテクチャが多様なデータセットに一般化できる程度はどの程度か？
• RQ4 注意マップは、生成画像に埋め込まれたデータセット固有のアーティファクト（例：透かしやID）の記憶を明らかにできるか？
• RQ5 提案手法は、先行するGANおよびTransformerと比較して、計算コストを削減しつつSOTA性能を達成できるか？

主な発見

• StyleNATは、FFHQ-256で2.046の新SOTA FIDを達成し、StyleGAN-XLのような畳み込みモデルおよびTransformerベースのモデルを上回った。
• FFHQ-1024では、Transformerベースの GAN としての新SOTA FID 4.174 を記録し、優れたグローバル特徴モデリングを示した。
• 性能が優れているにもかかわらず、StyleGAN-XLと比較して28%のパラメータ削減と56%の推論スループット向上を達成した。
• 注意マップの可視化により、特に低解像度（例：16x16）で、データセット固有のアーティファクト（透かしや画像ID）がモデルに記憶されていることが明らかになった。
• 注意マップに透かしやIDが現れる割合は、それらが視覚的に存在するのと強く相関しており、学習データへの過剰適合の兆候である。
• 高品質な画像サンプルを生成しているにもかかわらず、ChuchデータセットではFIDスコアが著しく低く、分布的アーティファクトへの過剰適合が一般化の欠如を示唆している。