Adam梯度下降优化算法：从数学原理到PyTorch实战

1次阅读

没有评论

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

在深度学习模型训练中，传统的 随机梯度下降（SGD）存在两个主要痛点：

学习率敏感性：固定学习率在平坦区域下降过慢，在陡峭区域又容易震荡
稀疏梯度问题：某些特征更新频率极低时（如 NLP 中的罕见词），SGD 难以有效跟踪历史信息

比如在 ResNet 训练中，我们常看到这样的现象：

学习率设大了→损失函数剧烈震荡
学习率设小了→训练进度像蜗牛爬坡

Adam 的本质是 动量 + 自适应学习率，其更新规则分为四个关键步骤：

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

由于初始时刻 $m_0=v_0=0$ 会导致估计偏小：
$$\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}$$

下面我们实现一个带 warmup 和 gradient clipping 的 Adam：

import torch
import math

class CustomAdam(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0, 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

                # 梯度裁剪
                torch.nn.utils.clip_grad_norm_(p, max_norm=1.0)

                grad = p.grad.data
                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)

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

                # 更新一阶和二阶矩
                m.mul_(beta1).add_(grad, alpha=1-beta1)
                v.mul_(beta2).add_(grad.pow(2), alpha=1-beta2)

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

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

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

        return loss

在 MNIST 数据集上比较三种优化器的表现：

# 训练循环示例
for epoch in range(epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

SGD：收敛稳定但速度慢
RMSprop：初期快速但后期波动
Adam：兼具快速收敛和稳定性

参数	推荐值	作用说明
β₁	0.9	控制动量衰减率
β₂	0.999	控制自适应项衰减率
ε	1e-8	数值稳定性常数
batch_size	32-512	越大学习率可设越高
warmup	总 step 的 1%-5%	避免初期数值不稳定