从 SGD 到 AdamW 原理和代码解读

深度学习优化算法经历了 SGD -> SGDM -> NAG ->AdaGrad -> AdaDelta -> Adam -> Nadam -> AdamW 这样的发展历程。

接下来用一个框架来梳理所有的优化算法。

首先定义:待优化参数:$w$ , 目标函数:$f(x)$ , 初始学习率 $\alpha$

然后,开始进行迭代优化。在每个epoch $t$:

  1. 计算目标函数关于当前参数的梯度: $g_t = ∇f(w_t)$
  2. 根据历史梯度计算一阶动量和二阶动量: $m_t = \phi(g_1,g_2,…,g_t); V_t =\psi(g_1,g_2,…,g_t)$
  3. 计算当前时刻的下降梯度: $\eta = \alpha \cdot m_t/\sqrt V_t$
  4. 根据下降梯度进行更新: $w_{t+1} = w_t -\eta_t$

步骤3、4对于各个算法几乎都是一致的,主要的差别就体现在1和2上。

也就是计算一阶动量$m_t$ 和 二阶动量$V_t$时采用不同的套路。

此外在所有优化器代码里有一些函数作用是相通的:

共性的方法有:

  • add_param_group(param_group) : 把参数放进优化器中,这在Fine-tune预训练时可以使冻结层可训练,并随着训练的进行添加到优化器中。
  • load_state_dict(state_dict): 把优化器的状态加载进去。
  • state_dict():返回优化器状态,以dict形式
  • step(closure=None):优化一步参数
  • zero_grad(set_to_none=False):把所有的梯度值设为0

使用方法:

1
2
3
4
5
6
7
for input, target in dataset:
def closure():
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
return loss
optimizer.step(closure)

SGD

SGD没有动量的概念,也就是说:

代入步骤3,可以看到下降梯度就是最简单的

SGD最大的缺点就是下降速度慢,而且可能会在沟壑的两边持续震荡,停留在一个局部最优点。

SGD with Momentum

为了抑制震荡,SGDM认为梯度下降过程可以加入惯性。下坡的时候,如果发现是陡坡,那就利用惯性跑的快一点。

在SGD的基础上引入了一阶动量:

一阶动量就是各个时刻梯度方向的指数移动平均,约等于最近$1/(1-\beta_1)$个时刻的梯度向量和的平均值。

也就是说,t 时刻的下降方向,不仅由当前点的梯度方向决定,而且由此前累积的下降方向决定。

$\beta_1$的经验值为0.9,这意味着下降方向主要是此前累积的下降方向,并略微偏向当前时刻的下降方向。想象高速公路上汽车转弯,在高速向前的同时略微偏向,急转弯可是要出事的。

SGD with Nesterov Acceleration

SGD还有一个问题是困在局部最优的沟壑里震荡。想象一下你走到一个盆地,四周都是略高的小山,你觉得没有下坡的方向,那就只能呆在这里了。可是你如果爬上高地,就会方向外卖的世界还很广阔。

因此外卖不能停留在当前位置去观察未来的方向,而是要向前一步,多看一步,看远一些。

NAG全称Nesterov Accelerated Gradient,是在SGD、SGDM的基础上的进一步改进,改进点在于步骤1。

我们知道在时刻 t 的主要下降方向是由累积动量决定的,自己的梯度方向说了也不算。那与其看当前梯度方向,不如先看看如果跟着累积动量走了一步,那个时候再怎么走。

因此NAG在步骤1,不计算当前位置的梯度方向,而是计算如果按照累积动量走了一步,那个时候的下降方向:

然后用下一个点的梯度方向,与历史累积动量结合,计算步骤2中当前时刻的累加动量。

定义优化器:

1
CLASS torch.optim.SGD(params, lr=<required parameter>, momentum=0, dampening=0, weight_decay=0, nesterov=False)

参数:

  • params (iterable) – 优化器作用的模型参数。
  • lr (float) – learning rate,相当于是统一框架中的 $\alpha$
  • momentum (float, optional) – 动量参数。(默认值:0)
  • weight_decay (float, optional) – 权重衰减系数 weight decay (L2 penalty) (默认值:0)
  • dampening (float, optional) – dampening for momentum (默认值:0)
  • nesterov (bool, optional) – 允许 Nesterov momentum (默认值:False)

源码解读:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import torch
from .optimizer import Optimizer, required


