Min-SNR-γ

写在前面

我在做扩散模型训练时，AI 总是会给我推荐一个叫做 Min-SNR-$\gamma$ 的加权策略，最近我终于忍不住去拜读了一下相关论文。但是让我奇怪的是，提出这个看起来非常简单有效的训练技巧的论文的引用量居然只有三百多，相比 DL 领域动辄成千上万的引用数量而言，实在是有点让人费解。我实际读完之后也感觉这篇文章有点戛然而止，虎头蛇尾的意思，当然也可能是我水平不足，无法窥得其中的精髓。

Motivation

作者认为目前训练扩散模型普遍存在收敛速度慢的问题，究其原因，在于不同的扩散时间步（timestep, $t$）的优化方向不一致。而我们在训练扩散模型时，每次会随机采样一个 $t$，然后针对该 $t$ 计算梯度并执行梯度下降，此时就会导致其他 $t$ 的性能变差。

为了证明这件事，作者为此专门做了一个实验。将扩散时间步 $t\in [1,1000]$ 分桶，例如每 100 个为一桶。然后在一个训练好的模型上，只针对某个特定桶中的 $t$ 进行微调优化，然后查看不同 $t$ 上的 Loss 变化。结果如下图所示。

图中不同颜色的曲线代表了在不同的桶中进行针对性微调的后 Loss 的变化。可以看出来微调的部分 $t$ 的 Loss 明显下降，而其他部分的 Loss 则上升，且距离越远越明显。

为了解决这个问题，作者们提出了 Min-SNR-$\gamma$ 加权策略。

Method

Preliminary

我们回顾一下扩散模型的基本框架，包含两个方面，一个前向加噪过程，一个反向去噪过程。前向加噪过程中将训练数据表示为 $\mathbf x_0\sim P(\mathbf x_0)$，通过对 $\mathbf x_0$ 不断地添加噪声直到原始信息几乎消失，其数学描述为：

$$\begin{align} q(\mathbf x_t|\mathbf x_0)&=\mathcal N(\mathbf x_t;\alpha_t \mathbf x_0,\sigma_t^2 \mathbf I) \tag{1}\\ \mathbf x_t&=\alpha_t \mathbf x_0+\sigma_t \epsilon \tag{2} \end{align}$$

其中 $\epsilon \in \mathcal N(0,\mathbf I)$。

反向过程则是从噪声逐渐去噪直到恢复为原数据 $x_0$。但是反向过程无法直接求解，因此对其进行参数化，将其表述为：

$$\begin{align} p_\theta(\mathbf x_{t-1}|\mathbf x_t)=\mathcal N(\mathbf x_{t-1};\hat\mu_\theta(\mathbf x_t),\hat\Sigma_\theta(\mathbf x_t)) \tag{3} \end{align}$$

其中 $\hat\mu_\theta(\mathbf x_t), \hat\Sigma_\theta(\mathbf x_t)$ 由神经网络学习得到。但是 DDPM 的作者们发现 $\hat\Sigma_\theta$ 可直接设置为 $\sigma^2_t \mathbf I$，因此实际需要学习的只有 $\hat\mu_\theta$。同时他们发现直接预测噪声可以取得很好的效果，因此 DDPM 的 Loss 就变成了：

$$\begin{align} \mathcal L^t_{sample}(\theta)=\mathbb E_{\mathbf x_0,\epsilon}[||\epsilon-\hat\epsilon_\theta(\alpha_t\mathbf x_0+\sigma_t\epsilon)||^2_2]\tag{4} \end{align}$$

Multi-Task learning

前面我们提到，由于扩散模型在不同加噪步上的梯度不一致，导致了收敛速度慢。本文作者为了解决这一问题，将扩散模型的训练视为一个多任务学习，每一个 $t$ 代表了一个独立的学习目标。此时就变成了多任务优化问题，并且我们可以给出一个明确的定义，令 $\mathcal L^t(\theta)$ 表示不同 $t$ 下的损失，并且定义梯度更新方向为 $\delta\neq 0$ ，我们的优化目标就变成了

