Denoising diffusion models (DDMs) offer a flexible framework for sampling from high dimensional data distributions. DDMs generate a path of probability distributions interpolating between a reference Gaussian distribution and a data distribution by incrementally injecting noise into the data. To numerically simulate the sampling process, a discretisation schedule from the reference back towards clean data must be chosen. An appropriate discretisation schedule is crucial to obtain high quality samples. However, beyond hand crafted heuristics, a general method for choosing this schedule remains elusive. This paper presents a novel algorithm for adaptively selecting an optimal discretisation schedule with respect to a cost that we derive. Our cost measures the work done by the simulation procedure to transport samples from one point in the diffusion path to the next. Our method does not require hyperparameter tuning and adapts to the dynamics and geometry of the diffusion path. Our algorithm only involves the evaluation of the estimated Stein score, making it scalable to existing pre-trained models at inference time and online during training. We find that our learned schedule recovers performant schedules previously only discovered through manual search and obtains competitive FID scores on image datasets.

翻译：去噪扩散模型（DDMs）为从高维数据分布中采样提供了一个灵活的框架。DDMs通过在数据中逐步注入噪声，生成一条连接参考高斯分布与数据分布之间的概率分布插值路径。为了数值模拟该采样过程，必须选择一个从参考分布返回至干净数据的离散化调度方案。合适的离散化调度对于获得高质量样本至关重要。然而，除了手工设计的启发式方法外，选择该调度方案的通用方法仍然难以捉摸。本文提出了一种新颖算法，用于根据我们推导出的成本函数自适应地选择最优离散化调度。我们的成本函数衡量了模拟过程将样本从扩散路径中的一点传输至下一点所做的功。该方法无需超参数调优，并能自适应扩散路径的动力学与几何特性。该算法仅涉及对估计的斯坦分数（Stein score）的评估，使其能够在推理时扩展到现有预训练模型，并在训练期间在线运行。我们发现，通过学习得到的调度方案能够恢复先前仅通过手动搜索发现的性能优良的调度，并在图像数据集上获得具有竞争力的FID分数。