扩散模型的设计空间

在讨论扩散模型时，我们常常会产生一个想法，那就是一般情况下扩散模型的前向扩散过程是人为给定的，当前向扩散过程给定了，其对应的反向扩散过程也就被定义了，那么这个我们人为设定的扩散过程是不是最优的那个过程呢？假如不是的话，我们又该如何在避免发生优化过程中的“耦合”(改变一个因素的同时实际上影响到了多个因素)的情况下有的放矢地去对这个扩散过程进行优化呢？这篇来自英伟达团队的《Elucidating the Design Space of Diffusion-Based Generative Models》就比较完善地解答了这个问题。该文章的核心思想就是尽可能将模型训练以及各种参数选取的过程解耦合，让模型的训练过程变成一个与参数选取无关的纯黑盒过程，而推理过程则与参数选取相关，可以认为各种参数选择的差异实际上是我们对连续的ODE过程选取了不同的离散化策略【但是这里的分析并不是一般化的参数设置，所以关于这一部分仍然有研究空间】

现有扩散模型的形式化统一

设原有的数据分布空间为$p({\bf x})$，其标准差为$\sigma_d$，考虑通过向数据添加独立同分布的高斯噪声（标准差为$\sigma$）得到的平滑分布族$p({\bf x},\sigma)$。扩散模型的核心思想是随机采样一个噪声图像${\bf x}_N\sim\mathcal{N}({\bf 0},\sigma_{\max}^2{\bf I})$，逐步对其去噪变为图像序列${\bf x}_i$，其中${\bf x}_i\sim p({\bf x}_i;\sigma_i)$。【其中$\sigma_i$一般下尺度会随着$i$不断增大】【这是本文对于扩散模型的认知，但是我感觉这貌似更像是SMLD的理念，而不是DDPM的理念】

在该文中，其将扩散过程形式化为

$$\begin{align} d\hat {\bf x}(t)=\sqrt{2\dot\sigma(t)\sigma(t)}d{\bf w}(t)\\ {\bf x}(t)=s(t)\hat{\bf x}(t) \end{align}$$

这样的话，加噪过程和压缩真值的过程就被解耦了，加入的噪声都被实实在在地加到原始真值上了【这个加噪过程和宋飏提出的SMLD形式是完全一致的，即可以认为在$t$时在真值上加入的噪声总量的标准差为$\sigma(t)$】，而不是被$s(t)$压缩后的隐变量空间。进一步，我们可以将${\bf x}(t)$的SDE写出来

$$\begin{align} d{\bf x}(t)=\frac{\dot s(t)}{s(t)}{\bf x}(t)dt+s(t)\sqrt{2\dot\sigma(t)\sigma(t)}d{\bf w}(t) \end{align}$$

ODE采样

考虑到前面博客我们提到的前向SDE和反向ODE的对应关系，我们可以推知反向ODE为

$$\begin{align} d{\bf x}(t)=\left[\frac{\dot s(t)}{s(t)}{\bf x}(t)-s^2(t)\dot\sigma(t)\sigma(t)\nabla_{ {\bf x}(t) }\log q_t({\bf x}(t))\right]dt \end{align}$$

这和EDM原文中的结果还有一点差异，就是梯度的形式，EDM中是关于$\hat{\bf x}(t)$的分布的梯度，这二者的关系我们可以比较容易推导出来

$$\begin{align} q({\bf x}(t))d{\bf x}(t)=\hat q(\hat{\bf x}(t))d\hat{\bf x}(t)\\ q({\bf x}(t))=\hat q(\hat{\bf x}(t))\frac{1}{s(t)^n} \end{align}$$

代入对数函数并求梯度后，实际上就变成了

$$\begin{align} d{\bf x}(t)=\left[\frac{\dot s(t)}{s(t)}{\bf x}(t)-s^2(t)\dot\sigma(t)\sigma(t)\nabla_{ {\bf x}(t) }\log \hat q_t\left(\frac{ {\bf x}(t)}{s(t)}\right)\right]dt \end{align}$$