$$\begin{align} \mathcal L^t(\theta+\delta) \leq \mathcal L^t(\theta), \forall t \in \{1,...,T\}\tag{5} \end{align}$$

由于我们的目标是确定 $\delta$ 的取值，因此需要将其显式表达出来，所以我们对公式 5 的左侧进行一阶泰勒展开，得到

$$\begin{align} \mathcal L^t(\theta+\delta) \approx \mathcal L^t(\theta) + \langle \nabla_\theta \mathcal L^t(\theta), \delta \rangle \tag{6} \end{align}$$

此时公式 5 可以代换为

$$\begin{align} \langle \nabla_\theta \mathcal L^t(\theta), \delta \rangle \leq 0, \forall t \in \{1,...,T\}\tag{7} \end{align}$$

此时我们知道了，满足公式 7 的要求的更新方向就是可以照顾到所有 $t$ 的更新方向，那么 $\delta$ 到底怎么定义呢？我们不可能在高维空间中直接随机的去一个一个尝试，因此我们要想办法将公式 7 转换为一个可优化的问题。此处我们不按照原论文正文的表述，而是直接根据附录的推导过程来理解。

首先为了简化表示，我们令 $\mu=-\delta, g_t=\nabla_\theta \mathcal L^t(\theta)$，其中 $\mu$ 表示我们寻找的梯度，$g_t$ 表示不同 $t$ 下的当前梯度。此时我们的目标是找到一个 $\mu$ 使其满足

$$\begin{align} \langle\mu,g_t\rangle\geq 0 ,\forall t \in \{1,...,T\}\tag{8} \end{align}$$

这个公式可以等价为

$$\begin{align} \min_t \langle\mu,g_t\rangle\geq 0, \forall t \in \{1,...,T\}\tag{9} \end{align}$$

即对于内积最小 $t$ 而言，其内积都大于等于零。但是实际上我们的目标并不仅仅是满足这个公式，因为我们最终要找到梯度 $\mu$ 进行梯度下降，我们理所当然的希望能够获得一个较快的下降速度，所以我们还希望梯度在满足要求的情况下尽可能大。因此我们额外加一层要求，得到下式

$$\begin{align} \max_\mu\min_t \langle\mu,g_t\rangle \tag{10} \end{align}$$

但是这个目标有一个问题，即找到的 $\mu$ 可能过大，进而导致梯度下降不稳定，所以我们加一个正则项来限制 $\mu$ 的大小。

$$\begin{align} \max_\mu \left[\min_t \langle\mu,g_t\rangle-\frac{1}{2}||\mu||^2\right] \tag{11} \end{align}$$

可能有人注意到了从公式 10 开始我省略了大于等于零的条件。这是因为在公式 11 中令 $\mu=0$ 时目标函数等于 0，所以最大值不可能小于 0。如果最大值大于 0，说明存在一个可以让所有任务同时下降的方向；如果最大值等于 0，则说明此时已经是一个 Pareto 驻点。

由于 $g_t$ 是一组离散取值，我们用一组权重做一下恒等变换

$$\begin{align} \min_t\langle\mu,g_t\rangle=\min_{w_t}\sum_{t=1}^T w_t\langle\mu,g_t\rangle,\{w_t\geq0,\sum_{t=1}^T w_t=1\} \tag{12} \end{align}$$

将其代回公式 11

$$\begin{align} \max_\mu \left[\min_t \langle\mu,g_t\rangle-\frac{1}{2}||\mu||^2\right] &=\max_\mu\left[\min_{w_t}\sum_{t=1}^T w_t\langle\mu,g_t\rangle-\frac{1}{2}||\mu||^2\right]\\ &=\max_\mu\left[\min_{w_t}\left\langle\mu,\sum_{t=1}^T w_tg_t\right\rangle-\frac{1}{2}||\mu||^2\right] \tag{13} \end{align}$$

进一步化简，令 $G_w=\sum_{t=1}^T w_tg_t$。由于目标函数关于 $w_t$ 是线性的、关于 $\mu$ 是凹的，因此在满足强对偶条件时可以交换 $\min$ 和 $\max$ 的顺序：

