Adam优化器在神经网络训练中的原理与实践：从数学推导到PyTorch实现

1次阅读

没有评论

共计 2715 个字符，预计需要花费 7 分钟才能阅读完成。

传统随机梯度下降（SGD）虽然简单直接，但在处理非凸优化问题时存在明显局限性：

所有参数使用相同的学习率，无法适应不同特征的重要性差异
在损失函数曲面存在局部极小值或鞍点时容易陷入停滞
需要手动调整学习率衰减策略，增加了调参成本

这些问题在深层神经网络中尤为突出。2015 年 Kingma 和 Ba 提出的 Adam 优化器，通过结合动量（Momentum）和 RMSProp 的优点，实现了自适应学习率调整，显著提升了训练效率和模型性能。

Adam 的核心思想是维护两个移动平均量：梯度的一阶矩估计（动量）和二阶矩估计（未中心化的方差）。更新规则可分解为以下步骤：

计算梯度：
$$ g_t = \nabla_\theta J(\theta_t) $$
更新一阶矩估计（类似动量）：
$$ m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t $$
更新二阶矩估计（自适应学习率）：
$$ v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2 $$
偏差校正（应对初始零偏置）：
$$ \hat{m}_t = \frac{m_t}{1-\beta_1^t} $$
$$ \hat{v}_t = \frac{v_t}{1-\beta_2^t} $$
参数更新：
$$ \theta_{t+1} = \theta_t – \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t $$

其中 $\beta_1$ 控制动量衰减率（通常 0.9），$\beta_2$ 控制方差衰减率（通常 0.999），$\epsilon$ 是数值稳定性常数（1e-8）。

以下是自定义 Adam 优化器的完整实现，包含学习率 warmup 和梯度裁剪：

import torch
from torch.optim import Optimizer

class CustomAdam(Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0, amsgrad=False, warmup_steps=4000):
        defaults = dict(lr=lr, betas=betas, eps=eps,
                        weight_decay=weight_decay, amsgrad=amsgrad)
        super().__init__(params, defaults)
        self.warmup_steps = warmup_steps
        self.step_count = 0

    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()

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

                # 梯度裁剪（防止梯度爆炸）grad = torch.nn.utils.clip_grad_norm_(p, max_norm=1.0)

                # 参数状态初始化
                state = self.state[p]
                if len(state) == 0:
                    state['step'] = 0
                    state['m'] = torch.zeros_like(p.data)
                    state['v'] = torch.zeros_like(p.data)
                    if group['amsgrad']:
                        state['v_max'] = torch.zeros_like(p.data)

                m, v = state['m'], state['v']
                beta1, beta2 = group['betas']
                state['step'] += 1

                # 学习率 warmup
                lr = group['lr'] * min(state['step'] ** (-0.5), 
                                     state['step'] * self.warmup_steps ** (-1.5))

                # 核心更新逻辑
                m.mul_(beta1).add_(1 - beta1, grad)
                v.mul_(beta2).addcmul_(1 - beta2, grad, grad)

                # 偏差校正
                m_hat = m / (1 - beta1 ** state['step'])
                v_hat = v / (1 - beta2 ** state['step'])

                # 参数更新
                denom = v_hat.sqrt().add_(group['eps'])
                p.data.addcdiv_(-lr, m_hat, denom)

                # 权重衰减
                if group['weight_decay'] > 0:
                    p.data.add_(-group['lr'] * group['weight_decay'], p.data)

        self.step_count += 1
        return loss

关键实现细节：