论文

sDMD:用 Sinkhorn 散度改进扩散模型蒸馏

sDMD: Sinkhorn Distribution Matching Distillation

精选理由

一篇蒸馏方法论文,把 DMD 加上 Sinkhorn 最优传输修正,Qwen-Image 和 Wan2.1 上 2-4 步就能出图,多样性还比 DMD2 好。

论文提出 sDMD,针对 DMD 蒸馏方法中 reverse KL 导致的模式坍缩和低密度区域教师评分不可靠问题,引入基于 Sinkhorn 散度的最优传输修正项,并根据判别器输出自适应激活该修正。在 FLUX.1-dev、Qwen-Image-20B 和 Wan2.1-14B 上仅需 2-4 步采样,sDMD 在 GenEval、DPG、WISE 和 VBench 上超过蒸馏基线,样本多样性优于 DMD2。

原文 · arXiv cs.LG

sDMD: Sinkhorn Distribution Matching Distillation

Distribution Matching Distillation (DMD) distills multi-step diffusion models into few-step generators by minimizing reverse KL divergence. Despite its strong performance, DMD guides samples by local density alone, which introduces two limitations: teacher scores can be unreliable in regions of low teacher density, and the mode-seeking behavior of reverse KL can collapse the student onto a subset of the teacher modes. We first connect DMD to Wasserstein gradient flow, showing that its gradient follows from a fixed-point formulation with integral KL (IKL) as the energy functional. This perspective motivates augmenting IKL with an optimal transport term that addresses these limitations by guiding outlier samples toward the target and promoting mode coverage. We realize this term using Sinkhorn divergence estimated from finite batches. In practice, however, errors in the estimated Sinkhorn velocity can perturb in-domain samples for which IKL guidance is already accurate. To this end, we introduce sDMD, which adaptively activates the Sinkhorn correction based on discriminator outputs, focusing on outlier samples where IKL may be unreliable. With 2--4 sampling steps on FLUX.1-dev, Qwen-Image-20B, and Wan2.1-14B, sDMD outperforms distillation baselines on GenEval, DPG, WISE, and VBench while improving sample diversity over DMD2.