如何让GDN2更稳定一些？

我们知道，DeltaNet的核心是$\mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{k}}^{\top}$形式的更新矩阵（参考《线性注意力简史：从模仿、创新到反哺》），在$\eta \leq 2$并对$\mathbf{\mathit{k}}$施加L2 Normalize时，它的连乘积具备良好的数值稳定性（参考《为什么DeltaNet要加L2 Normalize？》和[《DeltaNet的核心逆矩阵的元素总是在[-1, 1]内》](https://spaces.ac.cn/archives/11563)），这是它能稳定训练的关键，后续工作GDN、KDA也沿用了这一设置。

后来，GDN2试图将更新矩阵一般化为$\mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top}$，类似的选择也出现在RWKV7中。然而，这种一般化可能会带来一些新的不稳定问题，本文对此做一个简单分析。

问题描述[#](https://spaces.ac.cn/archives/11945#%E9%97%AE%E9%A2%98%E6%8F%8F%E8%BF%B0)

GDN2的一般形式为

$$ (\text{1}) \mathbf{\mathit{S}}_{t} = \mathbf{\mathit{S}}_{t - 1} \mathbf{\mathit{D}}_{t} \left(\right. \mathbf{\mathit{I}} - \eta_{t} \mathbf{\mathit{k}}_{t} \mathbf{\mathit{j}}_{t}^{\top} \left.\right) + \eta_{t} \mathbf{\mathit{v}}_{t} \mathbf{\mathit{k}}_{t}^{\top} $$

其中$\mathbf{\mathit{k}}_{t} , \mathbf{\mathit{j}}_{t} \in \mathbb{R}^{d}$，$\eta_{t} \geq 0$，$\mathbf{\mathit{D}}_{t}$是对角矩阵，对角线元素通常限制在$\left[\right. 0 , 1 \left]\right.$内，代表遗忘矩阵。KDA对应于$\mathbf{\mathit{j}}_{t} = \mathbf{\mathit{k}}_{t}$这一特例，而GDN进一步对应于$\mathbf{\mathit{D}}_{t} = \gamma_{t} \mathbf{\mathit{I}}$这一特例。 当我们将递归展开时，$\mathbf{\mathit{S}}_{t}$会出现$\mathbf{\mathit{D}}_{t} \left(\right. \mathbf{\mathit{I}} - \eta_{t} \mathbf{\mathit{k}}_{t} \mathbf{\mathit{j}}_{t}^{\top} \left.\right)$的连乘积，所以为了保证数值稳定性，我们应当希望

$$ (\text{2}) \parallel \mathbf{\mathit{D}}_{t} \left(\right. \mathbf{\mathit{I}} - \eta_{t} \mathbf{\mathit{k}}_{t} \mathbf{\mathit{j}}_{t}^{\top} \left.\right) \parallel_{2} \leq 1 $$

其中$\parallel \cdot \parallel_{2}$是谱范数。由于$\mathbf{\mathit{D}}_{t}$是对角矩阵且对角元已经在$\left[\right. 0 , 1 \left]\right.$内，所以它满足$\parallel \mathbf{\mathit{D}}_{t} \parallel_{2} \leq 1$，因此上式的一个充分条件是

$$ (\text{3}) \parallel \mathbf{\mathit{I}} - \eta_{t} \mathbf{\mathit{k}}_{t} \mathbf{\mathit{j}}_{t}^{\top} \parallel_{2} \leq 1 $$

很不幸的是，接下来我们将会看到，当$\mathbf{\mathit{k}}_{t} , \mathbf{\mathit{j}}_{t}$不共线时（即GDN2希望拓展的区间），上式是不可能成立的。

谱之范数[#](https://spaces.ac.cn/archives/11945#%E8%B0%B1%E4%B9%8B%E8%8C%83%E6%95%B0)

简单起见，下面的讨论省略下标$t$，并且假设$\mathbf{\mathit{k}} , \mathbf{\mathit{j}}$都已经归一化（否则总可以将$\parallel \mathbf{\mathit{k}} \parallel , \parallel \mathbf{\mathit{j}} \parallel$吸收到$\eta$中），我们有

$$ (\text{4}) \parallel \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top} \parallel_{2} = \frac{\eta + \sqrt{\eta^{2} - 4 \eta \mathbf{\mathit{k}}^{\top} \mathbf{\mathit{j}} + 4}}{2} $$

这个结论大概是这篇文章最繁琐的地方，所以我们将它的证明放到最后。由于$\mathbf{\mathit{k}} , \mathbf{\mathit{j}}$已经归一化，所以有$\mathbf{\mathit{k}}^{\top} \mathbf{\mathit{j}} \in \left[\right. - 1 , 1 \left]\right.$，而当$\mathbf{\mathit{k}}^{\top} \mathbf{\mathit{j}} < 1$且$\eta > 0$时，我们有

$$ (\text{5}) \frac{\eta + \sqrt{\eta^{2} - 4 \eta \mathbf{\mathit{k}}^{\top} \mathbf{\mathit{j}} + 4}}{2} > \frac{\eta + \sqrt{\eta^{2} - 4 \eta + 4}}{2} = \frac{\eta + \left|\right. \eta - 2 \left|\right.}{2} = max \left(\right. \eta , 2 \left.\right) - 1 \geq 1 $$

也就是说只要$\mathbf{\mathit{k}} \neq \mathbf{\mathit{j}}$且$\eta > 0$，那不管怎么调，$\mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top}$的谱范数都严格大于1！更准确地有

