Locality-aware Parallel Decoding for Efficient Autoregressive Image Generation
Locality-aware Parallel Decoding for Efficient Autoregressive Image Generation
基本情報
- arXiv ID: 2507.01957v1 (https://arxiv.org/abs/2507.01957)
- 著者: Zhuoyang Zhang、Luke J. Huang(共同筆頭著者)、Chengyue Wu、Shang Yang、Kelly Peng、Yao Lu、Song Han
- 所属: MIT、NVIDIA、First Intelligence
- 投稿日: 2025年7月2日
- カテゴリ: cs.CV
簡単に説明すると
この論文は、画像を生成するAIの処理速度を3.4倍以上向上させる「Locality-aware Parallel Decoding (LPD)」という新しい手法を提案しています。
従来の画像生成AIは、画像を左上から右下へ1つずつのパッチ(小領域)を順番に生成していました。
しかしLPDは、複数のパッチを同時に生成できる仕組みを導入しています。
例えば、家の絵を描くとき、通常は屋根から順番に描いていきますが、LPDでは屋根の一部と窓、ドアなどを同時に描けるようになります。
ただし、同時に描く部分は互いに離れた場所を選ぶことで、描画の一貫性を保っています。
この手法により、256×256サイズの画像生成で必要なステップ数を256から20に、512×512では1024から48に削減し、品質を維持しながら3.4倍以上の高速化を実現しています。
1. 研究概要
1.1 背景と動機
大規模言語モデル(LLMs)で成功を収めた自己回帰モデリングは、視覚生成タスクにも広く適用されています。
視覚生成における自己回帰モデルは、画像を1次元のトークン列に変換し、各トークンを順次予測する方式を採用しています。
この手法は、言語モデリングとの高い互換性により、統一的なマルチモーダル生成システムの構築において重要な役割を果たしています。
しかし、既存の自己回帰型の視覚生成手法には2つの主要な課題があります。
第一に、1ステップごとに1トークンしか生成できないため、メモリバウンドなワークロードとなり、生成ステップ数に比例して遅延が増大します。
第二に、高解像度の画像生成では、トークン数が解像度の2乗に比例して増加するため、計算コストが急激に上昇します。
最近の研究では、並列化による高速化が試みられていますが、4倍程度の高速化が限界でした。
非自己回帰的なMaskGITのような手法も存在しますが、双方向コンテキストのための完全なアテンション計算が必要で、効率性に課題があります。
1.2 主要な貢献
本研究では、以下の3つの主要な貢献をしています。
- Flexible Parallelized Autoregressive Modelingという新しいアーキテクチャを提案する。
これにより任意の生成順序と並列度を実現可能にする - Locality-aware Generation Order Scheduleという新しいスケジューリング手法を設計する。
これは空間的局所性を考慮して並列生成グループを形成し、文脈サポートを最大化しながら相互依存性を最小化する - ImageNetデータセットでの評価により、生成ステップ数を12.8倍から21.3倍削減することを実証する。
具体的には256×256で256→20ステップ、512×512で1024→48ステップへの削減を達成
2. 提案手法
2.1 手法の概要
LPDは、2つの核心的な技術要素から構成されています。
第一に、Flexible Parallelized Autoregressive Modelingは、従来のデコーダ専用アーキテクチャを拡張し、任意の生成順序と並列度をサポートします。
これは、学習可能な位置クエリトークンを使用して、目標位置でのトークン生成を誘導することで実現されます。
位置クエリトークンは、ターゲット位置の位置埋め込みを共有された学習可能な埋め込みに加えることで構築されます。
第二に、Locality-aware Generation Order Scheduleは、画像生成における強い空間的局所性を活用します。
このスケジューリングは、2つの原則に基づいています。
すなわち、ターゲット位置は既存のコンテキストに空間的に近いことで強い条件付けを確保し、同じ並列ステップで予測されるトークンは空間的に離れていることで相互依存性を削減します。
2.2 技術的詳細
Flexible Parallelized Autoregressive Modelingの核心は、コンテキスト表現とトークン生成の役割を分離することにあります。
従来のモデルでは、各トークンが将来のトークンへのコンテキスト提供と次トークンの生成の両方を担っていました。
提案手法では、以前に生成されたトークンがコンテキストを提供し、位置クエリトークンが並列生成を駆動します。
訓練時には、特殊な訓練用アテンションマスクを使用します。
このマスクは2つのアテンションパターンを含みます。
Context Attentionは、後続のトークンがコンテキストトークンに因果的にアテンションできるようにします。
Query Attentionは、同じステップ内の位置クエリトークン間の相互可視性を確保し、後続のトークンがクエリトークンにアテンションすることを防ぎます。
推論時には、生成された画像トークンのエンコーディングと位置クエリトークンによるデコーディングを交互に実行します。
これらの操作は、特殊な推論用アテンションマスクを使用して単一ステップに融合できます。
Locality-aware Generation Order Scheduleは、LlamaGenモデルの50,000枚の画像生成時のアテンション分析に基づいています。
分析により、トークンのアテンションは近傍領域に集中し、距離とともに急速に減衰することが判明しました。
この知見に基づき、各ステップで生成するトークンを選択するアルゴリズムを設計しました。
2.3 新規性
既存手法との主な違いは、以下の点にあります。
第一に、提案するFlexible Parallelized Autoregressive Modelingは、並列生成されるトークン間の相互可視性を保証します。
SARやARPGなどの既存手法では、並列生成されるトークンが互いに独立していましたが、本手法では全ての同時予測ターゲット位置間の可視性を確保しています。
第二に、KVキャッシュに生成されたトークンのみを保存し、位置クエリトークンは保存しません。
RandARなどの手法では、位置指示トークンもKVキャッシュに保存する必要があり、メモリ消費が2倍になっていました。
第三に、空間的局所性を明示的に考慮した生成順序スケジューリングを導入しています。
従来の手法は固定された並列順序を使用していたため、並列化と生成品質が制限されていました。
3. 実験結果
3.1 実験設定
実験は、ImageNetデータセットのクラス条件付き画像生成タスクで実施されました。
256×256と512×512の2つの解像度で評価しました。
モデルは3つのサイズで訓練されました。
LPD-L(337Mパラメータ)、LPD-XL(752Mパラメータ)、LPD-XXL(1.4Bパラメータ)です。
全てのモデルは標準的なデコーダ専用トランスフォーマーアーキテクチャを使用し、LlamaGenトークナイザーを採用しています。
評価指標として、50,000枚の生成画像に対するFID(Fréchet Inception Distance)を主要指標としました。
IS(Inception Score)、Precision、Recallも報告しています。
効率性の評価は、NVIDIA A100 GPUでBFloat16精度を使用して実施しました。
3.2 主要な結果
ImageNet 256×256での実験結果では、LPDは生成ステップ数を256から20に削減しながら、品質を維持することに成功しました。
LPD-XLモデルは、わずか20ステップでFID 2.10を達成しました。
これはARPGと比較して3.2倍のステップ数削減と4.2倍の低遅延を実現しています。
32ステップに増やすとFID 1.92となり、ARPG-XXLと同等の品質を達成しながら、3.4倍の低遅延を実現しました。
ImageNet 512×512では、生成ステップ数を1024から48に削減し、21.3倍のステップ数削減を達成しました。
LPD-XXLモデルは48ステップでFID 2.21を記録し、従来手法と比べて21.3倍少ないステップ数で同等の品質を達成しました。
遅延の測定結果では、LPD-XXLは256×256画像生成で0.37秒、512×512で0.82秒を達成しました。
これは、ラスター順序の自己回帰モデルと比較して、それぞれ10.5倍、19.9倍の高速化に相当します。
3.3 既存手法との比較
他の並列化型の自己回帰モデルとの比較では、LPDが全ての面で優れた性能を示しました。
ARPGと比較すると、同等のモデルサイズにおいて、より少ないステップ数で高品質な生成を実現しています。
例えば、LPD-XL(32ステップ)はARPG-XL(64ステップ)よりも優れたFIDを達成しました。
具体的には、半分のステップ数でより良い画質を実現しています。
PARやNARなどの固定並列順序を使用する手法と比較して、LPDの柔軟な生成順序により大幅な品質向上を実現しました。
特に、高並列度での品質劣化が抑制されています。
非自己回帰的なMaskGITと比較しても、LPDは必要なアテンション計算量を削減することで、同等の品質を維持しながら高速化を達成しています。
4. 実用性評価
4.1 実装の容易性
LPDは標準的なトランスフォーマーアーキテクチャに基づいており、実装は比較的容易です。
位置クエリトークンの導入と特殊なアテンションマスクの実装が主な変更点となります。
生成順序スケジュールは事前に計算して保存できるため、推論時の追加遅延はありません。
論文ではPyTorchによる実装例も提供されており、既存のフレームワークへの統合も容易です。
4.2 計算効率
メモリ効率の面では、KVキャッシュに生成されたトークンのみを保存するため、従来手法と同等のメモリ使用量を維持しています。
位置クエリトークンは推論時、KVキャッシュへの保存が不要なため、追加のメモリオーバーヘッドはありません。
計算効率では、並列生成により大幅な高速化を実現しています。
バッチサイズ1での遅延測定では、256×256画像で10.5倍、512×512画像で19.9倍の高速化を達成しました。
スループットも同様に向上し、大規模な画像生成タスクでの実用性を示しています。
4.3 応用可能性
LPDは任意の生成順序をサポートするため、様々な画像編集タスクに自然に拡張できます。
実験では、インペインティング(画像の一部を修復)、アウトペインティング(画像を拡張)、クラス条件付き編集などのゼロショット画像編集タスクを実証しました。
これらのタスクでは、編集対象外の領域のトークンを事前にKVキャッシュへ入力します。
その後、編集領域をランダムな順序で生成します。
また、提案手法は他の視覚生成タスクやマルチモーダル生成システムへの統合も期待できます。
特に、フラットなトークン表現を維持しているため、CLIPやDINOなどの既存の視覚認識基盤モデルとの互換性も保たれています。
5. まとめと所感
5.1 論文の意義
この研究は、自己回帰型の画像生成の効率性を10倍以上向上させる重要な進歩を示しています。
特に、品質を維持しながら10倍以上の高速化を実現したことは、実用的な観点から非常に価値があります。
理論的には、空間的局所性を明示的に考慮した生成順序の設計という新しい視点を提供しています。
アテンション分析に基づいた原理的なアプローチは、他の生成タスクへの応用も期待できます。
実用的には、高解像度の画像生成の計算コストを最大21.3倍削減し、リアルタイムアプリケーションへの道を開いています。
また、ゼロショット編集能力により、単一のモデルで多様なタスクに対応できる汎用性も示しています。
5.2 今後の展望
今後の改善点として、以下が考えられます。
第一に、より大規模なモデルや高解像度(1024×1024以上)での評価が期待されます。
現在の実験は512×512までですが、さらなる高解像度での効率性向上の可能性があります。
第二に、動画生成への拡張が興味深い研究方向です。
時間的な局所性も考慮することで、現在の画像生成と同様の高速化を動画生成でも実現できると考えられます。
第三に、他のモダリティ(テキスト、音声など)との統合により、より汎用的なマルチモーダル生成システムの構築が期待されます。
提案手法の柔軟性は、このような統合において有利に働く可能性があります。