Switti: Designing Scale-Wise Transformers for Text-to-Image Synthesis

テキスト条件付き拡散モデル（DM）[4, 56, 15, 59, 51, 92, 3]は、テキストから画像（T2I）生成のデファクトスタンダードとなっている。 その印象的な性能にもかかわらず、DMの周知の制限として、逐次推論が遅いことが挙げられ、これがリアルタイムや大規模な生成タスクの妨げとなっている。 公開されている最先端のT2I拡散モデルの大半[4, 51, 15, 58, 92, 39]は、VAE [33]の潜在空間で動作し、高解像度画像のより効率的なサンプリングを可能にしている。 しかし、これらのモデルでも潜在空間において$20{-}50$の拡散ステップを必要とする。

拡散蒸留法[49, 67, 46, 60, 61, 85, 84, 32]は、拡散ステップ数をわずか$2{-}4$に削減するための最も有望な方向性である。 DMD2 [84]やADD [60]などの現在の最先端アプローチは、$4$ステップで強力な生成性能を示しており、実画像に対する追加の敵対的訓練のおかげで、画質の面で教師モデルの性能を上回る可能性さえある。

自己回帰（AR）モデルは、画像生成のための有望な代替パラダイムであり、主に3つのグループに分類される：次トークン予測 [68, 14, 87, 38]、次スケール予測 [72, 47, 69]、およびマスクされた自己回帰モデル [41, 16]である。

次トークン予測ARモデルは、GPT様の因果的トランスフォーマー [53, 54, 13, 73] に類似しており、ラスター順序（左から右、上から下）などの走査戦略を用いてトークンごとに画像を生成する。 トークンは通常、VQ-VAEベースの離散画像トークナイザー [74, 14, 86, 38] を用いて取得される。 VQ-VAEは画像を低解像度の2D潜在空間にマッピングし、各潜在「ピクセル」を学習された語彙の要素に割り当てる。

マスクされた自己回帰画像モデリング（MAR） [41, 16] は、マスクされた画像生成モデル [6, 40, 7] を拡張し、1つのステップでランダムな順序で複数のマスクされたトークンを予測する。 特筆すべきは、MARが連続的なトークンで動作し、訓練には拡散損失を使用し、トークンのサンプリングには軽量なトークン単位の拡散モデルを使用することである。 Fluid [16] はこのアプローチをT2I生成に適用し、そのスケーリング挙動を探究している。

VAR [72] によって導入された次スケール予測ARモデリングは、画像を異なる解像度のスケールの連続として表現する。 次トークン予測やマスクされたARモデリングとは異なり、スケール単位のトランスフォーマーは、以前に生成された低解像度のスケールに注意を払いながら、より高い解像度のすべてのトークンを並行して予測する。

画像をスケールの連続として表現するために、VAR [72] は階層的VQ-VAEを使用し、残差量子化（RQ） [38] を用いて段階的に構築された異なる解像度（スケール）の潜在変数のピラミッドに画像をマッピングする。 以下では、このVAEモデルを RQ-VAE と呼ぶ。 RQ-VAEの各潜在変数は、単層VQ-VAEと同様に、すべてのスケールで共有される語彙からの離散トークンのセットに関連付けられる。

サンプリング時、スケール単位のARモデル $\theta$ は反復的にスケールごとに画像トークンを予測し、以下のように定式化される：

ここで、$s_{i}$ は現在のスケールのRQ-VAEトークンを表し、$N$ はスケールの総数、$c$ は条件付け情報である。 このモデルは、図 5（左）に示すように、ブロック単位の因果的注意マスクを持つトランスフォーマー [54] である。 VAR [72] は DiT [50] からトランスフォーマーアーキテクチャを採用している。

最近の研究では、次スケール予測モデルをT2I生成に適用している [47, 90, 69]。 STAR [47] は VARから事前学習されたRQ-VAEモデルを使用し、テキスト条件付けを効果的に処理するようにジェネレータを修正している。 本稿執筆時点でSTARはまだリリースされていないが、我々はSTARをベースラインアーキテクチャとして考え、そこから段階的に提案モデルである Switti に向けて進展させている。