$$\begin{align} \max_\mu\left[\min_{w_t}\langle\mu,G_w\rangle-\frac{1}{2}||\mu||^2\right]&=\max_\mu\min_{w_t}\left[\langle\mu,G_w\rangle-\frac{1}{2}||\mu||^2\right]\\ &=\min_{w_t}\max_\mu\left[\langle\mu,G_w\rangle-\frac{1}{2}||\mu||^2\right] \tag{14} \end{align}$$

对于公式 14 而言，其中 $\max(\cdot)$ 部分可以通过求导来计算，当 $\mu=G_w$ 时，$\max$ 成立，取得最大值。则公式 14 可以化简为

$$\begin{align} \min_{w_t}\max_\mu\left[\langle\mu,G_w\rangle-\frac{1}{2}||\mu||^2\right]&=\min_{w_t}\left[\langle G_w,G_w\rangle-\frac{1}{2}||G_w||^2\right]\\ &=\min_{w_t}\frac{1}{2}||G_w||^2\\ &=\min_{w_t}\frac{1}{2}\left\{\left\|\sum_{t=1}^T w_t\nabla_\theta\mathcal L^t(\theta)\right\|^2\;\middle|\;w_t\geq0,\sum_{t=1}^T w_t=1\right\} \tag{15} \end{align}$$

省略没有影响的分数项，就可以得到论文中的优化目标。

此时我们的目标是求解公式 15，从而得到最优加权方案。但是这里还有一点小问题，此时加权权重 $w_t$ 只有一个概率分布约束，但是在实际求解中可能得到极端值，即某个 $t$ 的 $w_t$ 非常接近 1 而导致其他权重几乎为零。这会导致训练不稳定，并不是我们想要的结果。因此我们加入一个正则项来约束一下。

$$\min_{w_t\in\Delta}\left\{\left\|\sum_{t=1}^T w_t\nabla_\theta\mathcal L^t(\theta)\right\|_2^2+\lambda\sum_{t=1}^T w_t^2\right\}, \Delta=\left\{w_t\geq 0,\sum_{t=1}^T w_t=1 \right\} \tag{16}$$

针对这个优化目标，我们有两个现成的方法，分别是 Frank-Wolfe 算法和 UGD（Unconstrained Gradient Descent），这两个方法都是迭代优化算法，这种迭代优化的特性导致其求解效率非常低下，每次梯度下降都要计算所有 $t$ 的解，并且其基于 mini-batch 得到的解也不够稳定。

这里我实在不想展开这两个算法了，我们直接看看作者提供的解法吧。

Min-SNR-γ

作者为了避免上述两种算法的缺点，干脆直接设定了一套固定的加权策略。需要注意，原论文为了统一不同的预测目标，首先是在 pred-$\mathbf x_0$ 的设置下讨论权重，而不是 pred-noise。作者测试了五种不同的加权策略，分别是：

1. Constant weighting：$w_t=1$，所有 $t$ 平等对待。
2. SNR weighting：$w_t=SNR(t)$，其中 $SNR(t)=\frac{\alpha_t^2}{\sigma_t^2}$。
3. Max-SNR-$\gamma$：$w_t=\max\{SNR(t),\gamma\}$，给 SNR 设置一个下限。
4. Min-SNR-$\gamma$：$w_t=\min\{SNR(t),\gamma\}$，给 SNR 设置一个上限。
5. UGD weighting：每次迭代都根据公式 16 重新求解 $w_t$。

我这里没有展开 UGD，因为它在每次梯度下降时都会改变，并且在这里只是作为一个 baseline 使用。

随着 $t$ 逐渐接近 0，噪声越来越小，$SNR(t)$ 会迅速增大。因此原始的 SNR weighting 会把大量权重集中到低噪声的时间步上。Max-SNR-$\gamma$ 只是避免高噪声时间步的权重太接近 0，并没有处理低噪声时间步权重过大的问题。Min-SNR-$\gamma$ 做的事情则恰好相反：当 $SNR(t)>\gamma$ 时，直接把权重截断为 $\gamma$，从而避免训练被少数高 SNR 时间步支配。论文默认令 $\gamma=5$。

