跳转至

DiffRatio: Training One-Step Diffusion Models Without Teacher Supervision

会议: NeurIPS 2026
arXiv: 2502.08005
领域: 图像生成
关键词: 扩散模型、单步生成、密度比估计、变分分数蒸馏、推断期并行扩展

一句话总结

DiffRatio 抛弃了传统分数蒸馏中依赖预训练教师网络打分的双网络模式,直接通过时间步条件化的轻量密度比估计器计算学生分布与真实数据分布的对数比梯度,消除了双重分数估计偏差,在 CIFAR-10 与 ImageNet 上实现了领先的无教师单步扩散生成与推断期并行重要性重采样扩展。

研究背景与动机

基于扩散的生成模型通过正向扩散加噪与反向去噪过程建模复杂数据分布,在图像、视频与 3D 生成等领域取得了巨大成功。然而,标准高斯去噪扩散模型在采样时通常需要数十至上千步连续函数评估(NFEs),极大地限制了其实际部署与交互式生成效率。尽管微分方程求解器与自适应协方差估计等加速算法能够将采样步数压至 10 到 20 步,但要在极限的单步(1-step)或少步采样下保持高保真度,仍需依赖蒸馏(Distillation)方法。其中,以变分分数蒸馏(Variational Score Distillation, VSD)为代表的分数蒸馏范式将预训练多步教师扩散模型蒸馏为单步隐变量生成器,成为实现极速生成的关键技术路线。

然而,现有的两阶段分数蒸馏流水线存在严重的梯度估计偏差累积问题。在第一阶段,预训练教师网络在真实数据分布上学习分数函数,受限于网络容量与优化困难,教师分数必然携带估计误差;在第二阶段,学生生成器不仅依赖这一不完美的教师网络提供外部分数监督,还需额外维护一个随训练动态更新的学生分数网络来逼近自身分布的分数。最终施加给生成器的蒸馏梯度依赖于学生分数与教师分数的差值,导致双重估计偏差在此累积并放大;即使学生分数网络被充分训练,教师监督带来的固有系统偏差也无法消除。同时,蒸馏过程中必须同时维护并频繁评估教师模型、学生打分网络以及生成器本身,显存与计算开销巨大。

针对上述矛盾,本文提出了一种无需教师持续监督的全新训练视角:既然扩散散度(DiKL / DiJS)的梯度本质上取决于学生分布与数据分布的分数差,而分数差恒等于两者对数密度比对空间坐标的梯度,便无需将两者的分数割裂估计再相减。核心 idea:将单步扩散蒸馏转化为扩散时间步上的分类概率比估计,通过轻量级判别器直接拟合对数密度比的梯度,彻底摆脱教师在线监督与双网络打分误差,并天然赋予密度比在推断阶段作为无外挂奖励的并行验证重采样能力。

方法详解

整体框架

DiffRatio 的整体流程分为两阶段:首先在目标数据集上通过标准去噪分数匹配(DSM)预训练一个扩散网络,并固定其时间步参数 \(t=t_{\text{init}}\) 作为单步生成器的权重初始化;随后进入核心的无教师蒸馏阶段,完全冻结/丢弃多步教师监督,转而训练一个随扩散时间条件化的紧凑二分类判别网络 \(c_\eta(x_t, t)\)。在每一次训练迭代中,单步生成器从隐空间采样输出一批生成样本 \(x_0\),并将其扩散加噪到随机选取的噪声水平 \(t\);同时从真实数据分布中抽取样本并加入同分布噪声。轻量分类器学习区分带噪真实样本与带噪生成样本,进而输出贝叶斯后验概率以直接给出对数密度比。生成器接收由该密度比对输入坐标求导得到的解析梯度反向传播完成单步参数更新。在推断阶段,训练收敛的密度比分类器还可在无需外部 Reward 模型的情况下,对多个并行生成的候选样本执行重要性重采样,实现完全并行的质量扩展。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["标准 DSM 预训练权重<br/>固定时间步 t=t_init 初始化生成器"] --> B["时间条件化扩散加噪<br/>生成分布 q_θ(x_t) 与真实数据 p_d(x_t)"]
    B --> C["扩散密度比估计器<br/>轻量分类网络拟合贝叶斯后验"]
    C --> D["解析对数比梯度反传<br/>根据 DiKL / DiJS / DiRM 准则优化生成器"]
    D --> E["推断期并行扩展<br/>基于单步候选池的重要性重采样"]

