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

📄 arXiv: 2512.05251v1 📥 PDF

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

分类: stat.ML, cs.LG

发布日期: 2025-12-04

💡 一句话要点

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

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

关键词: 扩散模型 单步采样 自蒸馏 确定性流 ELBO估计 体积一致性 贝叶斯推断

📋 核心要点

  1. 现有采样算法计算成本高，需要多次迭代才能生成高质量样本。
  2. 通过学习步长条件ODE，使用自蒸馏和确定性流，实现单步采样。
  3. 在合成和贝叶斯基准测试中，显著减少计算量，同时保持样本质量和ELBO估计。

📝 摘要（中文）

从非归一化目标分布中采样是机器学习和统计学中的一项基本但具有挑战性的任务。现有的采样算法通常需要许多迭代步骤才能生成高质量的样本，导致计算成本高昂。本文提出了一种单步扩散采样器，它学习一个步长条件ODE，使得一个大的步长通过状态空间一致性损失重现许多小步长的轨迹。此外，本文还表明，由于常见的离散积分器产生不匹配的前向/后向转移核，扩散采样器中的标准ELBO估计在少步长情况下会退化。受此分析的启发，本文推导了一种用于ELBO估计的确定性流(DF)重要性权重，而无需后向核。为了校准DF，本文引入了一种体积一致性正则化，该正则化使沿流的累积体积变化在不同步长分辨率上对齐。因此，本文提出的采样器仅需一步或几步即可实现采样和稳定的证据估计。在具有挑战性的合成和贝叶斯基准测试中，它以数量级更少的网络评估实现了具有竞争力的样本质量，同时保持了稳健的ELBO估计。

🔬 方法详解

问题定义：现有的扩散模型采样方法通常需要大量的迭代步骤，导致计算成本非常高昂。尤其是在需要快速采样的场景下，例如实时应用或需要大量样本进行统计推断时，这种高计算成本就成为了一个瓶颈。此外，在少步采样的情况下，标准ELBO（Evidence Lower Bound）估计会退化，导致模型评估不准确。

核心思路：本文的核心思路是通过学习一个步长条件ODE（常微分方程），使得一个大的步长能够近似模拟多个小步长的轨迹。这通过一个状态空间一致性损失来实现，本质上是一种自蒸馏。同时，为了解决少步采样时ELBO估计退化的问题，论文引入了确定性流（Deterministic Flow）的重要性权重，避免了对后向核的依赖。

技术框架：整体框架包含以下几个主要部分：1) 学习一个步长条件ODE，该ODE描述了从噪声到数据的扩散过程。2) 使用状态空间一致性损失来训练该ODE，使得单步采样能够近似多步采样。3) 推导确定性流的重要性权重，用于在少步采样时进行准确的ELBO估计。4) 引入体积一致性正则化，用于校准确定性流，确保不同步长分辨率下体积变化的对齐。

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

关键设计：状态空间一致性损失用于确保单步采样与多步采样的轨迹一致。确定性流的重要性权重基于ODE的雅可比行列式计算，用于校正ELBO估计。体积一致性正则化通过比较不同步长分辨率下的体积变化来校准确定性流。具体的网络结构和参数设置取决于具体的应用场景，但通常会采用类似于U-Net的结构来学习步长条件ODE。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在合成和贝叶斯基准测试中，能够以数量级更少的网络评估次数，达到与现有方法具有竞争力的样本质量。例如，在某些基准测试中，该方法仅需一步采样即可达到与传统扩散模型数百步采样相当的性能，同时保持了稳健的ELBO估计，验证了该方法的有效性和优越性。

🎯 应用场景

该研究成果可广泛应用于需要快速采样的领域，例如贝叶斯推断、生成对抗网络（GAN）的训练、强化学习中的策略梯度估计等。通过减少采样所需的计算量，可以加速模型训练和推理过程，提高效率，并降低计算成本。此外，该方法在药物发现、材料设计等领域也具有潜在的应用价值。

📄 摘要（原文）

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 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.