优化器学习小结

2025-9-10

优化器的python实现

基本原理

基本构建类：torch.optim.Optimizer

# SGD 示例
optimizer_sgd = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Adam 示例
optimizer_adam = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))

zero_grad() 方法 任何 PyTorch 训练循环中的一个关键步骤是调用 optimizer.zero_grad()。在 PyTorch 中，反向传播 (loss.backward()) 期间计算的梯度默认是累积的；它们被加到每个参数的 .grad 属性上，而不是被覆盖。这种设计选择对于训练循环神经网络 (RNN) 或为模拟大批量训练实现梯度累积等场景非常有用。然而，对于标准训练，在计算当前迭代的梯度之前，必须清除上一次迭代的梯度。zero_grad() 方法通过遍历所有注册的参数并将其 .grad 属性设置为 None 或清零来完成此任务。因此，标准的训练工作流程是三个步骤的重复序列：

optimizer.zero_grad()：重置上一步的梯度。
loss.backward()：为当前的小批量数据计算梯度。
optimizer.step()：使用计算出的梯度更新模型的参数。

step() 方法 step() 方法是优化器的核心，特定算法的数学更新规则在此执行。调用此方法时，它会遍历每个参数组，并对组内的每个参数执行优化步骤。它访问参数计算出的梯度（存储在 p.grad 中），并使用它以及任何存储的状态（如动量缓冲）和超参数（如学习率）来更新参数的值 (p.data)。这正是优化理论的抽象数学公式被转化为张量上的具体计算操作的地方。

本质上，优化器的不同在于step方法中，对于每个张量的grad属性取值的更新方式。

状态管理

现代优化器是有状态的，这意味着它们在训练迭代中维护信息以指导未来的更新。这种状态管理以及动态控制超参数的能力是通过一系列复杂的特性来处理的。

参数组 param_groups 是高级训练方案的一个关键特性。传递给优化器构造函数的参数列表在内部存储为字典列表，每个字典代表一个参数组。这种结构允许为模型的不同部分设置不同的超参数。一个在迁移学习中的常见用例是为大型预训练主干模型的参数应用较小的学习率，而为新初始化的分类器头使用较大的学习率，从而防止预训练的特征在训练早期被破坏性地改变。同样，可以从权重衰减中排除某些参数，如偏置项。

状态字典 (state_dict) 优化器维护着在训练过程中演变的状态。对于带动量的 SGD，这个状态是每个参数的“速度”或“动量缓冲”。对于 Adam，它包括一阶和二阶矩估计（m 和 v）。这个状态对算法的功能至关重要，必须被保存才能正确地恢复训练。state_dict() 方法返回一个包含此信息的字典，该字典通常将内部参数 ID 映射到其对应的状态张量。相反，load_state_dict() 允许将此状态加载回优化器，确保训练过程的无缝继续。

学习率调度器 (lr_scheduler) 学习率可以说是最重要的超参数，并且在整个训练过程中很少保持不变。学习率调度器是根据预定计划调整学习率的实用工具。常见的策略包括 StepLR（在指定周期按因子衰减学习率）和 CosineAnnealingLR（遵循余弦曲线平滑地降低学习率）。这些调度器通过在每个步骤或周期修改其 param_groups 内的 lr 值与优化器交互，为控制收敛动态提供了强大的机制。

Example

# 模仿pytorch结构的simpleSGD
import torch

class SimpleSGD:
    def __init__(self, params, lr=0.01):
        """
        初始化优化器。
        :param params: 一个可迭代的待优化参数。
        :param lr: 学习率。
        """
        self.params = list(params)
        self.lr = lr

    def zero_grad(self):
        """清除所有被优化参数的梯度。"""
        for p in self.params:
            if p.grad is not None:
                p.grad.detach_() # 原地操作
                p.grad.zero_()

    def step(self):
        """执行单次优化步骤。"""
        with torch.no_grad(): # 内部操作不应被 autograd 跟踪
            for p in self.params:
                if p.grad is None:
                    continue
                # 核心 SGD 更新规则
                p.data -= self.lr * p.grad.data

梯度估计，学习率，收敛轨迹

优化器设计的本质在于对数据点（样本）提供可靠的梯度估计方法，在数据集和模型上获得鲁棒的收敛轨迹。

前LLM时代的核心优化器状态：速度量 — 动量

概述

标准的动量更新规则引入了一个“速度”向量 $v_t$ ，它累积了过去梯度的指数衰减移动平均。设 $\theta_t$ 为时间步 t 的模型参数， $g_t = \nabla_{\theta}L(\theta_{t-1})$ 为损失函数 L 相对于前一时刻参数的梯度，η 为学习率， $β \in [0,1)$ 为动量系数。更新方程为：

