深度学习里的流形优化 II——Hyperball May Not Be a Free Lunch

Too Long, Don't Read:

在本篇博客中我们将通过一个有趣的消融实验验证权重衰减本质上只是显式控制模型的梯度范数，在这个过程中权重衰减等价于对学习率进行调控。除此之外我们围绕这个进行了一些有趣的讨论。

1.引言

在上一部分的工作之中，我们简单地介绍了权重衰减作为一种控制模型参数的策略的正则化方法：在现代优化学习的过程中，我们一般发现模型更新时计算得到的梯度范数变小会随着两个因素而改变：1）随着优化的进行，损失平面逐渐变得平坦导致梯度更小；2）因模型范数的变大，导致梯度范数的变小。在这个过程中我们认为第二个因素中，模型范数的变大在正常的优化过程中是一定会出现的现象，那么梯度范数会因为模型范数的增长而导致梯度过早的变小导致模型更新量的减小。

那么权重衰减的意义是通过每一步减小模型范数来增加梯度范数，它甚至会让梯度范数扩大至比一开始更新还要大的状态，那么这一个启发式的思想来保证了更新过程中梯度不会出现过早衰减而导致模型更新变慢的现象。

但是目前依然存在一个问题，权重衰减真的只是通过增加梯度范数来提高模型的优化性能吗？在这里存在一个很强的几何观点：权重衰减本身是将模型重投影到一个半径更小的球上，是否可能是在这个更小的球面上进行更新来得到更好的优化动态；还是只是增大了更新步长就获得了更好的优化动态？

在本blog中我们设计了一个有趣的消融实验去讨论这个过程。

2.对齐实验的设置

我们考虑两个不同的优化动态，一个是没有权重衰减的muon$w_t^1$，一个是有权重衰减的muonwd $w_t^2$ 。在这个过程中我们一般能观测到在初期更新时，$w^1$ 的val loss一开始会低于 $w^2$，在更新一段时间后 $w^2$ 的val loss会加速衰减，变得比 $w^1$ 更低。那么在这个过程中，我们通过一种标量化学习率对齐的策略来观测，两个不同的更新动态能否通过只修改学习率的策略进行对齐呢？

事实上是可以的，我们通过有效学习率的思想:

\[\eta_{eff}=\frac{ \eta }{\|w\|_F}\]

将每一步更新看作更新整个模型的百分比看作更新步长，那么我们可以通过先运行 $w^2$ 的实验，并且记录下来 $w^2$ 实验中每一步更新的有效学习率 $ \eta_{eff,t}^2=\frac{\eta_t}{|w^2_t|}_F $ 。并在运行 $ w^1 $ 的更新过程中，通过只进行对学习率进行标量乘法来让两者的有效学习率做对齐。

这个实验过程中，我们固定了muon和muonwd两者之间的每次更新模型参数的百分比，其中 $w^1$ 在一个不断膨胀的球面上进行更新； $w^2$ 在一个存在回缩的球面上进行更新。我们希望观测这两者的val loss之间是否存在直接联系。这个实验留给读者自行进行，不放出具体的实验图像，只讨论结果。

通过这个有趣的实验，我们能发现muon和muonwd的损失曲线在进行有效学习率对齐后，几乎能呈现出完全重合的过程，这作为一个有利的证据证明了一点，模型在一个大球面上更新还是在一个小球面上更新，只要相对更新步长相同，那么方向是不会有很大差异的。

注：这个观点是trivial的，也是不trivial的。首先如果你从尺度不变性的思想出发，它确实似乎是合理的，但是事实上目前的大部分架构，因为激活函数的非齐次性，你在不进行十分精细的操作下是很难直接得到尺度不变性的结论，故这个观点并没有很直接。

通过一个这样的消融实验，我们迈出了透明优化重要一步：权重衰减确实是一种学习率控制的策略，并且在控制学习率后，哪怕模型范数出现了较大的波动，但它不会对任何更新动态发生改变。那么更进一步，我们去用这个思想讨论Hyperball风格优化器。

