[論文レビュー] SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers

Enze Xie, Jun Song Chen+9 more|arXiv (Cornell University)|Oct 14, 2024

Advanced Vision and Imaging被引用数 4

ひとこと要約

Sanaは、32倍の深さの圧縮自己符号化器、DiTアーキテクチャにおける線形自己注意機構、文脈内命令微調整を施したデコーダー専用言語モデル（Gemma）によるテキスト符号化、および推論ステップを削減するフローベースのDPMソルバーを用いることで、画質が最先端でありながら、前例のない高速さで4096×4096解像度の画像を生成する、極めて効率的なテキストから画像への拡散フレームワークを導入する。Sana-0.6Bは、FLUXなどの最先端モデルと比較して100倍速く推論を実行でき、モデルサイズは20倍小さく、16GBのラップトップGPUでも動作可能である。

ABSTRACT

We introduce Sana, a text-to-image framework that can efficiently generate images up to 4096$ imes$4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU. Core designs include: (1) Deep compression autoencoder: unlike traditional AEs, which compress images only 8$ imes$, we trained an AE that can compress images 32$ imes$, effectively reducing the number of latent tokens. (2) Linear DiT: we replace all vanilla attention in DiT with linear attention, which is more efficient at high resolutions without sacrificing quality. (3) Decoder-only text encoder: we replaced T5 with modern decoder-only small LLM as the text encoder and designed complex human instruction with in-context learning to enhance the image-text alignment. (4) Efficient training and sampling: we propose Flow-DPM-Solver to reduce sampling steps, with efficient caption labeling and selection to accelerate convergence. As a result, Sana-0.6B is very competitive with modern giant diffusion model (e.g. Flux-12B), being 20 times smaller and 100+ times faster in measured throughput. Moreover, Sana-0.6B can be deployed on a 16GB laptop GPU, taking less than 1 second to generate a 1024$ imes$1024 resolution image. Sana enables content creation at low cost. Code and model will be publicly released.

研究の動機と目的

• 4K解像度における、既存の高解像度テキストから画像への拡散モデルの高い計算コストと遅い推論速度を解決すること。
• ラップトップの16GB GPUメモリを備えた一般消費者向けハードウェアでも、高精細な4K画像の効率的かつリアルタイムの推論を可能にすること。
• 画像品質やテキスト-画像アライメントを損なわずに、モデルサイズと推論遅延を低減すること。
• 従来のテキスト符号化器（例：CLIP、T5）の制限を克服し、デコーダー専用LLMの指示従いおよび推論能力を活用すること。
• 最小限のハードウェア要件で、超高解像度画像合成を可能にするスケーラブルなエンドツーエンドのトレーニングおよび推論パイプラインを開発すること。

提案手法

• 画像の空間的次元を32倍圧縮する深さ圧縮自己符号化器（AE-F32）を提案し、従来の8倍圧縮自己符号化器と比較して、潜在トークン数を16倍削減した。
• DiT内のすべての標準自己注意モジュールを線形自己注意機構に置き換え、高解像度生成における計算複雑度をO(N²)からO(N)に低減した。
• MLPブロックに統合された3×3のディープワイズ畳み込みであるMix-FFNを導入し、局所的トークン集約を可能にし、位置エンコーディングの必要性を排除した（NoPE）。
• T5に代えて、小規模なデコーダー専用言語モデル（Gemma）をテキスト符号化器として採用し、複雑な人間の指示（CHI）プロンプトとコンテキスト内学習を用いて、指示従いの精度と画像-テキストアライメントを向上させた。
• クリップスコアに基づく動的キャプション選択を備えたマルチVLMキャプションパイプラインを設計し、トレーニングデータの品質と収束性を向上させた。
• フローベースのサンプリング手法としてFlow-DPM-Solverを提案し、推論ステップを28–50から14–20に削減しながら、生成品質を向上させ、遅延を低減した。

実験結果

リサーチクエスチョン

• RQ1 コンactなアーキテクチャを用いて、高品質かつ低遅延で4K（4096×4096）解像度の画像生成を実現できるか？
• RQ2 標準自己注意機構をDiTで置き換えることで、品質劣化を伴わず、高解像度推論を効率的に行えるか？
• RQ3 適切に複雑な人間の指示でプロンプトされた小規模なデコーダー専用LLM（例：Gemma）が、T5やCLIPと比較してテキスト-画像アライメントで優れるか？
• RQ4 高度な自己符号化、効率的自己注意、最適化されたサンプリングの組み合わせにより、最先端モデルと比較して推論時間を100倍以上短縮できるか？
• RQ5 16GBのラップトップGPUに高解像度拡散モデルをデプロイし、1024×1024画像の推論を1秒未塔で実行することは可能か？

主な発見

• Sana-0.6Bは、16GBのラップトップGPU上で1024×1024画像を1秒未塔で生成し、リアルタイム推論の能力を示した。
• Sana-0.6Bは、A100 GPU上で4096×4096解像度の画像生成において、FLUXと比較して106倍速く推論を実行し、遅延を469秒から9.6秒に短縮した。
• Sana-0.6Bは、1024×1024画像生成においてFIDが5.81、CLIPスコアが28.36を達成し、SDXL や PixArt-Σ と比較して両方の指標で優れた性能を示した。
• Gemma-2Bをテキスト符号化器として使用した場合、FIDが5.9、CLIPスコアが26.8を達成し、同等または低い遅延でT5-LargeやT5-XXLと同等またはそれを上回る性能を示した。
• Flow-DPM-Solverにより、サンプリングステップを28–50から14–20に削減し、推論速度を向上させながら、生成品質を維持または向上させた。
• Mix-FFNを用いたNoPE（位置エンコーディングなし）設計により、位置エンコーディングなしで効果的なモデリングが可能となり、アーキテクチャを単純化し、トレーニングの安定性を向上させた。