分类： AI

在当前大规模语言模型（LLM）的参数高效微调方法中，LoRA（Low-Rank Adaptation）无疑是一个重要的手段。此前，我们在《梯度视角下的LoRA：简介、分析、猜测及推广》中曾简单讨论过LoRA。今天，我们来学习LoRA的一个新结论：给LoRA的两个矩阵分配不同的学习率，LoRA的效果还能进一步提升。这一结论出自最近的论文《LoRA+: Efficient Low Rank Adaptation of Large Models》（下称“LoRA+”）。

为什么调整学习率显得如此重要？

乍一看，给不同的矩阵分配不同的学习率似乎并没有什么特别之处，因为配置不同的学习率相当于引入了新的超参数，而通常来说，只要引入并精调超参数，模型性能都会有所提升。然而，“LoRA+”的特别之处在于，它从理论角度肯定了这一必要性，并且明确指出最优解必然是右矩阵的学习率大于左矩阵的学习率。

LoRA基本原理回顾

假设预训练参数为

，如果使用全量参数微调，那么增量也是一个 ( n \times m ) 矩阵。为了降低参数量，LoRA将更新量约束为低秩矩阵，即设 ( W = W_0 + AB )，其中

。在训练时，只更新 ( A ) 和 ( B )。

为了便于理解，我们假设层的输入是

，层的运算为

。由于“LoRA+”的结论跟预训练权重无关，不失一般性可以设

，那么层运算简化为

。

结论简析

“LoRA+”的结论是：为了使LoRA的效果尽可能接近最优，权重 ( B ) 的学习率应该要大于权重 ( A ) 的学习率。

这个结论的核心在于，虽然 ( A ) 和 ( B ) 在表面上是对称的，但实际上它们有着固有的不对称性。即使将 ( A ) 或 ( B ) 全零初始化，结论依然成立。

数值稳定性与贡献相当假设

1. 数值稳定性

数值稳定性意味着 ( X )、 ( XA )、 ( XAB ) 的每个分量都应该是 ( O(1) ) 级别的，而不依赖于网络宽度 ( n, m )。虽然 ( XA ) 是中间变量，但它的数值稳定性对梯度稳定性至关重要。

为了实现数值稳定性，假设 ( X ) 是 ( O(1) )，那么 ( A ) 和 ( B ) 应该分别用 ( 1/n ) 和 ( 1/r ) 的方差初始化。

2. 贡献相当

为了使LoRA最优， ( A ) 和 ( B ) 两个矩阵对效果应该有同等程度的贡献。我们可以通过损失函数 ( L ) 的变化来衡量 ( A ) 和 ( B ) 的贡献：

在Adam优化器中，如果两个参数的学习率之比是 ( \lambda )，那么经过长期训练后，这两个参数的数量级之比也是

。因此，如果 ，那么 ( B ) 和 ( A ) 的数量级之比也是 ，与我们的期望一致。

快速推导

通过矩阵求导，我们可以得到：

结合数值稳定性和贡献相当的假设，可以推导出：

总结

在这篇文章中，我们介绍并推导了“LoRA+”的结果，它支持LoRA的两个低秩矩阵 ( A ) 和 ( B ) 存在固有的不对称性。不管将哪个矩阵全零初始化，都应该将 ( B ) 的学习率设置得大于 ( A )，以达到更优的效果。这一理论不仅在实践中有效，也为我们提供了一个经典的理论指导训练的范例。

引言

今天我们分享的是一篇名为《Score Identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation》的新论文。该论文探讨了如何更快更好地蒸馏扩散模型。

即便没有做过蒸馏，大家可能也能猜到蒸馏的常规步骤：随机采样大量输入，然后用扩散模型生成相应结果作为输出，用这些输入输出作为训练数据对，来监督训练一个新模型。然而，众所周知，作为教师的原始扩散模型通常需要多步（比如1000步）迭代才能生成高质量输出，所以且不论中间训练细节如何，该方案的一个显著缺点是生成训练数据太费时费力。此外，蒸馏之后的学生模型通常或多或少都有效果损失。

