重新思考学习率与Batch Size（二）：平均场

上文《重新思考学习率与Batch Size（一）：现状》末尾我们说到，对于SignSGD、SoftSignSGD等$\tilde{\boldsymbol{\varphi}}_B$非线性依赖于$\tilde{\boldsymbol{g}}_B$的情形，计算过程的心智负担相当沉重，并且面临难以推广的困境。为此，笔者投入了一些精力去尝试简化其中的推导，万幸有些许收获，其中的关键思路便是本文的主题——平均场。

平均场是物理中常见的近似计算方法，它没有固定的形式，但大体思想就是将求平均移到函数之内。事实上，在《为什么Adam的Update RMS是0.2？》中我们就已经窥见过平均场的魅力，而在这篇文章中，我们再来见识它在计算SignSGD/SoftSignSGD的学习率规律上的奇效。

方法大意 #

沿着上文的记号，对于SignSGD我们有$\newcommand{sign}{\mathop{\text{sign}}}\tilde{\boldsymbol{\varphi}}_B=\sign(\tilde{\boldsymbol{g}}_B)$，我们需要先计算$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$和$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$，继而可以算出
\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\eta^* \approx \frac{\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}\boldsymbol{g}}{\tr(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]\boldsymbol{H})}\label{eq:eta-opt}\end{equation}
其中$\boldsymbol{g}$是梯度，$\boldsymbol{H}$是Hessian矩阵。根据假设，随机变量$\tilde{\boldsymbol{g}}_B$的均值为$\boldsymbol{g}$，协方差矩阵为$\boldsymbol{\Sigma}/B$，我们主要关心的是$\eta^*$与Batch Size $B$的关系。由于$\sign$是Element-wise的运算，因此我们可以从单个标量出发进行尝试。平均场方法，源于笔者某天突然发现的一个可能成立的近似关系
\begin{equation}\mathbb{E}[\sign(\tilde{g}_B)] = \mathbb{E}\bigg[\frac{\tilde{g}_B}{\sqrt{\tilde{g}_B^2}}\bigg]\approx \frac{\mathbb{E}[\tilde{g}_B]}{\sqrt{\mathbb{E}[\tilde{g}_B^2}]} = \frac{g}{\sqrt{g^2 + \sigma^2/B}}\end{equation}
看过《当Batch Size增大时，学习率该如何随之变化？》的读者，应该能惊奇地发现，这个只需一行就能快速推导出来的结果，跟原文中一大通假设和近似得出来的结果，只差一个无关紧要的常数$\pi/2$！这个事实让笔者意识到，平均场近似或许对学习率与Batch Size的关系完全够用了。

基于平均场的推导有诸多好处。首先是假设少，原始推导至少包含三个假设：分量独立、正态分布、$\text{erf}(x)$用$x/\sqrt{x^2+c}$近似，但是平均场近似可以去掉分布形式的假设，只需要假设它自身是可用的就行。然后是计算简单，上面我们一行就完成了计算，而原始推导即便诸多假设之下计算也是复杂得多。

计算过程 #

