Adam梯度下降原理详解：从数学推导到PyTorch实战

1次阅读

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

梯度下降是深度学习模型训练的基石，它通过迭代调整参数来最小化损失函数。传统 SGD 面临学习率难以选择、收敛速度慢的问题，而 Adam 通过结合动量法和自适应学习率，成为目前最受欢迎的优化器之一。本文将带你从数学原理到代码实现，全面掌握 Adam 的工作机制。

Adam 首先继承了动量法的思想，通过指数加权平均积累历史梯度：
$$ m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t $$
其中 $g_t$ 是当前梯度，$\beta_1$ 通常取 0.9。这相当于给梯度增加了惯性，在持续方向上加速。

同时引入 RMSprop 对梯度幅度的自适应调整：
$$ v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2 $$
$\beta_2$ 通常取 0.999，$v_t$ 反映了梯度平方的指数衰减平均。

由于初始时刻 $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 \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} $$
其中 $\eta$ 是基础学习率，$\epsilon$（默认 1e-8）防止除零。

import torch
import torch.nn as nn
from torchvision import datasets, transforms

# 定义简化版 Adam 实现
class MyAdam:
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):
        self.params = list(params)
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.m = [torch.zeros_like(p) for p in self.params]
        self.v = [torch.zeros_like(p) for p in self.params]
        self.t = 0

    def step(self):
        self.t += 1
        for i, p in enumerate(self.params):
            if p.grad is None:
                continue
            # 更新一阶矩估计
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * p.grad
            # 更新二阶矩估计
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * p.grad**2
            # 偏差修正
            m_hat = self.m[i] / (1 - self.beta1**self.t)
            v_hat = self.v[i] / (1 - self.beta2**self.t)
            # 参数更新
            p.data -= self.lr * m_hat / (torch.sqrt(v_hat) + self.eps)

# MNIST 训练示例
def train_mnist():
    transform = transforms.Compose([transforms.ToTensor(), 
                                   transforms.Normalize((0.1307,), (0.3081,))])
    train_set = datasets.MNIST("./data", train=True, download=True, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

    model = nn.Sequential(nn.Linear(784, 128),
        nn.ReLU(),
        nn.Linear(128, 10)
    )

    # 比较原生 Adam 和我们的实现
    optimizer1 = torch.optim.Adam(model.parameters(), lr=0.001)
    optimizer2 = MyAdam(model.parameters(), lr=0.001)

    for epoch in range(5):
        for batch_idx, (data, target) in enumerate(train_loader):
            data = data.view(data.size(0), -1)
            output = model(data)
            loss = nn.functional.cross_entropy(output, target)

            optimizer1.zero_grad()
            loss.backward()
            optimizer1.step()

在 MNIST 数据集上测试不同学习率效果：

lr=0.01：前几轮快速下降，后期可能震荡
lr=0.001（推荐初始值）：稳定收敛
lr=0.0001：收敛速度过慢

$\beta_1$ 越大（接近 1），动量效应越强，适合平稳梯度方向
$\beta_2$ 越大，对梯度幅度变化越不敏感
推荐默认值 $(\beta_1=0.9, \beta_2=0.999)$ 在大多数情况表现良好

稀疏梯度场景 ：
当梯度非常稀疏时（如 NLP 任务），可以适当增大 $\epsilon$
考虑使用 Adam 变体如 AdamW（解耦权重衰减）

早停策略 ：

best_loss = float("inf")
patience = 3
counter = 0

for epoch in range(100):
    train_loss = train_one_epoch()
    val_loss = validate()

    if val_loss < best_loss:
        best_loss = val_loss
        counter = 0
        torch.save(model.state_dict(), "best_model.pt")
    else:
        counter += 1
        if counter >= patience:
            print("Early stopping")
            break