v_t=\beta v_{t−1}+g_t

θ_t=θ_{t−1}−\eta v_t

这里， $v_0$ 初始化为零。速度 $v_t$ 维持了过去梯度方向的“记忆”。当连续的梯度指向相似的方向时，它们对 $v_t$ 的贡献会累积，导致在该方向上步长更大。相反，当梯度振荡时，它们的贡献倾向于相互抵消，从而抑制更新并防止不稳定的跳跃。另一种常见的公式将速度定义为真正的指数移动平均 (EMA) ：

v_t=\beta v_{t−1}+(1−\beta)g_t

Nesterov 加速度

Nesterov 加速梯度 (NAG) 是对标准动量的一种改进，在实践中通常能提供更好的收敛速度。NAG 的关键洞见在于，它不是在当前位置计算梯度，而是在一个近似的未来位置——即沿着当前动量方向的“前瞻”一步——计算梯度 1 。这使得优化器能够预测其前进方向并更快地修正路线。如果动量即将把更新带过一个最小值，那么在前瞻位置的梯度将指回最小值，提供一个修正项来减缓更新。

二阶矩 — 自适应学习率

自适应学习率方法通过为每个参数维护一个基于该参数梯度历史进行调整的逐参数学习率来自动化此过程。

AdaGrad：累积平方梯度

AdaGrad 是这条路线的起点。它为每个参数累积历史梯度的平方和，并用其平方根来缩放学习率：

G_t = G_{t-1} + g_t \odot g_t

\theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{G_t} + \epsilon} \odot g_t

其中 $\odot$ 表示逐元素乘法， $\epsilon$ 是防止除零的小常数。直觉上，频繁获得大梯度的参数学习率被压低，而稀疏出现、梯度罕见的参数保留较大的步长——这使 AdaGrad 在稀疏特征场景（如词嵌入）中表现出色。

它的致命缺陷同样来自 $G_t$ ：平方和单调递增，分母只增不减，学习率随训练不可逆地衰减到接近零。在深度网络的长程训练中，优化往往在远未收敛时就”熄火”了。

RMSProp：会遗忘的二阶矩

RMSProp 的修正非常直接：把累积和换成指数移动平均，让旧的梯度信息按几何速率被遗忘：

v_t = \beta_2 v_{t-1} + (1-\beta_2)\, g_t \odot g_t

\theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{v_t} + \epsilon} \odot g_t

$v_t$ 现在估计的是近期梯度幅度的平方（可以理解为对 $\mathbb{E}[g^2]$ 的滑动估计），而非全部历史的总量。学习率不再单调衰减，优化器可以适应损失面局部曲率的变化：进入平坦区域后 $v_t$ 变小、步长回升；进入陡峭区域后 $v_t$ 变大、步长收缩。

谁是 Adam：一阶矩与二阶矩的合流

到这里，两条线索已经齐备：动量维护梯度的一阶矩（方向的记忆），RMSProp 维护梯度的二阶矩（幅度的标定）。Adam（Adaptive Moment Estimation）就是把两者装进同一个优化器：

m_t = \beta_1 m_{t-1} + (1-\beta_1)\, g_t

v_t = \beta_2 v_{t-1} + (1-\beta_2)\, g_t \odot g_t

\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \qquad \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 = 0.9$ ， $\beta_2 = 0.999$ ， $\epsilon = 10^{-8}$ 。

偏差修正从哪里来 $m_t$ 和 $v_t$ 都从零初始化，训练初期它们严重偏向零。以 $v_t$ 为例，展开递推可得 $v_t = (1-\beta_2)\sum_{i=1}^{t}\beta_2^{t-i}\, g_i^2$ 。若假设 $\mathbb{E}[g_i^2]$ 近似平稳，则

\mathbb{E}[v_t] \approx \mathbb{E}[g^2]\,(1-\beta_2^t)

系数 $(1-\beta_2^t)$ 正是初始化偏差的来源： $t$ 很小时它远小于 1（ $\beta_2=0.999$ 时尤其明显），除以它即得无偏估计 $\hat{v}_t$ 。没有这个修正，训练初期 $\sqrt{\hat v_t}$ 被低估，步长会被异常放大。

