DMD: One-step Diffusion with Distribution Matching Distillation

Authors: Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Frédo Durand, William T. Freeman, Taesung Park Affiliations: Massachusetts Institute of Technology, Adobe Research Venue: CVPR 2024

代码核对说明：原始 DMD 论文项目页公开了论文、视频与示例，但没有同步放出原始 DMD 官方训练仓库。因此本文的“代码映射 / 伪代码”以社区复现 devrimcavusoglu/dmd 为主，并用作者后续官方仓库 tianweiy/DMD2 辅助说明后续改进；两者不能混同为“原始 DMD 官方实现”。

1. Motivation (研究动机)

Diffusion Model 的核心优势是生成质量高、训练稳定，但它的采样过程本质上是一个多步迭代去噪过程：从高斯噪声开始，经过几十到几百次网络前向，才得到最终图像。这种高质量与高延迟的组合，使它在交互式应用、实时创作工具和低延迟部署中受到明显限制。

论文聚焦的不是“把采样器从 1000 步减到 20 步”，而是更激进的问题：能否把一个已经训练好的 diffusion teacher 直接蒸馏成 one-step generator，并且尽量不损失质量。此前工作大多沿着 knowledge distillation 的路线前进，即让 student 去拟合 teacher 的 noise-to-image 映射，或者在 ODE / flow 轨迹上保持一致性。但这类方法都有一个共同难点：

它们把问题建模为 逐样本、逐轨迹的对应学习；
而 diffusion 的映射本身是高维、复杂、强非线性的；
当 student 只允许 1 步时，这个映射过于难以被精确拟合。

因此，DMD 的切入点非常关键：不再强迫 student 复现 teacher 的逐点路径，而是只要求它生成的样本分布与 teacher 的样本分布一致。 这相当于把问题从“函数逼近”改写成“分布匹配”。

这件事值得研究有三个原因：

速度收益极大：如果 one-step generator 成立，理论上可把 50/100/500 步采样压缩成 1 步前向。
更符合生成任务目标：用户在意的是“生成图像看起来是否像 teacher 产生的分布”，而不是某个噪声是否必须映射到某一张固定图像。
为后续 few-step / adversarial diffusion distillation 奠定基础：DMD 后续直接催生了 DMD2 - Improved Distribution Matching Distillation for Fast Image Synthesis 这条工作线。

2. Idea (核心思想)

DMD 的核心思想可以概括为一句话：用两个 diffusion score estimator 的差，近似 one-step generator 的分布匹配梯度。

具体来说：

从点对点蒸馏改为分布对分布蒸馏
目标不再是让 $G_{θ} (z)$ 逼近 teacher 对同一个 $z$ 的输出，而是让 generator 的分布 $p_{fake}$ 接近 teacher / real 分布 $p_{real}$ 。
用 diffusion model 近似 real / fake 分布的 score
KL 散度关于 generator 的梯度可以写成 real score 和 fake score 的差。由于 diffusion model 本身就在学习带噪分布的 score，因此可以自然承担这个角色。
在线学习 fake critic
real score 对应的是冻结的 pretrained diffusion model；fake score 则由另一个持续在 generator 输出上训练的 diffusion model $μ_{fake}^{ϕ}$ 估计，因为 $p_{fake}$ 会随着训练不断变化。
再加一个 regression loss 抗 mode collapse
仅靠分布匹配会面临 GAN 类似的问题：优化可能偏向高密度模式、导致 mode dropping。于是作者额外构造了一个小规模的 noise-image paired dataset，用 LPIPS 约束 $G_{θ} (z)$ 与 teacher deterministic sampler 的输出在大结构上保持一致。

与 GAN 的相似点在于：都在做 real/fake 分布对齐；但与 GAN 的关键差异在于：DMD 的 critic 不是从零学习的判别器，而是继承了 pretrained diffusion model 的强先验，优化目标也来自 score difference，而非标准 adversarial minimax。

3. Method (方法)

3.1 整体框架

Figure 2 解读：这张图完整展示了 DMD 的三路训练信号。左侧是 one-step generator $G_{θ}$ ，输入随机 latent / noise $z$ ，一次前向得到假图像 $x$ 。中间的 distribution matching 分支会对 $x$ 注入随机噪声得到 $x_{t}$ ，然后分别送入冻结的 real score model $μ_{real}$ 与在线更新的 fake score model $μ_{fake}$ ，两者给出的“去噪方向差”近似为让图像“更真实、更不假”的梯度。下方是 regression 分支：从预先离线生成的 noise-image paired dataset 中取 $(z_{ref}, y_{ref})$ ，约束 $G_{θ} (z_{ref})$ 与 teacher 输出 $y_{ref}$ 在感知空间中接近。最终 generator 同时接收 distribution-level realism 信号与 structure-preserving regularization。