这一节我们将利用平均场近似，给出SignSGD完整的计算过程。首先是均值$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$，其实上一节的计算其实已经差不多完整了，这里只需要补充少许细节。我们用分量写法：
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]_i = \mathbb{E}[\sign((\tilde{g}_B)_i)] = \mathbb{E}\bigg[\frac{(\tilde{g}_B)_i}{\sqrt{(\tilde{g}_B)_i^2}}\bigg]\approx \frac{\mathbb{E}[(\tilde{g}_B)_i]}{\sqrt{\mathbb{E}[(\tilde{g}_B)_i^2]}} = \frac{g_i}{\sqrt{g_i^2 + \sigma_i^2/B}} = \frac{\sign(g_i)}{\sqrt{1 + (\sigma_i^2/g_i^2)/B}}\end{equation}
其中$\sigma_i^2 = \boldsymbol{\Sigma}_{i,i}$。由于我们最终主要关心$\eta^*$与$B$的关系，这两者都是标量的，所以这里我们再用一次平均场近似，将与$B$有关的分母部分以标量形式分离出来：
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]_i \approx \frac{\sign(g_i)}{\sqrt{1 + (\sigma_i^2/g_i^2)/B}} \approx \frac{\sign(g_i)}{\sqrt{1 + \mathcal{B}_{\text{simple}}/B}} \triangleq \mu_i\end{equation}
这里的$\mathcal{B}_{\text{simple}}$就是上一篇文章的$\mathcal{B}_{\text{simple}} = \tr(\boldsymbol{\Sigma})/\boldsymbol{g}^{\top}\boldsymbol{g}$，它又等于$\mathbb{E}[\sigma_i^2]/\mathbb{E}[g_i^2]$（这个$\mathbb{E}$是对下标$i$取平均），也就是说，它是将原本跟下标$i$有关的$\sigma_i^2/g_i^2$，替换成跟下标无关的某种平均值$\mathbb{E}[\sigma_i^2]/\mathbb{E}[g_i^2]$。这样近似之后结果得以简化，但仍保留了关于$B$的函数形式。

然后是二阶矩$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$，这里我们重新引入分量独立假设以简化结果。不引入这个假设其实也可以计算，不过结果会复杂一些，并且也需要另外的假设来简化计算，所以还不如直接引入独立假设。在独立假设之下，$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}$分$i\neq j$和$i=j$两部分计算，当$i\neq j$时，
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j} = \mathbb{E}[(\tilde{\varphi}_B)_i(\tilde{\varphi}_B)_j] = \mathbb{E}[(\tilde{\varphi}_B)_i]\mathbb{E}[(\tilde{\varphi}_B)_j] \approx \mu_i \mu_j\end{equation}
当$i=j$时就更简单了，因为$\sign$的平方必然是1，所以它的期望自然也是1。因此，总的结果可以简写成$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}\approx \mu_i\mu_j + \delta_{i,j}(1 - \mu_i\mu_j)$。

反常现象 #

将上述计算结果代入到式$\eqref{eq:eta-opt}$，我们得到
\begin{equation}\eta^* \approx \frac{\sum_i |g_i|}{\frac{1}{\beta}\sum_i H_{i,i} + \beta\sum_{i\neq j} H_{i,j}\sign(g_i g_j)}\label{eq:eta-opt-sign}\end{equation}
其中$\beta = (1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}$。注意$\beta$关于$B$是单调递增的，并且$\beta\in(0,1)$，所以$\beta$可以看成是标准化的Batch Size。然而，关于$\beta$却并不总是单调的，所以这就可能会出现“Batch Size增大，学习率反而应该减小”的反常行为，原论文称之为“Surge现象”。

让我们一步步来理解。当$B\ll \mathcal{B}_{\text{simple}}$时，有$\beta\approx \sqrt{B/\mathcal{B}_{\text{simple}}}$，此时$\beta \ll 1$，那么式$\eqref{eq:eta-opt-sign}$分母中$1/\beta$项将主导，于是有
\begin{equation}\eta^* \approx \frac{\sum_i |g_i|}{\sum_i H_{i,i}}\beta \approx \frac{\sum_i |g_i|}{\sum_i H_{i,i}}\sqrt{B/\mathcal{B}_{\text{simple}}}\propto \sqrt{B}\end{equation}
这表明SignSGD的学习率在小Batch Size时适用于平方根缩放。由于我们在分析时要假定Hessian矩阵的正定性，所以必然有$\sum_i H_{i,i} > 0$，那么当$\sum_{i\neq j} H_{i,j}\sign(g_i g_j) \leq 0$时，式$\eqref{eq:eta-opt-sign}$关于$\beta$始终是单调递增的，所以$\eta^*$关于$B$也是单调递增的，此时不存在反常表现。