[docs]class SGD(Optimizer):
r"""Implements stochastic gradient descent (optionally with momentum).

Nesterov momentum is based on the formula from
`On the importance of initialization and momentum in deep learning`__.

Args:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float): learning rate
momentum (float, optional): momentum factor (default: 0)
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
dampening (float, optional): dampening for momentum (default: 0)
nesterov (bool, optional): enables Nesterov momentum (default: False)

Example:
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> optimizer.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optimizer.step()

__ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf

.. note::
The implementation of SGD with Momentum/Nesterov subtly differs from
Sutskever et. al. and implementations in some other frameworks.

Considering the specific case of Momentum, the update can be written as

.. math::
\begin{aligned}
v_{t+1} & = \mu * v_{t} + g_{t+1}, \\
p_{t+1} & = p_{t} - \text{lr} * v_{t+1},
\end{aligned}

where :math:`p`, :math:`g`, :math:`v` and :math:`\mu` denote the
parameters, gradient, velocity, and momentum respectively.

This is in contrast to Sutskever et. al. and
other frameworks which employ an update of the form

.. math::
\begin{aligned}
v_{t+1} & = \mu * v_{t} + \text{lr} * g_{t+1}, \\
p_{t+1} & = p_{t} - v_{t+1}.
\end{aligned}

The Nesterov version is analogously modified.
"""

def __init__(self, params, lr=required, momentum=0, dampening=0,
weight_decay=0, nesterov=False):
if lr is not required and lr < 0.0:
raise ValueError("Invalid learning rate: {}".format(lr))
if momentum < 0.0:
raise ValueError("Invalid momentum value: {}".format(momentum))
if weight_decay < 0.0:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))

defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
weight_decay=weight_decay, nesterov=nesterov)
if nesterov and (momentum <= 0 or dampening != 0):
raise ValueError("Nesterov momentum requires a momentum and zero dampening")
super(SGD, self).__init__(params, defaults)

def __setstate__(self, state):
super(SGD, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('nesterov', False)

[docs] @torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step.

Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()

for group in self.param_groups:
weight_decay = group['weight_decay']
momentum = group['momentum']
dampening = group['dampening']
nesterov = group['nesterov']

for p in group['params']:
if p.grad is None:
continue
d_p = p.grad # 得到每个参数的梯度,也就是g_t
if weight_decay != 0: # 如果使用weight_decay的话,相当于目标函数上加上 L2正则
d_p = d_p.add(p, alpha=weight_decay)
if momentum != 0:
param_state = self.state[p]
if 'momentum_buffer' not in param_state:
buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
else:
buf = param_state['momentum_buffer']
# 计算动量,momentum参数beta_1一般取0.9,相当于之前的动量buf乘以0.9再加上此次梯度
# d_p乘以(1-beta_1)=0.1
buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
if nesterov:
# 如果通过nesterov方式更新参数,那么eta_t就相当于g_t+m_t*beta_1
d_p = d_p.add(buf, alpha=momentum)
else:
# 如果不通过nesterov方式更新参数,那么\eta_t就是相当于是上一步计算出的动量m_t
d_p = buf

p.add_(d_p, alpha=-group['lr']) #最后用学习率更新梯度

return loss

AdaGrad

此前我们都没有用到二阶动量。二阶动量的出现,才意味着“自适应学习率”优化算法时代的到来。

SGD及其变种以同样的学习率更新每个参数,但深度神经网络往往包含大量的参数,这些参数并不是总会用得到(想想大规模的embedding)。

对于偶尔更新的参数,我们了解的信息太少,希望能从每个偶然出现的样本身上多学习一些,即学习率大一些。

怎么样去度量历史更新频率呢?那就是二阶动量——该维度上,迄今为止所有梯度值的平方和:

我们在回顾一些步骤3中的下降梯度:

可以看出,此时实质上的学习率由$\alpha$ 变成了,$\alpha/\sqrt{V_t}$。

一般为了避免分母为0,会在分母上加一个小的平滑项。因此$\sqrt{V_t}$是恒大于0的,而且参数更新越频繁,二阶动量越大,学习率就越小。

这一方法在稀疏数据场景下表现非常好,但也存在一些问题:因为$\sqrt{V_t}$ 是单调递增的,会使学习率单调递减至0,可能会使训练过程提前结束,即便后续还有数据也无法学到必要的知识。

定义优化器:

1
CLASS torch.optim.Adagrad(params,lr=0.01,lr_decay=0,weight_decay=0,initial_accumulator_value=0,eps=1e-10)

参数:

  • params (iterable) – 优化器作用的模型参数。
  • lr (float) – learning rate – 相当于是统一框架中的 。
  • lr_decay(float,optional) – 学习率衰减 (默认值:0)
  • weight_decay (float, optional) – 权重衰减系数 weight decay (L2 penalty) (默认值:0)
  • eps(float,optional):防止分母为0的一个小数 (默认值:1e-10)

源码解读:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
[docs]class Adagrad(Optimizer):
"""Implements Adagrad algorithm.

It has been proposed in `Adaptive Subgradient Methods for Online Learning
and Stochastic Optimization`_.

Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-2)
lr_decay (float, optional): learning rate decay (default: 0)
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-10)

.. _Adaptive Subgradient Methods for Online Learning and Stochastic
Optimization: http://jmlr.org/papers/v12/duchi11a.html
"""

def __init__(self, params, lr=1e-2, lr_decay=0, weight_decay=0, initial_accumulator_value=0, eps=1e-10):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= lr_decay:
raise ValueError("Invalid lr_decay value: {}".format(lr_decay))
if not 0.0 <= weight_decay:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
if not 0.0 <= initial_accumulator_value:
raise ValueError("Invalid initial_accumulator_value value: {}".format(initial_accumulator_value))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))

