深度学习优化器算法巧思速览

这一篇博文想写很久了，一直没有下笔，核心原因也是有一些待办的思路在攻关验证。

我们先从一个核心的问题出发，

1. 为什么要研究优化器算法？

它的关联问题：训练为什么要调参，调的是什么参？

如果就这个问题去问各种大语言模型，它们能给出一堆的理由。

但就博主而言，答案只有一个：

干掉调参，解放生产力，榨干算力。

说到底就一个字"穷"。

在多年的研发生涯里，对调参这个事深恶痛绝，为什么辛辛苦苦架构出来的模型，一训练就崩，训练收敛慢到龟速，这严重影响了开发进度，并且增加了很多不可抗力的消耗。

我相信有很多业内同行，都有这种痛，训练了很久，效果依旧很差，泛化能力也不行，然后就开始苦恼，为什么自己没有足够的钱，足够的算力。

明明自己很好的思路，戛然而止，退而求其次。

早年间，博主经常半夜醒来，看训练的损失曲线，生怕训崩。就算没有训崩，自己花费了大量时间精力，却没有很好的回报。

一次又一次，是很打击信心的。

在付出了大量时间和人民币之后，博主终于从泥潭里爬出来了，时光荏苒，这个困扰我九年的问题，画上句号了。

那大语言模型是怎么回答这个问题的。

核心就一句话：

"没有新优化器，下一代模型根本训不起来。"

从理论上看，它是在解决一个尚未被完全理解的复杂高维优化问题，充满挑战与机遇。

解决基础性训练难题------让模型"能学"
从工程上看，它是降低AI研发成本、推动技术普及的关键杠杆。

追求极致的效率与效益------让模型"快学"且"省学"
从性能上看，它是提升模型最终准确性、鲁棒性和泛化能力的决定性因素。

提升模型的终极性能------让模型"学好"

最终达到，拓展AI的技术边界------让"不可能"成为"可能"

当然就这个问题，大家可以自行去追问各家的大语言模型，给出的结论大同小异。

2. 那博主为什么要写这篇博文？

最基本的还是希望抛砖引玉，希望能有更多的同行在力大砖飞，烧钱的当下，不要放弃底层算法的研究。

同时为更多的深度学习小白提供一个新的视角，学习并应用深度学习，温故而知新。

3. 那什么是优化器算法？

优化器算法是驱动机器学习模型学习的"引擎"。它的核心任务是：在训练过程中，根据损失函数计算出的梯度（即方向），以某种策略更新模型的参数，从而最小化损失函数。

可以将训练过程想象成在复杂地形中寻找最低点：

损失函数：代表地形的高度。
模型参数：代表我们在地形中的位置。
梯度：代表我们脚下最陡峭的下坡方向。
优化器：就是那个决定"往哪个方向走、走多大步、以及是否要考虑之前的惯性"的导航策略。

Adam (Adaptive Moment Estimation)

思想：目前最流行和默认的优化器之一。它结合了Momentum和RMSProp的优点。
- 它计算梯度的一阶矩（均值，提供动量）和二阶矩（未中心化的方差，用于自适应调整学习率）。
- 然后对这两个矩进行偏差校正，使其在训练初期不那么偏向于0。
优点：
- 通常收敛速度快。
- 对超参数的选择相对鲁棒（默认参数通常就能工作得很好）。
- 能处理噪声和稀疏梯度。

如果把Adam的一阶矩和二阶矩去掉，它就蜕变为SGD。

而随机梯度下降(朴素SGD)是一种优化算法，通过随机选取单个样本来近似梯度，从而迭代更新模型参数，收敛至最小值。

换句话说，朴素SGD是一个没有应用任何先验补充的野蛮人，较于Adam的平滑学习而言，它就像一只无头苍蝇，到处乱撞，也不知道该撞多少次才能收敛至最小值。

4. Adam相较于朴素SGD，它做了哪些改进？

引入动量缓冲m，也就是一阶矩，指数加权平滑梯度，它积累了历史梯度的方向趋势。使得朴素SGD的动荡趋于平稳平滑。
引入自适应步长v，也就是二阶矩，指数加权平均的平方，它积累了历史梯度平方的值趋势。

最终以 grad = m / sqrt(v) 作为目标梯度进行更新。

对于动量一阶矩，基本没啥好说的，就是求历史平均梯度，使得训练平稳。

核心还是自适应步长v，对于频繁更新、梯度大的参数，其二阶矩估计值大，因此实际更新步长会被调小（除以一个大数），避免"步子太大"而越过最优点。对于不频繁更新、梯度小的参数，则给予更大的相对步长，鼓励其更新。

所以Adam能加速较于朴素SGD训练收敛，二阶矩功不可没。