当$\sum_{i\neq j} H_{i,j}\sign(g_i g_j) > 0$时，根据基本不等式我们可以得出式$\eqref{eq:eta-opt-sign}$分母存在一个最小值点
\begin{equation}\beta^* = \sqrt{\frac{\sum_i H_{i,i}}{\sum_{i\neq j} H_{i,j}\sign(g_i g_j)}}\end{equation}
注意$\beta\in(0, 1)$，所以还有一个附加条件$\beta^*\in(0, 1)$，此时$\eta^*$关于$B$就不再是单调递增，而是先增后减，存在一个临界Batch Size，超过这个临界Batch Size后学习率反而应该降低，这便是“Surge现象”。

原因反思 #

为什么会出现Surge现象这种反常行为呢？事实上，这是优化器本身的假设与我们的分析方法不完全相容的体现。具体来说，我们为了估计最优学习率，将Loss的增量展开到了二阶近似，并假设了Hessian矩阵的正定性。在这些设定之下，最优更新量应该是牛顿法，即$\boldsymbol{H}^{-1}\boldsymbol{g}$。

在牛顿法视角下，不同优化器实际上是对Hessian矩阵的不同假设，比如SGD对应于假设$\boldsymbol{H}=\eta_{\max}^{-1} \boldsymbol{I}$，而SignSGD则对应于假设$\newcommand{diag}{\mathop{\text{diag}}}\boldsymbol{H}=\eta_{\max}^{-1} \diag(|\boldsymbol{g}|)$，当然实际训练我们只能将$\boldsymbol{g}$替代为$\tilde{\boldsymbol{g}}_B$。Surge现象实际体现了$B\to\infty$时，SignSGD所假设的Hessian矩阵与实际Hessian矩阵的偏离程度在变大。

我们知道，如今的LLM模型参数都是以亿起步的，不管是完整的Hessian矩阵还是协方差矩阵，其计算都是近乎不可能的事情，这也是我们计算二阶矩$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$时要引入独立假设的原因之一，此时协方差矩阵就只是一个对角阵，估算才是可行的。Hessian矩阵也是类似，我们往往只能对特定结构的Hessian矩阵进行计算。

例如，代入$\boldsymbol{H}=\eta_{\max}^{-1} \diag(|\boldsymbol{g}|)$到式$\eqref{eq:eta-opt-sign}$可得$\eta^*\approx \eta_{\max} \beta = \eta_{\max} / \sqrt{1 + \mathcal{B}_{\text{simple}}/B}$，这个形式就很简洁了，并且没有反常行为。这是否意味着Surge现象不会出现了？并不是，Surge现象是客观存在的，这里更多的是想说：当我们在实验中观察到Surge现象时，也许首要的事情并不是修正$\eta^*$的变化规律，而应该是要更换优化器了。

损失变化 #

有了$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$和$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$，我们还可以像上一篇文章一样计算$\overline{\Delta\mathcal{L}}$，特别有意思的是，它跟SGD的结果具有相同的格式
\begin{equation}\overline{\Delta\mathcal{L}} = \mathcal{L}(\boldsymbol{w}) - \mathbb{E}[\mathcal{L}(\boldsymbol{w} - \eta^*\tilde{\boldsymbol{g}}_B)] \approx \frac{(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}\boldsymbol{g})^2}{2\tr(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]\boldsymbol{H})}\approx \frac{\Delta\mathcal{L}_{\max}}{1 + \mathcal{B}_{\text{noise}}/B}\end{equation}
其中
\begin{equation}\Delta\mathcal{L}_{\max} = \frac{\frac{1}{2}(\sum_i |g_i|)^2}{\sum_i H_{i,i} + \sum_{i\neq j} H_{i,j}\sign(g_i g_j)},\quad \mathcal{B}_{\text{noise}} = \frac{\mathcal{B}_{\text{simple}}\sum_i H_{i,i}}{\sum_i H_{i,i} + \sum_{i\neq j} H_{i,j}\sign(g_i g_j)}\end{equation}
注意这里是保留了完整Hessian矩阵的，所以结果其实颇为有趣——尽管学习率$\eta^*$可能会出现Surge现象，但损失函数的平均增量并没有这个现象，它关于$B$始终是单调递增的，并且还保持跟SGD相同的形式，这意味着我们可以推导出相同的“训练数据量-训练步数”关系：
\begin{equation}\left(\frac{S}{S_{\min}} - 1\right)\left(\frac{E}{E_{\min}} - 1\right) = 1\end{equation}
一个更值得思考的问题是，为什么SGD和SignSGD的更新量截然不同，包括学习率$\eta^*$的表现也有明显差异，但$\overline{\Delta\mathcal{L}}$关于$B$的关系却有着相同的形式。这单纯就只是巧合，还是有更深刻的原理在背后支撑？

