One-Step Diffusion Samplers via Self-Distillation and Deterministic Flow

📄 arXiv: 2512.05251 📥 PDF

作者: Pascal Jutras-Dube, Jiaru Zhang, Ziran Wang, Ruqi Zhang

分类: stat.ML, cs.LG

发布日期: 2026-02-28

💡 一句话要点

提出基于自蒸馏和确定性流的单步扩散采样器，加速采样并稳定证据估计。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 单步采样 自蒸馏 确定性流 ELBO估计

📋 核心要点

  1. 现有的采样算法通常需要多次迭代才能生成高质量样本，计算成本高昂，限制了其应用。
  2. 本文提出单步扩散采样器，通过自蒸馏学习步长条件ODE，实现一步采样模拟多步采样的效果，加速采样过程。
  3. 实验表明，该方法在保证样本质量的同时，显著减少了网络评估次数，并保持了稳健的ELBO估计。

📝 摘要（中文）

本文提出了一种单步扩散采样器，通过学习步长条件ODE，利用状态空间一致性损失，使得一个大步能够复现多个小步的轨迹。研究表明，由于常见的离散积分器导致前向/后向转移核不匹配，扩散采样器中的标准ELBO估计在少步情况下性能下降。为此，本文推导了一种无需后向核的确定性流(DF)重要性权重用于ELBO估计。为了校准DF，引入了体积一致性正则化，以对齐不同步长分辨率下沿流的累积体积变化。因此，所提出的采样器仅需一步或几步即可实现快速采样和稳定的证据估计。在具有挑战性的合成和贝叶斯基准测试中，它以数量级更少的网络评估次数实现了具有竞争力的样本质量，同时保持了稳健的ELBO估计。

🔬 方法详解

问题定义：论文旨在解决从非归一化目标分布中进行高效采样的问题。现有扩散模型采样方法通常需要大量的迭代步骤，导致计算成本高，限制了其在实际应用中的可行性。此外，在少步采样的情况下，由于离散积分器带来的前向和后向转移核不匹配，导致ELBO估计的性能下降。

核心思路：论文的核心思路是学习一个步长条件ODE，使得仅需一步即可生成高质量的样本，从而显著减少采样所需的计算量。通过自蒸馏的方式，让单步采样器学习多步采样的轨迹，保证采样质量。同时，针对少步采样中ELBO估计不准确的问题，提出了基于确定性流(DF)的重要性权重，并引入体积一致性正则化来校准DF。

技术框架：该方法主要包含以下几个模块：1) 步长条件ODE学习模块，通过状态空间一致性损失，学习一步采样模拟多步采样的轨迹。2) 确定性流(DF)重要性权重估计模块，用于在少步采样情况下进行准确的ELBO估计。3) 体积一致性正则化模块，用于校准DF，保证不同步长分辨率下沿流的累积体积变化一致。整体流程是，首先训练步长条件ODE，然后使用DF重要性权重进行ELBO估计，并通过体积一致性正则化进行校准。

关键创新：论文的关键创新在于：1) 提出了单步扩散采样器，显著减少了采样所需的计算量。2) 提出了基于确定性流(DF)的重要性权重，解决了少步采样中ELBO估计不准确的问题。3) 引入了体积一致性正则化，用于校准DF，保证了ELBO估计的准确性。与现有方法相比，该方法能够在保证样本质量的同时，显著减少网络评估次数，并保持稳健的ELBO估计。

关键设计：状态空间一致性损失用于训练步长条件ODE，鼓励单步采样器学习多步采样的轨迹。确定性流(DF)重要性权重的计算基于前向过程，避免了对后向核的依赖。体积一致性正则化通过最小化不同步长分辨率下沿流的累积体积变化差异来校准DF。具体的网络结构和参数设置取决于具体的应用场景和数据集。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在合成数据集和贝叶斯基准测试中，仅需一步或几步即可达到与传统多步扩散模型相当的样本质量，同时网络评估次数减少了几个数量级。此外，该方法在少步采样情况下，能够保持稳健的ELBO估计，优于传统的ELBO估计方法。

🎯 应用场景

该研究成果可广泛应用于需要从复杂分布中进行采样的机器学习和统计学领域，例如贝叶斯推断、生成模型、强化学习等。其快速采样的特性使其在资源受限或对实时性要求高的场景下具有重要价值，例如移动设备上的图像生成、机器人控制等。未来，该方法有望进一步扩展到高维数据和更复杂的分布。

📄 摘要（原文）

Sampling from unnormalized target distributions is a fundamental yet challenging task in machine learning and statistics. Existing sampling algorithms typically require many iterative steps to produce high-quality samples, leading to high computational costs. We introduce one-step diffusion samplers which learn a step-conditioned ODE so that one large step reproduces the trajectory of many small ones via a state-space consistency loss. We further show that standard ELBO estimates in diffusion samplers degrade in the few-step regime because common discrete integrators yield mismatched forward/backward transition kernels. Motivated by this analysis, we derive a deterministic-flow (DF) importance weight for ELBO estimation without a backward kernel. To calibrate DF, we introduce a volume-consistency regularization that aligns the accumulated volume change along the flow across step resolutions. Our proposed sampler therefore achieves both fast sampling and stable evidence estimate in only one or few steps. Across challenging synthetic and Bayesian benchmarks, it achieves competitive sample quality with orders-of-magnitude fewer network evaluations while maintaining robust ELBO estimates.