LAMB：结合层级自适应学习率与Adam.

paper：Large Batch Optimization for Deep Learning: Training BERT in 76 minutes

LARS是一种常用的对大型神经网络进行大批量训练的方法，它通过设置分层自适应的学习率，能够在几分钟内在ImageNet上训练ResNet模型。然而LARS在训练BERT等自注意力模型时表现比较差，这说明它在不同任务中的表现不一致。

作者首先提出了一种原则性的学习率分层自适应策略，并提出了一种新的分层自适应大批量优化技术LAMB，通过将批量大小增加到TPU v3的内存限制，实现了使用 $76$ 分钟训练BERT模型。

1. 学习率分层自适应

在大批量训练中，参数更新通常是分层进行的。对于网络第 $i \in [h]$ 层的参数，通用的参数更新形式为：

x_{t + 1}^{(i)} = x_{t}^{(i)} - η_{t} \frac{ϕ (| | x_{t}^{(i)} | |)}{| | g_{t}^{(i)} | |} g_{t}^{(i)}

$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} = β_{1} m_{t - 1} + (1 - β_{1}) (g_{t} + λ x_{t})

$m_t = \beta_1 m_{t-1} + (1-\beta_1)(g_t+\lambda x_t)$

x_{t + 1}^{(i)} = x_{t}^{(i)} - η_{t} \frac{ϕ (| | x_{t}^{(i)} | |)}{| | m_{t}^{(i)} | |} m_{t}^{(i)}

$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} = β_{1} m_{t - 1} + (1 - β_{1}) g_{t} v_{t} = β_{2} v_{t - 1} + (1 - β_{2}) g_{t}^{2} m_{t} = \frac{m_{t}}{1 - β_{1}^{t}}, v_{t} = \frac{v_{t}}{1 - β_{2}^{t}} r_{t} = \frac{m_{t}}{\sqrt{v_{t}} + ϵ} x_{t + 1}^{(i)} = x_{t}^{(i)} - η_{t} \frac{ϕ (| | x_{t}^{(i)} | |)}{| | r_{t}^{(i)} + λ x_{t}^{(i)} | |} (r_{t}^{(i)} + λ x_{t}^{(i)})

$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)})$