概要
Wan2.2-Lightningなどで使われている蒸留（高速化）技術の一つ。
教師モデルの生成多様性をある程度保ったままステップ数を削減できる。
DMD1(通称)
One-step Diffusion with Distribution Matching Distillation
https://arxiv.org/abs/2311.18828
旧来の蒸留では訓練時に教師モデルの生成経路と1対1で比較していた為、生徒モデルは模倣しやすいパターンに偏りがちだった。
DMDでは生徒モデルが教師モデルの生成分布をどれだけカバーできているかを計測・フィードバックして表現力を向上させる。
生成分布の広がりをどのように計測するか
非蒸留のモデル同士であれば、モデルの予測値は「入力画像をどう弄れば条件に合った画像になりそうか」の期待値（対数確率密度の勾配）であるため、いくらかの予測値を使うことで生成分布の形を比較することができる。
しかしながらDMDで蒸留された生徒モデルは1ステップで直接画像を生成するよう訓練されるので、教師モデルと直接比較できない。
そこでDMDでは、教師モデルとも生徒モデルとも異なる第三の"fake"拡散モデルを生徒モデルの生成結果で訓練し、このfakeモデルから近似的に生徒モデルの生成分布を計測している。
DMD2(通称)
Improved Distribution Matching Distillation for Fast Image Synthesis
https://arxiv.org/html/2405.14867v2
DMDを改良し、より高品質なモデルをもっと効率的に訓練できるようにしたもの。
- 品質担保のため2ステップ以上のステップ数に対応。
- 訓練時に生徒モデルの生成結果に基づく中間状態の画像を入力として予測させる訓練を行っている。
通常のモデルは予測対象にランダムノイズを乗せたものを入力として学習しており、生成時はステップ数が少ないと学習時に想定しなかったような中間潜在画像を受け取ることになるが、そのミスマッチを改善。 - 教師モデル以上の品質を狙うため、教師データセットから直接GANの識別器を構築し、その識別器が教師画像群と見分けられない画像を生成するよう訓練している。
- DMD1では訓練の安定のためにある程度の数のサンプルについて教師モデルの出力と1対1で同じ入力ノイズから近い出力が出るか比較していたが廃止。
訓練が不安定だった原因について、生成分布の比較のためのfakeモデルが十分生徒モデルを近似できてないためと考えfakeモデルの訓練を手厚くして改善。
Fake-Score-network-Free DMD
Distribution Matching Distillation without Fake Score Network
https://arxiv.org/abs/2605.19256
オリジナルのDMDは拡散モデルをターゲットとした蒸留だったが、
Flow Matchingモデルをターゲットとして改良することで分布マッチングのためのfakeモデルを不要としたもの。
教師モデルなしでの学習も可能。
Flow Matchingは「起点画像（t2iならノイズ）から終点画像への直線的な流れ（透過合成みたいなもの）」を学習しているため、
理想的には点Xtのベクトル場が求まればその終点X0からXt側への逆写像も一意に求まる。
DMDの目標は生成結果であるX0の分布を教師モデルに合わせることなので、教師モデルのベクトル場をそのまま模倣させてしまえば良い。
ただしそれをそのままやると教師モデルのコピーにしかならずステップ数を削減できない。
そこでConsistency Models損失を組み合わせることで早期ステップから一貫して画像らしい画像に向かって予測するよう強制している。
(拡散モデルにしろFMモデルにしろ非蒸留のモデルは高ノイズ域では砂嵐の高周波をべったり塗りつぶすだけで多様性はほぼ出ない為、
教師モデルの曲率の高い高ノイズ域についてはそこまで模倣させなくても別に問題ないらしい？)
直線パスのFMであれば、理想的には起点画像と終点画像が同じであれば予測値は常に同じになるので、
生徒モデル自身による予測精度の高い時刻tでの予測値を使えば事前学習済みの教師モデルも不要になる。