从模块关系看，DMD 包含三类网络：

Base / real score model $μ_{real}$ ：冻结的 teacher diffusion model；
Generator $G_{θ}$ ：目标 one-step generator；
Fake score model $μ_{fake}$ ：初始化自 teacher，并持续在 generator 输出上训练。

其训练数据流可以写成：

z: [B, C, H, W] or [B, 4, 64, 64]
  -> G_theta
  -> x_fake
  -> q_t(x_t | x_fake)
  -> {mu_real, mu_fake}
  -> {pred_real, pred_fake}
  -> distribution matching gradient
  -> update G_theta
 
(z_ref, y_ref) from paired dataset
  -> G_theta(z_ref) = x_ref
  -> LPIPS(x_ref, y_ref)
  -> regularize G_theta
 
x_fake (stopgrad)
  -> q_t(x_t | x_fake)
  -> mu_fake
  -> denoising loss
  -> update mu_fake

3.2 Base model、one-step generator 与带噪 score

论文假设已经给定一个 pretrained diffusion denoiser $μ_{base} (x_{t}, t)$ 。作者在推导中使用 mean-prediction 形式的 diffusion 模型；若实际底座是 $ϵ$ -prediction，只需按 Appendix H 的变量变换即可。

生成器的设计非常直接：

结构上继承 teacher 的 denoiser 主体；
移除 time-conditioning 的角色，使其变成固定噪声等级上的 one-step generator；
初始化时直接复制 teacher 权重，使

G_{θ} (z) = μ_{base} (z, T - 1), \forall z .

这一步很重要，因为它让 one-step generator 从一开始就不是随机初始化，而是站在“teacher 在最大噪声处的预测”这个合理起点上。

对任意生成样本 $x = G_{θ} (z)$ ，论文通过 forward diffusion 得到带噪样本：

q_{t} (x_{t} ∣ x) \sim N (α_{t} x, σ_{t}^{2} I) .

真实分布的 score 用冻结的 pretrained diffusion model 给出：

s_{real} (x_{t}, t) = - \frac{x _{t} - α _{t} μ _{base} ( x _{t} , t )}{σ _{t}^{2}} .

假分布的 score 则由在线学习的 fake critic 给出：

s_{fake} (x_{t}, t) = - \frac{x _{t} - α _{t} μ _{fake}^{ϕ} ( x _{t} , t )}{σ _{t}^{2}} .

这里的关键点是：DMD 并不直接比较 $μ_{base}$ 和 $G_{θ}$ 的输出，而是把两者都搬到“带噪分布的 score 场”里去比较。

Figure 4 解读：图 (a) 说明如果直接在原始 real/fake 分布上做 score difference，两个分布可能根本不重叠，于是某些位置上 real score 或 fake score 没有良好定义。图 (b) 展示了加噪后的效果：当样本被 diffusion 到较高噪声层级后，两个分布会在整个空间中充分重叠，此时 score difference 才成为一个 everywhere-defined 的训练信号。DMD 的“先加噪、再做分布匹配”本质上就是为了解决这个可定义性问题。

3.3 Distribution Matching Loss

论文的目标是最小化 fake 分布与 real 分布之间的 KL 散度：

D_{K L} (p_{fake} ∥ p_{real}) = E_{z \sim N (0, I), x = G_{θ} (z)} [- lo g p_{real} (x) + lo g p_{fake} (x)] .

对 generator 参数求梯度，有：

\nabla_{θ} D_{K L} = E_{z, x = G_{θ} (z)} [- (s_{real} (x) - s_{fake} (x)) \nabla_{θ} G_{θ} (z)] .

问题在于：这个式子仍然需要原始分布上的 score，而 diffusion model 实际上只提供带噪分布的 score 近似。因此作者把梯度改写到带噪样本 $x_{t}$ 上，得到最终近似：

\nabla_{θ} D_{K L} ≃ E_{z, t, x, x_{t}} [w_{t} α_{t} (s_{fake} (x_{t}, t) - s_{real} (x_{t}, t)) \nabla_{θ} G_{θ} (z)],

其中 $t \sim U (T_{m i n}, T_{m a x})$ ，论文设置为 $T_{m i n} = 0.02 T$ 、 $T_{m a x} = 0.98 T$ 。

为了让不同噪声水平的梯度幅度更稳定，论文引入了时间相关的权重：