$$ (\text{6}) \parallel \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top} \parallel_{2} \leq 1 \Longleftrightarrow \mathbf{\mathit{k}} = \mathbf{\mathit{j}} \textrm{ } 且 \textrm{ } \eta \leq 2 \left(\right. 或 \textrm{ } \eta = 0 \left.\right) $$

但如果$\mathbf{\mathit{k}} = \mathbf{\mathit{j}}$，那么就退化为KDA了，失去了GDN2的推广意义。而在$\mathbf{\mathit{k}} \neq \mathbf{\mathit{j}}$时，$\parallel \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top} \parallel_{2}$一定大于1，这意味着GDN2的稳定性没有保障。该如何破局呢？

保证策略[#](https://spaces.ac.cn/archives/11945#%E4%BF%9D%E8%AF%81%E7%AD%96%E7%95%A5)

庆幸的是，我们还有遗忘门矩阵$\mathbf{\mathit{D}}$，由于它被限制在$\left[\right. 0 , 1 \left]\right.$内，它们乘积$\mathbf{\mathit{D}} \left(\right. \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top} \left.\right)$的谱范数大概率会小于1，于是实际训练并不一定失败。然而，这仅仅是概率上的，并非理论上的保证。

最理想的办法是直接求出$\mathbf{\mathit{D}} \left(\right. \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top} \left.\right)$的谱范数，然后对它做谱裁剪之类的操作，然而对于一般的$\mathbf{\mathit{D}}$，该矩阵的谱范数没有通用的闭式解。

好在$\mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top}$的谱范数有解析解，我们可以退而求其次，利用它的谱范数，对$\mathbf{\mathit{D}}$做裁剪操作：

$$ (\text{7}) \mathbf{\mathit{D}} \left(\right. \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top} \left.\right) \rightarrow min \left(\right. \mathbf{\mathit{D}} , 1 / \rho \left.\right) \left(\right. \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top} \left.\right) $$

其中$\rho = \parallel \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top} \parallel_{2}$。这样一来，我们只会修改$\mathbf{\mathit{D}}$中大于$1 / \rho$的对角元，由于实践中$\eta$通常不大，所以平均而言$1 / \rho$不会明显小于1，那么改动面不会很广，即以尽可能小和简洁的修改，保证了GDN2的稳定性。 值得指出的是，$\rho$和$min \left(\right. \mathbf{\mathit{D}} , 1 / \rho \left.\right)$都不过是Token-wise的运算，可以独立在RNN的算子之外，并且$\rho$也有闭式解，那么经过融合后，基本不会影响效率，因此算得上是一个心智负担极小的改动了。

证明过程[#](https://spaces.ac.cn/archives/11945#%E8%AF%81%E6%98%8E%E8%BF%87%E7%A8%8B)

最后补上式$(\text{4})$的证明。不失一般性，假设$\mathbf{\mathit{k}} , \mathbf{\mathit{j}}$不共线，将它们正交化为$\mathbf{\mathit{u}}_{1} = \mathbf{\mathit{j}}$、$\mathbf{\mathit{u}}_{2} = \frac{\mathbf{\mathit{k}} - c \mathbf{\mathit{j}}}{\sqrt{1 - c^{2}}}$，其中$c = \mathbf{\mathit{k}}^{\top} \mathbf{\mathit{j}}$。记$\mathbf{\mathit{U}}_{1} = \left[\right. \mathbf{\mathit{u}}_{1} , \mathbf{\mathit{u}}_{2} \left]\right.$，将它扩充成完整正交基$\mathbf{\mathit{U}} = \left[\right. \mathbf{\mathit{U}}_{1} , \mathbf{\mathit{U}}_{2} \left]\right.$，记$\mathbf{\mathit{A}} = \mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top}$。

