Adam梯度下降优化实战：解决深度学习训练中的震荡与收敛问题

1次阅读

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

在深度学习模型训练过程中，优化算法的选择直接影响模型的收敛速度和最终性能。Adam（Adaptive Moment Estimation）作为自适应学习率优化算法的代表，因其优异的性能成为许多开发者的首选。然而，实际应用中我们常常会遇到训练过程震荡、收敛不稳定等问题。本文将深入剖析 Adam 的工作原理，分享实战中的调参技巧，并通过代码示例展示如何稳定训练过程。

在深度学习任务中，我们通常需要优化高维非凸函数。传统的随机梯度下降（SGD）虽然简单，但在面对这类问题时存在明显不足：

对所有参数使用相同的学习率，无法适应不同参数的重要性差异
在损失函数曲面平坦区域进展缓慢，而在陡峭区域容易震荡
需要手动调整学习率，对超参数选择敏感

Adam 算法结合了动量法（Momentum）和 RMSprop 的优点，通过计算梯度的一阶矩估计（均值）和二阶矩估计（未中心化的方差）来为不同参数设计独立的自适应学习率。其核心公式如下：

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

其中，$m_t$ 和 $v_t$ 分别是一阶和二阶矩估计，$\beta_1$ 和 $\beta_2$ 是衰减率超参数，$\epsilon$ 是为数值稳定性添加的小常数。

以下是 PyTorch 中实现 AdamW 优化器并加入学习率 warmup 的代码示例。AdamW 是 Adam 的变种，对权重衰减进行了更合理的处理：

import torch
import math

class AdamWarmup(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0.01, warmup_steps=4000):
        defaults = dict(lr=lr, betas=betas, eps=eps,
                        weight_decay=weight_decay, warmup_steps=warmup_steps)
        super().__init__(params, defaults)

    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 = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError('Adam does not support sparse gradients')

                state = self.state[p]

                # 初始化状态
                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

                # 衰减一阶和二阶矩估计
                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().add_(group['eps'])  # 添加 epsilon 防止除零

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

                p.data.addcdiv_(exp_avg, denom, value=-lr)
                # 应用权重衰减（AdamW 的方式）p.data.mul_(1 - lr * group['weight_decay'])

        return loss

实践中发现，学习率 (η) 和 batch size(B)存在以下经验关系：

$\eta = \eta_{base} \times \frac{B}{B_{base}}$

其中 $\eta_{base}$ 和 $B_{base}$ 是基准学习率和 batch size。当 batch size 增大时，应相应增大学习率。

梯度爆炸：当 $\beta_2$ 设置过小（如 <0.9）时，二阶矩估计可能无法及时更新，导致梯度爆炸。建议保持 $\beta_2$ 在 0.9-0.999 之间。
学习率震荡：在训练初期，由于二阶矩估计 $v_t$ 较小，可能导致学习率 $\eta/\sqrt{v_t}$ 过大。解决方案包括：
使用学习率 warmup
提高 $\epsilon$ 值（如从 1e- 8 提高到 1e-6）
权重衰减处理：原始 Adam 中的权重衰减实现存在问题，推荐使用 AdamW 变种。