技術概要：FlashAR

問題定義

大規模な自己回帰（AR）モデルは、視覚データを離散トークン系列として扱うことで、画像生成において最先端の性能を達成しています。しかし、標準的な AR モデルはラスタースキャン復号に依存しており、トークンを厳密に左から右、上から下へと逐次的に生成します。この逐次的な依存関係により、推論遅延は画像解像度に対して二次的に（$O(HW)$）スケールし、高解像度生成が著しく遅くなるだけでなく、現代の GPU 並列性の効果的な活用を阻害しています。

既存の加速戦略は重大なトレードオフに直面しています：

1. パラダイムの再設計：トークン予測順序を変更する方法（VAR、NAR など）やマルチスケール生成を採用する手法は、多くの場合ゼロから再学習を必要とし、膨大な既存の事前学習済みラスタ-AR モデルのレポジトリと互換性がありません。
2. 学習後適応：離散拡散適応などのアプローチは、元の予測目的を変更し、学習と推論の間にギャップを生み出し、因果モデルに固有の微細な構造的整合性を劣化させる可能性があります。

本論文が取り組む中心的な課題は以下の通りです：「事前学習済みラスタースキャン AR モデルを、学習済み生成事前分布を保持し、学習後のオーバーヘッドを最小化しながら、高並列生成器へと変換することは可能か？」

手法：FlashAR

FlashAR は、事前学習済みラスタースキャン AR モデルを双方向次トークン予測構造を用いて並列生成器に変換する、軽量な学習後適応フレームワークです。このフレームワークは、以下の 3 つの中核技術コンポーネントを導入します：

1. 双ヘッド復号のための中間ブランチ化

最終層に新しい予測ヘッドを接続する（これは元の水平ラスタースキャン目的にバイアスがかかる）のではなく、FlashAR は中間層（$m$）でネットワークを分岐させます。

• アーキテクチャ：事前学習済みのバックボーンは、共有トランク（$F_{0:m-1}$）と 2 つの並列パスに分割されます。
  • 水平ヘッド：行方向予測のために、元の上位層（$F_{m:L-1}$）を再利用します。
  • 垂直ヘッド：列方向予測のために、上位層のクローンとして初期化された学習可能なブロック（$\tilde{F}_{m:L-1}$）を利用します。
• 根拠：中間特徴量はより豊かな意味的抽象化を保持し、最終層表現ほど特定の左から右への方向に特化していません。ここで分岐することで、垂直パスは必要な文脈にアクセスしつつ、クリティカルパスの深さを延長することなく、実行時の効率性を維持できます。

2. 学習可能な融合ゲート

水平予測と垂直予測は相補的な依存関係を捉えますが、その相対的な重要性は空間的位置（例えば、水平エッジ対垂直エッジ）によって異なります。静的な平均化では不十分であり、文脈の衝突を引き起こす可能性があります。

• メカニズム：FlashAR は、2 つの先行状態（$h^H_{p,q-1}$ および $h^V_{p-1,q}$）の文脈に基づいて予測を動的に結合する学習可能な融合ゲートを採用します。
• 数式：融合されたログオッズ $z_{p,q}$ は、$g_{p,q} z^H_{p,q-1} + (1 - g_{p,q}) z^V_{p-1,q}$ として計算されます。ここでゲート $g_{p,q}$ は MLP とシグモイド関数を通じて導出されます。これにより、モデルは方向性の手がかりを適応的に重み付けできます。

3. 2 段階安定化適応パイプライン

急速な収束を確保し、事前学習済み多様体の破滅的忘却を防ぐため、FlashAR は 2 段階の学習戦略を使用します：

• ステージ 1（ブランチ初期化）：バックボーンと水平ヘッドは固定されます。垂直ヘッドと融合ゲートのみが学習されます。これにより、確立された水平事前分布を乱すことなく、新しい垂直パスが予測能力を迅速に獲得できます。
• ステージ 2（共同微調整）：ネットワーク全体（バックボーン、両方のヘッド、ゲート）が共同で微調整され、新しい対角並列復号パラダイムに対して直交する空間表現を調和させます。

4. ハードウェア対応推論

理論上の並列性をウォールクロック時間の高速化に変換するため、著者らは以下を実装しています：

• FlexAttention：密なゼロパディングによるメモリ断片化を回避するため、高密度な 2 次元近接マスクをオンザフライで動的にコンパイルします。
• バッチ化 KV キャッシュ：対角ステップですべてのトークンを同時に更新し、カーネル起動のオーバーヘッドを均質化します。

主要な貢献

1. FlashAR フレームワーク：ゼロから事前学習するコストなしに、事前学習済みラスタ-AR モデルを並列生成器に変換する学習後適応手法。
2. 双ヘッド中間ブランチ化：中間層で分岐することにより、水平事前分布の保持と直交する垂直予測の導入という緊張関係を解決する新しいアーキテクチャ。
3. 動的融合：水平信号と垂直信号を適応的に結合することで、空間的不整合を解決する学習可能なゲート。
4. 効率的な学習パイプライン：最小限のデータ要件で高速な収束を確保する 2 段階の適応プロセス。

実験結果

著者らは、LlamaGen（ImageNet 256×256）およびEmu3.5（テキストから画像 512×512）において FlashAR を評価しました。

• 高速化：FlashAR は、標準的な AR ベースラインと比較して、512×512 画像生成において最大22.9 倍のウォールクロック高速化を達成しました。LlamaGen においては、447.2 画像/秒のスループットを達成し（標準 AR は 117.9）、大幅な向上を示しました。
• データ効率：この適応には、元の学習データのわずか**0.05%**のみが必要です（例：Emu3.5 では 8 万トークン対、事前学習トークン 1500 億）。
• 品質の維持：
  • LlamaGen：FlashAR-L は FID 3.16、IS 289.0 を達成し、BlockDiffusion（FID 4.55）を上回るとともに、ゼロから学習した NAR-L をも Inception Score で凌駕しました。
  • Emu3.5：FlashAR は復号ステップを 1024 から 63 に削減しました。GenEval スコアはわずかに低下するのみ（80.48 → 80.29）で、「色」および「位置」の指標においてベースラインを上回りました。
• 収束：FlashAR は BlockDiffusion よりも著しく速く収束し、BlockDiffusion の 75 エポックと比較して 25 エポックで優れた FID を達成しました。

意義と主張

本論文は、FlashAR が並列復号の効率性と事前学習済み自己回帰モデルの忠実性の間のギャップを成功裏に埋めたと主張しています。元の学習目的への修正を最小限に抑えることで、FlashAR は既存の大規模モデルが持つ強力な生成能力を保持します。著者らは、このアプローチが、ゼロから新しい並列アーキテクチャを学習する際の禁止的なコストや、目的変更型適応手法でしばしば見られる品質の劣化を伴わずに、リアルタイムかつ高解像度の画像生成に対する実用的な解決策を提供すると強調しています。

この研究は、モデルアーキテクチャが中間ブランチ化と動的融合を通じて、事前学習済みモデルの固有のバイアスを慎重に尊重するように適応されれば、対角並列復号が自己回帰生成のための実現可能かつ効率的なパラダイムであることを示しています。