#

校正亚当 (raDAM) 优化器

我们已经在 PyTorch 中实现了它，作为我们的 AmsGrad 实现的扩展，因此只需要实施修改即可。

在训练的初始阶段，Adam optimizer 有时会收敛到糟糕的局部最佳值；尤其是在训练变形金刚时。研究使用热身来应对这种情况；对于最初的训练步骤（热身阶段），他们使用较低的学习率。本文将问题确定为训练初始阶段自适应学习率的高方差，并使用新的校正术语来减少方差。

本文还评估了两种方差缩减机制：adam-2K：仅计算前 2k 步长的自适应学习率（ $v_{t}$ 在 Adam 中），而不更改参数或计算动量（ $m_{t}$ )。adam-eps：Adam 很大 $ϵ \approx 1 0^{- 4}$ 。

纠正了亚当

让 $σ (g_{1}, ..., g_{t})$ 和 $ψ (g_{1}, ..., g_{t})$ 成为计算动量和自适应学习速率的函数。对亚当来说，他们是

σ (g_{1}, ..., g_{t}) ψ (g_{1}, ..., g_{t}) = \frac{( 1 - β _{1} ) \sum _{i = 1}^{t} β _{1} ^{t - i} g _{i}}{1 - β _{1} ^{t}} = \frac{1 - β _{2} ^{t}}{( 1 - β _{2} ) \sum _{i = 1}^{t} β _{2} ^{t - i} g _{i}^{2}}

指数移动平均线作为简单移动平均线

指数移动平均线的分布可以近似为简单移动平均线。

p (\frac{( 1 - β _{2} ) \sum _{i = 1}^{t} β _{2} ^{t - i} g _{i}^{2}}{1 - β _{2} ^{t}}) \approx p (\frac{\sum _{i = 1}^{f (t, β_{2})} g _{t + 1 - i}^{2}}{f ( t , β _{2} )})

这里我们取最后一个 $f (t, β_{2})$ 梯度的简单移动平均线。 $f (t, β_{2})$ 满足以下条件，

\frac{( 1 - β _{2} ) \sum _{i = 1}^{t} β _{2} ^{t - i} \cdot i}{1 - β _{2} ^{t}} = \frac{\sum _{i = 1}^{f (t, β_{2})} ( t + 1 - i )}{f ( t , β _{2} )}

这给了， $f (t, β_{2}) = \frac{2}{1 - β _{2}} - 1 - \frac{2 t β _{2} ^{t}}{1 - β _{2} ^{t}}$

缩放反向卡方

从上面看，我们有 $p (ψ^{2} (g_{1}, ..., g_{t})) \approx p (\frac{\sum _{i = 1}^{f (t, β_{2})} g _{t + 1 - i}^{2}}{f ( t , β _{2} )})$ 哪里 $g_{i} \sim N (0, σ^{2})$ 。请注意， $s i g ma$ 这里是标准差，与动 $σ (.)$ 量不同。

缩放逆卡方是 $p$ 正态分布均值的逆平方分布。 $p (\frac{\sum _{i = 1}^{f (t, β_{2})} g _{t + 1 - i}^{2}}{f ( t , β _{2} )}) \sim Scale-inv X^{2} (ρ, \frac{1}{σ ^{2}})$ 在哪里 $ρ = f (t, β_{2})$ 。

整改

他们证明了随时间变化的变化 $ψ (.)$ 而 $ρ$ 降低 $ψ^{2} (.) \sim Scale-inv X^{2} (ρ, \frac{1}{σ ^{2}})$ 。

因此，方差最小化 $ρ$ 为最大值 $ρ_{\infty} = \frac{2}{1 - β _{2}} - 1$ 。让最小方差为 $C_{va r}$

为了确保自适应学习率 $ψ (.)$ 具有一致的方差，我们使用以下方法校正方差 $r$

r = \frac{C _{va r}}{Va r [ ψ ( . ) ]}

近似值 $Va r [ψ (.)]$