关键设计

1. 扩散密度比估计与统一梯度估计器:消解双重分数估计偏差

传统 VSD 方法分别训练学生打分网络 \(s_{\psi_2}^{q_\theta}(x_t, t)\) 与教师打分网络 \(s_{\psi_1}^{p_d}(x_t, t)\),并利用两者的差值作为扩散散度的梯度信号: $\(\nabla_{\theta}\mathrm{DiKL} \approx \sum_{t=1}^T w(t) \int q_\theta(x_t) \left( s_{\psi_2}^{q_\theta}(x_t,t) - s_{\psi_1}^{p_d}(x_t,t) \right) \frac{\partial x_t}{\partial \theta} dx_t\)$ 两组参数相互独立的打分网络会导致各自的逼近误差直接叠加。DiffRatio 利用数学恒等式 \(\nabla_{x_t}\log q_\theta(x_t) - \nabla_{x_t}\log p_d(x_t) = \nabla_{x_t}\log\frac{q_\theta(x_t)}{p_d(x_t)}\),通过在加噪时间步 \(t\) 上引入平衡标签的二分类问题(\(y=1\) 为真实数据,\(y=0\) 为生成器输出),建立贝叶斯后验分类器 \(c_\eta(x_t, t) \approx p(y=1|x_t)\)。密度比直接转化为对数几率: $\(\nabla_{x_t}\log\frac{q_\theta(x_t)}{p_d(x_t)} = \nabla_{x_t} \operatorname{logit}(1 - c^*(x_t, t))\)$ 生成器的梯度更新直接改写为单个分类网络输出 logit 的空间导数。该设计不仅将必须同时在线维护的网络数量由三个压缩为两个,避免了两个分数估计网络的误差累积,而且消除了教师模型的固定偏差传递,在大样本极限下具备渐近一致性保证。

2. 泛化散度目标族构建:从 DiKL 到 DiJS 与真实度最大化

在确定了密度比分类器 \(c_\eta\) 之后,扩散生成的目标不再受限于逆向 KL 散度。通过改变对分类器预测值的变换形式,DiffRatio 将训练目标拓展为一个可灵活选择的散度族。具体而言,采用对数变换代替 logit 变换,可以推导出扩散 Jensen-Shannon 散度(DiJS)的梯度估计器: $\(\nabla_\theta \mathrm{DiJS}(q_\theta(x_0) \| p_d(x_0)) \approx \sum_{t=1}^T w(t) \int q_\theta(x_t) \nabla_{x_t} \log (1 - c_\eta(x_t, t)) \frac{\partial x_t}{\partial \theta} dx_t\)$ 这一形式在各个时间尺度上最小化生成样本被判别为假的概率。更进一步,本文还提出了扩散真实度最大化准则(DiRM),直接最大化生成样本被判定为真实的对数似然: $\(\nabla_\theta \mathrm{DiRM}(\theta) \approx - \sum_{t=1}^T w(t) \int q_\theta(x_t) \nabla_{x_t} \log c_\eta(x_t, t) \frac{\partial x_t}{\partial \theta} dx_t\)$ 与传统 GAN 相比,DiffRatio 摆脱了判别器与生成器之间脆弱且极易震荡的极小极大博弈动态,不需要梯度惩罚(Gradient Penalty)或谱归一化(Spectral Normalization)等启发式稳定技巧,即可在不同噪声尺度下平稳收敛。

3. 轻量级 U-Net 编码器判别结构与多层级特征初始化

为了提高训练时的计算效率与显存利用率,DiffRatio 的分类器 \(c_\eta(x_t, t)\) 并不复用完整的教师 U-Net 架构,而是仅采用原 U-Net 骨干的 Encoder(下采样)半侧结构,并在末端外接 Sigmoid 激活层映射为标量概率。该分类网络参数量仅为标准打分网络的约一半,大幅降低了每轮前向反向传播的 FLOPs。与此同时,在高维复杂图像建模场景中,直接从随机初始化训练单步生成器容易发生模式坍塌。消融研究揭示,单纯依靠端到端泛函映射(Functional Mapping)不足以抑制坍塌,关键在于预训练扩散模型在加噪去噪过程中所学习到的丰富多层级表征特征(Feature Space Richness)。因此,DiffRatio 仅在初始化阶段将生成器加载为预训练 U-Net 并固定初始时间戳 \(t_{\text{init}}=2.5\),此后完全摆脱教师监督。

