Soft Diffusion:谷歌新框架从通用扩散中正确调度、学习和采样

选自arXiv

机器之心编译

编辑：蛋酱

近来，扩散模型成为 AI 领域的研究热点。谷歌研究院和 UT-Austin 的研究者在最新的一项研究中充分考虑了「损坏」过程，并提出了一个用于更通用损坏过程的扩散模型设计框架。

我们知道，基于分数的模型和去噪扩散概率模型（DDPM）是两类强大的生成模型，它们通过反转扩散过程来产生样本。这两类模型已经在 Yang Song 等研究者的论文《Score-based generative modeling through stochastic differential equations》中统一到了单一的框架下，并被广泛地称为扩散模型。

目前，扩散模型在包括图像、音频、视频生成以及解决逆问题等一系列应用中取得了巨大的成功。Tero Karras 等研究者在论文《Elucidating the design space of diffusionbased generative models》中对扩散模型的设计空间进行了分析，并确定了 3 个阶段，分别为 i) 选择噪声水平的调度，ii) 选择网络参数化（每个参数化生成一个不同的损失函数），iii) 设计采样算法。

近日，在谷歌研究院和 UT-Austin 合作的一篇 arXiv 论文《Soft Diffusion: Score Matching for General Corruptions》中，几位研究者认为扩散模型仍有一个重要的步骤：损坏（corrupt）。一般来说，损坏是一个添加不同幅度噪声的过程，对于 DDMP 还需要重缩放。虽然有人尝试使用不同的分布来进行扩散，但仍缺乏一个通用的框架。因此，研究者提出了一个用于更通用损坏过程的扩散模型设计框架。

具体地，他们提出了一个名为 Soft Score Matching 的新训练目标和一种新颖的采样方法 Momentum Sampler。理论结果表明，对于满足正则条件的损坏过程，Soft Score MatchIng 能够学习它们的分数（即似然梯度），扩散必须将任何图像转换为具有非零似然的任何图像。

在实验部分，研究者在 CelebA 以及 CIFAR-10 上训练模型，其中在 CelebA 上训练的模型实现了线性扩散模型的 SOTA FID 分数——1.85。同时与使用原版高斯去噪扩散训练的模型相比，研究者训练的模型速度显著更快。

论文地址：https://arxiv.org/pdf/2209.05442.pdf

方法概览

通常来说，扩散模型通过反转逐渐增加噪声的损坏过程来生成图像。研究者展示了如何学习对涉及线性确定性退化和随机加性噪声的扩散进行反转。

具体地，研究者展示了使用更通用损坏模型训练扩散模型的框架，包含有三个部分，分别为新的训练目标 Soft Score Matching、新颖采样方法 Momentum Sampler 和损坏机制的调度。

首先来看训练目标 Soft Score Matching，这个名字的灵感来自于软过滤，是一种摄影术语，指的是去除精细细节的过滤器。它以一种可证明的方式学习常规线性损坏过程的分数，还在网络中合并入了过滤过程，并训练模型来预测损坏后与扩散观察相匹配的图像。

只要扩散将非零概率指定为任何干净、损坏的图像对，则该训练目标可以证明学习到了分数。另外，当损坏中存在加性噪声时，这一条件总是可以得到满足。

具体地，研究者探究了如下形式的损坏过程。

在过程中，研究者发现噪声在实证（即更好的结果）和理论（即为了学习分数）这两方面都很重要。这也成为了其与反转确定性损坏的并发工作 Cold Diffusion 的关键区别。

其次是采样方法 Momentum Sampling。研究者证明，采样器的选择对生成样本质量具有显著影响。他们提出了 Momentum Sampler，用于反转通用线性损坏过程。该采样器使用了不同扩散水平的损坏的凸组合，并受到了优化中动量方法的启发。

这一采样方法受到了上文 Yang Song 等人论文提出的扩散模型连续公式化的启发。Momentum Sampler 的算法如下所示。

下图直观展示了不同采样方法对生成样本质量的影响。图左使用 Naive Sampler 采样的图像似乎有重复且缺少细节，而图右 Momentum Sampler 显著提升了采样质量和 FID 分数。

最后是调度。即使退化的类型是预定义的（如模糊），决定在每个扩散步骤中损坏多少并非易事。研究者提出一个原则性工具来指导损坏过程的设计。为了找到调度，他们将沿路径分布之间的 Wasserstein 距离最小化。直观地讲，研究者希望从完全损坏的分布平稳过渡到干净的分布。

实验结果

研究者在 CelebA-64 和 CIFAR-10 上评估了提出的方法，这两个数据集都是图像生成的标准基线。实验的主要目的是了解损坏类型的作用。

研究者首先尝试使用模糊和低幅噪声进行损坏。结果表明，他们提出的模型在 CelebA 上实现了 SOTA 结果，即 FID 分数为 1.85，超越了所有其他仅添加噪声以及可能重缩放图像的方法。此外在 CIFAR-10 上获得的 FID 分数为 4.64，虽未达到 SOTA 但也具有竞争力。

此外，在 CIFAR-10 和 CelebA 数据集上，研究者的方法在另一项指标采样时间上也表现更好。另一个额外的好处是具有显著的计算优势。与图像生成去噪方法相比，去模糊（几乎没有噪声）似乎是一种更有效的操纵。

下图展示了 FID 分数如何随着函数评估数量（Number of Function Evaluations, NFE）而变。从结果可以看到，在 CIFAR-10 和 CelebA 数据集上，研究者的模型可以使用明显更少的步骤来获得与标准高斯去噪扩散模型相同或更好的质量。