[論文レビュー] DiffiT: Diffusion Vision Transformers for Image Generation

Ali Hatamizadeh, Jiaming Song+3 more|arXiv (Cornell University)|Dec 4, 2023

Generative Adversarial Networks and Image Synthesis被引用数 7

ひとこと要約

本稿では、従来のU-Netアーキテクチャに代わり階層的トランスフォーマーエンコーダーデコーダーを採用する視覚トランスフォーマーに基づく拡散モデルDiffiTを提案する。時間依存自己注意（TMSA）機構を導入し、ノイズ除去の各段階に応じて注意を動的に適応させる。これにより、潜在空間におけるImageNet-256で1.73、ImageNet-512で2.67という最先端のFIDスコアを達成した。

ABSTRACT

Diffusion models with their powerful expressivity and high sample quality have achieved State-Of-The-Art (SOTA) performance in the generative domain. The pioneering Vision Transformer (ViT) has also demonstrated strong modeling capabilities and scalability, especially for recognition tasks. In this paper, we study the effectiveness of ViTs in diffusion-based generative learning and propose a new model denoted as Diffusion Vision Transformers (DiffiT). Specifically, we propose a methodology for finegrained control of the denoising process and introduce the Time-dependant Multihead Self Attention (TMSA) mechanism. DiffiT is surprisingly effective in generating high-fidelity images with significantly better parameter efficiency. We also propose latent and image space DiffiT models and show SOTA performance on a variety of class-conditional and unconditional synthesis tasks at different resolutions. The Latent DiffiT model achieves a new SOTA FID score of 1.73 on ImageNet256 dataset while having 19.85%, 16.88% less parameters than other Transformer-based diffusion models such as MDT and DiT,respectively. Code: https://github.com/NVlabs/DiffiT

研究の動機と目的

• 拡散ベースの画像生成における視覚トランスフォーマーの有効性を調査し、畳み込み型U-Netアーキテクチャを越えることを目的とする。
• 既存の拡散モデルにおける細粒度の時間的モデリングの欠如に取り組み、注意機構が異なるノイズ除去段階に適応しない問題を解決することを目的とする。
• 画像および潜在空間における生成の両方に対して、階層的トランスフォーマーに基づくアーキテクチャを導入することで、ノイズ除去ネットワークの設計パターンを統一することを目的とする。
• FIDスコアを指標として複数のベンチマークで最先端の画像生成品質を達成することを目的とする。

提案手法

• 時間ステップ情報を自己注意のクエリ、キー、バリューの投影に直接統合する、新規の時間依存自己注意（TMSA）モジュールを提案。これにより、ノイズ除去段階に応じた動的適応が可能になる。
• 特徴の階層を維持し、スケール間での特徴の効果的な精練を可能にする、階層的かつU字型のトランスフォーマーエンコーダーデコーダー構造（DiffiT）を設計。
• 潜在空間にTMSAモジュールを適用する高解像度画像生成モデルである潜在DiffiTを導入。これにより、効率的かつ高精度な合成が可能になる。
• 時間トークンを学習可能な投影を介して自己注意機構に統合し、現在のノイズ除去ステップに応じて注意の挙動を制御できるようにする。
• 相対的位置バイアスを時間トークンと組み合わせることで、空間的なインダクティブバイアスを保持しつつ、時間的適応性を実現する。
• 条件付けのためのアダプティブレイヤーナーミャライゼーション（AdaLN）を避ける。代わりに、時間依存性を直接注意機構に埋め込む。

Figure 2 : Overview of the image-space DiffiT model. Downsample and Upsample denote convolutional downsampling and upsampling layers, respectively. Please see the supplementary materials for more information regarding the DiffiT architecture.

実験結果

リサーチクエスチョン

• RQ1 視覚トランスフォーマーは、サンプル品質を維持または向上させながら、拡散ベースの画像生成においてU-Netアーキテクチャを効果的に置き換えることができるか？
• RQ2 拡散モデルにおける自己注意機構を、異なるノイズ除去タイムステップに適応させることで、進化する画像構造をよりよく捉えることができるか？
• RQ3 時間依存行動を、単純な位置埋め込みや正規化層ではなく、注意機構そのものに直接統合することで、どのような影響が生じるか？
• RQ4 階層的設計を備えた統一されたトランスフォーマー基盤アーキテクチャは、画像空間および潜在空間両方の生成において、従来のU-NetおよびDiTスタイルのモデルを上回る性能を発揮するか？
• RQ5 提案された時間依存自己注意機構は、CIFAR10、FFHQ-64、ImageNet-256/512といった多様なベンチマークで、FIDスコアにどの程度の向上効果をもたらすか？

主な発見

• 提案された時間依存自己注意（TMSA）モジュールは、標準的な自己注意に置き換えることで、VE設定ではCIFAR10でFIDを0.28、VP設定では0.25低減した。
• 潜在DiffiTは、ImageNet-256データセットで1.73という新たな最先端FIDスコアを達成し、DDPM++やMDTといった先行手法を上回った。
• アブレーションスタディの結果、TMSAは代替手法よりも顕著に優れていた。TMSAから時間依存性を除去するとFIDは3.97に上昇し、時間トークンをMLP層にのみ適用した場合でもFIDは3.81に留まった。
• 質的注意マップの可視化では、TMSAが細粒度のオブジェクトディテールを捉え、ノイズ除去ステップに応じて動的に適応していることが示された。TMSAを搭載しないモデルとは対照的であった。
• パowerコサインスケジュールを用いた分類器フリーのガイドランスにより、ImageNet-256ではFIDが1.73に改善され、最適なガイドランススケールは4.6であった。これより高いスケールでは性能が劣化した。
• エンコーダーとデコーダーブロックを両方備えた完全なDiffiTアーキテクチャは、CIFAR10でFID1.95を達成し、高精度な生成に完全なU字型設計の重要性が示された。

Figure 3 : The DiffiT Transformer block applies linear projection to spatial and time-embedding tokens before combining them together to form query, key, and value vectors for each token. These vectors are then used to compute multi-head self-attention activations, followed by two linear layers.