作者将这些固定权重代入公式 16 的优化目标进行比较。结果当然还是每次动态求解的 UGD 最接近 Pareto 最优，但是 Min-SNR-$\gamma$ 已经非常接近 UGD，并且明显优于 Constant、SNR 和 Max-SNR-$\gamma$。换句话说，Min-SNR-$\gamma$ 用一个只和噪声调度有关的静态权重，近似替代了一个依赖当前梯度、每一步都要重新求解的动态权重。

不同预测目标下的权重

这里还有一个很容易弄混的问题。上面给出的 $w_t=\min\{SNR(t),\gamma\}$ 是 pred-$\mathbf x_0$ 下的权重，如果模型预测的是噪声 $\epsilon$，不能原封不动地照搬这个公式。

根据公式 2，我们可以在 $\mathbf x_0$ 和 $\epsilon$ 两种预测目标之间进行转换：

$$\begin{align} \hat\epsilon_\theta(\mathbf x_t)&=\frac{\mathbf x_t-\alpha_t\hat{\mathbf x}_\theta(\mathbf x_t)}{\sigma_t}\\ \left\|\epsilon-\hat\epsilon_\theta(\mathbf x_t)\right\|_2^2 &=\frac{\alpha_t^2}{\sigma_t^2}\left\|\mathbf x_0-\hat{\mathbf x}_\theta(\mathbf x_t)\right\|_2^2\\ &=SNR(t)\left\|\mathbf x_0-\hat{\mathbf x}_\theta(\mathbf x_t)\right\|_2^2 \tag{17} \end{align}$$

所以，不加权的 pred-noise Loss，本身就等价于使用 SNR weighting 的 pred-$\mathbf x_0$ Loss。这也解释了为什么前面会说 SNR weighting 在 pred-noise 下和 Constant weighting 等价。这里的“等价”并不是说 $SNR(t)=1$，而是两种预测目标之间的变换已经偷偷带入了一个 SNR 系数。

如果我们希望在统一的 pred-$\mathbf x_0$ 空间中施加权重 $W(t)$，那么 pred-noise 下的实际权重就要除以 $SNR(t)$。同理，在本文采用的 variance-preserving 设置下，论文附录证明了 pred-$\mathbf v$ 的 Loss 等价于 $(SNR(t)+1)$ 倍的 pred-$\mathbf x_0$ Loss。因此 Min-SNR-$\gamma$ 在三种预测目标下的实际权重分别是：

预测目标	实际乘在 MSE Loss 上的权重
$\mathbf x_0$	$\min\{SNR(t),\gamma\}$
$\epsilon$	$\dfrac{\min\{SNR(t),\gamma\}}{SNR(t)}=\min\left\{1,\dfrac{\gamma}{SNR(t)}\right\}$
$\mathbf v$	$\dfrac{\min\{SNR(t),\gamma\}}{SNR(t)+1}$

对于最常见的 pred-noise 而言，这个公式就非常直观了：当噪声较大、$SNR(t)\leq\gamma$ 时，权重为 1；当噪声很小、$SNR(t)>\gamma$ 时，权重变成 $\gamma/SNR(t)$。也就是说，它并没有增强高噪声时间步，而是在主动降低那些过于简单的低噪声时间步的影响。

Experiments

论文主要在 CelebA 和 ImageNet 上进行实验，网络结构同时包含 ViT 和 UNet。消融实验默认使用 ImageNet $256\times256$ 的 latent space、ViT-B、cosine noise scheduler 和 1000 个扩散时间步，Min-SNR 的默认截断值为 $\gamma=5$。

实验结果中最醒目的结论是，Min-SNR-$5$ 达到 FID=10 所需的迭代次数只有其他加权策略的约 $1/3.4$，因此论文声称获得了 3.4 倍的收敛加速。这里说的是达到同一 FID 所需的训练迭代数，而不是严格测量得到的单步训练速度。不过 Min-SNR 只多出几个逐元素运算，几乎不会增加单步训练开销，所以迭代数的降低基本可以反映实际训练时间的降低。