並行して行われた研究であるHART [69] は、わずか $0.7$B パラメータの軽量なT2Iスケール単位ARモデルを提案している。 これは主に、VARの離散的RQ-VAEの限界に対処するために、連続的な誤差残差をモデル化する追加の拡散モデルを導入することで、ハイブリッドモデルを実現している：スケール単位のARモデルと再構成された潜在変数を精緻化するための拡散モデルの組み合わせである。 対照的に、我々はVAR [72] から事前学習されたRQ-VAEを使用し、$512{\times}512$ 解像度用に微調整しながら、スケール単位の生成トランスフォーマーの設計にのみ焦点を当てている。 我々のスケール単位の生成モデル設計とHARTのハイブリッドトークン化を組み合わせることは、将来の研究の有望な方向性となり得る。

MAR [41] も自己回帰と拡散の事前分布を組み合わせたハイブリッドモデルと考えることができる。 Disco-diff [83] は次トークン予測トランスフォーマーによって生成された離散トークンに条件付けられた拡散モデルを使用する。 DART [19] は非マルコフ的拡散モデルのバックボーンとしてARトランスフォーマーを導入している。 これらの研究とは対照的に、Switti は拡散事前分布を一切使用していない。

出発点として、我々はVAR [72]とSTAR [47]に密接に従った基本的なテキスト条件付きアーキテクチャを設計する。 スケール方向の自己回帰型テキスト画像生成パイプラインは、主に3つの要素で構成される：画像トークナイザーとしてのRQ-VAE [38]、事前学習済みテキストエンコーダー [55]、そしてスケール方向ブロック方向の因果的トランスフォーマー [72]である。

我々のモデルは、VAR [72]から事前学習済みRQ-VAEを採用しており、これは画像を$512{\times}512$スケールで表現する。 セクション 4.3で議論されるように、我々は$512{\times}512$解像度でわずかに調整を行う。

結果として得られるモデルの強力な画像-テキスト整合性を確保するために、我々はT2I拡散モデリングの文献 [51, 15] に従い、2つのテキストエンコーダーを採用する：CLIP ViT-L [55] とOpenCLIP ViT-bigG [27]である。 各モデルから抽出されたテキスト埋め込みはチャネル軸に沿って連結される。

基本的なトランスフォーマーアーキテクチャはVAR [72]から採用されており、我々は各トランスフォーマーブロックの自己注意層とフィードフォワードネットワーク（FFN）の間にクロス注意 [75] 層を組み込む。 OpenCLIP ViT-bigGからのプールされた埋め込みは、適応型層正規化（AdaLN） [81] を介してトランスフォーマーブロックに伝播される。

我々はまた、SDXL [51] に従って、生成された画像における意図しないオブジェクトのクロッピングを緩和するために、クロッピングパラメータに対する条件付けを組み込む。 具体的には、中心クロップ座標 $c_{top},c_{left}$ をフーリエ特徴埋め込みに変換し、それらを連結する。 次に、得られたベクトルをOpenCLIP ViT-bigGの隠れサイズに線形層を介してマッピングし、プールされた埋め込みに加算する。

層正規化（LN）層 [1] は注意層とFFNブロックの入力に適用される。 RMSNorm [89] 層はクエリ-キー（QK）正規化に使用される。 我々はまた、正規化された2D回転位置埋め込み（RoPE） [47, 22] を使用し、これによりモデルのより高い解像度への適応が速くなる。 FFNブロックはGeLU活性化関数 [21] を使用する。 基本アーキテクチャの可視化は付録 Aにある。

ここでは、$d{=}20$個のトランスフォーマーブロックを持つ基本モデルの学習性能を分析し、モデルの安定性と収束性を向上させる修正点を紹介する。

