【Diffusion勉強会】DDT: Decoupled Diffusion Transformer

DEEP LEARNING JP [DL Papers] DDT: Decoupled Diffusion Transformer Kai Yamashita, Matsuo Lab http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 Title: DDT: Decoupled Diffusion Transformer Authors: Shuai Wang, Zhi Tian, Weilin Huang, Limin Wang TL;DR DDTは従来のDiffusion Transformerが同じモジュールで担っていた「低周波の意味抽出」と「高周波の詳細復元」を、条 件エンコーダと速度デコーダに分離したエンコーダ-デコーダ構造に置き換えることで最適化ジレンマを解消した。 ImageNet 256×256/512×512でFID 1.31/1.28のSOTAを達成し、訓練は約4倍の効率化へ、推論も自己条件を動的計画で共有 することで最大3倍高速化している。 2

Background • Diffusion-Models – データを”破壊”する順過程とデータを”復元”する逆過程からなる確率過程に基づく生成モデル のダイアグラム – 順過程では確率微分方程式(SDE)に従って元のデータにノイズを与える – 逆過程ではスコア関数を学習し，reverse SDEを解くことでデータを生成できる – 従来の生成モデル(VAE, GANなど)に比べて極めて高品質かつ多様なデータの生成を可能とす るため，今日の生成モデルのデファクトスタンダートとなる – SDEではなくODEを扱う枠組みとして，Flow-Matchingがある 3

Background • Diffusion-Transformer(DiT) – Diffusion Modelのために提案されたVision-TransformerをベースとしたNeural Network Architecture – 入力となるノイズデータをViT風にパッチ分割し，AdaLN-Zeroによってクラス埋め込みや時 刻埋め込みを注入 – 従来主流だったU-Netベースのモデルよりも性能において上回るが，収束が遅いことが問題視 されている 4

What’s the bottleneck of DiT? • DiTはノイズ入力をエンコードして低周波成分(意味的成分)を抽出するのと同じモ ジュール内で高周波成分(詳細成分)をデコードする仕組みを採用 • これによって最適化ジレンマが発生する – 低周波の意味情報を効果的にエンコードするには，入力に含まれる高周波成分を抑制，減衰さ せる必要がある – 同じモジュール内で高周波成分をデコードしようとすると，高周波成分の減衰が必要なエンコ ードの要件と高周波成分を復元するデコードの条件が矛盾する • そもそもDiTは低周波のエンコードに問題を抱えている – 特にノイズが多いステップ(reverse processの初期部分)で課題を抱えており，これは計算リソ ースを初期部分に多く割り当てると改善することが実験で確認されていることより示唆される – 従来のモデルがピクセル情報のみに依存していることもこの制約の一要因 5

