#

Rectified Adam (RAdam) optimizer

This implementation is based on the official implementation of the paper On the Variance of the Adaptive Learning Rate and Beyond.

We have implemented it in PyTorch as an extension to our AMSGrad implementation thus requiring only the modifications to be implemented.

Adam optimizer sometimes converges to a bad local optima during the initial stages of the training; especially when training transformers. Researches use warmups to counter this; for the the initial training steps (warm-up stage) they use a low learning rate. This paper identifies the problem to be the high variance of adaptive learning rate during initial stages of training, and counters it using a new rectification term to reduce variance.

The paper also evaluates two variance reduction mechanisms: Adam-2k: Only compute the adaptive learning rate ( $v_{t}$ in Adam) during the first 2k steps, without changing parameters or calculating momentum ( $m_{t}$ ). Adam-eps: Adam with large $ϵ \approx 1 0^{- 4}$ .

Rectified Adam

Let $σ (g_{1}, ..., g_{t})$ and $ψ (g_{1}, ..., g_{t})$ be the functions to calculate momentum and adaptive learning rate. For Adam, they are

σ (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}}

Exponential moving average as simple moving average

The distribution of exponential moving average can be approximated as a simple moving average.

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

Here we are taking the simple moving average of the last $f (t, β_{2})$ gradients. $f (t, β_{2})$ satisfies the following,

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

which gives, $f (t, β_{2}) = \frac{2}{1 - β _{2}} - 1 - \frac{2 t β _{2} ^{t}}{1 - β _{2} ^{t}}$

Scaled inverse chi-squared

From above we have $p (ψ^{2} (g_{1}, ..., g_{t})) \approx p (\frac{\sum _{i = 1}^{f (t, β_{2})} g _{t + 1 - i}^{2}}{f ( t , β _{2} )})$ where $g_{i} \sim N (0, σ^{2})$ . Note that $s i g ma$ here is the standard deviation and different from $σ (.)$ for momentum.

Scaled inverse chi-squared is the distribution of squared inverse of mean of $p$ normal distributions. $p (\frac{\sum _{i = 1}^{f (t, β_{2})} g _{t + 1 - i}^{2}}{f ( t , β _{2} )}) \sim Scale-inv X^{2} (ρ, \frac{1}{σ ^{2}})$ where $ρ = f (t, β_{2})$ .

Rectification

They prove that variance of $ψ (.)$ decreases with $ρ$ when $ψ^{2} (.) \sim Scale-inv X^{2} (ρ, \frac{1}{σ ^{2}})$ .

Therefore the variance is minimized at maximal $ρ$ which is $ρ_{\infty} = \frac{2}{1 - β _{2}} - 1$ . Let the minimum variance be $C_{va r}$

In order to ensure that the adaptive learning rate $ψ (.)$ has consistent variance, we rectify the variance with $r$

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

Approximating $Va r [ψ (.)]$

They estimate $Va r [ψ (.)] \approx \frac{Va r [ ψ ^{2} ( . )]}{4 E [ ψ ^{2} ( . )}$ based on first order expansion of $ψ^{2} (.)$ 🤪 I didn't get how it was derived.

From $Scale-inv X^{2}$ distribution we have,

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

which gives, $Va r [ψ (.)] \approx \frac{ρ}{2 ( ρ - 2 ) ( ρ - 4 ) σ ^{2}}$

Rectification term

We have

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

where $C_{va r}$ is $Va r [ψ (.)]$ for $ρ_{\infty}$ . Lt $ρ$ and step $t$ be $ρ_{t}$ , and $r_{t}$ be the rectification term at 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}}

This gives,

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

#

Rectified Adam Optimizer

This class extends from AMSAdam optimizer defined in amsadam.py .

148class RAdam(AMSGrad):

#

Initialize the optimizer

params is the list of parameters
lr is the learning rate $α$
betas is a tuple of ( $β_{1}$ , $β_{2}$ )
eps is $\overset{ϵ}{^}$ or $ϵ$ based on optimized_update
weight_decay is an instance of class WeightDecay defined in __init__.py
optimized_update is a flag whether to optimize the bias correction of the second moment by doing it after adding $ϵ$
amsgrad is a flag indicating whether to use AMSGrad or fallback to plain Adam
degenerate_to_sgd whether to use sgd when the rectification term $r_{t}$ is intractable.
defaults is a dictionary of default for group values. This is useful when you want to extend the class 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)

#

Take an update step for a given parameter tensor

state is the optimizer state of the parameter (tensor)
group stores optimizer attributes of the parameter group
grad is the current gradient tensor $g_{t}$ for the parameter $θ_{t - 1}$
param is the parameter tensor $θ_{t - 1}$

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

#

Calculate weight decay

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

#

Get $m_{t}$ and $v_{t}$ ; i.e. $σ (.)$ and $ψ (.)$ without bias correction

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

#

Calculate $t$ the number of optimizer steps

195        state['step'] += 1

#

Perform RAdam update

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

#

Calculate rectification term $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}$ is tractable when $ρ_{t} >= 4$ . We are being a little more conservative since it's an approximated value

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

#

Do the RAdam parameter update

state is the optimizer state of the parameter (tensor)
group stores optimizer attributes of the parameter group
param is the parameter tensor $θ_{t - 1}$
m and v are the uncorrected first and second moments $m_{t}$ and $v_{t}$ ; i.e. $σ (.)$ and $ψ (.)$ without bias correction

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):

#

Get $β_{1}$ and $β_{2}$

235        beta1, beta2 = group['betas']

#

Bias correction term for $\overset{m}{^}_{t}$ , $1 - β_{1}^{t}$

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

#

Bias correction term for $\overset{v}{^}_{t}$ , $1 - β_{2}^{t}$

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

#

Get learning rate

244        lr = self.get_lr(state, group)

#

If $r_{t}$ is intractable

247        if r is not None:

#

Whether to optimize the computation by combining scalar computations

249            if self.optimized_update:

#

Denominator $v_{t} + \overset{ϵ}{^}$

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

#

Step size $α r_{t} * \frac{1 - β _{2} ^{t}}{1 - β _{1} ^{t}}$

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

#

Update parameters $θ_{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)

#

Computation without optimization

258            else:

#

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

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

#

Step size $\frac{α r _{t}}{1 - β _{1} ^{t}}$

262                step_size = lr * r / bias_correction1

#

Update parameters $θ_{t} \leftarrow θ_{t - 1} - α r_{t} \cdot \frac{m ^ _{t}}{v ^ _{t} + ϵ}$

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

#

If $r_{t}$ is intractable do a SGD with momentum

268        elif self.degenerated_to_sgd:

#

Step size $\frac{α}{1 - β _{1} ^{t}}$

270            step_size = lr / bias_correction1

#

Update parameters $θ_{t} \leftarrow θ_{t - 1} - α \cdot \overset{m}{^}_{t}$

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

#

Plot $r_{t}$ against $t$ for various $β_{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()