3.时变权重衰减的优化器——Hyperball

Hyperball是一个在现代优化器之中表现出超强潜力的优化器，它主要包括两个方面：1）对更新量进行时变缩放；2）对模型更新进行时变缩放。这两部耦合起来的优化器，并且因为Hyperball风格优化器本质上是一种缩放优化器，它可以直接在其他优化器上即插即用。在Muon上进行Hyperball风格设计的MuonH在很多环境中起到了有效的加速效应。其更新过程数学表达为：

\[W_{t+1}=\frac{\|W_t\|_F}{\left\|W_t-\eta u_t \frac{\|W_t\|_F}{\|u_t\|_F}\right\|_F}W_t - \frac{\eta\|W_t\|_F}{\left\|W_t-\eta u_t \frac{\|W_t\|_F}{\|u_t\|_F}\right\|_F} u_t \frac{\|W_t\|_F}{\|u_t\|_F}.\]

这个式子可以更好说明Hyperball风格优化器本质上是一个时变缩放控制器。那么我们能否通过第二章的实验策略通过对齐有效学习率的方法来做到对Hyperball风格优化器更新动态的模仿，甚至让Muonwd和MuonH实现互相对齐呢？

在这里我们依照第二章的实验设置，然而在数值计算上存在微小改动对齐角速度而非对齐有效学习率，但两者之间数值差异不大。我们将MuonWD对齐至MuonH，同时也将MuonH对齐至MuonWD实现一个互相对齐的实验，具体实验结果如下图所示：

我们看到了在这个过程中确实实现了较好的对齐，那么我们成功的让权重衰减风格优化器，Hyperball风格优化器通过仅通过学习率调度对齐的思想实现了不同优化器的对齐。那么在这个过程中我们可以大胆下一个过于粗暴的结论无论是权重衰减，还是Hyperball风格优化器，它们都没有对更新的方向产生影响，只是内置了一种学习率调度策略。它和全新的预条件处理是完全不同的两种优化策略方向。

那么在这个过程中，我们从Hyperball的一个有趣现象出发，我们会发现Hyperball风格优化器相较于传统优化器， 它的收敛往往出现于迭代步数的晚期，但会迅速超越无hyperball风格优化器，这个现象完全等同于权重衰减过程中，无权重衰减的模型会比有权重衰减的模型收敛更快的现象，它们的本质原因是相同的。

那么为什么会收敛到更好的解呢？河谷景观是一个启发式的解决思想：模型初期的步长更大时，虽然Loss没有明显改善，但是其迭代至更好的损失曲面上。Hyperball风格优化器也可以看作这一种更新过程，通过初期较大的有效学习率来保证模型更新到更好的损失曲面。

然而，这也带来了一个更加直接的问题，我们已经将权重衰减和Hyperball风格优化器进行一个严格的对照来描述哪怕在尺度不变性不成立的情况下，它们也能保证好在不同的模型范数下有着相同的更新方向。然而河谷景观到底是不是一个可靠的思想呢？这为我们带来了一个问题，也是我们之后研究的方向为什么模型初期的步长更大就能走到更好的损失曲面？为什么权重衰减往往会在后期反超？我们需要一种更加定性的结论去思考这个过程。

4.致谢

本篇Blog的主要作者为:Yihao Xiao，和Jialong Sun。除此之外我们感谢为我们提供重要贡献的的合作者:Zitian Gao; Zeming Wei;Chutian Wang;Ran Tao; Jiaye Teng;Bryan Dai。最后向供有价值思想的,并我们提供有益讨论的同伴们致谢:Huaqing Zhang;Tian Xie;Guangyu Chen;Huanran Chen;Ganzhao Yuan; Zhenjie Zhou; Kaiyue Wen.

我们要特别的感谢Kaiyue Wen在学习率调度，优化器策略上做出的卓越贡献，他的贡献成为了本篇Blog的相关理论基础。