w_{t} = \frac{σ _{t}^{2}}{α _{t}} \cdot \frac{CS}{∥ μ _{base} ( x _{t} , t ) - x ∥ _{1}},

其中 $C$ 是通道数， $S$ 是空间位置数。直观理解：

当噪声水平很高时，raw gradient 容易失衡；
用 $∥ μ_{base} (x_{t}, t) - x ∥_{1}$ 做归一化，可以让不同 $t$ 的训练信号落在更可控的量级上；
这也是 DMD 相比直接照搬 DreamFusion / ProlificDreamer 权重更稳定的关键之一。

与此同时，fake critic 通过标准 denoising loss 在线更新：

L_{denoise}^{ϕ} = ∥ μ_{fake}^{ϕ} (x_{t}, t) - x_{0} ∥_{2}^{2},

并沿用 base diffusion model 的 timestep weighting 策略。

3.4 Regression Loss 与最终目标

Figure 3 解读：这张二维 toy figure 解释了为什么 distribution matching 还不够。左图是初始 fake / real 分布；图 (a) 只使用 real score 时，所有 fake sample 都会被吸向最近的真实 mode，出现典型 mode collapse；图 (b) 使用 real+fake score 后，覆盖范围扩大，但仍可能错过某些 mode；图 (c) 再加 regression loss 后，模型既保留了分布匹配带来的 realism，又能维持 mode coverage。它说明 DMD 的 regression 分支不是“锦上添花”，而是 one-step 训练稳定性的必要结构件。

paired dataset 的构造方式是：

采样随机噪声 $z$ ；
用 teacher diffusion model 的 deterministic sampler 生成对应图像 $y$ ；
得到一组 $(z, y)$ 对；
在训练时只需偶尔从这个小型 paired set 中采样即可。

回归损失定义为：

L_{reg} = E_{(z, y) \sim D} ℓ (G_{θ} (z), y),

其中距离函数 $ℓ$ 使用 LPIPS。最终 generator 的训练目标为：

L_{G} = D_{K L} + λ_{reg} L_{reg},

除特别说明外，论文默认 $λ_{reg} = 0.25$ 。

3.5 Classifier-Free Guidance (CFG) 的适配

对于 text-to-image teacher（如 Stable Diffusion v1.5），DMD 的适配思路是：

用 guided teacher 生成 paired dataset，以便 regression loss 对齐的是实际部署时的 guided sampling 输出；
在计算 distribution matching gradient 时，把 real score 替换为 guided model 对应的 mean prediction；
fake score 的定义不变；
训练时使用固定 guidance scale，这也是本文只支持固定 CFG 强度的重要原因。

这意味着原始 DMD 的 text-to-image 版本虽然已经很强，但并不天然支持 inference-time 任意改 guidance scale。这是它与后续更灵活的 few-step distillation 方法的一个明显差别。

3.6 基于论文 Appendix + 复现代码整理的伪代码

下列伪代码以论文 Appendix B / Eq. (6-9) 为主，并对照 devrimcavusoglu/dmd 的 CIFAR-10 复现实现整理；若与论文公式有轻微常数项差异，以后文“Paper vs Code Differences”表为准。

3.6.1 构造 paired dataset

def build_paired_dataset(mu_base, num_pairs, solver, cond):
    # z: [B, C, H, W] for pixel-space EDM
    # z: [B, 4, 64, 64] for latent SD v1.5
    pairs = []
    for _ in range(num_pairs):
        z = sample_gaussian_latent(cond)           # z: [B, Cz, Hz, Wz]
        y = deterministic_sample(mu_base, z, solver, cond)
        # y: [B, C, H, W] or decoded image batch
        pairs.append((z, y))
    return pairs

3.6.2 Distribution Matching Loss

def distribution_matching_loss(mu_real, mu_fake, x, num_steps=1000, class_ids=None):
    # x: [B, C, H, W], fake image from generator
    t_min = int(0.02 * num_steps)
    t_max = int(0.98 * num_steps)
    timestep = randint(t_min, t_max, shape=[x.shape[0]])   # [B]
    noisy_x, sigma_t = forward_diffusion(x, timestep)      # noisy_x: [B, C, H, W]
 
    with no_grad():
        pred_fake = mu_fake(noisy_x, sigma_t, class_ids)   # [B, C, H, W]
        pred_real = mu_real(noisy_x, sigma_t, class_ids)   # [B, C, H, W]
 
    weighting = abs(x - pred_real).mean(dim=[1, 2, 3], keepdim=True)  # [B, 1, 1, 1]
    grad = (pred_fake - pred_real) / weighting                         # [B, C, H, W]
    target = stopgrad(x - grad)                                        # [B, C, H, W]
    loss = 0.5 * mse_loss(x, target)
    return loss

