【DL輪読会】Diffusion Adversarial Post-Training for One-Step Video Generation

Diffusion Adversarial Post-Training for One-Step Video Generation Daiki Miyake, Matsuo Lab, M1 1

書誌情報 • タイトル: Diffusion Adversarial Post-Training for One-Step Video Generation • 著者: Shanchuan Lin, Xin Xia, Yuxi Ren, Ceyuan Yang, Xuefeng Xiao, Lu Jiang (ByteDance) • arXivリンク: https://arxiv.org/abs/2501.08316 • デモページリンク: https://seaweed-apt.com/ 2

• https://arxiv.org/abs/2501.08316
• https://seaweed-apt.com/

概要 • 動画拡散モデルの蒸留 – 拡散モデルでの生成は多くの反復計算が必要だが，その回数を少なくしたい • 既存研究と比較して – より高解像度の動画が対象 – 実データのみを学習に使う (合成データを使わない) 3

背景 • 既存の動画生成モデルの蒸留手法の欠点 – 512x512・16fpsと，解像度が低い – 最低でも4 stepの反復計算を必要とする • 提案手法の目的 – 1280x720・24fpsと，高解像度での生成 – 1 stepで高品質に生成 4

既存研究 • (画像)拡散モデルの蒸留 – 決定論的手法 ノイズから教師モデルによる生成サンプルへのマッピングを学習する手法 (拡散モデルがODEを解いているとみなすと，ノイズと生成サンプルは一対一対応する) 最適化が進みにくく，複数stepでの生成が必要 – 分布的手法 教師モデルによる生成データの分布を近似する手法 データ生成が高コスト，かつ教師モデルの品質が天井になる 複数stepでの生成が必要 • 上記手法の動画拡散モデルへの応用 – 低解像度，複数stepでの生成が必要 5

提案手法 • 学習済みの動画拡散モデルを追加学習し，1 stepでの生成を目指す (Adversarial Post-Training: APT) • 追加学習の段階では，実データのみを使う 6

提案手法 • Discriminatorを導入し，GANの枠組みで学習する Discriminator Generator (拡散モデル) z: ノイズ c: プロンプト 7

提案手法 • Generator – 学習済みの拡散モデルで初期化 (ただし出力がサンプルになるよう調整) – 敵対的学習の前に，Consistency Modelと 同様の手法で少し蒸留を行う 8

提案手法 • Discriminator – Generatorと同様のアーキテクチャを採用， 元の拡散モデルで初期化 – 学習可能なQueryベクトルとのCross Attention層を追加 それらの出力の線形和をとることで スカラーの出力を得る – ランダムに選んだタイムステップを shift関数で変換したものを入力 9

提案手法 • GANの学習を安定化させるためにR1正則化がよく使われる • 勾配計算のためには2次勾配の計算が必要 PyTorch FSDP, Gradient Checkpointing, FlashAttentionを用いた大 規模学習では，そもそも2次勾配の計算に対応していない • 以下の計算でR1正則化を近似する 10

実験 • 学習設定 – 最初に画像のみで学習 (128~256 H100, batch-size=9062) – 次に動画のみで学習 (1024 H100, batch-size=2048) • 比較手法 – 蒸留された大規模モデルとして，FLUX-Schnell, SD3.5-Turbo, SDXL-DMD2, SDXL-Hyper, SDXL-Lightning, SDXL-Nitro, SDXL-Turboらと比較 11

定性評価 • 元の拡散モデルと比較して，より自然な画像を生成する傾向がある 12

定性評価 • 既存の蒸留モデルと比較しても，より精細な画像を生成できている 13

定性評価 14

定量評価 • 視覚的観点(質や色)，構造的観点，テキスト整合性の3観点でuser preferenceを，蒸留前のモデルと比較 • 既存モデルは蒸留により悪化，提案手法の視覚的観点のみ向上 15

• 動画では，提案手法は元モデルより構造的観点で悪化 16

定量評価 • 提案手法と既存モデルとのpreferenceの比較 • テキスト整合性が既存モデルより弱い 17

Ablation • 正則化がないと，生成画像にグリッド上のアーティファクトが生成さ れる • Discriminatorのパラメータ数をGeneratorよりも減らすと性能悪化 • Cross Attention層を最後の層だけにすると構造的観点で悪化 • バッチサイズが小さいとモード崩壊を起こす • など 18

まとめ • 拡散モデルを高解像度・1 stepで蒸留するAPTを提案 • Generatorと同形のDiscriminatorを導入し，近似したR1正則化によっ て学習を安定 • 基本的には既存モデルよりも良い性能だが，構造的観点・テキスト整 合性については劣化がみられる 19