LAMB：结合层级自适应学习率与Adam.
LARS是一种常用的对大型神经网络进行大批量训练的方法，它通过设置分层自适应的学习率，能够在几分钟内在ImageNet上训练ResNet模型。然而LARS在训练BERT等自注意力模型时表现比较差，这说明它在不同任务中的表现不一致。
作者首先提出了一种原则性的学习率分层自适应策略，并提出了一种新的分层自适应大批量优化技术LAMB，通过将批量大小增加到TPU v3的内存限制，实现了使用$76$分钟训练BERT模型。
1. 学习率分层自适应
在大批量训练中，参数更新通常是分层进行的。对于网络第$i \in [h]$层的参数，通用的参数更新形式为：
\[x_{t+1}^{(i)} = x_{t}^{(i)} -\eta_t \frac{\phi(|| x_{t}^{(i)}||)}{||g_{t}^{(i)}||} g_{t}^{(i)}\]其中梯度$g_t^{(i)}$是由基本优化算法$\mathcal{A}$计算得到的。$\phi(|| x_{t}^{(i)}||)$是指通过参数的l2范数对学习率进行调整。更新梯度也通过l2范数标准化为$\frac{g_{t}^{(i)}}{||g_{t}^{(i)}||}$。此时参数更新量与梯度大小无关，不会受梯度爆炸/消失的影响。尽管梯度更新的方向有所偏差，对结果不会产生明显的影响。
在实践中函数$\phi(\cdot)$可以简单地设置为$\phi(z)=\min(\max(z,\gamma_l),\gamma_u)$或$\phi(z)=z$。后者对应的自适应学习率即为$\frac{|| x_{t}^{(i)}||}{||g_{t}^{(i)}||}$。
2. LARS 与 LAMB
LARS是动量算法与层级自适应学习率的结合。其中动量采用滑动平均累计，并额外引入权重衰减：
\[m_t = \beta_1 m_{t-1} + (1-\beta_1)(g_t+\lambda x_t)\] \[x_{t+1}^{(i)} = x_{t}^{(i)} -\eta_t \frac{\phi(|| x_{t}^{(i)}||)}{||m_{t}^{(i)}||} m_{t}^{(i)}\]LAMB是Adam与层级自适应学习率的结合。LAMB的自适应表现在两个方面，使用梯度二阶矩的平方根进行的归一化和学习率的分层归一化。
\[m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t \\ v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2 \\ m_t = \frac{m_t}{1-\beta_1^t}, \quad v_t = \frac{v_t}{1-\beta_2^t} \\ r_t = \frac{m_t}{\sqrt{v_t}+\epsilon} \\ x_{t+1}^{(i)} = x_{t}^{(i)} -\eta_t \frac{\phi(|| x_{t}^{(i)}||)}{||r_{t}^{(i)}+\lambda x_t^{(i)}||} (r_{t}^{(i)}+\lambda x_t^{(i)})\]
3. 实验分析
作者展示了训练BERT的实验结果。通过将批量扩大到$32$k(受限于TPU内存)，可以在$76$分钟内训练完成，并且获得与通常训练相当的准确率。
在训练时设置的部分超参数如下：