3.6.3 Fake score model 的去噪更新

def denoising_loss(mu_fake, x, timestep, sigma_data, class_ids=None):
    # x: [B, C, H, W], stop-gradient fake image
    x_t, sigma_t = forward_diffusion(stopgrad(x), timestep)            # [B, C, H, W]
    pred_fake = mu_fake(x_t, sigma_t, class_ids)                       # [B, C, H, W]
    weight = 1.0 / sigma_t**2 + 1.0 / sigma_data**2                    # [B]
    weight = weight[:, None, None, None]                               # [B, 1, 1, 1]
    loss = mean(weight * (pred_fake - stopgrad(x)) ** 2)
    return loss

3.6.4 一次完整训练步

def train_step(generator, mu_real, mu_fake, paired_batch, class_ids=None):
    # paired_batch = (z_ref, y_ref)
    # z, z_ref: [B, C, H, W] (EDM) or [B, 4, 64, 64] (SD latent)
    # y_ref: [B, C, H, W]
    z_ref, y_ref = paired_batch
 
    z = randn_like(y_ref)                                              # [B, C, H, W]
    g_sigma = fixed_generator_sigma_at_T_minus_1(batch=z.shape[0])     # [B]
    z = z * g_sigma[0]
    z_ref = z_ref * g_sigma[0]
 
    x = generator(z, g_sigma, class_ids)                               # [B, C, H, W]
    x_ref = generator(z_ref, g_sigma, class_ids)                       # [B, C, H, W]
 
    loss_dm = distribution_matching_loss(mu_real, mu_fake, x, class_ids=class_ids)
    loss_reg = lpips(resize_to_224(x_ref), resize_to_224(y_ref))
    loss_g = loss_dm + lambda_reg * loss_reg
    update(generator, loss_g, grad_clip=10.0)
 
    t = randint(1, 1000, shape=[x.shape[0]])                           # [B]
    loss_d = denoising_loss(mu_fake, x, t, sigma_data=mu_fake.sigma_data, class_ids=class_ids)
    update(mu_fake, loss_d, grad_clip=10.0)
 
    return {
        "loss_g": loss_g,
        "loss_dm": loss_dm,
        "loss_reg": loss_reg,
        "loss_denoise": loss_d,
    }

3.6.5 推理

def one_step_inference(generator, cond):
    # z: [B, C, H, W] or [B, 4, 64, 64]
    z = sample_gaussian_latent(cond)
    g_sigma = fixed_generator_sigma_at_T_minus_1(batch=z.shape[0])
    x = generator(z * g_sigma[0], g_sigma, cond)   # one forward pass only
    return x

3.7 实现指南：顺序、接口、数据流与验证

3.7.1 推荐实现顺序

先加载 teacher 并验证 deterministic sampler：确保可以稳定生成 paired dataset；
实现 one-step generator 初始化：从 teacher 复制权重，固定到 $T - 1$ 的输入噪声级；
实现 forward diffusion 与 DM loss：这是核心梯度来源；
实现 fake critic 的 denoising update：保证 $μ_{fake}$ 能跟上 generator 分布漂移；
再接入 regression loss：最后用 LPIPS 稳住大结构与 mode coverage；
最后接入 CFG / latent-space 细节：只在 pixel-space 跑通后再迁移到 SD v1.5。

3.7.2 关键接口约定

模块	输入	输出	接口要点
`generator`	`z: [B, C, H, W]` 或 latent `z: [B, 4, 64, 64]`，固定 `g_sigma`	`x_fake: [B, C, H, W]` / latent	只执行一步前向
`forward_diffusion`	`x_fake, timestep/noise_level`	`x_t, sigma_t`	必须与 teacher 的 noise schedule 对齐
`mu_real`	`x_t, t/sigma_t, cond`	去噪预测 `pred_real`	冻结参数
`mu_fake`	`x_t, t/sigma_t, cond`	去噪预测 `pred_fake`	在线更新
`paired_dataset`	`z_ref`	`y_ref`	由 deterministic teacher sampler 离线生成
`lpips_branch`	`x_ref, y_ref`	标量回归损失	实现中通常先 resize 到 224

3.7.3 最小数据流检查

x, pred_real, pred_fake, x_t 的 shape 必须完全一致；
weighting_factor 必须严格为正，且不能出现 NaN / Inf；
初始阶段若 mu_fake 仍接近 mu_real，则 DM 梯度不应异常爆炸；
推理时必须确认只调用 generator 一次；
text-to-image 版本需核对 latent decode 只在 regression / eval 路径上发生，而不是每次 DM loss 都 decode。