有没有方法能一次性解决这两个缺点呢？这就是上述论文试图要解决的问题。

思路简介

论文将所提方案称为“Score Identity Distillation（SiD）”，基于几个恒等式来设计和推导了整个框架。实际上，它的设计思想与几个恒等式并没有直接联系，其次几个恒等式都是已知的公式而不是新的，所以这个名字显得相当随意。

本文将其称之为“重现江湖”，是因为SiD的思路跟之前在《从去噪自编码器到生成模型》介绍过的论文《Learning Generative Models using Denoising Density Estimators》（简称“DDE”）几乎一模一样，甚至最终形式也有五六分相似。只不过当时扩散模型还未露头角，所以DDE是将其作为一种新的生成模型提出的，在当时反而显得非常小众。而在扩散模型流行的今天，它可以重新表述为一种扩散模型的蒸馏方法，因为它需要一个训练好的去噪自编码器——这正好是扩散模型的核心。

接下来笔者用自己的思路去介绍SiD。

初级形式

假设我们有一个在目标数据集训练好的教师扩散模型 ( \epsilon_{\phi^}(x_t, t) )，它需要多步采样才能生成高质量图片。我们的目标是训练一个单步采样的学生模型 ( x = g_{\theta}(z) )，即一个类似GAN的生成器，输入指定噪声 ( z ) 就可以直接生成符合要求的图像。如果我们有很多的 ( (z, x) ) 对，那么直接监督训练就可以了（当然损失函数和其他细节还需要进一步确定，读者可以自行参考相关工作），但如果没有呢？肯定不是不能训，因为就算没有 ( \epsilon_{\phi^}(x_t, t) ) 也能训，比如GAN，所以关键是怎么借助已经训练好的扩散模型提供更好的信号。

SiD及前作DDE使用了一个看上去很绕但是也很聪明的思路：

如果 ( g_{\theta}(z) ) 产生的数据分布跟目标分布很相似，那么拿 ( g_{\theta}(z) ) 生成的数据集去训练一个扩散模型 ( \epsilon_{\psi^}(x_t, t) ) 的话，它也应该跟 ( \epsilon_{\phi^}(x_t, t) ) 很相似？

这个思路的聪明之处在于，它绕开了对教师模型生成样本的需求，也不需要训练教师模型的真实样本，因为“拿 ( g_{\theta}(z) ) 生成的数据集去训练一个扩散模型”只需要学生模型 ( g_{\theta}(z) ) 生成的数据（简称“学生数据”），而 ( g_{\theta}(z) ) 是一个单步模型，用它来生成数据时间上比较友好。

当然，这还只是思路，将其转换为实际可行的训练方案还有一段路要走。

方法与公式

扩散模型回顾

我们采用《生成扩散模型漫谈（三）：DDPM = 贝叶斯 + 去噪》的形式，对输入 ( x_0 ) 进行加噪：
[ x_t = \bar{\alpha}_t x_0 + \bar{\beta}_t \epsilon, \quad \epsilon \sim N(0, I) ]

训练 ( \epsilon_{\phi^}(x_t, t) ) 的方式则是去噪： [ \phi^ = \arg\min_{\phi} E_{x_0 \sim \tilde{p}(x_0), \epsilon \sim N(0, I)} \left[ | \epsilon_{\phi}(\bar{\alpha}_t x_0 + \bar{\beta}_t \epsilon, t) – \epsilon |^2 \right] ]

同样地，如果我们想用 ( g_{\theta}(z) ) 的学生数据训练一个扩散模型，那么训练目标是：
[ \psi^* = \arg\min_{\psi} E_{z, \epsilon \sim N(0, I)} \left[ | \epsilon_{\psi}(x_t^{(g)}, t) – \epsilon |^2 \right] ]

