[論文レビュー] Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think

Sihyun Yu, Sangkyung Kwak+5 more|arXiv (Cornell University)|Oct 9, 2024

Reinforcement Learning in Robotics被引用数 4

ひとこと要約

本論文は、DINOv2などの事前学習済み自己教師付きエンコーダーから得られる高品質な外部視覚特徴と、ノイズの多い潜在表現を一致させるシンプルな正則化手法であるRepresentation Alignment (REPA) を提案する。この外部特徴を拡散モデルの初期層に蒸留することで、REPAは学習を17.5倍以上高速化し、分類器フリーのガイドランスを使用した際のSOTA FID=1.42を達成し、学習効率と生成品質の両方を顕著に向上させる。

ABSTRACT

Recent studies have shown that the denoising process in (generative) diffusion models can induce meaningful (discriminative) representations inside the model, though the quality of these representations still lags behind those learned through recent self-supervised learning methods. We argue that one main bottleneck in training large-scale diffusion models for generation lies in effectively learning these representations. Moreover, training can be made easier by incorporating high-quality external visual representations, rather than relying solely on the diffusion models to learn them independently. We study this by introducing a straightforward regularization called REPresentation Alignment (REPA), which aligns the projections of noisy input hidden states in denoising networks with clean image representations obtained from external, pretrained visual encoders. The results are striking: our simple strategy yields significant improvements in both training efficiency and generation quality when applied to popular diffusion and flow-based transformers, such as DiTs and SiTs. For instance, our method can speed up SiT training by over 17.5$ imes$, matching the performance (without classifier-free guidance) of a SiT-XL model trained for 7M steps in less than 400K steps. In terms of final generation quality, our approach achieves state-of-the-art results of FID=1.42 using classifier-free guidance with the guidance interval.

研究の動機と目的

• 大規模な拡散変換器における内部表現学習の質の低さに起因する、遅く非効率な学習の課題に対処すること。
• 高品質な外部視覚特徴を組み込むことで、拡散モデルの学習効率と生成品質が顕著に向上するかどうかを調査すること。
• ノイズの多い拡散入力と、綺麗な画像の事前学習済み視覚エンコーダーとの間の入力不一致を、シンプルで効果的な正則化戦略によって克服すること。
• DINOv2のような強力な自己教師付きエンコーダーと表現一致を図ることで、収束が速まり、生成性能が優れていることを示すこと。
• アーキテクチャの変更なしに、実用的で即座に利用可能な方法を提供すること。

提案手法

• REPAは、拡散変換器のノイズの多い潜在隠れ状態の投影を、クリーン画像に適用された事前学習済み視覚エンコーダー（例：DINOv2）の対応する特徴と一致させる正則化損失を導入する。
• この手法は、拡散モデルの隠れ状態を外部視覚エンコーダーの特徴空間と同じ空間にマップするための線形プロキシヘッドを使用し、特徴マッチングを可能にする。
• 一致損失は学習中に計算され、拡散変換器の初期層にのみ適用されるため、計算オーバーヘッドを最小限に抑える。
• このアプローチは、DiT や SiT アーキテクチャを含む、拡散およびフローベースの変換器と互換性があり、VAE埋め込み入力を使用する潜在拡散フレームワークとも連携可能である。
• 外部エンコーダーのファインチューニングや、拡散モデルのアーキテクチャの変更を必要としないため、容易に展開可能である。
• 学習プロセスは分類器フリーのガイドランスを活用し、FIDとユーザースタディを用いた標準的なテキストから画像へのベンチマークで評価される。

Figure 1: Representation alignment makes diffusion transformer training significantly easier. Our framework, REPA, explicitly aligns the diffusion model representation with powerful pretrained visual representation through a simple regularization. Notably, model training becomes significantly more e

実験結果

リサーチクエスチョン

• RQ1 高品質な外部視覚特徴と拡散変換器表現を一致させることで、学習収束が顕著に高速化されるか？
• RQ2 DINOv2のような事前学習済み自己教師付きエンコーダーと表現一致を図ることで、標準的な拡散学習と比較して生成画像の品質が向上するか？
• RQ3 拡散モデル特徴と外部視覚特徴の一致は、学習中にどのように変化するか？また、明示的な正則化によって加速可能か？
• RQ4 提案されたREPA手法は、DiT や SiT といった異なる拡散変換器アーキテクチャにおいても有効か？
• RQ5 REPAは、ベースラインモデルと比較してはるかに少ない学習ステップでSOTAの生成品質を達成できるか？

主な発見

• REPAは、SiT-XL/2の性能を700万ステップのベースラインに一致させるために必要な学習ステップ数を40万ステップ未満にまで削減し、収束速度を17.5倍速めた。
• 分類器フリーのガイドランス（ガイドランススケール4.0）を用いた際、SOTAのFIDスコア1.42を達成し、従来の手法を上回った。
• REPAを用いた場合、下流の分類タスクにおける線形プローブ精度が顕著に向上し、より良い意味的表現学習が実現していることが示された。
• DINOv2との表現一致は、長期間の学習と大規模なモデルスケールに伴い、段階的に向上し、外部の教師信号の価値を裏付けた。
• ノイズの多い入力でも、表現一致損失は拡散モデル特徴と強力な自己教師付きエンコーダー特徴の間の意味的ギャップを効果的に低減した。
• 定性的な結果では、『消防車』『ラップトップ』『スペースシャトル』『サンゴ礁』など多様なクラスにおいて、高精細で多様性に富み、プロンプトに整合した生成が得られた。

(a) Semantic gap: Linear probing