3.7.4 最小 smoke test

# smoke test: batch=2 on CIFAR-10-like tensor space
z = torch.randn(2, 3, 32, 32, device=device)
y = torch.randn(2, 3, 32, 32, device=device).clamp(-1, 1)
class_ids = torch.tensor([0, 1], device=device)
 
x = generator(z, g_sigma, class_ids)
assert x.shape == (2, 3, 32, 32)
 
loss_dm = distribution_matching_loss(mu_real, mu_fake, x, class_ids=class_ids)
loss_reg = lpips(resize_to_224(x), resize_to_224(y))
assert torch.isfinite(loss_dm)
assert torch.isfinite(loss_reg)

3.8 Code-to-paper mapping 与 Paper-vs-Code 差异

3.8.1 代码映射表（以 `devrimcavusoglu/dmd` 为主）

论文概念	代码位置	关键类 / 函数	说明
Distribution Matching Loss (Eq. 7/8)	`dmd/loss.py`	`DistributionMatchingLoss.forward`	采样 $t$ 、forward diffusion、real/fake 预测、weighting、stop-grad MSE 近似
Generator 总损失	`dmd/loss.py`	`GeneratorLoss.forward`	`D_KL + lambda_reg * LPIPS`
Fake critic 去噪损失 (Eq. 6)	`dmd/loss.py`	`DenoisingLoss.forward`	用 EDM 风格的 weighted MSE 更新 `mu_fake`
forward diffusion / fixed generator sigma	`dmd/modeling_utils.py`	`forward_diffusion`, `get_fixed_generator_sigma`	体现“生成器固定在 $T - 1$ 噪声级”
训练主循环	`dmd/training/training_loop.py`	`train_one_epoch`	先更新 generator，再更新 `mu_fake`
训练入口 / 超参数	`dmd/train.py`	`run`, `train`	数据加载、优化器、FID 评测、checkpoint
paired dataset 读取	`dmd/dataset/cifar_pairs.py`	`CIFARPairs`	从离线保存的 `(image, latent)` 对中读取
单步推理	`dmd/generate.py`	`DMDGenerator` / `EDMGenerator` 路径	固定 generator sigma，一次前向生成
后续官方改进版	`tianweiy/DMD2/main/`	`train_sd.py`, `sd_unified_model.py`, `edm_unified_model.py`	说明作者后续如何把 DMD 扩展到更强 few-step / GAN-augmented 训练

3.8.2 Paper vs Code Differences

方面	论文表述	代码侧观察	影响
原始 DMD 官方代码	论文 / 项目页未附官方训练仓库	公开可用的是后续官方 `DMD2`，原始 DMD 常用社区复现	阅读实现细节时必须区分“论文原始方法”与“后续改进版 / 社区版”
Generator 形式	“base denoiser without time-conditioning”	社区复现中仍保留 EDM 风格接口，但通过 `get_fixed_generator_sigma()` 固定到 $T - 1$ 噪声级	工程上更容易复用 teacher 结构；概念上等价于 one-step 固定噪声生成器
Eq. 8 权重	公式写成 $\frac{σ _{t}^{2}}{α _{t}} \frac{CS}{∥ μ _{ba se} ( x _{t} , t ) - x ∥ _{1}}$	社区复现直接用 `abs(x - pred_real).mean(...)` 做归一化，再把常数吸收到实现里	常数项和参数化差异来自 EDM/mean-prediction 具体实现，不改变核心“按噪声级归一化”的思想
Fake critic 权重	“沿用 base diffusion 的 timestep weighting”	社区复现显式用 EDM 权重 `1 / sigma_t^2 + 1 / sigma_data^2`	更接近 EDM 训练习惯，便于复现 CIFAR-10 结果
时间采样	论文用 $t \sim U (T_{m i n}, T_{m a x})$	社区代码对 DM loss 用 `[0.02T, 0.98T)`，对 denoising update 用 `[1, T)` 避免 $σ = 0$	这是合理的数值稳定性处理
text-to-image 实现	论文给出 CFG 适配原则	原始 DMD 无官方发布代码；官方 DMD2 则在 `main/train_sd.py` 等文件里实现更成熟的 SD/SDXL 训练框架	复现 T2I DMD 时，很多工程细节需要从 DMD2 反推

4. Experimental Setup (实验设置)

4.1 数据集与任务

论文覆盖了四类设置：

