深度学习优化器分析

---

前言

文章目标：
1、理解优化器的在训练中的作用
2、以adam优化器为例分析，优化器功能的优劣
3、对adam优化器展开内存分析
4、分析adam优化器的pytorch实现

---

一、优化器是什么？

一言以蔽之，优化器指导梯度下降方向与步长，目标是调整的权重更快、更准确的使目标函数达到最优值（最大值或在最小值）。

以一个最小化的典型训练脚步为例：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的线性模型
class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        # 初始化权重和偏置，这里使用matmul作为线性变换
        self.weight = nn.Parameter(torch.randn(4, 4))
        self.bias = nn.Parameter(torch.randn(4))

    def forward(self, x):
        # 使用matmul进行前向传播
        return torch.matmul(x, self.weight) + self.bias

# 实例化模型
model = LinearModel()

# 定义损失函数，这里使用均方误差
criterion = nn.MSELoss()

# 定义优化器，这里使用Adam
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 准备一些假数据进行演示，这里使用均匀分布的随机数
inputs = torch.randn(4, 4)
targets = torch.randn(4)

# 训练循环
num_epochs = 10
for epoch in range(num_epochs):
    # 前向传播
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    # 清零梯度
    optimizer.zero_grad()
    # 反向传播
    loss.backward()
    # 更新权重
    optimizer.step()
    # 打印损失信息
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 训练结束
print("Training complete")

训练过程大致分为以下步骤循环：
1、前向传播
2、计算前向传播的误差（loss）
3、反向传播（以loss值为指导，计算梯度）
4、优化器通过梯度更新权重

一般我们使用梯度下降法更新参数：
因此一个最基础的权重更新公式，其中r为学习率：
W(t+1) = W(t) - ε * g(t)
公式中唯一与优化器相关的参数就是学习率ε,一个标量。实际上述公式表达的是随机梯度下降法（Stochastic Gradient Descent，SGD）。随机指的是一次梯度更新使用的是一个batch的数据集，而非完整的数据集。

二、优化器算法的演进

优化器算法的演进围绕着学习率ε与梯度值g(t)两个方面不断改进。

1.SGD算法—梯度值方向的改进

在使用SGD算法时，我们有时会遇到震荡的问题（每次迭代后梯度变化过于剧烈），导致模型收敛过慢。因此引入动量（Momentum）的概念，动量是数学领域的概念，在物理领域的动量可以被理解为惯性。

在优化器中引入动量的好处是，可以抵消梯度中那些变化剧烈的分量，从而减小振荡、加快收敛速度。

• 未引入动量时，权重更新过程中振荡的向极值逼近
  * 引入动量后，向极值逼近的振荡减小
  加入动量后的SGD算法的权重更新公式(SGD-M)为：
  W(t+1) = W(t) + V(t+1)
  其中V(t+1) = μV(t ) - ε * g(t+1)表示携带动量后需要更新的梯度； 该方法可以减小更早时刻梯度对当前梯度的影响，μ参数通常取值为0.9。
  动量的存在带来两个影响，（1）振荡的权重更新被平滑了；（2）非振荡的权重更新的被加速了。

[拓展–继续演进]

但是第二点影响有个副作用：由于动量的存在，权重的更新首前一个step的gard的影响，导致它在最后的收敛阶段不停震荡，从而浪费很多时间。

针对这个问题，NAG(Nesterov’s Accelerated Gradient)算法被提出的，它的基本思想是在梯度更新之前往前看一步，让优化器可以预知未来的情况，从而对现在的梯度进行调整。
NAG算法中的
W(t+1) = W(t) + V(t + 1)
V(t + 1) = μV(t) - ε * g (W(t) + μ * V(t) )
根据当前梯度方向向前走一步之后再计算出来的梯度, 再用当前梯度进行修正。通过这种方式，在接近最优解时，Nesterov优化器比标准动量SGD算法有更快的收敛速度。因为这种“前瞻性”的操作能帮助优化器避免走得太远，就像是提前刹车，所以在接近最优解时更加稳定。On the importance of initialization and momentum in deep learning

动量法首先计算当前的梯度值（图中的小的蓝色向量），然后在更新的累积梯度（大的蓝色向量）方向上前进一大步，Nesterov加速梯度下降法NAG首先在先前累积梯度（棕色的向量）方向上前进一大步，计算梯度值，然后做一个修正（绿色的向量）。

pytorch源码实现：注意这里实现的是简化版本的NAG 算法，有助于降低计算量，推导过程参考

2.AdaGrad – 优化学习率方向

学习率代表了优化的步长，如果设定过大，可能会导致参数在最优点附近震荡，如果设置过小，则会导致模型收敛速度过慢。所以通常情况下我们会选择一种让学习率逐步减小的策略，例如余弦退火算法、StepLr策略等等。

但是即使是这样，我们还有一个问题没有解决，那就是不同参数应该有不同的学习率。

那是不是可以让模型自动地调整学习率呢？
AdaGrad方法就可以实现这个愿望。

AdaGrad的公式为：
$W(t+1)=W(t)-\frac{lr}{\sqrt{S(t)}+\epsilon }g(t)$

其中，$S(t)=S(t-1)+g^2(t)$代表着历史中所有梯度的平方和；ε是一个极小量，防止除零， lr是学习率本身。

这样做的好处是可以减小之前震荡过大的参数的学习率，增大更新速率较慢的参数的学习率，从而让模型可以更快收敛。

