[論文レビュー] Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

Patrick Esser, Sumith Kulal+15 more|arXiv (Cornell University)|Mar 5, 2024

Computer Graphics and Visualization Techniques被引用数 84

ひとこと要約

本論文はノイズスケールをバイアスしたサンプリングを導入し、別モダリティ重みを持つMM-DiTバックボーンを用いたマルチモーダルテキスト-画像トランスフォーマーと、最先端拡散モデルと比較して競合または上回る性能を示す拡張スケーリング研究により、 rectified flow モデルの高解像度画像合成を改善する。

ABSTRACT

Diffusion models create data from noise by inverting the forward paths of data towards noise and have emerged as a powerful generative modeling technique for high-dimensional, perceptual data such as images and videos. Rectified flow is a recent generative model formulation that connects data and noise in a straight line. Despite its better theoretical properties and conceptual simplicity, it is not yet decisively established as standard practice. In this work, we improve existing noise sampling techniques for training rectified flow models by biasing them towards perceptually relevant scales. Through a large-scale study, we demonstrate the superior performance of this approach compared to established diffusion formulations for high-resolution text-to-image synthesis. Additionally, we present a novel transformer-based architecture for text-to-image generation that uses separate weights for the two modalities and enables a bidirectional flow of information between image and text tokens, improving text comprehension, typography, and human preference ratings. We demonstrate that this architecture follows predictable scaling trends and correlates lower validation loss to improved text-to-image synthesis as measured by various metrics and human evaluations. Our largest models outperform state-of-the-art models, and we will make our experimental data, code, and model weights publicly available.

研究の動機と目的

• 知覚的に関連する内容へノイズスケールをバイアスすることにより、高解像度画像合成のための rectified flow の定式化を動機づけ、改善する。
• テキストと画像トークン間の双方向情報フローを可能にする、スケーラブルなマルチモーダルトランスフォーマーのバックボーンを開発し、テキスト-to-画像生成を改善する。
• データセットとサンプリング設定を横断して拡散と rectified-flow の variantes を体系的に比較し、優れた訓練・サンプリング戦略を特定する。
• 提案モデルの8Bパラメータまでのスケーリング挙動を実証し、検証損失と画像-テキスト評価指標との相関を評価する。

提案手法

• rectified flow モデルでノイズスケールを再重み付けし、知覚的に関連するタイムステップを重視する訓練をバイアスし、間置きタイムステップを強調するために重み付け可能なノイズ予測目的 (L_w) を生み出す。
• RF、EDM、LDM風スケジュールを含むバリアントを比較し、logit-normal、mode-based、CosMap のタイムステップ分布など、特注の SNR サンプラーを用いる。
• MM-DiT を導入、画像ストリームとテキストストリームのための2つの別々の重みセットを持つマルチモーダル拡散バックボーンで、クロスアテンションおよびMLP処理中の双方向相互作用を可能にする。
• 高解像度データでの事前学習とファインチューニングを QK 正規化で安定化させ、bf16 精度での訓練を可能にしつつ、再構成を改善するため latent チャンネルを拡張（d=16）する。
• 改良されたオートエンコーダ（潜在空間 d=16）、CogVLM による合成キャプションと元のキャプションを50/50で混合、そしてテキスト条件付き画像生成のためのモーダリティ毎にスケーラブルな拡散バックボーンを利用する。

実験結果

リサーチクエスチョン

• RQ1 中間で知覚的に関連するスケールへタイムステップサンプリングをバイアスすることは、標準的な拡散 formulationsと比べて高解像度画像合成における rectified flow の性能を改善するか？
• RQ2 画像とテキストトークンの別々のストリームを持つMM-DiTのマルチモーダルバックボーンは、従来の DiT、CrossDiT、UViT と比べてテキスト-to-画像生成で優れているか？
• RQ3 rectified-flow ベースのモデルにおけるスケーリング傾向はどのように現れ、検証損失の低下が自動評価および人間評価を含むテキスト-to-画像性能の向上へどのように結びつくか？
• RQ4 データ前処理とキャプション拡張（合成キャプションと元キャプションの併用）が、GenEvalスタイルの指標に対して大規模 T2I モデルでどのような影響を与えるか？
• RQ5 高解像度のファインチューニングに不可欠な訓練安定化技術（QK 正規化、混合精度ファインチューニング、さまざまなアスペクト比への位置エンコーディング）は何か？

主な発見

• 中間のタイムステップを強調するノイズサンプリング戦略（例: rf/lognorm(0.00, 1.00)）は、CLIPとFIDの指標で強い性能を発揮し、しばしば最先端の拡散モデルを上回るまたは同等となる。
• 適切に設計されたタイムステップサンプリングを用いた rectified flow のバリアントは、いくつかの設定でLDM-LinearおよびEDMのベースラインを上回り、特に低サンプリングステップで優れる。
• テキストと画像モダリティ別に別々の重みを持つ MM-DiT マルチモーダルバックボーンは、検証損失、CLIP、FIDで CC12M において vanilla DiT、CrossDiT、UViT を大幅に上回る。
• オートエンコーダの潜在チャネルを d=16 に拡大すると再構成指標が改善され、より良いスケーリングを支援する。容量の増加は画像品質の向上と相関する。
• 合成キャプション（CogVLM 生成）と元のキャプションを50/50で混ぜると GenEval スコアが改善され、合成キャプションが訓練データを効果的に増強できることを示唆する。
• QK正規化による訓練安定化は bf16 混合精度での安定したファインチューニングを可能にし、高解像度スケーリングとより良い注意機構の挙動を促進する。
• 8Bパラメータまでのスケーリング実験は、検証損失が低いほど自動評価と人間評価の両方でテキスト-to-画像性能が向上することを示す。