CIFAR-10
- 32×32，包含 conditional 与 unconditional 两种设置；
- 主要用于快速验证 DMD 的机制与消融。
ImageNet-64×64
- class-conditional；
- 是论文主打的 one-step class-conditional 图像生成 benchmark。
LAION-Aesthetic 6.25+
- 大约 3M 图文样本；
- 用于 distill Stable Diffusion v1.5，guidance scale = 3；
- 评测用 zero-shot MS COCO-30K。
LAION-Aesthetic 6+
- 大约 12M 图文样本；
- guidance scale = 8；
- 目标更偏向视觉质量而非 FID 最优化。

4.2 Baseline 与评测指标

对比基线 包括：

GAN：BigGAN-deep、StyleGAN-T、GigaGAN 等；
多步 diffusion sampler：ADM、DPM++、UniPC；
few-step / one-step diffusion distillation：Progressive Distillation、TRACT、Consistency Model、Diff-Instruct、InstaFlow、LCM-LoRA、UFOGen。

指标：

FID：衡量生成分布与真实分布的距离；
CLIP Score：用于 text-to-image 对齐度评估；
Latency：batch size = 1 的单样本推理时间。

评测协议：

CIFAR-10 / ImageNet：各生成 50,000 张图，用 EDM 官方评测代码算 FID；
COCO-30K：随机采样 30,000 个 COCO 2014 val prompts，生成 512×512 图像后下采样到 256×256，和 40,504 张真实图像计算 FID，并用 OpenCLIP-G 计算 CLIP score。

4.3 训练配置

设置	Teacher / Backbone	Pair 数据构造	优化器 / 训练细节	训练规模
CIFAR-10 conditional	EDM `edm-cifar10-32x32-cond-vp`	100k noise-image pairs，Heun 18 steps	AdamW, lr=5e-5, wd=0.01, betas=(0.9,0.999), LPIPS(VGG), $λ_{re g} = 0.25$ , grad clip=10	7 GPU，总 batch 392，300k iter
CIFAR-10 unconditional	EDM `edm-cifar10-32x32-uncond-vp`	500k pairs，Heun 18 steps	AdamW，同上，但 $λ_{re g} = 0.5$	同量级设置
ImageNet-64	EDM `edm-imagenet-64x64-cond-adm`	25k pairs，Heun 256 steps	AdamW, lr=2e-6, wd=0.01, betas=(0.9,0.999), mixed precision, $λ_{re g} = 0.25$ , grad clip=10	7 GPU，总 batch 336，350k iter
LAION-Aesthetic 6.25+	Stable Diffusion v1.5	500k pairs，PNDM 50 steps, guidance=3	AdamW, lr=1e-5, wd=0.01, betas=(0.9,0.999), tiny VAE decode for reg, LPIPS(VGG), mixed precision, grad checkpoint	72 GPU，总 DM batch 2304，reg batch 1152，20k iter
LAION-Aesthetic 6+	Stable Diffusion v1.5	12M pairs，PNDM 50 steps, guidance=8	AdamW, lr=1e-5, wd=0.01, betas=(0.9,0.999), mixed precision, grad checkpoint；作者记录了多轮 schedule 调整	约 80×A100，训练两周，最佳版本累计 165k iter

补充几点实现细节：

CIFAR / ImageNet 的噪声区间采用 $σ_{m i n} = 0.002, σ_{m a x} = 80$ ，离散为 1000 bins；
LPIPS 前通常把图像上采样到 224×224；
distribution matching loss 与 fake score denoising loss 的权重都设为 1；
对高 guidance T2I，作者在 Table 5 里展示了逐步把 regression pair 数量、regression 权重、VAE 类型和 max DM step 调到更优配置的过程。

5. Experimental Results (实验结果)

5.1 ImageNet-64×64：one-step 逼近 teacher

Figure T1 解读：这是论文 Table 1 的裁剪图，核心信息是 DMD 在 1 次前向 的约束下把 ImageNet-64×64 的 FID 做到 2.62。它不仅显著优于已有 one-step / few-step 蒸馏方法，也把与 teacher EDM（2.32）的差距压到 0.3 左右，说明“分布匹配 + 小规模 paired regression”确实能把 one-step 模型拉到接近多步 diffusion teacher 的质量。

Method	Fwd Pass	FID ↓
BigGAN-deep	1	4.06
ADM	250	2.07
Progressive Distillation	1	15.39
DFNO	1	7.83
BOOT	1	16.30
TRACT	1	7.43
Meng et al.	1	7.54
Diff-Instruct	1	5.57
Consistency Model	1	6.20
DMD	1	2.62
EDM (Teacher)	512	2.32

如果把 Appendix E 的 CIFAR-10 结果也一并看：