4. 推断期并行重要性重采样扩展

传统基于评分或对抗蒸馏的方法在训练结束后,辅助网络(学生评分网络或判别器)通常被彻底遗弃。而 DiffRatio 学习到的时间条件分类器在最小噪声水平 \(t_{\text{min}}\) 下保留了对生成分布与真实分布之间微小残差的判决能力。由于目标分布满足 \(p_d(x_0) = q_\theta(x_0) \frac{p_d(x_0)}{q_\theta(x_0)}\),系统可以在推断期以单步生成器独立并行生成 \(M\) 个候选样本 \(x_0^{(1:M)} \sim q_\theta(x_0)\),并利用学到的密度比权重计算离散重要性概率: $\(w(x_0^{(i)}) \approx \frac{c_\eta(x_{t_{\text{min}}}^{(i)}, t_{\text{min}})}{1 - c_\eta(x_{t_{\text{min}}}^{(i)}, t_{\text{min}})}, \quad i \sim \mathrm{Cat}\left( \frac{w(x_0^{(1)})}{\sum_j w(x_0^{(j)})}, \dots, \frac{w(x_0^{(M)})}{\sum_j w(x_0^{(j)})} \right)\)$ 这一采样过程使得单步扩散模型首次拥有了类似于大语言模型的“Test-time Compute Scaling”机制。更为关键的是,由于所有候选样本生成与判别器评分均在时间维度相互独立,重采样计算可以 100% 并行执行,无需任何额外的自回归或串行扩散步骤。

损失函数 / 训练策略