对于与$\mathbf{\mathit{k}} , \mathbf{\mathit{j}}$都正交的向量$\mathbf{\mathit{x}}$，有$\mathbf{\mathit{A}} \mathbf{\mathit{x}} = \mathbf{\mathit{x}}$和$\mathbf{\mathit{x}}^{\top} \mathbf{\mathit{A}} = \mathbf{\mathit{x}}^{\top}$，那么$\mathbf{\mathit{A}} \mathbf{\mathit{U}}_{2} = \mathbf{\mathit{U}}_{2}$和$\mathbf{\mathit{U}}_{2}^{\top} \mathbf{\mathit{A}} = \mathbf{\mathit{U}}_{2}^{\top}$，于是

$$ (\text{8}) \mathbf{\mathit{U}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{U}} = \left(\right. \begin{matrix}\mathbf{\mathit{U}}_{1}^{\top} \\ \mathbf{\mathit{U}}_{2}^{\top}\end{matrix} \left.\right) \mathbf{\mathit{A}} \left(\right. \begin{matrix}\mathbf{\mathit{U}}_{1} & \mathbf{\mathit{U}}_{2}\end{matrix} \left.\right) = \left(\right. \begin{matrix}\mathbf{\mathit{U}}_{1}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{U}}_{1} & \mathbf{\mathit{U}}_{1}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{U}}_{2} \\ \mathbf{\mathit{U}}_{2}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{U}}_{1} & \mathbf{\mathit{U}}_{2}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{U}}_{2}\end{matrix} \left.\right) = \left(\right. \begin{matrix}\mathbf{\mathit{U}}_{1}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{U}}_{1} & 0 \\ 0 & \mathbf{\mathit{I}}_{d - 2}\end{matrix} \left.\right) $$

由于正交变换不改变奇异值，所以$\mathbf{\mathit{A}}$的非平凡奇异值只能来自矩阵$\mathbf{\mathit{U}}_{1}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{U}}_{1}$，直接算出来得

$$ (\text{9}) \mathbf{\mathit{B}} \triangleq \mathbf{\mathit{U}}_{1}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{U}}_{1} = \left(\right. \begin{matrix}1 - \eta c & 0 \\ - \eta \sqrt{1 - c^{2}} & 1\end{matrix} \left.\right) $$

$2 \times 2$矩阵的奇异值，有比较巧妙的直接计算方法。设两个奇异值为$\sigma_{1} \geq \sigma_{2}$，利用奇异值与F范数和行列式的关系可得

$$ (\text{10}) \underbrace{\sigma_{1}^{2} + \sigma_{2}^{2} = \parallel \mathbf{\mathit{B}} \parallel_{F}^{2} = \eta^{2} - 2 \eta c + 2 , \sigma_{1} \sigma_{2} = \left|\right. det \left(\right. \mathbf{\mathit{B}} \left.\right) \left|\right. = \left|\right. 1 - \eta c \left|\right.} \\ \Downarrow \\ (\text{11}) \sigma_{1} \pm \sigma_{2} = \sqrt{\sigma_{1}^{2} + \sigma_{2}^{2} \pm 2 \sigma_{1} \sigma_{2}} = \sqrt{\eta^{2} - 2 \eta c + 2 \pm 2 \left|\right. 1 - \eta c \left|\right.} $$

我们主要关心$\sigma_{1}$，它等于$\pm$两项的平均，此时可以去掉绝对值（反正正负都要算），最后化简得

$$ (\text{12}) \sigma_{1} = \frac{\eta + \sqrt{\eta^{2} - 4 \eta c + 4}}{2} $$

前面已经证明了它不小于1，所以$\sigma_{1}$就是谱范数。

文章小结[#](https://spaces.ac.cn/archives/11945#%E6%96%87%E7%AB%A0%E5%B0%8F%E7%BB%93)

本文分析了GDN2将更新矩阵从$\mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{k}}^{\top}$推广到$\mathbf{\mathit{I}} - \eta \mathbf{\mathit{k}} \mathbf{\mathit{j}}^{\top}$后可能出现的稳定性问题，并讨论了相应的应对策略。

_转载到请包括本文地址：[https://spaces.ac.cn/archives/11945](https://spaces.ac.cn/archives/11945 "如何让GDN2更稳定一些？")_

_更详细的转载事宜请参考：_[《科学空间FAQ》](https://spaces.ac.cn/archives/6508#%E6%96%87%E7%AB%A0%E5%A6%82%E4%BD%95%E8%BD%AC%E8%BD%BD/%E5%BC%95%E7%94%A8 "《科学空间FAQ》")