DMD-conditional 在 CIFAR-10 上达到 2.66 FID；
unconditional DMD 为 3.77 FID；
虽仍落后于 EDM teacher（1.84），但对 one-step 模型来说已非常有竞争力。

5.2 消融：Distribution Matching 与 Regression 缺一不可

Figure 5 解读：上半部分对比“有 / 无 distribution matching”时的生成结果。没有 distribution matching 时，样本虽然还能维持局部结构，但 realism 和整体语义完整性明显下降。下半部分对比“有 / 无 regression loss”，去掉 regression 后模型更容易 mode collapse，论文特地用重复出现的灰色汽车示例说明样本多样性显著受损。这个图和 Figure 3 一起说明：DMD 不是单一损失就能稳定工作，而是两个分支共同决定了 realism 与 diversity 的平衡。

Figure T2 解读：这是论文 Table 2 的裁剪图。左表是损失项消融，右表是不同 sample weighting 的比较。可以直接看出：去掉 distribution matching 后，CIFAR-10 / ImageNet FID 都明显恶化；去掉 regression loss 后虽然有分布驱动，但训练会更不稳定且更容易 mode dropping。右表则说明作者提出的 Eq. (8) 权重明显好于直接使用 $σ_{t} / α_{t}$ 或 $σ_{t}^{3} / α_{t}$ 的做法。

关键消融结果如下：

Training loss	CIFAR FID ↓	ImageNet FID ↓
w/o Distribution Matching	3.82	9.21
w/o Regression Loss	5.58	5.61
DMD	2.66	2.62

Sample weighting	CIFAR FID ↓
$σ_{t} / α_{t}$	3.60
$σ_{t}^{3} / α_{t}$	3.71
Eq. (8) (Ours)	2.66

5.3 Zero-shot Text-to-Image：接近 SD v1.5，速度约 30×

Figure 1 解读：论文开头的 teaser 很有代表性。作者把 DMD 的 one-step 输出和 Stable Diffusion 50-step 输出混在一起，让读者盲猜哪张来自哪个模型。这个展示方式想表达的不是“DMD 每张都更好”，而是：在大量常规 prompt 下，one-step 结果已经足够接近 teacher，使得肉眼区分明显变难。对一篇 acceleration 论文来说，这比单纯报一个 latency 数字更有说服力。

Figure 6 解读：这是与 InstaFlow、LCM 和 Stable Diffusion 的定性比较。左列 DMD 为 1 step / 90ms，中间方法虽然也快，但细节、材质、prompt 对齐度通常不如 DMD；最右侧的 50-step SD 仍然最稳，但 DMD 已经把差距压到了一个“速度优势远大于视觉劣势”的范围内。也就是说，DMD 证明了 one-step T2I 并非只能生成“看起来像草图”的结果，而是能进入可用的高质量区间。

在 low-guidance（guidance=3）设置下，MS COCO-30K 的主结果如下：

Method	Resolution	Latency ↓	FID ↓
DPM++ (4 step)	512	0.26s	22.36
UniPC (4 step)	512	0.26s	19.57
LCM-LoRA (4 step)	512	0.19s	23.62
InstaFlow-0.9B	512	0.09s	13.10
UFOGen	512	0.09s	12.78
DMD	512	0.09s	11.49
SD v1.5 (Teacher)	512	2.59s	8.78

结论非常清楚：

与 teacher 相比，DMD 的 FID 只差 2.71；
但延迟从 2.59s 降到 0.09s；
在作者报告中，FP16 下可达 20 FPS。

5.4 High-guidance 结果：DMD 仍然领先同类加速方法

论文还专门训练了一个 guidance scale = 8 的模型，用于更贴近高视觉质量部署场景。Table 4 的关键结果如下：

Method	Latency ↓	FID ↓	CLIP Score ↑
DPM++ (4 step)	0.26s	22.44	0.309
UniPC (4 step)	0.26s	23.30	0.308
LCM-LoRA (1 step)	0.09s	77.90	0.238
LCM-LoRA (2 step)	0.12s	24.28	0.294
LCM-LoRA (4 step)	0.19s	23.62	0.297
DMD	0.09s	14.93	0.320
SD v1.5 (Teacher)	2.59s	13.45	0.322

这组结果说明：

DMD 不只在“对 FID 友好的 guidance=3”设定下有效；
在更常见的高 guidance 视觉质量设定里，它仍显著优于同类 one/few-step baseline；
但它与 teacher 的 gap 也会更明显，说明高 guidance regime 下的蒸馏难度更高。

5.5 局限性、实现代价与整体结论