作者还将不同时间步上的未加权 MSE 单独拿出来比较。在 pred-$\mathbf x_0$ 的设置下，Constant weighting 在高噪声时间步上表现较好，却不擅长低噪声时间步；SNR weighting 则刚好相反，几乎把注意力都放在低噪声部分。Min-SNR-$\gamma$ 在各个噪声区间的 Loss 都下降得更快，这与前面“平衡不同时间步”的说法是吻合的。

更重要的是，这个效果并不局限于 pred-$\mathbf x_0$ 或 ViT。论文分别测试了 $\mathbf x_0$、$\epsilon$ 和 $\mathbf v$ 三种预测目标，Min-SNR-$\gamma$ 在三种设置中都收敛得更快。换成 UNet 后也能观察到同样的趋势，论文给出的 FID 如下：

训练迭代数	200K	400K	600K	800K	1M
Baseline ($\mathbf x_0$)	25.93	15.41	11.54	9.52	8.33
Min-SNR-$5$ ($\mathbf x_0$)	7.99	5.34	4.69	4.41	4.28
Baseline ($\epsilon$)	8.55	5.43	4.64	4.35	4.21
Min-SNR-$5$ ($\epsilon$)	7.32	4.98	4.48	4.24	4.14

可以看到，Min-SNR 对 pred-$\mathbf x_0$ 的提升极其明显；对于原本就隐含了 SNR weighting 的 pred-noise，提升则相对温和，尤其是训练到后期以后，两者的差距已经很小。这一点其实比单独强调 3.4 倍加速更值得注意，因为 Min-SNR 的收益显然和预测目标有关，并不是在所有配置下都能凭空加速三倍。

作者也测试了 $\gamma\in\{1,5,10,20\}$。不同设置的 FID 虽然存在波动，但是总体比较稳定，其中 $\gamma=5$ 在大部分设置下取得了最好或者接近最好的结果，因此被设为默认值。最终，论文还报告了 CelebA-64 上 1.60、ImageNet-64 上 2.28，以及 ImageNet-256 上 2.06 的 FID。至少从实验结果而言，这个简单的加权策略确实是有效的。

实际使用

Min-SNR-$\gamma$ 的实现非常简单，下面给出一个示意代码：

import torch.nn.functional as F

snr = alpha_t.square() / sigma_t.square().clamp_min(1e-8)
min_snr = snr.clamp(max=gamma)

if prediction_type == "x0":
    weight = min_snr
elif prediction_type == "epsilon":
    weight = min_snr / snr.clamp_min(1e-8)
elif prediction_type == "v":
    weight = min_snr / (snr + 1)

per_sample_loss = F.mse_loss(prediction, target, reduction="none")
per_sample_loss = per_sample_loss.flatten(1).mean(1)
loss = (weight * per_sample_loss).mean()

如果使用 DDPM 中更常见的 $\bar\alpha_t$ 记法，那么 $\alpha_t^2=\bar\alpha_t$，$\sigma_t^2=1-\bar\alpha_t$，因此

$$SNR(t)=\frac{\bar\alpha_t}{1-\bar\alpha_t}. \tag{18}$$

需要注意，权重应该先乘在每个样本各自的 MSE 上，然后再对 batch 求平均。如果先把整个 batch 的 Loss 平均成一个标量，就无法再按照每个样本采到的 $t$ 进行加权了。

一点个人看法

读到这里，应该就能理解我在开头所说的虎头蛇尾是什么意思了。论文前半部分花了大量篇幅将扩散模型训练描述为一个多任务优化问题，又从共同下降方向一路推导到 Pareto 最优权重。按照这个故事继续发展下去，我们理应得到一个从公式 16 推导出来的近似解。但是 Min-SNR-$\gamma$ 并不是从公式 16 推出来的，作者只是在几个手工设计的静态权重中进行比较，然后发现它的目标函数值比较接近 UGD。