其中 ( x_t^{(g)} = \bar{\alpha}t g{\theta}(z) + \bar{\beta}t \epsilon )，是由学生数据加噪后的样本，其分布记为 ( p{\theta}(x_t^{(g)}) )。

学生模型的学习

我们可以通过最小化教师模型和学生模型生成的数据分布差异来学习学生模型：
[ \theta^* = \arg\min_{\theta} E_{z, \epsilon \sim N(0, I)} \left[ | \epsilon_{\phi^}(x_t^{(g)}, t) – \epsilon_{\psi^}(x_t^{(g)}, t) |^2 \right] ]

注意这个优化依赖于 ( \theta )，所以当 ( \theta ) 通过上式发生改变时，( \psi^* ) 的值也随之改变，因此需要交替优化，类似GAN一样。

点睛之笔

上述方法存在理论与实践之间的gap，主要体现在两个问题：

1. 理论上要求先求出上式的最优解，然后才去优化，但实际上从训练成本考虑，我们并没有将它训练到最优就去优化了；
2. 理论上 ( \psi^* ) 随 ( \theta ) 而变，即应该写成 ( \psi^(\theta) )，从而在优化时应该多出一项 ( \psi^(\theta) ) 对 ( \theta ) 的梯度，但实际上在优化时我们都只当 ( \psi^* ) 是常数。

SiD的核心贡献是通过恒等变换，尽量消除优化目标对 ( \psi^* ) 的依赖，从而有效缓解上述两个问题。

恒等变换

我们具体来看做了什么恒等变换。我们先来看去噪目标：
[ E_{x_0 \sim \tilde{p}(x_0), x_t \sim p(x_t | x_0)} \left[ | \epsilon_{\phi}(x_t, t) + \bar{\beta}t \nabla{x_t} \log p(x_t | x_0) |^2 \right] ]

根据得分匹配相关结果，上述目标的最优解是 ( \epsilon_{\phi^}(x_t, t) = -\bar{\beta}t \nabla{x_t} \log p(x_t) )。同理，学生模型的训练目标的最优解是 ( \epsilon_{\psi^}(x_t^{(g)}, t) = -\bar{\beta}t \nabla{x_t^{(g)}} \log p_{\theta}(x_t^{(g)}) )。

此时我们有：
[ E_{z, \epsilon \sim N(0, I)} \left[ | \epsilon_{\phi^}(x_t^{(g)}, t) – \epsilon_{\psi^}(x_t^{(g)}, t) |^2 \right] ]

通过恒等变换，我们可以将上式化简为：
[ E_{z, \epsilon \sim N(0, I)} \left[ \langle \epsilon_{\phi^}(x_t^{(g)}, t) – \epsilon_{\psi^}(x_t^{(g)}, t), \epsilon_{\phi^*}(x_t^{(g)}, t) – \epsilon \rangle \right] ]

这就是SiD的核心结果，能够高效地实现蒸馏。

其他细节

1. 论文的推导默认了 ( \bar{\alpha}_t = 1 )。
2. 论文的结果是以 ( \mu(\bar{x}_t) = x_t – \bar{\beta}_t \epsilon(x_t, t) / \bar{\alpha}_t ) 为标准给出的，这与扩散模型常见的表述方式不同。
3. SiD最终取了上式的相反数作为额外的损失函数，加权到改进的损失函数上，以取得更优的蒸馏效果。

延伸思考

对于式(3)和式(4)的交替优化，有不少读者可能会想到，但SiD的精彩之处是提出了恒等变换，使得训练更加稳定高效。

文章小结

在这篇文章中，我们介绍了一种新的将扩散模型蒸馏为单步生成模型的方案，其思想可以追溯到前两年的利用去噪自编码器训练生成模型的工作，它不需要获得教师模型的真实训练集，也不需要迭代教师模型来生成样本对，而引入了类似GAN的交替训练，同时提出了关键的恒等变换来稳定训练过程，整个方法有颇多值得学习之处。