一般规律 #

依旧是从平均场近似出发，笔者得到了一个倾向于后者的答案。不管是$\eta^*$还是$\overline{\Delta\mathcal{L}}$，核心难度都是计算$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$和$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$，所以我们的目标是探寻两者的统一计算规律。

我们一般地设$\tilde{\boldsymbol{\varphi}}_B=\tilde{\boldsymbol{H}}{}_B^{-1}\tilde{\boldsymbol{g}}_B$，$\tilde{\boldsymbol{H}}_B$是某个半正定矩阵，那么我们可以写出
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] = \mathbb{E}[\tilde{\boldsymbol{H}}{}_B^{-1}\tilde{\boldsymbol{g}}_B]\approx \underbrace{\mathbb{E}[\tilde{\boldsymbol{H}}_B]^{-1}}_{\text{记为}\hat{\boldsymbol{H}}{}^{-1}}\mathbb{E}[\tilde{\boldsymbol{g}}_B] = \hat{\boldsymbol{H}}{}^{-1}\boldsymbol{g}\end{equation}
以及
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}] = \mathbb{E}[\tilde{\boldsymbol{H}}{}_B^{-1}\tilde{\boldsymbol{g}}_B\tilde{\boldsymbol{g}}_B^{\top}\tilde{\boldsymbol{H}}{}_B^{-1}]\approx \mathbb{E}[\tilde{\boldsymbol{H}}_B]^{-1}\mathbb{E}[\tilde{\boldsymbol{g}}_B\tilde{\boldsymbol{g}}_B^{\top}]\mathbb{E}[\tilde{\boldsymbol{H}}_B]^{-1} = \hat{\boldsymbol{H}}{}^{-1}(\boldsymbol{g}\boldsymbol{g}^{\top} + \boldsymbol{\Sigma}/B)\hat{\boldsymbol{H}}{}^{-1}
\end{equation}
代入$\overline{\Delta\mathcal{L}}$的表达式，我们得到
\begin{equation}\overline{\Delta\mathcal{L}} \approx \frac{1}{2}\frac{(\boldsymbol{g}^{\top}\hat{\boldsymbol{H}}{}^{-1}\boldsymbol{g})^2}{\boldsymbol{g}^{\top}\hat{\boldsymbol{H}}{}^{-1}\boldsymbol{H}\hat{\boldsymbol{H}}{}^{-1}\boldsymbol{g} + \tr(\boldsymbol{\Sigma}\hat{\boldsymbol{H}}{}^{-1}\boldsymbol{H}\hat{\boldsymbol{H}}{}^{-1})/B}\end{equation}
注意上式关于$\hat{\boldsymbol{H}}$是齐次的，如果我们假设$\hat{\boldsymbol{H}}$与$B$的关系可以单独分理出一个标量形式如$\hat{\boldsymbol{H}}\approx f(B) \boldsymbol{G}$，其中$f(B)$是$B$的标量函数，$\boldsymbol{G}$跟$B$不明显相关，那么分子分母是可以同时把$f(B)$约掉的，最终关于$B$的关系，可以整理成如下形式
\begin{equation}\overline{\Delta\mathcal{L}} \approx \frac{\Delta\mathcal{L}_{\max}}{1 + \mathcal{B}_{\text{noise}}/B}\end{equation}
这就证明了$\overline{\Delta\mathcal{L}}$关于$B$具有相同的渐近规律，其核心是关于$\hat{\boldsymbol{H}}$的齐次性。相比之下，$\eta^*$就没有这么统一的结果，因为它关于$\hat{\boldsymbol{H}}$并不是齐次的。