Related Work on Efficient DiT Training • 計算コストの改善 – Linear AttentionやSparse Attentionなどによる計算量削減 • Resampling 手法 – Log-Norm SamplingやLoss Reweightingなどによる学習ダイナミクスの安定 化 • 表現学習による強化 – 外部の視覚的な機能バイアスを導入することで学習効率を改善(後述 6

REPA Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think – DINOv2やCLIPなどのPretrainedのImage Encoderの特徴量と，DiTのEncoder出力との間で AlignmentすることでDiTの学習速度が驚異的に高速化 7

Preliminary Analysis • 解析対象は簡単のため，Linear-based Flow-Matchingを用いる • 上図に示されるとおり，拡散モデルはスペクトルの低周波→高周波へと自己回帰的に細部を洗練 していく，これは前述のジレンマを起こす要因になる 8

Preliminary Analysis • Lemma – Linear-Based Flow-Matchngのノイズスケジューラにおいて時刻tで保持される最大周波数 𝑓𝑚𝑎𝑥 (𝑡)は，クリーンデータの最大周波数を𝐾𝑓𝑟𝑒𝑞 とすると，以下の式で下界が与えられる – すなわち，tが大きく(ノイズが少なく)なるほど， • ノイズが少なくなるため，意味情報のエンコードは容易になる • 最悪のシナリオを想定すると，タイムステップtにおいてタイムステップsに進むためには 𝑡 2 残差周波数𝑓𝑚𝑎𝑥 𝑠 − 𝑓𝑚𝑎𝑥 (𝑡)を復元しなければならないが， 1−𝑡 が急激に増加するため ，後段ステップほど大変なタスクを解くことになるため難しくなる 仮説： もし初期ステップに計算リソースを多く割いて性能が改善するのであれば，DiTは低周波のエンコー ドに問題を抱えていることが示唆される，逆に後段ステップに計算リソースを割いて性能が改善する のであれば，高周波成分のデコードに問題を抱えていることがわかる 9

Preminary Analysis • 前スライドの仮説の検証 • • • SiT-XL/2モデルを用いて2nd-orderのAdam-likeのsolverを用いて，計算リソースを初期ステップに多 めに割り当てた場合の性能の変化を検証 均等に割り当てるよりも初期ステップに割り当てた方が性能が向上したため，DiTは低周波のエンコー ドに問題を抱えていることがわかる よって本研究では，Encoder-Decoderを分離することで，エンコードの性能を向上させることを考える 10

Proposed Method • Decoupled Diffusion Transformer – DiTをCondition EncoderとVelocity Decoderに分離 – Condition Encoderは低周波成分の抽出をノイズ入力，クラスラベル，タイムステップをもと にして行う – Velocity Decoderはノイズ入力と，Condition Encoderの出力をもとに高周波成分の速度を出 力する • Condition Encoder – 既存のDiT/SiTをベースモデルとして，AttentionブロックとFFNブロックを交互に配置，長い残差接続は廃止 – REPAを参考にし，Encoderの各層の出力とDINOv2の特徴量とのAllignmentをとる • Velocity Decoder – Encoderと同様の構造だが，入力はノイズ入力およびEncoder出力とタイムステップ(クラスラベルなし) – Decoderの学習は通常のFlow-MatchingのLossと同様に学習 11

Proposed Method • Sampling Acceleration – Condition Encoderの出力は，REPAによる表現アライメント学習とデコーダへの自己条件づけ により，タイムステップ間で高い類似性を持つ – 従って，エンコーダ出力を隣接ステップに再利用することでエンコーダの計算量を節約し，推 論を高速化できる – すなわち，タイムステップtが再計算タイムステップ集合Φに含まれるならEncoder出力を計算 し，それ以外のタイムステップでは前ステップの出力を再利用 – Φを，一定間隔ごとに取るのが最もナイーブな方法だが，動的計画法を用いて Φを構成するこ とも考えられる．𝑧𝑡 の類似度行列𝑆 ∈ ℝ𝑁 × 𝑁 (Nはタイムステップ数)で与えられる以下のコスト 関数を最小化する Φを動的計画法で求めることもできる 12

Experiment • ImageNet-256x256での実験 – ベースラインは，SiT, DiTやREPAなど – 近年の改良手法(SwiGLU, RoPE, RMSNormとlog-norm sampling)を既存DiT/SiTに導入し， Improved-Baselineとしてこれも検証 13

Results • • • • 256x256, 512x512両解像度において，DDTの性能がベースラインを上回る 元のDiT/SiTよりも圧倒的に速い学習速度で性能が向上している REPAによってDiT/SiTの学習速度は向上するが，DDTはさらに高速 400 epochほどで使用するVAEの理論FID上限1.20に迫る 14

Results with Different Model Size • • Improved手法(SwiGLU, RoPE etc.)を導入することでベースライン の性能は向上する Improved手法を導入した場合においてもDDTは性能を押し上げる 15

Evaluation of Sampling Acceleration • • • Uniformな割り当てよりも動的計画法を用いたもののほうが性能が良い 共有率を上げてもFID劣化は軽微である Encoder出力の類似度の相関Mapより，最小の類似度も0.8程であり，タ イムステップ間でEncoder出力は類似していることがわかる 16

Ablation • • • • Encoder およびDecoderのLayer数を色々変更して比較 意外にも20 Encoder Layer / 4 Decoder Layerという極端な比率が最良であった モデルサイズが大きいほどエンコーダ肥大化が有効になる Encoder-Decoder分離により3x3 Convでも奮闘するが，依然Attetion + MLPが最強 17

Conclusion • 本論文では従来のDiTが抱えていた「低周波セマンティクスの抽出」と「高周波ディティールの復 元」を同じモジュールで同時に学習しなければならない最適化ジレンマを再考し，この問題を根 本的に解決するDecoupled Diffusion Transformer(DDT)を提案 • 具体的には 「低周波セマンティクスの抽出」を担うCondition Encoderおよび 高周波ディティー ルの復元」を担うVelocity DecoderにDiTを分離 • モデルスケールを拡大するほどエンコーダをデコーダより大きくする方が性能向上に寄与するこ とを見出した • 実験では，DDTが少ない学習回数で従来手法を上回る高品質な画像生成を達成 • さらに，エンコーダとデコーダを分離したことでエンコーダ出力をタイムステップ間で再利用す ることによる推論高速化を導入でき，提案した動的計画アルゴリズムを用いることで画像劣化を 最小限に抑えたまま推論を高速化できることを示した • すなわち，本研究は学習効率だけでなく，推論効率の改善にも直結するものである 18