其中$\hat q$是关于$\hat{\bf x}$的分布【即加噪真值，这样的话，在模型训练过程中就无需引入关于$s(t)$的信息了】。我们可以用模型拟合分数项

$$\mathbb{E}_{ {\bf y}\sim p_{data}}\mathbb{E}_{ {\bf n}\sim\mathcal{N}({\bf 0},\sigma^2{\bf I})}||D({\bf y}+{\bf n};\sigma)-{\bf y}||_2^2$$ $$\nabla_{ {\bf x}(t) }\log \hat q_t\left(\frac{ {\bf x}(t)}{s(t)}\right)=\frac{D\left(\frac{ {\bf x}(t)}{s(t)};\sigma(t)\right)-\frac{ {\bf x}(t)}{s(t)} }{\sigma(t)^2}$$

SDE采样

考虑到前面博客我们提到的前向SDE和反向SDE的对应关系，我们也可以推知反向SDE为

$$d{\bf x}(t)=\left\{\frac{\dot s(t)}{s(t)}{\bf x}(t)-2s^2(t)\dot\sigma(t)\sigma(t)\nabla_{ {\bf x}(t) }\log \hat q_t\left(\frac{ {\bf x}(t)}{s(t)}\right)\right\}dt+s(t)\sqrt{2\dot\sigma(t)\sigma(t)}d\bar{\bf w}(t)$$

在本文中，其将该形式做了重新组织

$$\begin{align} d{\bf x}(t)=\Bigg\{\frac{\dot s(t)}{s(t)}{\bf x}(t)-s^2(t)\dot\sigma(t)\sigma(t)\nabla_{ {\bf x}(t) }\log \hat q_t\left(\frac{ {\bf x}(t)}{s(t)}\right)\Bigg\}dt\\-\Bigg\{s^2(t)\dot\sigma(t)\sigma(t)\nabla_{ {\bf x}(t) }\log \hat q_t\left(\frac{ {\bf x}(t)}{s(t)}\right)\Bigg\}dt+s(t)\sqrt{2\dot\sigma(t)\sigma(t)}d\bar{\bf w}(t) \end{align}$$

其将$s^2(t)\frac{\dot\sigma(t)}{\sigma(t)}$视为系数$\beta(t)$，那么形式就变成了

$$\begin{align} d{\bf x}(t)=\underbrace{\Bigg\{\frac{\dot s(t)}{s(t)}{\bf x}(t)-s^2(t)\dot\sigma(t)\sigma(t)\nabla_{ {\bf x}(t) }\log \hat q_t\left(\frac{ {\bf x}(t)}{s(t)}\right)\Bigg\}dt}_{ODE项} \\\underbrace{\underbrace{-\Bigg\{\beta(t)\sigma^2(t)\nabla_{ {\bf x}(t) }\log \hat q_t\left(\frac{ {\bf x}(t)}{s(t)}\right)\Bigg\}dt}_{确定性噪声衰减}+\underbrace{\sqrt{\beta(t)}\sigma(t)d\bar{\bf w}(t)}_{不确定性噪声注入} }_{朗之万SDE} \end{align}$$

文章认为，确定性噪声衰减项和不确定噪声的注入这两项中的随机性实际彼此抵消了【额外进行了了噪声消除，但是增加了新的噪声】，使得SDE过程在某个时刻的方差和ODE过程的方差是一致的，因此$\beta(t)$实际上表示原始噪声被新噪声替换的相对速率。

这一视角揭示了随机性在实际中的价值：隐式的朗之万扩散将样本驱动到当前时间步的期望边缘分布，主动修正先前采样步骤中的误差。然而，用离散SDE求解器步长近似朗之万项本身也会引入误差。先前的研究表明非零$\beta(t)$具有积极作用，但据我们所知，宋飏等人中$\beta$的隐式选择并不具备特殊性质。因此，最优随机性程度应通过实验确定。

