快速扩散模型数值采样，近似平均方向求解器

deardao2024-04-27 11:45

浙大提出近似平均方向求解器，快速扩散模型数值采样

扩散模型的生成轨迹几乎位于图像空间的一个二维子空间中！于是，通过直接学习生成轨迹的平均方向，我们可以进一步减小离散误差。

论文题目：

Fast ODE-based Sampling for Diiffusion Models in Around 5 Steps

论文地址：

项目地址：

给定噪声输入，扩散模型利用评分函数（score function）进行迭代去噪实现图像生成、音频生成、视频生成、文生图等任务。这个过程可以解释为离散求解某个随机微分方程（SDE），或是求解其对应的概率流常微分方程（PF-ODE）。

在扩散模型加速采样上，现有方法主要分为两大类。一是设计更快的数值求解器，在增加步长同时保持较小的离散误差。这些方法成功地将采样步数从 1000 减少至 20 以下。二是基于知识蒸馏，在数据分布和预先指定的噪声分布之间建立一一映射。如此，训练好的学生模型只需一步就可以实现高质量生成。

本文结合两类方法的长处，在保持较低训练开销的同时，利用蒸馏得到的知识，进一步加速扩散模型采样。

-- 提出 AMED-Solver，一种新的扩散模型 ODE 求解器，其通过学习近似平均方向来最小化离散误差。

-- 提出 AMED-Plugin，一个可以应用于各种扩散模型 ODE 求解器的插件，仅引入较小的训练开销和可忽略的采样开销。

我们称离散求解过程中中间输出构成的集合为扩散模型的采样轨迹。

扩散模型的采样轨迹虽然位于一个维度非常高（通常上千甚至上万）的空间，但其可以近似地用两个主成分来刻画，也就是说，这条轨迹几乎处在一个二维子空间上！

利用 EDM 提供的 Heun Solver $2$ 进行 40 步离散采样得到 1000 条轨迹，对每条轨迹做主成分分析并提取部分主成分重构轨迹，得到的相对重构误差非常小，重构轨迹也能解释原轨迹 99% 以上的方差。