[拓展：继续演进]

AdaGrad算法存在一个明显的问题，就是它要考虑所有的历史数据，这样就会导致越来越大，学习率的分母变大了，也就会导致学习率越来越小，模型收敛的速度也就会变慢了。

解决这个问题的方法还是采用指数加权移动平均法，减小更早的梯度对当前学习率的影响，而AdaGrad+指数加权移动平均法就是RMSProp方法了，其公式为：

$W(t+1)=W(t)-\frac{lr}{\sqrt{S(t)}+\epsilon }g(t)$

其中，$S(t)=\alpha S(t-1)+(1-\alpha )g^2(t)$代表着历史中所有梯度的平方和， α代表指数加权移动平均法的参数。

3、Adam算法(Adaptive Moment Estimation 自适应矩估计)–同时优化梯度与学习率

SGD-M + AdaGrad 结合起来的方法就是Adam算法论文，公式如下
$W(t+1)=W(t)-\frac{lr}{\sqrt{S(t)}+\epsilon }V(t)$

其中：$V(t)=\beta 1\ast V(t-1)+(1-\beta 1)\ast g(t)$, $S(t)=\beta 2\ast S(t-1)+(1-\beta 2)\ast g^2(t)$, 将其带入权重更新公式为：

$W(t+1)=W(t)-\frac{lr}{\sqrt{\beta 2\ast S(t-1)+(1-\beta 2)\ast g^2(t)}+\epsilon }\ast (\beta 1\ast V(t-1)+(1-\beta 1)\ast g(t))$

注意：这里公式中的表示与前文表示一致，V(t)是与梯度g(t)呈现一阶相关性的变量， S（t）是与g(t)呈现二阶相关性的变量。
但是这种描述与论文原文描述不一致，为了方便看原文算法，将符号标识进行一次转换W->θ（权重）、V->m(一阶矩)、S->V(二阶矩)，lr->α(学习率)

变化后的公式如下：

$\theta (t+1)=\theta (t)-\frac{\alpha }{\sqrt{\beta 2\ast V(t-1)+(1-\beta 2)\ast g^2(t)}+\epsilon }\ast (\beta 1\ast m(t-1)+(1-\beta 1)\ast g(t))$

m(t)实现梯度方向的控制，受超参β 1控制，属于一阶矩估计，也有称一阶动量（数学：均值，物理：前后方向的惯性）
V(t)实现动态学习率，受超参β 2控制，属于二阶矩估计，也有称二阶动量(数学：方差、物理：旋转方向的惯性，转动惯量 )

论文中给出权重的更新算法

pytorch实现算法

torch实际的代码

def _single_tensor_adam(
	...
):
	...
    for i, param in enumerate(params):
        grad = grads[i] if not maximize else -grads[i]
        exp_avg = exp_avgs[i]
        exp_avg_sq = exp_avg_sqs[i]
        step_t = state_steps[i]

        # update step
        step_t += 1

        if weight_decay != 0: # 惩罚项： 防止过拟合
            grad = grad.add(param, alpha=weight_decay)

		...

        # Decay the first and second moment running average coefficient
        exp_avg.lerp_(grad, 1 - beta1) # m(t) = beta1 * m(t - 1) + (1 - beta1) * grad
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2) # v(t) = beta2 * v(t - 1) + (1 - beta2) * grad^2

        if capturable or differentiable:
            ...
        else:
            step = _get_value(step_t)

            bias_correction1 = 1 - beta1**step # (1 - beta1^t)
            bias_correction2 = 1 - beta2**step # (1 - beta2^t)

            step_size = lr / bias_correction1  # 动态步长

            bias_correction2_sqrt = bias_correction2**0.5 # 开根号

            if amsgrad: # 类似RMSProp方法
                # Maintains the maximum of all 2nd moment running avg. till now
                torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])

                # Use the max. for normalizing running avg. of gradient
                denom = (max_exp_avg_sqs[i].sqrt() / bias_correction2_sqrt).add_(eps)
            else:
                denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps) # denom = v(t)^0.5 / (1 - beta2^t)^0.5 + eps

            param.addcdiv_(exp_avg, denom, value=-step_size) # param = param - step_size * m(t) / denom

		...

adam算法的内存占用分析

deepspeed在ZeRo论文中如下分析
模型状态（model states）: 模型参数（fp16）、模型梯度（fp16）和Adam状态（fp32的模型参数备份，fp32的momentum和fp32的variance）。
那么寻找下Adam状态占用的内存：
首先在adam优化器代码中我们很轻易的找到momentum和variance， dtype确实是float32。

 exp_avg = exp_avgs[i]
 exp_avg_sq = exp_avg_sqs[i]

fp32的模型参数备份在哪？
经过一番查找发现混合精度下的fp32模型参数备份内存并非是优化器持有，而是由apex或者amp产生。apex的例子。

[拓展：继续演进]

adam算法中遗漏了NAG的算法，因此便有集成Nesterov 的Nadam算法，重点在于梯度更新的公式
$g(t)=g(w(t)-\alpha \ast m(t-1)/\sqrt{V(t)})$

---

总结

本文以pytorch视角总结了优化器的算法，以学习率与梯度更新方向两点为优化方向，介绍了主流优化算法的演进，最后以adam算法为例打开pytorch的优化器实现，并尝试分析优化器带来的内存占用。