我々は、セクション 4.1で説明した$256{\times}256$画像-テキストデータセットを用いて、混合精度BF16/FP32で$150$K回のイテレーションでモデルを学習させる。 詳細な学習および評価のセットアップは付録 Bに記載している。 学習中、我々は活性化ノルムおよびFID [24]、CLIPスコア [23]、PickScore [34]などの標準的な指標を追跡する。

まず、我々は学習中の安定性の問題を観察し、大規模実験での最終的な発散や最適以下の性能につながることを確認した。 我々の調査により、この問題の根源はトランスフォーマーの活性化ノルムの急速な成長にあることが明らかになった。これは図 3（青）に示されている。 最後のトランスフォーマーブロックの活性化ノルムは学習イテレーションを通じて成長し、$10^{16}$という極めて大きな値に達する。

したがって、学習を安定させるための最初のステップは、学習中にモデルヘッドをFP32にキャストすることである。 我々は、これが活性化ノルムを大幅に減少させる重要な技術的詳細であり、図 4（オレンジ）に示すように、はるかに良い収束をもたらすことを発見した。 しかし、この手法は問題を完全に解決するものではない。なぜなら、活性化ノルムは依然として成長し続け、学習の終わりには$10^{4}$という高い値に達するからである。

学習中の活性化ノルムの成長をさらに抑えるために、我々は「サンドイッチ」のような正規化 [12, 92]を採用し、活性化ノルムを適切な範囲に保つ。 具体的には、各注意機構とフィードフォワードブロックの直後に追加の正規化層を挿入し、効率性のためにLN層をRMSNormに置き換える。 図 3に示すように、この修正によって学習中の活性化ノルムの成長がさらに抑制され、図 4（緑）に示すように、わずかにモデルの性能が向上する。

最後に、言語トランスフォーマーの標準的な実践 [73, 13]に従い、FFNブロックのGELU [21]活性化をSwiGLU [63]活性化に置き換える。これにより、学習可能なゲーティングメカニズムを通じて情報の流れを動的に制御することができる。 SwiGLUはLLMでは一般的な選択肢であるが、我々は最終的な性能にはほとんど影響を与えないことに気付いた。

全体として、我々は図 2に記述したアーキテクチャのトランスフォーマーブロックを示し、提案されたアーキテクチャを持つスケール単位のARモデルをSwitti (AR)と呼ぶ。

以下では、事前学習されたSwitti (AR)の異なるスケールにおける注意マップを分析する。 具体的には、$i\geq j$となる各スケールのペア$s_{i}$と$s_{j}$に対して、$s_{i}$から$s_{j}$への平均注意スコアを$\overline{A}_{s_{i},s_{j}}=\frac{1}{|s_{i}|}\sum_{m\in s_{i},n\in s_{j}}A_{mn}$という式を用いて計算する。 ここで、$A_{mn}$は元の自己注意マップを表し、$|s_{i}|$はスケール$s_{i}$におけるトークン数を示す。 テキストプロンプトのバッチにわたって平均化された可視化結果を図6に示す。 画像トークン間の自己注意は主に現在のスケールに集中しており、先行するスケールに対しては著しく弱いことが観察できる。

この注意の挙動に基づき、我々は自己注意層が現在のスケールのトークンにのみ注意を向けるように注意マスクを更新することを提案する（図5（右）参照）。 これは、トランスフォーマーがもはやブロック単位で因果的でなくなることを意味し、より安価な注意操作によって効率的なサンプリングが可能になり、キー値（KV）キャッシュの必要性も排除される。 興味深いことに、この修正はCLIPスコアとPickScoreの観点からパフォーマンスをわずかに向上させることも図4に示されている。 総じて、我々は提案する非因果的トランスフォーマーアーキテクチャを持つスケール単位のモデルをSwittiと呼ぶ。

最後に、我々は異なるモデルスケールにおけるテキスト条件付けの効果を検証する。 具体的には、クロスアテンションマップとサンプリング中にテキストプロンプトを切り替えた際のモデルの挙動を分析する。