有效分析 #

看到这里，想必大家都已经对平均场方法有所了解，它的主要特点就是计算简单，或者更本质上说，平均场就是挑简单的、能计算的方向去计算，这就导致了它极大的灵活性。灵活性在很多时候也是一种缺点，它意味着我们很难掌握下一步的规律。

至于要解释为什么这样做是有效的，那就更难了，只能具体问题具体分析，甚至有可能具体问题也很难分析下去。笔者的感觉，平均场方法是三分计算、三分幸运、三分直觉，再加上一分的玄学。当然，尝试一下是没问题的，我们就以前面SignSGD的计算为例，尝试做一下分析。

很明显，SignSGD最核心的计算是$\mathbb{E}[\sign(x)]$，我们记$\mathbb{E}[x]=\mu,\mathbb{E}[x^2]=\mu^2 + \sigma^2$，然后写出
\begin{equation}\sign(x) = \frac{x}{\sqrt{x^2}} = \frac{x}{\sqrt{\mu^2 + \sigma^2 + (x^2 - \mu^2 - \sigma^2)}}\end{equation}
假设$x^2 - \mu^2 - \sigma^2$是小量，我们做泰勒展开
\begin{equation}\sign(x) = \frac{x}{\sqrt{\mu^2 + \sigma^2}} - \frac{1}{2}\frac{x(x^2 - \mu^2 - \sigma^2)}{（\mu^2 + \sigma^2)^{3/2}} + \frac{3}{8}\frac{x(x^2 - \mu^2 - \sigma^2)^2}{（\mu^2 + \sigma^2)^{5/2}}-\cdots \end{equation}
现在分母都跟$x$无关的，分子是关于$x$的多项式，所以两边求期望，第一项便是平均场近似的结果$\mu/\sqrt{\mu^2 + \sigma^2}$。为了观察平均场近似的合理性，我们计算第二项
\begin{equation}\frac{1}{2}\frac{\mathbb{E}[x(x^2 - \mu^2 - \sigma^2)]}{（\mu^2 + \sigma^2)^{3/2}} = \frac{1}{2}\frac{\mathbb{E}[x^3] - (\mu^3 + \mu\sigma^2)}{（\mu^2 + \sigma^2)^{3/2}} \end{equation}
这涉及到了$\mathbb{E}[x^3]$，这是一个新的统计量，它是平均场误差的关键因素。我们可以拿正态分布$\mathcal{N}(x;\mu,\sigma^2)$来感知一下，此时$\mathbb{E}[x^3]=\mu^3 + 3\mu\sigma^2$，代入上式的
\begin{equation}\frac{\mu\sigma^2}{（\mu^2 + \sigma^2)^{3/2}} = \frac{\sigma^2/\mu^2}{（1 + \sigma^2/\mu^2)^{3/2}}\end{equation}
右端是一个有界的式子，最大值在$\sigma^2/\mu^2=2$取到，结果是$2/3^{3/2}=0.3849\cdots$。这表明平均场近似的误差极可能是有限的，并且误差项随着$\sigma\to 0$和$\sigma\to\infty$都趋于0，这些都一定程度上体现了平均场近似的可用性。

广义近似 #