他们 $Va r [ψ (.)] \approx \frac{Va r [ ψ ^{2} ( . )]}{4 E [ ψ ^{2} ( . )}$ 根据一阶扩张估计 $ψ^{2} (.)$ 🤪 我不明白它是如何得出的。

从 $Scale-inv X^{2}$ 分发来看，

E [ψ^{2} (.)] Va r [ψ^{2} (.)] = \frac{ρ / σ ^{2}}{ρ - 2} = \frac{2 ρ / σ ^{4}}{( ρ - 2 ) ^{2} ( ρ - 2 )}

这给了， $Va r [ψ (.)] \approx \frac{ρ}{2 ( ρ - 2 ) ( ρ - 4 ) σ ^{2}}$

整改期限

我们有

r Va r [ψ (.)] = \frac{C _{va r}}{Va r [ ψ ( . ) ]} \approx \frac{ρ}{2 ( ρ - 2 ) ( ρ - 4 ) σ ^{2}}

在 $C_{va r}$ 哪 $Va r [ψ (.)]$ 里 $ρ_{\infty}$ 。Lt $ρ$ and step $t$ be $ρ_{t}$ ，然后 $r_{t}$ 成为 step 的整改期限 $t$ 。

C_{va r} Va r [ψ (g_{1}, ..., g_{t})] \approx \frac{ρ _{\infty}}{2 ( ρ _{\infty} - 2 ) ( ρ _{\infty} - 4 ) σ ^{2}} \approx \frac{ρ _{t}}{2 ( ρ _{t} - 2 ) ( ρ _{t} - 4 ) σ ^{2}}

这给了，

r_{t} = \frac{( ρ _{t} - 2 ) ( ρ _{t} - 4 ) ρ _{\infty}}{( ρ _{\infty} - 2 ) ( ρ _{\infty} - 4 ) ρ _{t}}

139import math
140from typing import Dict, Optional
141
142import torch
143
144from labml_nn.optimizers import WeightDecay
145from labml_nn.optimizers.amsgrad import AMSGrad

#

纠正亚当优化器

这个类是从中定义的 AmsadAM 优化器扩展而来的amsadam.py 。

148class RAdam(AMSGrad):

#

初始化优化器

params 是参数列表
lr 是学习率 $α$
betas 是 ( $β_{1}$ , $β_{2}$ ) 的元组
eps 是 $\overset{ϵ}{^}$ 或 $ϵ$ 基于optimized_update
weight_decay 是在中WeightDecay 定义的类的实例 __init__.py
optimized_update 是一个标志，是否在添加后通过这样做来优化第二个时刻的偏差校正 $ϵ$
amsgrad 是一个标志，指示是使用 AmsGrad 还是回退到普通的 Adam
degenerate_to_sgd 纠正术语 $r_{t}$ 难以处理时是否使用 sgd。
defaults 是组值的默认字典。当你想扩展类时，这很有用RAdam 。

155    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
156                 weight_decay: WeightDecay = WeightDecay(),
157                 optimized_update: bool = True,
158                 amsgrad=False,
159                 degenerated_to_sgd=True, defaults=None):

#

175        self.degenerated_to_sgd = degenerated_to_sgd
176        super().__init__(params, lr, betas, eps, weight_decay, optimized_update, amsgrad, defaults)

#

对给定参数张量执行更新步骤

state 是参数（张量）的优化器状态
group 存储参数组的优化程序属性
grad 是参数的当前梯 $g_{t}$ 度张量 $θ_{t - 1}$
param 是参数张量 $θ_{t - 1}$

178    def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.nn.Parameter):

#

计算体重衰减

189        grad = self.weight_decay(param, grad, group)

#

Get $m_{t}$ an $v_{t}$ d; 即 $ψ (.)$ 不 $σ (.)$ 进行偏差校正

192        m, v = self.get_mv(state, group, grad)

#

$t$ 计算优化器步数

195        state['step'] += 1

#

执行 raDAM 更新

198        self.r_adam_update(state, group, param, m, v)

#

计算整改期限 $r_{t}$

200    @staticmethod
201    def calc_rectification_term(beta2: float, step: int) -> Optional[float]:

#

$β_{2}^{t}$

207        beta2_t = beta2 ** step

#

$ρ_{\infty} = \frac{2}{1 - β _{2}} - 1$

209        rho_inf = 2 / (1 - beta2) - 1

#

$ρ_{t} = \frac{2}{1 - β _{2}} - 1 - \frac{2 t β _{2} ^{t}}{1 - β _{2} ^{t}}$

