Distribution Matching Distillation without Fake Score Network¶
会议: NeurIPS 2026
arXiv: 2605.19256
OpenReview: https://openreview.net/forum?id=b7mPyu8rgW
领域: 模型压缩
关键词: 扩散模型蒸馏, 概率流映射, 分布匹配, 假样本速度代理, 无辅助评分网络
一句话总结¶
针对少步生成中分布匹配蒸馏（DMD）依赖辅助网络估计生成分布分数的显存与优化开销，本文提出 FSF-DMD，利用流映射生成器固有的端点伪速度作为无需辅助网络的假样本速度代理，结合流映射一致的反向模拟与自教师机制，在少步生成与从头训练中均取得优于传统 DMD2 的效果。
研究背景与动机¶
扩散模型与流匹配模型通过多步概率流常微分方程（PF-ODE）数值求解展现了强大的生成能力，但其高昂的采样延迟促使学术界转向少步乃至单步生成模型。当前主流方案主要沿着两条互补范式展开：基于前向散度的流映射（Flow-Map）与一致性模型，以真实数据计算前向轨迹积分，具备良好的模式覆盖特性但容易在低概率区域散布密度；基于反向散度的分布匹配蒸馏（DMD），通过生成样本在加噪分布上的真实分数与假样本分数之差引导模式收缩，有效提升生成质量。近期研究通常联合优化流映射目标与 DMD 目标，以兼顾模式覆盖与模式寻求优势。
然而，DMD 的核心依赖于一个辅助假样本分数（或假样本速度）估计网络，用于在训练过程中动态追踪正在快速演化的生成分布。在模型训练中，维护这样一个额外的网络不仅带来了显著的显存压力与计算开销，更导致脆弱的双时间尺度更新策略（TTUR）或复杂的分布对齐启发式规则，因为辅助网络往往无法在有限的更新步数内准确跟上生成器的激进更新。这种生成器与辅助跟踪器之间的动态失配，构成了基于 DMD 的少步蒸馏管线中的主要工程瓶颈。
本文敏锐地指出，在生成器本身具备流映射结构的体系下，显式的假样本评分网络并非必不可少。流映射天然具备时间可逆性与单射性质，其端点伪速度（endpoint pseudo-velocity）在概率流 ODE 耦合下严格坍缩为假样本后验速度。核心 idea：利用流映射生成器的端点伪速度直接作为假样本速度的解析代理，彻底摒弃辅助假样本评分网络，实现轻量、自同步且高保真的无需显式跟踪器的分布匹配蒸馏。
方法详解¶
整体框架¶
FSF-DMD 的总体管线围绕流映射生成器展开，将流映射的前向一致性蒸馏目标 \(\mathcal{L}_{\mathrm{CD}}\) 与无需辅助网络的分布匹配目标 \(\mathcal{L}_{\mathrm{FSF\text{-}DMD}}\) 联合优化。在每次训练迭代中，流映射生成器接收高斯噪声并输出少步生成样本；随后样本经过前向独立扰动生成加噪样本，由教师网络计算真实速度场，同时由当前流映射生成器自身评估端点伪速度作为假样本速度代理；二者构建速度残差后通过 stop-gradient 机制反馈优化生成器。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入噪声与条件样本"] --> B["阶段1：流映射少步生成<br/>多步欧拉累加输出样本"]
B --> C["阶段2：一致性蒸馏锚定<br/>JVP 保证流映射结构完备"]
B --> D["加噪扰动生成样本<br/>独立扰动构建中间态"]
D --> E["阶段3：生成器端点伪速度代理<br/>无需辅助网络直接计算假样本速度"]
D --> F["冻结教师网络<br/>评估真实速度场"]
E --> G["速度残差与动态缩放<br/>构建无需假网络匹配信号"]
F --> G
G --> H["阶段4：流映射一致反向模拟更新<br/>Stop-Gradient 轨迹级梯度回传"]
C --> I["联合梯度反向传播更新生成器"]
H --> I关键设计¶
1. 生成器端点伪速度代理：消除辅助评分网络的显式追踪开销
在传统 DMD 中，假样本速度网络 \(v_\psi\) 试图逼近生成分布的条件期望速度 \(v_{\mathrm{fake}}(\hat{x}_t;t) = \mathbb{E}_{\hat{x}|\hat{x}_t}[t^{-1}(\hat{x}_t - \hat{x})]\)。当扰动过程独立于生成器输入时，后验分布 \(p(\hat{x}|\hat{x}_t)\) 在全空间积分极其繁复，传统方法不得不引入辅助网络 \(v_\psi\) 经由流匹配迭代近似。而在流映射模型中，生成样本与加噪过程若沿概率流 ODE（PF-ODE）轨迹建立逆映射耦合 \(\hat{x}_t = f(\hat{x}; 0, t)\)，由于速度场的 Lipschitz 连续性保证了 ODE 解的单射性与唯一可逆性，后验分布严格退化为狄拉克脉冲：\(p(\hat{x}|\hat{x}_t) = \delta(\hat{x} - f(\hat{x}_t; t, 0))\)。于是后验期望直接得到解析解： $\(v_{\mathrm{fake}}(\hat{x}_t; t) = \frac{\hat{x}_t - f(\hat{x}_t; t, 0)}{t} = F_\theta(\hat{x}_t; t, 0)\)$ 其中 \(F_\theta(\hat{x}_t; t, 0)\) 即为流映射模型从时间 \(t\) 积分映射到端点 \(0\) 的伪速度。FSF-DMD 将该端点伪速度作为无偏且自洽的假样本速度替代物，无需任何外部分支或辅助网络，使假样本信号随着生成器参数演化实时同步更新，从根源上消除了双网络对抗步数不平衡与训练震荡。
2. 一致性蒸馏正则化锚定：保证端点伪速度结构自洽
端点伪速度 \(F_\theta(\hat{x}_t; t, 0)\) 之所以能严格充当假样本速度代理，前提是生成器函数 \(f_\theta(x_t; t, s) = x_t + (s-t)F_\theta(x_t; t, s)\) 在全时域内严格满足流映射积分规律。若直接仅优化分布匹配目标，生成器会迅速退化，丢失流映射流形结构。为此，FSF-DMD 引入无对抗判别器的一致性蒸馏损失 \(\mathcal{L}_{\mathrm{CD}}\) 作为几何锚定约束： $\(\mathcal{L}_{\mathrm{CD}} = \mathbb{E}_{x,z,t,s} \left[ w_{t,s} \left\| F_\theta(x_t; t, s) - \mathrm{sg}\left[ v_\Phi(x_t; t) + (s-t)\frac{d}{dt}F_\theta(x_t; t, s) \right] \right\|_2^2 \right]\)$ 该目标通过雅可比向量积（JVP）的有限差分计算时域全导数，确保伪速度网络对概率流 ODE 轨迹的忠实逼近。实验证明，去掉 \(\mathcal{L}_{\mathrm{CD}}\) 会导致端点代理严重偏离实际假样本流场，导致训练发散；而二者协同则提供了天然的真实数据正向模式覆盖信号，免去了传统 DMD2 必需的判别器 GAN 损失。
3. 流映射一致的反向模拟机制：消除少步生成的时序失配
在推断阶段，少步流映射生成遵循分段欧拉累加：\(\hat{x} = z - \frac{1}{N}\sum_{i=1}^N F_\theta(\hat{x}_{i/N}; \frac{i}{N}, \frac{i-1}{N})\)。传统 DMD2 的反向模拟基于去噪跳步采样，与流映射沿连续轨迹分段积分的形式并不兼容。FSF-DMD 专门针对流映射设计了一致的反向模拟目标，通过对中间前向状态施加 stop-gradient，将复合生成轨迹表达为可微端点等效平均伪速度 \(\tilde{F}_\theta(z; 1, 0) = \frac{1}{M}\sum_{i=1}^M F_\theta(\mathrm{sg}[\hat{x}_{i/M}]; \frac{i}{M}, \frac{i-1}{M})\)。构造出的 \(M\) 步少步分布匹配损失为： $\(\mathcal{L}_{\mathrm{FSF\text{-}DMD}}(\tilde{F}, \Phi) = \mathbb{E}_{z,t} \left[ w_t \left\| \tilde{F}_\theta(z; 1, 0) - \mathrm{sg}\left[ \tilde{F}_\theta(z; 1, 0) - \Delta_{\tilde{F},\Phi}(\hat{x}_t, t) \right] \right\|_2^2 \right]\)$ 其中速度残差 \(\Delta_{\tilde{F},\Phi}(\hat{x}_t, t) = F_\theta(\hat{x}_t; t, 0) - v_\Phi(\hat{x}_t; t)\)。结合动态权重 \(w_t = |\tilde{F}_\theta(z; 1, 0) - v_\Phi(\hat{x}_t; t)|^{-1}\) 与时间转移偏置采样，精确消除了少步自回归累加误差。
4. 自教师扩展机制：无教师环境下的从头自反向散度修正
由于 MeanFlow 形式的参数化在端点极限极限满足 \(F_\theta(x_t; t, t) \approx v_t^*(x_t)\)，生成器本身在瞬时速度极限下天然充当瞬时速度场评估器。FSF-DMD 据此打破了对离线教师网络 \(v_\Phi\) 的依赖，将外部教师速度替换为自身的对角线速度估算 \(F_\theta(\hat{x}_t; t, t)\)，形成内省式残差： $\(\Delta_F(\hat{x}_t, t) = F_\theta(\hat{x}_t; t, 0) - F_\theta(\hat{x}_t; t, t)\)$ 将一致性蒸馏替换为从头一致性训练 \(\mathcal{L}_{\mathrm{CT}}\) 后，模型无需任何预训练教师即可从随机初始化开始，在学习概率流流形的同时由内部反向散度驱动分布纠偏，突破了传统 DMD 只能用于蒸馏已有大模型的范式局限。
损失函数 / 训练策略¶
最终蒸馏目标结合了一致性蒸馏与无假样本网络分布匹配： $\(\mathcal{L}_{\mathrm{distill}} = \mathcal{L}_{\mathrm{CD}} + \lambda \mathcal{L}_{\mathrm{FSF\text{-}DMD}}(\tilde{F}, \Phi)\)$ 默认超参数设置：\(\lambda = 0.05\)，\(w_{t,s} = \cos(t)\)，时间转移参数 \(\gamma = 10.0\)。采用 NorMuon 优化器，学习率为 \(5\times 10^{-4}\)，批大小为 128，采用 BF16 精度。网络架构为 LightningDiT-B/1（133.8M 参数），通过有限差分近似步长 \(\epsilon = 0.005\) 计算 JVP。
实验关键数据¶
主实验¶
在 ImageNet-1K \(256\times 256\) 基准测试中，评估 2 步采样下的生成质量与训练效率。
| 方法 | 显存需求 (VRAM) | 训练耗时 (sec/step) | 收敛峰值步数 (Peak Step) | 2步生成 FID50K (↓) | Inception Score (↑) |
|---|---|---|---|---|---|
| Pretrained Flow Map (iMF-B/1) | - | - | - | 9.10 | 148.71 |
| Teacher Velocity Network (8步) | - | - | - | 8.54 | - |
| Consistency Distillation (Sabour et al.) | 47 GB | 0.23 | 36K | 6.06 | 188.90 |
| DMD2 (TTUR=1:5 + GAN) | 60 GB | 1.46 | 15K | 4.18 | 237.12 |
| DMD2 + IDA 动态参数混合 | 60 GB | 1.46 | 12K | 4.20 | 244.09 |
| DMD2 + \(\mathcal{L}_{\mathrm{CD}}\) (联合一致性) | 76 GB | 2.16 | 54K | 4.67 | 221.87 |
| FSF-DMD (本文, M=1 单步模拟) | 49 GB | 0.31 | 3K | 3.96 | 207.83 |
| FSF-DMD (本文, M=2 双步模拟) | 52 GB | 0.38 | 4K | 3.85 | 262.50 |
在放宽初始化条件至纯流匹配预训练（Flow-Matching Init）及从头训练（From-Scratch）的设置下，FSF-DMD 亦表现出色：
| 训练配置 / 初始化 | 方法 | 训练步数 | FID50K (↓) | Inception Score (↑) |
|---|---|---|---|---|
| Flow-Matching 初始化 | Consistency Distillation | 19K | 5.53 | 280.90 |
| Flow-Matching 初始化 | DMD2 (TTUR=1:5) | 3K | 5.53 | 287.33 |
| Flow-Matching 初始化 | DMD2 + \(\mathcal{L}_{\mathrm{CD}}\) | 33K | 4.90 | 262.03 |
| Flow-Matching 初始化 | FSF-DMD (本文) | 8K | 4.70 | 246.01 |
| 随机初始化 (从头训练) | iMF 基线 | 400K | 9.10 | 148.71 |
| 随机初始化 (从头训练) | iMF + DMD (TTUR=1:5) | 400K | 13.38 | 138.80 |
| 随机初始化 (从头训练) | iMF + TwinFlow | 400K | 12.17 | 144.26 |
| 随机初始化 (从头训练) | RCGM + TwinFlow | 400K | 22.58 | 106.04 |
| 随机初始化 (从头训练) | FSF-DMD (本文自教师) | 400K | 6.35 | 194.91 |
消融实验¶
| 蒸馏配置与组件消融 | FID50K | 说明 |
|---|---|---|
| 纯一致性蒸馏基线 (\(\mathcal{L}_{\mathrm{CD}}\)) | 6.06 | 缺少分布级反向散度修正 |
| + FSF-DMD 基础代理 (式 14, M=1) | 5.70 | 引入端点伪速度代理，显现分布修正增益 |
| + DMD2 风格反向模拟 (\(M=2\)) | 5.81 | 传统去噪步进与流映射欧拉积分失配 |
| + 流映射一致反向模拟 (式 22, \(M=2\)) | 5.34 | 保持流映射连续轨迹一致性，显著优于 DMD2 模拟 |
| + 动态损失加权与时间转移 (\(\lambda=0.05, \gamma=10\)) | 4.45 | 归一化自适应梯度权重使优化更平稳 |
| + 随机时间掩码 (\(s \leftarrow t\), 完整模型) | 3.85 | 锚定瞬时流速度，达到最佳 2 步生成水准 |
| 完整模型退化为单步模拟 (\(M=1\)) | 3.96 | 得益于流映射半群性质，退化极小 |
| DMD2 退化为单步模拟 | 4.78 | 传统 DMD2 缺乏半群连续性，单步模拟掉点严重 |
关键发现¶
- 收敛速度与显存开销：FSF-DMD 仅需 4K 训练步即可达到 3.85 FID，比 DMD2（15K 步）快近 4 倍；由于省去了辅助网络与鉴别器，显存开销从 DMD2 的 60GB 骤降至 49GB（单步模拟）/ 52GB（双步模拟），单步训练耗时从 1.46s 缩减至 0.38s，加速近 3.8 倍。
- 流映射半群性质优势：当模拟步数由 \(M=2\) 降至 \(M=1\) 时，FSF-DMD 的 FID 仅从 3.85 轻微上升至 3.96；相比之下，DMD2 从 4.18 剧烈恶化至 4.78。这表明流映射的内在半群性质 \(f(x_t; t, s) = f(f(x_t; t, r); r, s)\) 赋予了端点伪速度对单步估计的高度鲁棒性。
- 从头训练的稳定性：在无预训练教师的从头训练中，传统 DMD 联合训练甚至发生崩溃（FID 恶化至 13.38），而 TwinFlow 也仅取得 12.17；FSF-DMD 自教师机制无需生成器兼顾额外的反向速度头，将 FID 推进至 6.35，证明了假样本速度代理在纯正向流形构造过程中的自洽纠偏能力。
亮点与洞察¶
- 理论推导的极简闭式转化：发现概率流 ODE 耦合下粒子后验分布严格退化为点脉冲，使原本需要整网逼近的高维假样本速度解析坍缩为端点伪速度，设计极其精妙。
- 计算图完全同构与即时同步：摒弃额外的假样本网络意味着反向散度纠偏信号与生成器权重完全同源，免去了 TTUR 的双频更新调参，避免了跟踪滞后导致的模式坍塌。
- 可复用的少步架构方法论：该假样本速度代理范式不仅适用于扩散蒸馏，亦可无缝推广至连续时间一致性模型、大语言模型多步思维链蒸馏等任何具备常微分方程流形结构的生成架构中。
局限与展望¶
- 理论成立依赖流映射强假设：推导依赖于生成器结构具有良好的流映射性质及 ODE 速度场的 Lipschitz 连续性，若流映射前向训练未良好收敛，端点伪速度代理的质量将显著下滑。
- 评测范围尚待向十亿级文生图扩展：当前实验集中于 ImageNet-1K \(256\times 256\) 及 DiT-Base 架构，尚未在大规模十亿参数文生图模型（如 FLUX、SD3）上进行充分验证。
- 高阶求解器兼容性：当前反向模拟与前向生成均采用一阶欧拉法，未来值得探索将端点伪速度代理扩展至 Runge-Kutta、Heun 等高阶 ODE 数值积分格式。
相关工作与启发¶
- vs DMD / DMD2: DMD2 需引入辅助网络 \(v_\psi\) 经由流匹配目标追踪演化生成分布，并配合 GAN 损失与 TTUR 启发式平衡；本文证明在流映射架构下，辅助评分网络是冗余的，直接利用端点伪速度即可实现更优质量与更高算力利用率。
- vs TwinFlow: TwinFlow 虽亦省去了独立网络，但强行迫使生成器网络在负时间参数上学习第二任务（自对抗反向速度头），造成了生成容量的割裂；FSF-DMD 仅使用同一生成头求端点极限，容量全部专注于图像生成。
- vs rCM / SenseFlow: rCM 引入切空间归一化稳定训练但带来计算额外负担；SenseFlow 需借助隐式分布参数混合（IDA）与额外的图像基础模型判别器；FSF-DMD 保持端到端一致性蒸馏，结构更轻量统一。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙发现流映射在 PF-ODE 下假样本速度与端点伪速度的解析等价性，彻底去除了辅助评分网络。
- 实验充分度: ⭐⭐⭐⭐☆ 详尽对比了流映射初始化、流匹配初始化与从头训练三种严苛设定，指标全面，唯缺大尺寸文生图评估。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密清晰，动机由浅入深，算法伪代码与实验消融设计极为规范。
- 价值: ⭐⭐⭐⭐⭐ 极大简化了扩散少步蒸馏管线的工程复杂度与显存门槛，推动了流映射模型的实用化部署。