得分匹配网络训练

在监督式神经网络训练中，存在多种公认的最佳实践。例如，建议将输入和输出信号的幅度固定为（例如）单位方差，并避免每个样本的梯度幅度出现大幅波动。若直接训练神经网络来建模$D$将远非理想，因为输入和输出的方差都在随着噪声方差的扩大而逐渐扩大。

考虑到上述问题，此前研究的训练往往是训练一个$F_\theta$，使得$D_\theta({\bf x};\sigma)={\bf x}-\sigma F_\theta({\bf x};\sigma)$，但是这样处理也存在的一个问题是，$F_\theta({\bf x};\sigma)$的输出的误差会被放大$\sigma$倍，因此当$\sigma$较大时神经网络需要做十分精细的调整，否则会带来较大的误差隐患，这种时候看起来直接输出$D_\theta({\bf x};\sigma)$反而比较合理。本文受先前自适应混合信号与噪声的参数化方法启发，提出用$\sigma$依赖的跳跃连接来预条件神经网络，使其能够估计${\bf y},{\bf n}$或者二者的中间状态。$D_\theta$的表示如下所示

$$D_\theta({\bf x};\sigma)=c_{\text{skip} }(\sigma){\bf x}+c_{\text{out} }(\sigma)F_\theta(c_{\text{in} }(\sigma){\bf x};c_{\text{noise} }(\sigma))$$

其中$F_\theta$为待训练的神经网络，$c_{\text{skip} }(\sigma)$调节跳跃连接，$c_{\text{in} }(\sigma)$和$c_{\text{out} }(\sigma)$用于缩放输入和输出尺度，$c_{\text{noise} }(\sigma)$将噪声水平映射为$F_\theta$的条件输入【根据我们前面的分析，在噪声尺度较大时，$c_{\text{skip} }(\sigma)$会变小，这样后面的神经网络$F_\theta$将会承担更近似于$D_\theta$的拟合任务。$c_{\text{noise} }(\sigma)$一般是基于经验确定】。在训练时也引入一定的期望加权

$$\mathbb{E}_{\sigma,{\bf y},{\bf n} }[\lambda(\sigma)||D({\bf y}+{\bf n};\sigma)-{\bf y}||^2_2]$$

引入的$\lambda(\sigma)$相当于引入了一个噪声水平的概率分布$p_{\text{train} }(\sigma)$，在训练时，噪声水平$\sigma$从该分布中随机选取。再将上述式子写为与$F_\theta$相关的形式

$$\begin{align} \mathbb{E}_{\sigma,{\bf y},{\bf n} }[\lambda(\sigma)||c_{\text{skip} }(\sigma)({\bf y}+{\bf n})+c_{\text{out} }(\sigma)F_\theta(c_{\text{in} }(\sigma)({\bf y}+{\bf n});c_{\text{noise} }(\sigma))-{\bf y}||^2_2]\\ \mathbb{E}_{\sigma,{\bf y},{\bf n} }\left[\lambda(\sigma)c_{\text{out} }(\sigma)^2\left\Vert F_\theta(c_{\text{in} }(\sigma)({\bf y}+{\bf n});c_{\text{noise} }(\sigma))-\frac{ {\bf y} - c_{\text{skip} }(\sigma)({\bf y}+{\bf n})}{c_{\text{out} }(\sigma)}\right\Vert^2_2\right] \end{align}$$

文章认为上述训练目标中的有效权重实际上是$\lambda(\sigma)c_{\text{out} }(\sigma)^2$，因此为了平衡不同训练样本的损失权重，其令$\lambda(\sigma)=1/c_{\text{out} }(\sigma)^2$。

ODE采样的高阶方法

本文中的ODE采样方法与DDPM以及SMLD等文章中所用的方法不同，为一种高阶方法，下面参考苏神的博客做简单介绍。

【待完善】