211        rho = rho_inf - 2 * step * beta2_t / (1 - beta2_t)

#

$r_{t}$ 什么时候是可以处理 $ρ_{t} >= 4$ 的。我们稍微保守一点，因为它是近似值

215        if rho >= 5:

#

$r_{t} = \frac{( ρ _{t} - 2 ) ( ρ _{t} - 4 ) ρ _{\infty}}{( ρ _{\infty} - 2 ) ( ρ _{\infty} - 4 ) ρ _{t}}$

217            r2 = (rho - 4) / (rho_inf - 4) * (rho - 2) / rho * rho_inf / (rho_inf - 2)
218            return math.sqrt(r2)
219        else:
220            return None

#

是否更新 R adAM 参数

state 是参数（张量）的优化器状态
group 存储参数组的优化程序属性
param 是参数张量 $θ_{t - 1}$
m 和v 是未校正的第一个和第二个时刻 $m_{t}$ $v_{t}$ ；即 $σ (.)$ 和 $ψ (.)$ 没有偏差校正

222    def r_adam_update(self, state: Dict[str, any], group: Dict[str, any], param: torch.nn.Parameter,
223                      m: torch.Tensor, v: torch.Tensor):

#

获取 $β_{1}$ 和 $β_{2}$

235        beta1, beta2 = group['betas']

#

偏差校正术语 $\overset{m}{^}_{t}$ ， $1 - β_{1}^{t}$

237        bias_correction1 = 1 - beta1 ** state['step']

#

偏差校正术语 $\overset{v}{^}_{t}$ ， $1 - β_{2}^{t}$

239        bias_correction2 = 1 - beta2 ** state['step']
240
241        r = self.calc_rectification_term(beta2, state['step'])

#

获取学习率

244        lr = self.get_lr(state, group)

#

如果 $r_{t}$ 是棘手的

247        if r is not None:

#

是否通过组合标量计算来优化计算

249            if self.optimized_update:

#

分母 $v_{t} + \overset{ϵ}{^}$

251                denominator = v.sqrt().add_(group['eps'])

#

步长 $α r_{t} * \frac{1 - β _{2} ^{t}}{1 - β _{1} ^{t}}$

253                step_size = lr * math.sqrt(bias_correction2) * r / bias_correction1

#

更新参数 $θ_{t} \leftarrow θ_{t - 1} - α r_{t} \frac{1 - β _{2} ^{t}}{1 - β _{1} ^{t}} \cdot \frac{m _{t}}{v _{t} + ϵ ^}$

256                param.data.addcdiv_(m, denominator, value=-step_size)

#

无需优化的计算

258            else:

#

分母 $\frac{v _{t}}{1 - β _{2} ^{t}} + ϵ$

260                denominator = (v.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])

#

步长 $\frac{α r _{t}}{1 - β _{1} ^{t}}$

262                step_size = lr * r / bias_correction1

#

更新参数 $θ_{t} \leftarrow θ_{t - 1} - α r_{t} \cdot \frac{m ^ _{t}}{v ^ _{t} + ϵ}$

265                param.data.addcdiv_(m, denominator, value=-step_size)

#

如果 $r_{t}$ 难以解决，那就用势头做新加坡元

268        elif self.degenerated_to_sgd:

#

步长 $\frac{α}{1 - β _{1} ^{t}}$

270            step_size = lr / bias_correction1

#

更新参数 $θ_{t} \leftarrow θ_{t - 1} - α \cdot \overset{m}{^}_{t}$

273            param.data.add_(m, alpha=-step_size)

#

阴谋 $r_{t}$ $t$ 对抗各种 $β_{2}$

Plot of r_t

276def _test_rectification_term():

#

282    import matplotlib.pyplot as plt
283    import numpy as np
284
285    beta2 = [0.9999, 0.999, 0.99, 0.9, 0.8, 0.6, 0.5]
286    plt.plot(np.arange(1, 5_000), [[RAdam.calc_rectification_term(b, i) for b in beta2] for i in range(1, 5_000)])
287    plt.legend(beta2)
288    plt.title("Optimizer")
289    plt.show()
290
291
292if __name__ == '__main__':
293    _test_rectification_term()