docs/dmdx.md

# DMDX (Adversarial Distribution Matching) 蒸留

[← README に戻る](../README.md) — 関連: [PCM](pcm.md) / [DMD2](dmd2.md) / [蒸留全般](distillation.md)

## 論文

**"Adversarial Distribution Matching for Diffusion Distillation Towards Efficient Image and Video Synthesis"**
- arxiv: [2507.18569v1](https://arxiv.org/html/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_total` 2.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/](https://huggingface.co/darask0/anima-distill-loras/tree/main/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](distillation.md#41-phase-a-の詳細-postmortem5-層因果分析) の教訓を再確認)