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

1次阅读

没有评论

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

在神经网络训练中，随机梯度下降（SGD）长期以来是默认的优化器选择。然而，SGD 在处理非凸优化问题时存在明显不足：

固定学习率问题 ：所有参数使用相同的学习率，难以适应不同特征的更新需求
梯度震荡 ：在损失函数曲面陡峭区域容易产生剧烈震荡
收敛速度慢 ：需要精心调整学习率调度策略才能达到较好效果

这些痛点催生了自适应优化器的发展，其中 Adam（Adaptive Moment Estimation）因其卓越的性能成为当前最流行的选择之一。

Adam 的核心思想是结合动量（Momentum）和 RMSProp 的优点，通过计算梯度的一阶矩估计和二阶矩估计来实现参数的自适应更新。其算法流程如下：

计算当前时间步的梯度：
$$g_t = \nabla_\theta f_t(\theta_{t-1})$$
更新有偏一阶矩估计（动量项）：
$$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 = \theta_{t-1} – \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

其中 $\beta_1$ 和 $\beta_2$ 是衰减率超参数，控制着历史信息的保留程度。$\epsilon$ 是为数值稳定性添加的小常数。

下面我们实现一个完整的 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):
        """Performs a single optimization step"""
        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 = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError('Adam does not support sparse gradients')

                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(p, max_norm=1.0)

                state = self.state[p]

                # Initialize state
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p.data)
                    state['exp_avg_sq'] = torch.zeros_like(p.data)

                exp_avg, exp_avg_sq = state['exp_avg'], state['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)

                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])

                # Learning rate warmup
                warmup_factor = min(state['step'] / self.warmup_steps, 1.0)
                step_size = group['lr'] * warmup_factor / bias_correction1

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

                # Weight decay
                if group['weight_decay'] > 0:
                    p.data.add_(p.data, alpha=-group['lr'] * group['weight_decay'])

        self.step_count += 1
        return loss

在 CIFAR-10 数据集上使用 ResNet-18 架构进行测试，我们得到以下训练曲线：

图 1：不同优化器的训练损失对比