论文明确提到的局限性：

与 100 / 1000 步等更精细的 diffusion sampling 相比，one-step 仍有轻微质量差距；
模型上限受 teacher 限制，尤其在 text rendering、小人脸、小目标细节上继承了 SD v1.5 的缺点；
text-to-image 训练使用固定 guidance scale，缺乏推理时灵活调节能力。

从工程实现角度看，原始 DMD 还额外有三个成本：

必须离线构造 paired dataset：对大模型尤其昂贵；
需要同时维护 generator 与 fake critic：显存和训练复杂度都不低；
原始 T2I DMD 缺少官方代码：很多工程细节要靠论文 appendix 和后续 DMD2 - Improved Distribution Matching Distillation for Fast Image Synthesis 反推。

整体结论：

DMD 的真正贡献，不只是把 FID 做到了 2.62 / 11.49；
更重要的是它证明了：one-step diffusion distillation 不必执着于路径一致性，也可以从 distribution matching 的角度获得高质量生成器；
这种“real score - fake score + light paired regularization”的框架，直接启发了后续的 DMD2 - Improved Distribution Matching Distillation for Fast Image Synthesis，后者进一步去掉了 regression dataset 并引入 adversarial loss / two time-scale update，把质量继续往上推。

Paper Notes

探索

DMD: One-step Diffusion with Distribution Matching Distillation

DMD: One-step Diffusion with Distribution Matching Distillation

1. Motivation (研究动机)

2. Idea (核心思想)

3. Method (方法)

3.1 整体框架

3.2 Base model、one-step generator 与带噪 score

3.3 Distribution Matching Loss

3.4 Regression Loss 与最终目标

3.5 Classifier-Free Guidance (CFG) 的适配

3.6 基于论文 Appendix + 复现代码整理的伪代码

3.6.1 构造 paired dataset

3.6.2 Distribution Matching Loss

3.6.3 Fake score model 的去噪更新

3.6.4 一次完整训练步

3.6.5 推理

3.7 实现指南：顺序、接口、数据流与验证

3.7.1 推荐实现顺序

3.7.2 关键接口约定

3.7.3 最小数据流检查

3.7.4 最小 smoke test

3.8 Code-to-paper mapping 与 Paper-vs-Code 差异

3.8.1 代码映射表（以 `devrimcavusoglu/dmd` 为主）

3.8.2 Paper vs Code Differences

4. Experimental Setup (实验设置)

4.1 数据集与任务

4.2 Baseline 与评测指标

4.3 训练配置

5. Experimental Results (实验结果)

5.1 ImageNet-64×64：one-step 逼近 teacher

5.2 消融：Distribution Matching 与 Regression 缺一不可

5.3 Zero-shot Text-to-Image：接近 SD v1.5，速度约 30×

5.4 High-guidance 结果：DMD 仍然领先同类加速方法

5.5 局限性、实现代价与整体结论

目录

Paper Notes

探索

DMD: One-step Diffusion with Distribution Matching Distillation

DMD: One-step Diffusion with Distribution Matching Distillation

1. Motivation (研究动机)

2. Idea (核心思想)

3. Method (方法)

3.1 整体框架

3.2 Base model、one-step generator 与带噪 score

3.3 Distribution Matching Loss

3.4 Regression Loss 与最终目标

3.5 Classifier-Free Guidance (CFG) 的适配

3.6 基于论文 Appendix + 复现代码整理的伪代码

3.6.1 构造 paired dataset

3.6.2 Distribution Matching Loss

3.6.3 Fake score model 的去噪更新

3.6.4 一次完整训练步

3.6.5 推理

3.7 实现指南：顺序、接口、数据流与验证

3.7.1 推荐实现顺序

3.7.2 关键接口约定

3.7.3 最小数据流检查

3.7.4 最小 smoke test

3.8 Code-to-paper mapping 与 Paper-vs-Code 差异

3.8.1 代码映射表（以 devrimcavusoglu/dmd 为主）

3.8.2 Paper vs Code Differences

4. Experimental Setup (实验设置)

4.1 数据集与任务

4.2 Baseline 与评测指标

4.3 训练配置

5. Experimental Results (实验结果)

5.1 ImageNet-64×64：one-step 逼近 teacher

5.2 消融：Distribution Matching 与 Regression 缺一不可

5.3 Zero-shot Text-to-Image：接近 SD v1.5，速度约 30×

5.4 High-guidance 结果：DMD 仍然领先同类加速方法

5.5 局限性、实现代价与整体结论

目录

3.8.1 代码映射表（以 `devrimcavusoglu/dmd` 为主）