24/11/3 算法笔记 Adam优化器拆解

Adam 优化器是一种用于深度学习中的自适应学习率优化算法，它结合了两种其他流行的优化方法的优点：RMSprop 和 Momentum。简单来说，Adam 优化器使用了以下方法：

**指数加权移动平均（Exponentially Weighted Moving Average, EWMA）**：

Adam 维护了梯度的一阶矩估计（均值）和二阶矩估计（方差）的指数加权移动平均。
一阶矩估计（\(m_t\)）是过去梯度的加权平均，用于估计参数的期望方向。
二阶矩估计（\(v_t\)）是过去梯度平方的加权平均，用于估计参数更新的步长。

**自适应学习率**：

Adam 根据每个参数的二阶矩估计来调整学习率，使得每个参数都有自己的学习率，这是通过缩放梯度来实现的。
这种自适应性使得 Adam 在处理不同参数时更加灵活，能够加速收敛并提高模型性能。

**Momentum（动量）**：

Adam 引入了类似于 Momentum 优化器的动量项，它帮助优化器在优化过程中保持方向的一致性，并减少震荡。
动量项是通过将梯度的一阶矩估计乘以一个衰减因子（beta1）来实现的。

**RMSprop**：

Adam 从 RMSprop 继承了根据梯度的二阶矩估计来调整每个参数的学习率的思想。
这是通过将梯度除以二阶矩估计的平方根（加上一个小的常数 \(\epsilon\) 以保证数值稳定性）来实现的。

**Bias Correction（偏差校正）**：

由于使用了指数加权移动平均，一阶和二阶矩估计会随着时间而衰减，这可能导致偏差。Adam 通过偏差校正来调整这些估计，以获得更好的长期性能。

**AMSGrad（可选）**：

AMSGrad 是 Adam 的一个变种，它解决了在一些情况下 Adam 可能不会收敛的问题。
AMSGrad 通过维护一个最大二阶矩估计来调整学习率，而不是使用普通的二阶矩估计。

这些我们会在后面详细讲解

我们先来看下adam的源代码

复制代码

import torch
from . import functional as F
from .optimizer import Optimizer

