DMDX (Adversarial Distribution Matching) 蒸留
← README に戻る — 関連: PCM / DMD2 / 蒸留全般
論文
"Adversarial Distribution Matching for Diffusion Distillation Towards Efficient Image and Video Synthesis"
- arxiv: 2507.18569v1
- 著者: Yanzuo Lu et al. (Sun Yat-Sen University + ByteDance Seed Vision)
- 公開コード: なし (論文のみ)
核心アイデア
DMD2 の (real_score - fake_score) の 逆 KL 発散 はモード崩壊しがち (Phase B v2 と同症状)。DMDX は 学習可能 discriminator による hinge GAN (TVD = Total Variation Distance、対称) 最小化 に置換し、根本解決を狙う。
vs DMD2
| 観点 | DMD2 | DMDX (ADM) |
|---|---|---|
| 発散尺度 | 逆 KL (非対称、mean collapse 寄り) | TVD (対称) |
| 損失実装 | grad trick で score difference 直接最小化 | 学習可能 discriminator で hinge loss 敵対学習 |
| GAN 役割 | optional 正則化 | 本体損失 |
| 時刻情報 | t での score 比較のみ | t → t-Δt の ODE 進化中間点も考慮 |
| Initialization | MSE pretrain | 論文では ADP (= adversarial pre-training)、本実装は warm-start で代替 |
本実装の範囲
論文の ADM のみ移植、ADP は省略 (SAM 依存 pixel-space discriminator が重い、コスト効果が低い)。
代替: Civitai Anima Turbo を warm-start で「分布レベルの良い初期化」を確保。
実装: scripts/distill/{dmdx_loss,train_dmdx}.py
Modal: modal_app.py::train_dmdx_distill
ADM の構造
- Discriminator backbone = teacher MiniTrainDIT (frozen) — LADD の
AnimaLADDDiscriminatorを流用 - Trainable heads × 5 (blocks 2/8/14/20/26 を hook、各 spectral norm + BatchNormLocal)
- Cubic time schedule:
u ~ U(0,1), t = 1 - u**3(high-noise バイアス) - Δt evolution: teacher で 1-step Euler で
t → t-Δt進化、Δt = T/64 - Hinge loss: G =
mean(-D(fake))、D =relu(1-D(real)) + relu(1+D(fake))
訓練設定 (実走、2026-05-19)
| 項目 | 値 |
|---|---|
| Base | anima-base-v1.0.safetensors |
| Method | DMDX ADM-only |
| Student LoRA | wide (980 keys、single adapter) |
| LoRA rank | 32 |
| Disc heads | 5 (block 2/8/14/20/26)、各 ~3.4M = 17.1M total |
| Total outer | 5000 (1 outer = 2 disc + 1 gen) |
| n_critic_per_gen | 2 (DMD2 の 5 より少なめ、hinge GAN 系は 1:1〜2:1 が一般的) |
| n_student_steps | 4 |
| Teacher CFG | 4.5 |
| Student CFG | 1.0 |
| dt_ratio | 1/64 ≈ 0.0156 |
| recon_weight | 0.0 (pure ADM、LADD の Smooth-L1 anchor なし) |
| LR (gen / disc) | 5e-6 / 1e-5 |
| Warm-start | /models/loras/anima_turbo.safetensors (Civitai 公式 Turbo) |
| GPU | B200 |
| 訓練時間 | smoke ~35s (10 outer)、本番 ~5h 想定 (3.6s/outer) |
| コスト | smoke ~$0.5、本番 ~$30 想定 |
Smoke 結果 (10 outer、2026-05-19)
| outer | l_d_real | l_d_fake | l_d_total | l_g_adv | 評価 |
|---|---|---|---|---|---|
| 0 | 0.717 | 1.295 | 2.012 | -0.292 | 初期 hinge equilibrium |
| 5 | 0.724 | 1.277 | 2.001 | -0.276 | 安定継続 |
| 9 | 0.740 | 1.268 | 2.008 | -0.270 | smoke 完了 |
OOM / NaN / divergence なし。setup・warm-start (968/1016)・D init・ADM loss 全部正常。
本番訓練 (進行中、2026-05-19)
Step 500 検証
| outer | l_d_real | l_d_fake | l_d_total | l_g_adv |
|---|---|---|---|---|
| 0 | 0.717 | 1.295 | 2.012 | -0.292 |
| 90 | 0.774 | 1.227 | 2.001 | -0.218 |
| 340 | 0.938 | 1.034 | 1.972 | +0.048 (D 優勢) |
| 500 | 0.822 | 0.976 | 1.798 | -0.146 |
| 1000 | 0.933 | 0.971 | 1.904 | -0.093 |
GAN dynamics:
l_d_total2.012 → 1.798 → 1.904 (D 学習、後半は equilibrium 周辺で振動)l_g_adv-0.292 → +0.05 (一時 D 優勢) → -0.15 → -0.09 (G 持ち直し)- D と G の拮抗状態に到達、healthy GAN training
生成検証 (step 500/1500/2500/final、4-step / er_sde+simple)
step 500/1500/2500/final で 4 回 quick_check を実施 — すべて視覚的に同一。 warm-start (Civitai Anima Turbo) ベースラインとも視覚差なし:
| ckpt | 生成時間 (er_sde+simple) | 視覚的差異 |
|---|---|---|
| step 500 | 5.77s | warm-start とほぼ同一 |
| step 1500 | 6.07s | 同上 |
| step 2500 | — | 同上 |
| final (5000) | — | warm-start (Civitai Turbo) と区別困難 |
⚠ 重要な結論: warm-start dominance
5000 outer 訓練を完走したが、student LoRA は warm-start からほぼ移動しなかった:
- GAN dynamics は健全 (D vs G 振動、divergence / mode collapse なし)
- ただし重み update が保守的、独自の蒸留方向を獲得できず
- 出力品質は Civitai Anima Turbo と視覚的に同等
推測原因:
- hinge GAN の update が conservative (D が perfect 寄りで G の gradient saturate)
- warm-start anchor が強すぎる (Civitai Turbo が既に 4-step CFG=1 で機能)
- recon_weight=0 で soft anchor がなく、保守的振動に収束
配布
HF: darask0/anima-distill-loras/dmdx/
実質「Civitai Turbo + 軽い ADM 訓練」状態。研究目的・記録のため配布。 実用としては Anima Turbo を直接使うのと差異なし。
ComfyUI 推奨設定
LoraLoaderModelOnly:
lora_name: dmdx_student_final_comfy.safetensors
strength_model: 1.0
ModelSamplingAuraFlow:
shift: 3.0
KSampler:
steps: 4
cfg: 1.0
sampler_name: er_sde
scheduler: simple
期待品質 (paper 報告ベース)
SDXL 1-step での DMDX vs DMD2:
| metric | DMD2 | DMDX | 差 |
|---|---|---|---|
| CLIP Score | 35.22 | 35.26 | +0.04 |
| PickScore | 22.10 | 22.27 | +0.17 |
| HPSv2 | 27.45 | 27.70 | +0.25 |
| MPS | 10.69 | 11.20 | +0.51 |
| 多様性 LPIPS | 0.6715 | 0.7156 | +0.044 |
→ 多様性 +6.5% が最大の利点、画像品質は誤差レベル。Video (CogVideoX) では 8-step で 100-step 比 92-96% 加速の大きな成果。本実装は image only。
知見
DMDX の利点 (DMD2 比較で明確)
| 観点 | DMD2 + TrigFlow | DMDX (ADM) |
|---|---|---|
| critic 数 / generator update | 5 | 2 (-60%) |
| forwards / outer step | ~32 | ~23 (-28%) |
| LoRA adapter | 2 (student + fake_score、毎 forward で set_adapter() 切替必要) |
1 (student only) + 別 D heads |
| 並列訓練時の slowdown | ~1.7x (6.2 vs 3.7s/outer) | ほぼなし (3.6s/outer 維持) |
| 実装複雑度 | dual adapter PEFT 管理、leak 注意 | single adapter + 標準 D heads、シンプル |
仮に最終品質が DMD2 と同等以下でも、訓練効率の優位性は明確:
- generator update の頻度が高く (per outer)、G の学習機会が多い
- adapter switching overhead なし
- parallel training の安定性が高い
その他
- LADD discriminator の再利用: spectral norm + BatchNormLocal + multi-block hook の構造はそのまま流用可能、DMDX 専用 D を新規開発する必要なし
- cubic time schedule: shifted uniform より extreme な high-noise バイアス、
t_mean=0.0259-0.99の wide range (cubic の variance 大) recon_weight=0で開始: pure ADM の純粋性 verify。不安定なら LADD 流 Smooth-L1 anchor (>0) で stability boost 可能- billing limit 再発防止: 並列 launch 前に Modal spend cycle limit を確認 (DMD2 v1 が KeyboardInterrupt'd した先例あり)
- warm-start dominance の罠: 強い warm-start (Civitai Turbo) + 保守的な hinge update で student がほぼ動かない結果に。次回試すなら cold-start / lr_gen 引き上げ / recon_weight>0 のいずれかが必要
- GAN equilibrium ≠ 蒸留進行:
l_d_total振動と visual 変化は別軸。loss だけで判断せず、中間 ckpt 毎の visual 検証が必須 (distillation.md §4.1 layer 5 の教訓を再確認)