这个跳跃其实非常大。公式 16 的最优权重依赖当前模型、训练数据以及每个时间步的梯度，而 Min-SNR-$\gamma$ 只依赖预先设定的噪声调度。作者给出的解释是，不同任务的优化特征主要由噪声等级决定，并且训练一段时间后梯度会逐渐稳定，所以可以用固定权重近似。但是论文并没有进一步验证这两个假设，更没有说明为什么这个近似最后恰好应该是 $\min\{SNR(t),\gamma\}$。

其次，开头的分桶微调实验确实说明了不同时间步之间存在干扰，但是“存在梯度冲突”和“梯度冲突是收敛慢的主要原因”仍然是两件不同的事情。论文展示了 Min-SNR 能降低 Pareto 目标，也展示了它能加快 FID 收敛，但两者之间的因果关系并没有真正建立起来。如果能进一步展示训练过程中不同时间步梯度夹角的变化，或者在不同噪声调度下验证冲突程度与收敛速度的关系，整个故事会完整很多。

此外，论文所说的 3.4 倍加速来自一个特定模型达到 FID=10 所需的迭代次数，并不是任意网络、任意预测目标下都稳定存在的加速比例。从上面的 UNet 结果也能看到，Min-SNR 对 pred-$\mathbf x_0$ 的帮助远大于 pred-noise。至于最终的 SOTA FID，其中还混合了网络规模、采样器、classifier-free guidance 和训练时长等因素，因此我认为真正有说服力的仍然是同等设置下的消融实验，而不是最后那张 SOTA 表。

最后，$\gamma=5$ 虽然看起来很稳健，但论文对它的解释基本只有“实验效果不错”。为什么是 5？最合适的阈值是否应该随噪声调度、数据集或者模型大小变化？这些问题都没有答案。对于一篇试图用多任务优化解释训练加速的论文而言，最后停在一个经验超参数上，确实有一种刚刚把问题讲明白就结束了的感觉。

当然，这并不意味着 Min-SNR-$\gamma$ 没有价值。恰恰相反，作为一个工程技巧，它只需要几行代码，不增加可感知的计算开销，对 $\mathbf x_0$、$\epsilon$ 和 $\mathbf v$ 预测都适用，而且默认参数也比较稳定。它最有吸引力的地方并不是理论有多么完整，而是修改成本几乎为零，失败的风险又很低。这大概也是为什么每次询问 AI 如何改善扩散模型训练时，它都会被顺手推荐出来。

所以我最后的看法是：Min-SNR-$\gamma$ 是一个很好的训练技巧，但这篇论文只给出了一个颇有吸引力的解释框架，还没有真正完成这个解释。它在工程上解决了问题，却在理论故事最应该继续的地方停了下来。这种落差，可能就是我读完以后感到失望的真正原因。

Appendix

Equation 6

公式 6 是对公式 5 应用一阶泰勒展开得到，本人一直是数学苦手，在此也回顾一下泰勒展开的知识。 对于函数 $f(x)$ 而言，我们如果对 $x$ 施加一个极小变换 $\Delta x$ ，则相应的函数值变化可以近似表示为

$$f(x+\Delta x)\approx f(x)+f'(x)\Delta x$$

后面还要加一个余项 $O((\Delta x)^2)$ 才能成立等式。 但是在本文中由于 $\theta$ 是一个高维向量，此时 $\mathcal L^t(\theta)$ 是一个高维函数，其对应的梯度由对各个权重参数的偏导组成，即

$$\begin{align} \nabla_\theta \mathcal L^t(\theta)=\left\{ \frac{\partial \mathcal L^t}{\partial\theta_1},...,\frac{\partial \mathcal L^t}{\partial\theta_N} \right\},\theta=\{\theta_1,...,\theta_N\} \end{align}$$

此时两者变成了点乘的关系，也就是公式 6 的内积的形式。

Reference

[1] Hang, T., Gu, S., Li, C., Bao, J., Chen, D., Hu, H., … & Guo, B. (2023). Efficient diffusion training via min-snr weighting strategy. In Proceedings of the IEEE/CVF international conference on computer vision (pp. 7441-7451)