算法以交替单步(Single-step update)更新方式迭代执行。在每轮迭代中: 1. 按照幂律或预设加噪策略采样扩散时间 \(t' \in [t_{\text{min}}, t_{\text{max}}]\); 2. 生成器输出 \(x_0 = g_\theta(z)\) 并扩散加噪为 \(x_{t'}\),同时对真实批次样本加噪为 \(x'_{t'}\); 3. 以标准二分类交叉熵损失更新分类器 \(c_\eta\):\(\mathcal{L}_{\text{BCE}} = - \mathbb{E}_{p_d}[\log c_\eta(x'_{t'}, t')] - \mathbb{E}_{q_\theta}[\log(1 - c_\eta(x_{t'}, t'))]\); 4. 冻结分类器参数,根据 DiJS 目标(采用方差爆炸 VE 噪声权重 \(w(t) = \sigma_t^2\))将解析空间导数 \(\nabla_{x_t}\log(1 - c_\eta)\) 反向传播至生成器参数 \(\theta\); 5. 在 CIFAR-10 上使用批大小 64、学习率 \(5 \times 10^{-5}\);在 ImageNet(\(64\times 64\) 与 \(512\times 512\))上使用批大小 4096,无需任何判别器梯度惩罚或对抗截断技巧。

实验关键数据

主实验

在 CIFAR-10 无条件生成、ImageNet \(64\times 64\) 条件生成及 ImageNet \(512\times 512\) 高分辨率条件生成任务上,DiffRatio 与多步教师扩散模型、从头训练的一步模型以及各类教师监督蒸馏方案进行了系统对比。

数据集 方法分类 代表方法 采样步数 (NFE) FID (↓) Inception Score (↑)
CIFAR-10 (32×32) 教师模型 (多步) EDM (Karras et al.) 35 2.04 9.84
CIFAR-10 (32×32) 从头训练 (无教师) iCT-deep (Song & Dhariwal) 1 2.51 9.76
CIFAR-10 (32×32) 教师监督蒸馏 Progressive Distillation 1 8.34 8.69
CIFAR-10 (32×32) 教师监督蒸馏 Diff-Instruct (VSD 路线) 1 4.53 -
CIFAR-10 (32×32) 教师监督蒸馏 SiD (Zhou et al.) 1 2.03 10.02
CIFAR-10 (32×32) 教师监督蒸馏 SiDA (Zhou et al.) 1 1.52 10.32
CIFAR-10 (32×32) 本文方法 (无教师监督) DiffRatio-DiRM 1 4.87 9.85
CIFAR-10 (32×32) 本文方法 (无教师监督) DiffRatio-DiKL 1 3.81 9.90
CIFAR-10 (32×32) 本文方法 (无教师监督) DiffRatio-DiJS 1 2.39 9.93
ImageNet 64×64 教师模型 (多步) EDM (Karras et al.) 79 2.44 -
ImageNet 64×64 教师监督蒸馏 Diff-Instruct 1 5.57 -
ImageNet 64×64 教师监督蒸馏 SiD 1 2.02 -
ImageNet 64×64 教师监督蒸馏 DMD2 (w/ GAN) 1 1.51 -
ImageNet 64×64 本文方法 (无教师监督) DiffRatio-DiJS 1 1.54 -
ImageNet 512×512 教师模型 (多步) EDM2-M (w/ CFG) 126 2.01 -
ImageNet 512×512 教师监督蒸馏 SiD-M 1 2.06 -
ImageNet 512×512 教师监督蒸馏 SiDA-M 1 1.55 -
ImageNet 512×512 本文方法 (无教师监督) DiffRatio-DiJS-M 1 1.41 -

在 ImageNet \(512\times 512\) 极端高分辨率设置下,DiffRatio-DiJS-M 仅凭 1 步前向传播便取得了 1.41 的 FID,不仅大幅超越了耗费 126 步的教师模型 EDM2-M(FID 2.01),也显著优于包括 SiD-M(2.06)与 SiDA-M(1.55)在内的全部教师蒸馏基线。

消融实验与分析

1. 2D 高斯混合玩具实验中的分数差估计精度与生成质量

为直接定量证明密度比估计消除梯度偏差的能力,论文构建了由 40 个高斯成分组成的 2D 分布,计算解析真实分数差 \(\Delta s^* = \nabla \log q_\theta - \nabla \log p_d\) 并对比梯度精度与生成分布质量。

方法 估计机制 相对 L2 误差 (↓) 余弦相似度 (↑) Log-MMD (↓) Log-density (↑)
真实基准 (True) 解析真实分布 - 1.000 - -6.65
传统变分分数蒸馏 (VSD) 双网络 DSM 差值 (\(s^{q_\theta} - s^{p_d}\)) 显著偏高 显著偏低 -5.25 -10.88
DiffRatio (本文) 单分类器对数几率导数 (\(\nabla \operatorname{logit}\)) 显著偏低 全噪声谱稳定接近 1.0 -6.08 -6.79

2. 权重初始化策略与模式坍塌机制消融 (CIFAR-10)

分析不同权重初始化模式在 DiJS 准则训练单步生成器中的表现,解耦泛函映射与特征丰富度的作用。

初始化方式 预训练数据集与类别配置 生成表现 (FID) 训练状态
无初始化 (Random Init) - Collapsed 严重模式坍塌
单时间步 DSM 预训练 (Single-level) 完整 CIFAR-10 (固定单一 \(t_{\text{init}}\)) Collapsed 严重模式坍塌
多层级 DSM 预训练 (Multi-level) CIFAR-100 (排他 10 类) Collapsed 模式坍塌
多层级 DSM 预训练 (Multi-level) CIFAR-100 (排他 50 类) 6.20 稳定收敛,无坍塌
多层级 DSM 预训练 (Multi-level) CIFAR-100 (排他 90 类) 6.01 稳定收敛,表征多样
多层级 DSM 预训练 (Multi-level) 完整目标集 CIFAR-10 2.39 最佳生成质量

3. 推断期并行扩展 (Inference-Time Parallel Scaling) 评估

在 ImageNet \(512\times 512\) 上,利用训练好的轻量分类器在 \(t_{\text{min}}\) 处的估计密度比 \(w(x_0)\) 对 \(M\) 个并行样本执行重要性重采样(全部 NFE 均以并行方式执行,不增加串行延迟)。

候选样本数 \(M\) 串行延迟开销 累计并行 NFE ImageNet 512×512 FID (↓) 相对增益
1 (标准单步) 1 步 1 1.41 基线
10 1 步 \(10^\dagger\) 1.35 -0.06
20 1 步 \(20^\dagger\) 1.34 -0.07
50 1 步 \(50^\dagger\) 1.34 -0.07
100 1 步 \(100^\dagger\) 1.32 -0.09

关键发现

  • 消除双重偏差带来质的飞跃:在 2D 玩具实验中,VSD 的两个打分网络误差直接相加,使生成分布的 Log-density 恶化至 -10.88;而 DiffRatio 的单分类器架构将 Log-density 拉回至 -6.79,极度贴近真实数据上限(-6.65)。
  • DiJS 展现最佳稳定性与指标平衡:在散度族对比中,DiffRatio-DiJS 实现了 2.39 的 FID,明显优于 DiKL(3.81)与 DiRM(4.87),验证了对称的扩散 Jensen-Shannon 散度在稳定单步扩散梯度方向上的优越性。
  • 模式坍塌根源于表征缺失而非映射缺失:单时间步预训练拥有与多步预训练相同的初值空间映射,但依然发生模式坍塌;而利用 CIFAR-100 非重叠类别预训练的多尺度特征能够完全遏制坍塌。这明确证实扩散多尺度去噪特征的广度才是单步模型对抗坍塌的支柱。
  • 推断期无需额外奖励模型的并行 Scaling:随着候选样本数 \(M\) 从 1 提升至 100,FID 单调递减至 1.32,证明了训练收敛的密度比分类器作为内置 Verifier 的可靠性。

亮点与洞察

  • 将分数差重构为单网络密度比梯度:绕开了过去两阶段分数蒸馏必须同时学习学生打分与教师打分的固有思维定势,利用贝叶斯法则将分数差精确对应到二分类后验的对数几率空间导数,架构极简且数学优雅。
  • 非对抗训练实现极速稳定收敛:不同于常规 Diffusion-GAN 需要复杂的判别器正则化、渐进式网络结构或极其脆弱的超参微调,DiffRatio 不受极小极大博弈博弈面崩塌困扰,使用纯粹的二分类交叉熵即可平稳驱动生成器。
  • 重采样赋能零延迟测试期扩展:将训练阶段副产物(判别网络)变废为宝,在推断时无缝转变为无外挂打分器的分布校准器,候选样本生成与打分天然可跨 GPU 完全并行。

局限与展望

  • 依赖第一阶段基础扩散预训练权重:虽然在第二阶段蒸馏时摆脱了教师模型的在线监督与反向传播,但生成器初值仍需要预训练的多层级扩散特征来防止模式坍塌,尚无法做到严格意义上的“端到端从零冷启动单步训练”。
  • 极端流形维度退化风险:在高维图像空间中,当生成分布与真实分布重叠度极低时,分类器在微小噪声水平下的判别边界可能会过饱和,导致密度比梯度的有效动态范围受限。
  • 未来方向:探索无监督自编码潜空间内轻量表征预训练替代完整 DSM 预训练的可行性;将 DiffRatio 框架推广至文本到图像大模型(如 SDXL/FLUX)及视频生成扩散骨干的高保真少步蒸馏。

相关工作与启发

  • vs VSD (Variational Score Distillation) / Diff-Instruct: VSD 与 Diff-Instruct 依赖预训练教师网络打分并在线拟合学生网络打分,受制于两个网络的累积误差。DiffRatio 仅用单个轻量分类器直接学习分数差,消除了教师偏差累积,参数开销减半。
  • vs SiD / SiDA (Score Identity Distillation): SiD 系列需要教师网络权重、教师分数监督、学生分数网络三位一体,训练流程复杂。DiffRatio 完全摒弃在线教师打分网络,仅用单一二分类网络即可超越 SiDA 取得更优的 FID。
  • vs Diffusion-GAN: Diffusion-GAN 强依赖 StyleGAN 定制生成器骨干与复杂的对抗平衡技巧。DiffRatio 对网络架构完全无关,可直接在通用 EDM / EDM2 U-Net 上稳定运行,且无需额外梯度惩罚。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙利用对数比梯度将两阶段双网络打分压缩为单分类器密度比估计,理论洞察深刻。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D 玩具实验机理验证、CIFAR-10 及 ImageNet 512 高分辨率对比、初始化消融与推断扩展,数据详实。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑闭环,公式推导严谨,动机与设计高度契合。
  • 价值: ⭐⭐⭐⭐⭐ 为无教师监督单步扩散模型开辟了全新路径,推断并行扩展机制极具工业落地价值。