原本故事到这里，就接近完结了。

但在真实的场景下，虽然我们发现Adam不够好，但它的普及使得深度学习遍地开花。虽然仍然是需要调参，但是不像之前那么"玄学"了。

5. 后Adam家族时代，百家争鸣

由于这个话题展开，真的可以写一本书了。

所以本文的核心是"速览"，博主带着大家看一看这后Adam的各种巧思。

相关的算法实现，可以参考以下项目仓库：

PyTorch:
https://github.com/kozistr/pytorch_optimizer

TensorFlow/Keras:
https://github.com/NoteDance/optimizers

本文没有提及的其他算法，自行移步查阅。

5.1 砍Adam的显存

由于一阶矩m和二阶矩v都需要历史平滑，所以Adam至少要占用两倍的可训练模型参数。

这样一来，只要模型参数一大，那训练的时候 1+2 = 3 至少要存储三份权重。显存很快就不够用了。

所以，针对这个问题，我们开始磨刀霍霍向二阶矩v。

5.1.1 18年的Adafactor

[1804.04235v1] Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

社区比较知名的实现：

transformers/src/transformers/optimization.py at main · huggingface/transformers · GitHub

5.1.2 19年的SM3

[1901.11150] Memory-Efficient Adaptive Optimization

官方实现：
https://github.com/google-research/google-research/tree/master/sm3

Adafactor和SM3都是分解近似的做法。SM3的实现较为复杂，所以基本上没有被推广开来。所以很长一段时间都是Adafactor是主流。

但是Adafactor的实现稍微有些问题。

问题函数：

python 复制代码

@staticmethod
def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
    # copy from fairseq's adafactor implementation:
    # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505
    r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
    c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
    return torch.mul(r_factor, c_factor)

_approx_sq_grad 这个实现丢失了不少精度。

博主认为比较合理的实现，是把sqrt放到最后计算，精度会高些。

python 复制代码

@staticmethod
def _approx_sq_grad(row_exp_avg_sq, col_exp_avg_sq):
    row_factor = row_exp_avg_sq.unsqueeze(-1)
    row_factor = row_factor.mean(dim=-2, keepdim=True).div(row_factor)
    col_factor = col_exp_avg_sq.unsqueeze(-2)
    return row_factor.div(col_factor).sqrt_()

5.1.3 22年的Amos

[2210.11693]Amos: An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale

在Adafactor和SM3之后很长一段时间，砍优化器显存占用这个事情似乎被遗忘了。

直到Amos的出现，它进一步砍掉了v的显存占用，直接采用了平方均值，美其名曰"信息共享"。

显存不够用，又想保住精度，可以考虑采用Amos，当然它较之Adam还有不少改进点。

5.2 Adam二阶矩v为0的问题

导致v为0有很多原因，在模型训练的不同阶段，由于噪声也好，精度也好，会直接或者间接导致v为0。

前面提到 grad = m / sqrt(v)

早期Adam论文里的解决方案就是直接给v加上一个epsilon，一般设为1e-8，避免除以0。

而后续经过不少团队的实践发现这么做有点鲁莽。

然后就有人开始针对这个问题进行修改。

但是林林总总，都是把epsilon移来移去，例如梯度平方后就加上epsilon，再进行指数加权平均。

这个问题一直到了2024年，才有新的进展。

[2407.05872v2] Scaling Exponents Across Parameterizations and Optimizers

方法很简单，直接删除epsilon。

grad = m / sqrt(v) 改为 grad = atan2(m, sqrt(v))

从数值稳定的角度来说，atan2确实是稳定了许多，而且基本规避了一些特殊情况下训练跑崩，导致损失为nan的情况。

5.3 Adam的梯度长尾问题

这个很好理解，由于一阶矩m和二阶矩v都采用了指数平均，在不同程度上也是导致梯度长尾的诱因之一。

因为求平均值这个事，就跟奥运比赛打分一样，只用均值很不公平。去掉一个最高分，去掉一个最低分，然后再算平均相对合理一些。

求损失均值的时候一样存在，博主曾经设想过，也许求损失的中位数是一个可行的做法，但也有一定的局限性。

没有经过严格验证的求损失中位数思路的实现，仅供参考：

python 复制代码

def soft_median(losses, temperature=None):
    if temperature is None:
        temperature = max(0.1, 0.5 * losses.std())
    if losses.numel() % 2 == 0:
        losses = torch.cat([losses, losses.new_zeros(1)])
    x_sorted, _ = torch.sort(losses)
    n_loss = losses.shape[0]
    median_idx = (n_loss - 1) * 0.5
    idxs = torch.arange(n_loss, device=losses.device, dtype=losses.dtype)
    weights = torch.softmax(-torch.abs(idxs - median_idx) / temperature, dim=0)
    return torch.dot(weights, x_sorted)