之所以选择分析SignSGD，原因之一是我们通常用它作为Adam的理论近似。在《Adam的epsilon如何影响学习率的Scaling Law？》中，我们计算过一个理论上更好的近似SoftSignSGD，它考虑了$\epsilon$的影响。
\begin{equation}\sign(x)=\frac{x}{\sqrt{x^2}}\quad\to\quad\newcommand{softsign}{\mathop{\text{softsign}}}\softsign(x)=\frac{x}{\sqrt{x^2+\epsilon^2}}\end{equation}
此时$\tilde{\boldsymbol{\varphi}}_B = \softsign(\tilde{\boldsymbol{g}}_B)$。让我们直接进入主题
\begin{equation}\begin{aligned}
&\,\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]_i = \mathbb{E}[\softsign((\tilde{g}_B)_i)] = \mathbb{E}\bigg[\frac{(\tilde{g}_B)_i}{\sqrt{(\tilde{g}_B)_i^2 + \epsilon^2}}\bigg]\approx \frac{\mathbb{E}[(\tilde{g}_B)_i]}{\sqrt{\mathbb{E}[(\tilde{g}_B)_i^2]+ \epsilon^2}} \\[8pt]
=&\, \frac{g_i}{\sqrt{g_i^2 + \sigma_i^2/B + \epsilon^2}} = \frac{\softsign(g_i)}{\sqrt{1 + \sigma_i^2/(g_i^2 + \epsilon^2)/B}}\approx \frac{\softsign(g_i)}{\sqrt{1 + \mathcal{B}_{\text{simple}}/B}}\triangleq \nu_i\beta
\end{aligned}\end{equation}
这里的$\mathcal{B}_{\text{simple}}$有少许不同，它是$\tr(\boldsymbol{\Sigma})/(\boldsymbol{g}^{\top}\boldsymbol{g} + N\epsilon^2)$，其中$N$是模型总参数量，即$\boldsymbol{g}\in\mathbb{R}^N$；至于最后的$\nu_i=\softsign(g_i), \beta = (1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}$。接着计算$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$，在独立假设下当$i\neq j$时依旧可以分别求均值，因此有$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}=\nu_i \nu_j \beta^2$，所以只需要计算$i=j$的情形：
\begin{equation}\begin{aligned}
&\,\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,i} = \mathbb{E}[\softsign((\tilde{g}_B)_i)^2] = \mathbb{E}\bigg[\frac{(\tilde{g}_B)_i^2}{(\tilde{g}_B)_i^2 + \epsilon^2}\bigg]\approx \frac{\mathbb{E}[(\tilde{g}_B)_i^2]}{\mathbb{E}[(\tilde{g}_B)_i^2]+ \epsilon^2} \\[8pt]
=&\, \frac{g_i^2 + \sigma_i^2/B}{g_i^2 + \sigma_i^2/B + \epsilon^2} = 1 - \frac{1 - \softsign(g)^2}{1 + \sigma_i^2/(g_i^2 + \epsilon^2)/B}\approx 1 - \frac{1 - \softsign(g)^2}{1 + \mathcal{B}_{\text{simple}}/B}
\end{aligned}\end{equation}
这可以统一地写成$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}\approx \nu_i \nu_j\beta^2 + \delta_{i,j}(1-\beta^2)$，于是
\begin{equation}\eta^* \approx \frac{\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}\boldsymbol{g}}{\text{Tr}(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]\boldsymbol{H})} \approx \frac{\beta\sum_i \nu_i g_i}{\sum_i H_{i,i} + \beta^2(\sum_{i,j} \nu_i \nu_j H_{i,j} - \sum_i H_{i,i})}\end{equation}
上式除了$\beta$外，其余部份都跟$B$无关，因此我们已经得到$\eta^*$关于$B$的显式关系，形式跟SignSGD的大同小异。剩下的分析，可以参考《Adam的epsilon如何影响学习率的Scaling Law？》或者模仿前面的内容进行。

文章小结 #

这篇文章我们使用了平均场近似重新计算了SignSGD和SoftSignSGD的结论，大大简化了相关计算过程，并初步思考了这些计算的一般规律。

转载到请包括本文地址：https://www.kexue.fm/archives/11280

更详细的转载事宜请参考：《科学空间FAQ》