defaults = dict(lr=lr, lr_decay=lr_decay, eps=eps, weight_decay=weight_decay,
initial_accumulator_value=initial_accumulator_value)
super(Adagrad, self).__init__(params, defaults)

for group in self.param_groups:
for p in group['params']:
state = self.state[p]
state['step'] = 0
state['sum'] = torch.full_like(p, initial_accumulator_value, memory_format=torch.preserve_format)

def share_memory(self):
for group in self.param_groups:
for p in group['params']:
state = self.state[p]
state['sum'].share_memory_()

[docs] @torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step.

Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()

for group in self.param_groups:
params_with_grad = []
grads = []
state_sums = []
state_steps = []

for p in group['params']:
if p.grad is not None:
params_with_grad.append(p)
grads.append(p.grad)
state = self.state[p]
state_sums.append(state['sum'])
# update the steps for each param group update
state['step'] += 1
# record the step after step update
state_steps.append(state['step'])

F.adagrad(params_with_grad,
grads,
state_sums,
state_steps,
group['lr'],
group['weight_decay'],
group['lr_decay'],
group['eps'])

return loss

AdaDelta

由于AdaGrad 单调递减的学习率变化过于激进,我们考虑一个改变二阶动量计算方法的策略:不累积全部历史梯度,而只关注过去一段时间窗口的下降梯度。这也是就是AdaDelta名称中Delta的来历。

修改思路很简单,前面讲到,指数移动平均值大约就是过去一段时间的平均值,因此我们用这一方法来计算二阶累积动量:

接下来还是步骤3:

这就避免了二阶动量持续累积、导致训练过程提前结束的问题了。

RMSProp

定义优化器:

1
CLASS torch.optim.RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)

参数:

  • params (iterable) – 优化器作用的模型参数。
  • lr (float) – learning rate – 相当于是统一框架中的 $\alpha$。
  • momentum (float, optional) – 动量参数。(默认值:0)。
  • alpha(float,optional) – 平滑常数 (默认值:0.99)。
  • centered(bool,optional) – ifTrue, compute the centered RMSProp, the gradient is normalized by an estimation of its variance,就是这一项是 True 的话就把方差使用梯度作归一化。
  • weight_decay (float, optional) – 权重衰减系数 weight decay (L2 penalty) (默认值:0)
  • eps(float,optional):防止分母为0的一个小数 (默认值:1e-10)

源码解读:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import torch
from .optimizer import Optimizer


[docs]class RMSprop(Optimizer):
r"""Implements RMSprop algorithm.

Proposed by G. Hinton in his
`course <https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf>`_.

The centered version first appears in `Generating Sequences
With Recurrent Neural Networks <https://arxiv.org/pdf/1308.0850v5.pdf>`_.

The implementation here takes the square root of the gradient average before
adding epsilon (note that TensorFlow interchanges these two operations). The effective
learning rate is thus :math:`\alpha/(\sqrt{v} + \epsilon)` where :math:`\alpha`
is the scheduled learning rate and :math:`v` is the weighted moving average
of the squared gradient.

Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-2)
momentum (float, optional): momentum factor (default: 0)
alpha (float, optional): smoothing constant (default: 0.99)
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
centered (bool, optional) : if ``True``, compute the centered RMSProp,
the gradient is normalized by an estimation of its variance
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)

"""

def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= momentum:
raise ValueError("Invalid momentum value: {}".format(momentum))
if not 0.0 <= weight_decay:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
if not 0.0 <= alpha:
raise ValueError("Invalid alpha value: {}".format(alpha))

defaults = dict(lr=lr, momentum=momentum, alpha=alpha, eps=eps, centered=centered, weight_decay=weight_decay)
super(RMSprop, self).__init__(params, defaults)