一个有用的解读：信噪比 Adam 的有效更新量约为 $\eta \cdot \hat m_t / \sqrt{\hat v_t}$ 。当梯度方向稳定（ $|\mathbb{E}[g]|$ 相对 $\sqrt{\mathbb{E}[g^2]}$ 较大）时，比值接近 $\pm 1$ ，步长接近满额 $\eta$ ；当梯度噪声大、方向摇摆时， $\hat m_t$ 相互抵消而 $\hat v_t$ 不会，比值缩小，步长自动收敛保守。这也解释了 Adam 单步更新的量级天然近似有界于 $\eta$ ——学习率直接就是”每步最多走多远”的标尺，这是它对超参数相对不敏感、成为默认选择的重要原因。

AdamW：解耦的权重衰减

在 SGD 中，L2 正则（在损失上加 $\frac{\lambda}{2}\|\theta\|^2$ ）与权重衰减（更新时执行 $\theta \leftarrow (1-\eta\lambda)\theta$ ）是等价的。但在 Adam 中两者不再等价：L2 正则的梯度项 $\lambda\theta$ 会混入 $m_t$ 和 $v_t$ ，被 $\sqrt{\hat v_t}$ 除回去——梯度大的参数受到的正则化反而被稀释了，正则强度与梯度历史纠缠在一起。

AdamW 的做法是把权重衰减从梯度通道中拿出来，作为独立的一步直接作用在参数上：

\theta_t = \theta_{t-1} - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda\, \theta_{t-1} \right)

解耦之后， $\lambda$ 恢复了”以固定比例把权重拉向零”的本义，与学习率调度、自适应缩放各自独立可调。AdamW 因此成为 Transformer 时代事实上的标准优化器。

LLM 时代：内存、规模与新的更新几何

当模型进入十亿到万亿参数量级，优化器面临新的约束：Adam 需要为每个参数存两份状态（ $m$ 和 $v$ ），优化器状态的显存开销是参数本身的两倍。近年的工作大致沿三个方向展开。

压缩状态 Adafactor 注意到 $v$ 对于矩阵参数可以近似低秩分解：只存行和与列和两个向量，用外积重构二阶矩，将 $O(nm)$ 的状态压到 $O(n+m)$ 。它是 T5 等早期大模型预训练的主力。8-bit Adam 则直接对状态做分块量化，思路更工程化。

极简状态 Lion 通过程序搜索发现的更新规则只保留一份动量，且只取符号：

\theta_t = \theta_{t-1} - \eta\, \operatorname{sign}\!\big(\beta_1 m_{t-1} + (1-\beta_1)\, g_t\big) - \eta\lambda\,\theta_{t-1}

m_t = \beta_2 m_{t-1} + (1-\beta_2)\, g_t

符号函数让每个坐标的更新幅度恒为 $\eta$ ，二阶矩的”逐参数标定”被最粗粒度的归一化取代——内存减半，而在视觉与语言模型上常能追平甚至超过 AdamW（代价是对学习率和批量大小更敏感）。

重新设计更新的几何 Muon 代表了另一种思路：不再把参数看作一维向量的坐标集合，而是尊重权重矩阵的结构。它对动量矩阵做近似正交化（用 Newton–Schulz 迭代逼近 $U V^\top$ ，即把奇异值全部拉平为 1），使更新在所有奇异方向上均匀推进，而不是被最大奇异方向主导。在 2024–2025 年的多个 LLM 预训练实践（如 nanoGPT speedrun、Kimi 的 Moonlight）中，Muon 用于隐藏层矩阵、AdamW 用于嵌入和归一化参数的混合方案展现了可观的样本效率优势。这类方法提示：Adam 的逐坐标视角未必是深度网络优化的终点。

小结：优化器的统一视角

回望整条演化线，几乎所有一阶优化器都可以拆成同样的三个组件：

梯度的统计量：维护哪些关于梯度历史的量——一阶矩（动量/方向记忆）、二阶矩（幅度标定），还是它们的压缩近似。这决定了优化器的状态与内存开销。
更新规则：如何把统计量组合成参数增量——直接相加（SGD）、逐坐标相除（Adam）、取符号（Lion）、矩阵正交化（Muon）。这决定了更新的几何性质。
正则化通道：权重衰减是混入梯度（L2）还是独立作用（解耦衰减）。

SGD + momentum 是”只记方向”；AdaGrad/RMSProp 是”只标幅度”；Adam 两者兼备并做无偏修正；AdamW 把正则从自适应通道中解救出来；Adafactor/Lion 在状态开销上做减法；Muon 则质疑了逐坐标处理本身。所谓”谁是 Adam”——它不是某个孤立的发明，而是动量与自适应学习率两条线索在 2014 年的一次汇合，此后十年的优化器设计，大多是在这次汇合的框架内做加减法，直到最近才开始有人重画框架本身。

主题：优化器, 谁是Adam