class Adam(Optimizer):
    r"""Implements Adam algorithm.
    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
    The implementation of the L2 penalty follows changes proposed in
    `Decoupled Weight Decay Regularization`_.
    Args:
        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
    .. _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=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)

    @torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step.
        Args:
            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 = []
            exp_avgs = []
            exp_avg_sqs = []
            state_sums = []
            max_exp_avg_sqs = []
            state_steps = []
            for p in group['params']:
                if p.grad is not None:
                    params_with_grad.append(p)
                    if p.grad.is_sparse:
                        raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
                    grads.append(p.grad)
            # ... (rest of the step function implementation)

逐段拆解一下：

类定义和初始化方法 `init`

复制代码

class Adam(Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0, amsgrad=False):

params：要优化的参数或定义参数组的字典。
lr：学习率，默认为 1e-3。
betas：用于计算梯度和其平方的运行平均值的系数，默认为 (0.9, 0.999)。
eps：用于提高数值稳定性的项，默认为 1e-8。
weight_decay：权重衰减（L2 惩罚），默认为 0。
amsgrad：是否使用 AMSGrad 变体，默认为 False。

关于AMSGrad我们后续会讲解

参数验证

复制代码

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

参数验证确保所有传入的参数都在合理的范围内。例如，学习率（lr）应该是非负的，因为负的学习率没有意义，可能会导致优化过程中出现问题。
betas 参数应该在 0 和 1 之间，因为它们是用于计算梯度和梯度平方的指数移动平均的系数，超出这个范围的值可能会导致算法不稳定或不收敛。
eps（epsilon）值也应该是非负的，因为它用于提高数值稳定性，防止分母为零的情况发生。

默认参数字典

复制代码

        defaults = dict(lr=lr, betas=betas, eps=eps,
                        weight_decay=weight_decay, amsgrad=amsgrad)
        super(Adam, self).__init__(params, defaults)

状态设置方法 `setstate`

复制代码

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

在反序列化优化器时设置状态，并确保所有参数组都有 amsgrad 键，是为了在优化器被保存和加载后，能够保持 AMSGrad 变体的配置状态。这意味着，如果优化器在保存之前使用了 AMSGrad 变体，那么在加载优化器后，这个配置仍然会被保留，确保优化过程的连续性和一致性。

优化步骤方法 `step`

复制代码

    @torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step.
        Args:
            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()

step 方法执行单步优化。如果提供了 closure，则在启用梯度的情况下调用它，并获取损失值。

参数和梯度的准备

复制代码

        for group in self.param_groups:
            params_with_grad = []
            grads = []
            exp_avgs = []
            exp_avg_sqs = []
            state_sums = []
            max_exp_avg_sqs = []
            state_steps = []
            for p in group['params']:
                if p.grad is not None:
                    params_with_grad.append(p)
                    if p.grad.is_sparse:
                        raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
                    grads.append(p.grad)

            if p.grad.is_sparse:
                raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')

这行代码检查参数的梯度是否是稀疏的。Adam 优化器不支持稀疏梯度，如果发现梯度是稀疏的，会抛出一个运行时错误，建议使用 SparseAdam 优化器。

这段代码的主要目的是为每个参数组中的参数准备必要的数据，以便在后续步骤中进行参数更新。

完整的 `step` 方法实现

复制代码

                    exp_avgs.append(self.state[p]['exp_avg'])
                    exp_avg_sqs.append(self.state[p]['exp_avg_sq'])
                    state_sums.append(self.state[p]['step'])
                    if group['amsgrad']:
                        max_exp_avg_sqs.append(self.state[p]['max_exp_avg_sq'])
                    state_steps.append(self.state[p]['step'])

            for i, p in enumerate(params_with_grad):
                state = self.state[p]
                exp_avg, exp_avg_sq = exp_avgs[i], exp_avg_sqs[i]
                if group['amsgrad']:
                    max_exp_avg_sq = max_exp_avg_sqs[i]
                beta1, beta2 = group['betas']

                state['step'] = state['step'] + 1

                exp_avg.mul_(beta1).add_(grads[i], alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grads[i], grads[i], value=1 - beta2)
                if group['amsgrad']:
                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
                    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
                else:
                    denom = exp_avg_sq.sqrt().add_(group['eps'])

                step_size = group['lr']
                if group['weight_decay'] != 0:
                    step_size = step_size * (1 - beta2 * group['weight_decay'])

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

这是优化的核心，咱逐段讲解

准备状态数据

复制代码

exp_avgs.append(self.state[p]['exp_avg'])
exp_avg_sqs.append(self.state[p]['exp_avg_sq'])
state_sums.append(self.state[p]['step'])
if group['amsgrad']:
    max_exp_avg_sqs.append(self.state[p]['max_exp_avg_sq'])
state_steps.append(self.state[p]['step'])

这些行代码将每个参数的当前状态（一阶矩估计、二阶矩估计、步数）添加到之前初始化的列表中。如果启用了 AMSGrad，则还会添加最大二阶矩估计。

遍历有梯度的参数

复制代码

for i, p in enumerate(params_with_grad):

这行代码遍历有梯度的参数，并使用 enumerate 来同时获取参数的索引 i 和参数对象 p。

获取状态和超参数

复制代码

state = self.state[p]
exp_avg, exp_avg_sq = exp_avgs[i], exp_avg_sqs[i]
if group['amsgrad']:
    max_exp_avg_sq = max_exp_avg_sqs[i]
beta1, beta2 = group['betas']

这些行代码获取当前参数的状态，并从中提取一阶矩估计和二阶矩估计。如果启用了 AMSGrad，则还会获取最大二阶矩估计。同时，获取优化器的超参数 beta1 和 beta2。

AMSGrad是优化器的一个变种

更新步数

复制代码

state['step'] = state['step'] + 1

更新一阶矩估计和二阶矩估计

复制代码

exp_avg.mul_(beta1).add_(grads[i], alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(grads[i], grads[i], value=1 - beta2)

这两行代码更新一阶矩估计和二阶矩估计。mul_ 方法将当前估计乘以 beta1 或 beta2，add_ 方法将梯度（或梯度平方）乘以 (1 - beta1) 或 (1 - beta2) 后加到当前估计上。

AMSGrad 调整

复制代码

if group['amsgrad']:
    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
else:
    denom = exp_avg_sq.sqrt().add_(group['eps'])

调整学习率和更新参数

复制代码

step_size = group['lr']
if group['weight_decay'] != 0:
    step_size = step_size * (1 - beta2 * group['weight_decay'])
p.addcdiv_(exp_avg, denom, value=-step_size)

addcdiv_ 是一个 in-place 操作，它结合了加法、除法和乘法。这个操作通常用于执行参数更新，特别是在优化器中。

这个操作将计算：

参数 ← 参数 − 学习率 * 一阶矩估计/调整后的分母

更新规则

在每次迭代中，Adam 使用以下规则更新参数：

其中，θ是参数，mt 是一阶矩估计，vt 是二阶矩估计，ϵ是一个小常数，用于防止分母为零。

接下来我们来讲解AMSGrad

AMSGrad 是 Adam 的一个变种，它在更新规则中使用二阶矩估计的最大值，以确保即使梯度的方差减小，更新步长也不会变得过大。

二阶矩估计的概念

在优化算法的上下文中，二阶矩估计通常是指过去梯度平方的指数加权移动平均（Exponentially Weighted Moving Average, EWMA）。这个概念可以分解为以下几个部分：

梯度平方：首先计算每次迭代中参数的梯度，然后将这些梯度平方。
指数加权移动平均：对这些平方梯度值计算指数加权移动平均。这意味着给过去的平方梯度值赋予递减的权重，最近的梯度平方有更高的权重。
自适应学习率：二阶矩估计的结果用于调整每个参数的学习率。在 Adam 优化器中，每个参数的学习率会根据其二阶矩估计的平方根进行调整。

更新规则如下：

指数加权移动平均：

指数加权移动平均（Exponentially Weighted Moving Average，简称 EWMA）是一种用于估计数据集统计特性（如均值、方差）的时间序列方法。它给予最近的观测值更高的权重，而较早的观测值权重逐渐减小，权重按照指数衰减。

原理

EWMA 的核心思想是将最新的观测值与之前的估计值结合起来，以产生新的估计。这种方法特别适用于需要平滑数据或预测未来值的场景。

数学表达

参数解释

α（平滑参数）：这个参数控制了新数据对平均值的影响程度。α 值越大，新数据对平均值的影响越大，平滑效果越弱；α 值越小，平滑效果越强，但对新数据的反应越慢。

AMSGrad 与 Adam 的区别

最大二阶矩估计 ：AMSGrad 引入了最大二阶矩估计（max_exp_avg_sq），这是对 Adam 的改进。在 Adam 中，分母是直接使用二阶矩估计的平方根，而在 AMSGrad 中，分母使用最大二阶矩估计的平方根，这有助于解决 Adam 在某些情况下可能不会收敛的问题。

偏差矫正机制：

偏差的来源：
- 在 Adam 优化器中，一阶矩（mt）和二阶矩（vt）估计是基于指数加权移动平均（EWMA）计算的。由于 EWMA 的初始值为零，这会导致在优化过程的早期，这些估计值不能很好地反映梯度的真实分布，从而产生偏差。
偏差校正的方法：
- 为了解决这个问题，Adam 引入了偏差校正机制。具体来说，对于一阶矩和二阶矩的估计值，分别计算偏差校正因子，并用这些因子来调整估计值。
- 一阶矩的偏差校正因子为，二阶矩的偏差校正因子为，其中 t 是当前的迭代次数，β1和 β2是用于计算一阶和二阶矩估计的衰减率参数。
偏差校正的应用：
- 在每次迭代中，使用偏差校正因子对一阶矩和二阶矩的估计值进行调整，得到校正后的一阶矩（m^t）和二阶矩（v^t）：
- 然后，使用这些校正后的估计值来计算参数的更新步长。
偏差校正的目的：
- 偏差校正机制的目的是为了减少由于初始估计值偏差带来的影响，使得优化器在早期就能够以一个更加合理的学习率进行参数更新，从而提高模型的训练稳定性和收敛速度。
实际效果：
- 通过引入偏差校正机制，Adam 优化器能够在训练初期更准确地估计梯度的分布，避免了由于初始偏差导致的学习率过高或过低的问题，有助于模型更快地收敛到最优解。

24/11/3 算法笔记 Adam优化器拆解

类定义和初始化方法 __init__

参数验证

默认参数字典

状态设置方法 __setstate__

优化步骤方法 step

参数和梯度的准备

完整的 step 方法实现

准备状态数据

遍历有梯度的参数

获取状态和超参数

更新步数

更新一阶矩估计和二阶矩估计

AMSGrad 调整

调整学习率和更新参数

二阶矩估计的概念

原理

数学表达

参数解释

AMSGrad 与 Adam 的区别

类定义和初始化方法 `init`

状态设置方法 `setstate`

优化步骤方法 `step`

完整的 `step` 方法实现