def __setstate__(self, state):
super(RMSprop, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('momentum', 0)
group.setdefault('centered', False)

[docs] @torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step.

Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()

for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad
if grad.is_sparse:
raise RuntimeError('RMSprop does not support sparse gradients')
state = self.state[p]

# State initialization
if len(state) == 0:
state['step'] = 0
state['square_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
if group['momentum'] > 0:
state['momentum_buffer'] = torch.zeros_like(p, memory_format=torch.preserve_format)
if group['centered']:
state['grad_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)

square_avg = state['square_avg']
alpha = group['alpha']

state['step'] += 1

if group['weight_decay'] != 0:
grad = grad.add(p, alpha=group['weight_decay'])

square_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha)

if group['centered']:
grad_avg = state['grad_avg']
grad_avg.mul_(alpha).add_(grad, alpha=1 - alpha)
avg = square_avg.addcmul(grad_avg, grad_avg, value=-1).sqrt_().add_(group['eps']) # 计算当前步的动量
else:
avg = square_avg.sqrt().add_(group['eps'])

if group['momentum'] > 0:
buf = state['momentum_buffer']
buf.mul_(group['momentum']).addcdiv_(grad, avg)
p.add_(buf, alpha=-group['lr'])
else:
p.addcdiv_(grad, avg, value=-group['lr'])

return loss

RMSprop算是Adagrad的一种发展,和Adadelta的变体,效果趋于二者之间

Adam

谈到这里,Adam和Nadam的出现就很自然而然了——他们是前述方法的集大成者。

SGDM在SGD基础上增加了一阶动量,AdaGrad和AdaDelta在SGD基础上增加的二阶动量。

把一阶动量和二阶动量都用起来就是Adam了——Adaptive + Momentum

SGD的一阶动量:

加上AdaDelta的二阶动量:

优化算法里最常见的两个超参数$\beta_1,\beta_2$就都在这里了,前者是控制一阶动量,后者控制二阶动量。

Nadam

都说Adam是集大成者,但它遗漏了Nesterov,安装NAG的步骤1:

Nesterov+Adam = Nadam

定义优化器:

1
CLASS torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)

参数:

  • params (iterable) – 优化器作用的模型参数。
  • lr (float) – learning rate – 相当于是统一框架中的 。
  • betas(Tuple[float,float],optional) – coefficients used for computing running averages of gradient and its square ((默认值:(0.9, 0.999))
  • weight_decay (float, optional) – 权重衰减系数 weight decay (L2 penalty) (默认值:0)
  • eps(float,optional):防止分母为0的一个小数 (默认值:1e-10)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import math
import torch
from .optimizer import Optimizer


[docs]class Adam(Optimizer):
r"""Implements Adam algorithm.

It has been proposed in `Adam: A Method for Stochastic Optimization`_.

Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
amsgrad (boolean, optional): whether to use the AMSGrad variant of this
algorithm from the paper `On the Convergence of Adam and Beyond`_
(default: False)

.. _Adam\: A Method for Stochastic Optimization:
https://arxiv.org/abs/1412.6980
.. _On the Convergence of Adam and Beyond:
https://openreview.net/forum?id=ryQu7f-RZ
"""

def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
weight_decay=0, amsgrad=False):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
if not 0.0 <= betas[1] < 1.0:
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
if not 0.0 <= weight_decay:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
defaults = dict(lr=lr, betas=betas, eps=eps,
weight_decay=weight_decay, amsgrad=amsgrad)
super(Adam, self).__init__(params, defaults)

def __setstate__(self, state):
super(Adam, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('amsgrad', False)

[docs] @torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step.

Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()

for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad
if grad.is_sparse:
raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
amsgrad = group['amsgrad']

state = self.state[p]

# State initialization
if len(state) == 0:
state['step'] = 0
# Exponential moving average of gradient values
state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
# Exponential moving average of squared gradient values
state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
if amsgrad:
# Maintains max of all exp. moving avg. of sq. grad. values
state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
if amsgrad:
max_exp_avg_sq = state['max_exp_avg_sq']
beta1, beta2 = group['betas']

state['step'] += 1
bias_correction1 = 1 - beta1 ** state['step']
bias_correction2 = 1 - beta2 ** state['step']

if group['weight_decay'] != 0:
grad = grad.add(p, alpha=group['weight_decay'])

# Decay the first and second moment running average coefficient
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
if amsgrad:
# Maintains the maximum of all 2nd moment running avg. till now
torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
# Use the max. for normalizing running avg. of gradient
denom = (max_exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
else:
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])

step_size = group['lr'] / bias_correction1

p.addcdiv_(exp_avg, denom, value=-step_size)

return loss

AdamW

Adam的另一个改进版AdamW

AdamW就是Adam优化器加上L2正则,来限制参数值不可太大。

以往的L2正则是直接加在损失函数上,比如这样:加入正则, 损失函数就会变成:

所以在计算梯度$g_t$时要加上粉色的这一项。

但AdamW稍有不同,如下图所示,将正则加在了绿色位置。

至于为何这么做?直接摘录BERT里面的原话看看:

Just adding the square of the weights to the loss function is not the correct way of using L2 regularization/weight decay with Adam, since that will interact with the m and v parameters in strange ways. Instead we want to decay the weights in a manner that doesn’t interact with the m/v parameters. This is equivalent to adding the square of the weights to the loss with plain (non-momentum) SGD. Add weight decay at the end (fixed version).

意思s 如果直接将L2正则加到loss上去,由于Adam优化器的后续操作,该正则项将会与$m_t$和$v_t$产生奇怪的作用。因而,AdamW选择将L2正则项加在了Adam的$m_t$和$v_t$等参数被计算完之后,在于学习率$\eta$相乘之前,所以这也表明了weight_decay和L2正则虽目的一致、公式一致,但用法还是不同,二者有明显的区别。

以 PyTorch1.7.0 中的AdamW代码为例:

定义优化器:

1
CLASS torch.optim.AdamW(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, amsgrad=False)
  • params (iterable) – 优化器作用的模型参数。
  • lr (float) – learning rate – 相当于是统一框架中的 。
  • betas(Tuple[float,float],optional) – coefficients used for computing running averages of gradient and its square ((默认值:(0.9, 0.999))
  • weight_decay (float, optional) – 权重衰减系数 weight decay (L2 penalty) (默认值:0)
  • eps(float,optional):防止分母为0的一个小数 (默认值:1e-10)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import math
import torch
from .optimizer import Optimizer


[docs]class AdamW(Optimizer):
r"""Implements AdamW algorithm.

The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.

Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
weight_decay (float, optional): weight decay coefficient (default: 1e-2)
amsgrad (boolean, optional): whether to use the AMSGrad variant of this
algorithm from the paper `On the Convergence of Adam and Beyond`_
(default: False)

.. _Adam\: A Method for Stochastic Optimization:
https://arxiv.org/abs/1412.6980
.. _Decoupled Weight Decay Regularization:
https://arxiv.org/abs/1711.05101
.. _On the Convergence of Adam and Beyond:
https://openreview.net/forum?id=ryQu7f-RZ
"""

def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
weight_decay=1e-2, amsgrad=False):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
if not 0.0 <= betas[1] < 1.0:
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
if not 0.0 <= weight_decay:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
defaults = dict(lr=lr, betas=betas, eps=eps,
weight_decay=weight_decay, amsgrad=amsgrad)
super(AdamW, self).__init__(params, defaults)

def __setstate__(self, state):
super(AdamW, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('amsgrad', False)

[docs] @torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step.

Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()

for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue

# Perform stepweight decay
p.mul_(1 - group['lr'] * group['weight_decay'])

# Perform optimization step
grad = p.grad
if grad.is_sparse:
raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
amsgrad = group['amsgrad']

state = self.state[p]

# State initialization
if len(state) == 0:
state['step'] = 0
# Exponential moving average of gradient values
state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
# Exponential moving average of squared gradient values
state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
if amsgrad:
# Maintains max of all exp. moving avg. of sq. grad. values
state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
if amsgrad:
max_exp_avg_sq = state['max_exp_avg_sq']
beta1, beta2 = group['betas']

state['step'] += 1
bias_correction1 = 1 - beta1 ** state['step']
bias_correction2 = 1 - beta2 ** state['step']

# Decay the first and second moment running average coefficient
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
if amsgrad:
# Maintains the maximum of all 2nd moment running avg. till now
torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
# Use the max. for normalizing running avg. of gradient
denom = (max_exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
else:
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])

step_size = group['lr'] / bias_correction1

p.addcdiv_(exp_avg, denom, value=-step_size)

return loss

与Adam不一样的地方是:

Adam如果使用weight_decay的话,那么相当于目标函数上加了$1/2\gamma||\theta||^2$,所以相当于是梯度加上$\gamma\theta$故Adam使用了

grad = grad.add(p, alpha=group[‘weight_decay’])

而 AdamW 是 p.mul_(1 - group[‘lr’] * group[‘weight_decay’]) 直接让参数:

$\thetat =\theta{t-1}-\alpha\cdot\lambda\cdot\theta_{t-1}-\alpha\cdot\eta_t$