同样的，梯度在训练过程中变化很大，一些长尾样本带来的贡献就会被淹没掉。

带来的后果，不是过拟合，就是泛化差，能拿到次优解那是属于幸运儿了。

这个方向的研究多，也不多，因为很多长尾问题基本上不会考虑在优化器里解决，一般会采用损失加权惩罚的思路来缓解。

这篇论文可以帮助进一步理解梯度长尾问题。

[2201.05938v2] GradTail: Learning Long-Tailed Data Using Gradient-based Sample Weighting

当然它不是一个主流的方案和思路，主流的方案更多的是采用元学习之类的做法，局限性也比较大。

博主一直认为如果可以优雅解决长尾问题，那是新一轮的曙光。

5.4 Adam的过拟合问题

由于Adam本身的机制问题，

训练损失下降极快 → 模型迅速进入插值（interpolation）区域 → 参数范数容易膨胀 → 边界更复杂 → 泛化差。

当然长尾问题也是它导致过拟合的原因之一。

比较知名且使用广泛的方案是l2正则化，即权重衰减。

Adam 进化为 AdamW，也就是现在主流的优化器算法

它思路也是非常简单粗暴，在每次更新时，从权重中减去一个固定的比例（weight * weight_decay），是正则也是先验惩罚。

[1711.05101v3] Decoupled Weight Decay Regularization

权重衰减是一个很好的思路，但它带来了一个新的问题。衰减量设为多少才是合适的，也就是说，惩罚力度该如何界定。

衰减过大，学习收敛缓慢，衰减过小，没有起到作用。

随后Scheduled (Stable) Weight Decay也被提出，但是应用不广，鲜为人知。

[2011.11152] On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective

它的思路也很简单，通过汇总整个模型的参数信息，按照参数权重占比估算出每一层的衰减权重。

而有另一篇论文从另一个新颖的角度提出了一个方案。

[2103.06583] Preprint: Norm Loss: An efficient yet effective regularization method for deep neural networks

它的思路是在每次更新时，从权重中减去一个单元范数权重，可以近似看做是为权重衰减提供了范数先验。

5.5 学习率热身与梯度裁剪

在说到Adam过拟合的时候，我们很容易就发现了一个问题。

在不同的模型架构，训练的每个阶段，每层权重的值域是不一样的，而且这个值域随着训练的增加，也一直在变化。

由于这个核心问题的存在，训练早期梯度的波动就会很大，这个时候通常就需要学习率调参，或者在模型内部加入归一化层，目的尽可能快地把每一层的值域确立下来。

由此就引发出来学习率热身以及梯度裁剪相关的思考。

学习率热身相关的资料和论文也有很多，这里不展开细讲。

学习率规划热身的基本逻辑都是：

早期用极其小的学习率进行预热训练 → 中期慢慢地增大学习率 → 后期再固定学习率或者慢慢减少学习率

虽然很傻，但是确实有效。

21年的时候谷歌为了把归一化层删掉，就提出了自适应梯度裁剪方案。

[2102.06171] High-Performance Large-Scale Image Recognition Without Normalization

思路也很简单，根据每层梯度和权重的值域，按比例缩放当前的梯度。

25年终于有人想要把学习率预热删掉。

[2505.21910] Taming Transformer Without Using Learning Rate Warmup

思路跟Scheduled (Stable) Weight Decay很像，只不过这次是作用在学习率上罢了。

本质就是根据每层权重梯度比例算出来一个全局学习率的缩小率。由于每层的激活函数不一样，算出来一个全局缩小率，从逻辑上其实很牵强。

当然初次之外还有其他类似的思路，例如：

梯度范数化

[1711.02257v4] GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks

层范数化缩放

[1904.00962v5] Large Batch Optimization for Deep Learning: Training BERT in 76 minutes

梯度中心化

[2004.01461v2] Gradient Centralization: A New Optimization Technique for Deep Neural Networks

林林总总，大同小异。

博主根据自己的理解，也写了个梯度软裁剪，代码如下。

python 复制代码

@staticmethod
def _soft_clip(grad, var, epsilon=1e-12):
    dim = None if (r := var.dim()) <= 1 else tuple(range(1, r))
    var_norm = var.square().mean(dim=dim, keepdim=True).sqrt_().clamp_min_(epsilon)
    grad_norm = grad.square().mean(dim=dim, keepdim=True).sqrt_().clamp_min_(epsilon)
    clipped_norm = grad_norm.clamp_max(var_norm)
    return grad.mul_(clipped_norm / grad_norm)