LAMB优化器：让每层网络都迈出恰到好处的一步

可视化演示

展示编译器代码优化阶段如何通过DAG（有向无环图）识别并消除公共子表达式，减少冗余计算，优化基本块内的指令序列。

大批量训练难题

当指挥官的哨声只有一声，但队伍里有人背帐篷、有人背绳索、有人背氧气瓶，要求所有人迈出完全一样的步幅，这支队伍很快就会东倒西歪。深度学习里的大批量训练也有类似的尴尬：我们把 batch size 从 256 提高到 64K，本意是让每一步的梯度估计更平滑、更高效，但如果优化器给所有层一个统一学习率，某些层的参数步长会突然过大，另一些层却几乎不动。

这里面的核心矛盾是“层间尺度差异”。一个深层网络里，Embedding 层、Attention 层、FFN 层的参数范数可能相差好几个数量级。梯度经过多层反向传播后，各层的梯度范数更是参差不齐。Adam 这样的优化器虽然对每个参数都有自适应学习率，但它并没有考虑“这一层当前参数有多大”，于是当 batch 极大、学习率需要等比放大时，某一层就会成为最先“爆”的短板。

也正是这种不稳定性，让大批量训练常被形容为“悬崖边上跳舞”。想要不摔下去，要么小心翼翼地把学习率调得极小，让训练慢得难以接受；要么换一种能感知层间差异的步长策略。LAMB 就是为此而来的。

LAMB核心思想

这位经验丰富的向导不会只说一句“走”，而是先看一眼每个分队的装备总重量，再看一眼他们原本计划推出去的那一步有多大。如果一步迈得比装备重量还夸张，就把它收回来一些；如果一步明显太保守，就放大一些。在 LAMB 里，“装备总重量”是当前层的参数范数，而“计划推出去的一步”是 Adam 风格的原始更新方向。

用数学语言说，对第 \(i\) 层参数：

其中 \(\theta_t^{(i)}\) 是第 \(i\) 层的权重，\(\Delta_t^{(i)}\) 是由一阶矩和二阶矩计算出的 Adam 更新方向。这个比值 \(\displaystyle \frac{\|\theta_t^{(i)}\|}{\|\Delta_t^{(i)}\|}\) 就是“信任比率”。

这样一来，即使不同层重量悬殊，它们在每一次更新里走过的“相对距离”却能保持相近。登山队既能整体推进，又不会有人一脚踩空。

算法机制解析

想要真正驾驭 LAMB，不能只记住一个公式，还要看懂它和 Adam 共同使用的“仪表盘”。LAMB 保留了 Adam 对每个参数的动量估计和梯度二阶矩估计，只是多了一层“逐层裁判”。

首先，一阶矩 \(m\) 记录梯度方向的平滑趋势，二阶矩 \(v\) 记录梯度数值的平滑幅度。因为训练刚开始时这两项都偏小，还需要做偏差校正：

校正之后，LAMB 再计算该层的原始更新方向，并套上信任比率：

看到这里会发现，LAMB 并没有抛弃 Adam，而是在 Adam 的“脚下”按层加了可伸缩的弹簧：方向仍然由动量决定，但迈步大小由该层参数尺度说了算。

BERT预训练实践

当整套指挥系统第一次用在大规模探险队上，效果让所有人惊喜。BERT 这类预训练模型层数深、参数多，以往想缩短预训练时间，最直接的办法是加大 batch 并提高并行度，但 Adam 在超大批次下很难稳定收敛。LAMB 的出现，让 BERT 也能像滑雪一样顺着大批次赛道全力冲刺。

在 Google 当年的实验中，使用 LAMB 配合约 64K 的 batch size，BERT-Large 预训练在 1024 块 TPU 上大约 76 分钟就能完成，而传统方案动辄需要数天甚至数周。更关键的是，模型在 GLUE 等下游任务上的精度没有明显损失，有些任务甚至略有提升。这个结果让研究者意识到：阻碍超大批量训练的并非硬件，而是优化器缺乏层间感知能力。

从工程视角看，这相当于把每层参数的更新步长从“一刀切”变成“量体裁衣”。Embedding 层参数巨大，但更新方向可能很小；注意力层的参数范数小，却需要相对更大的更新。LAMB 通过信任比率让每一层都处在安全的前进区间，于是通信效率和大 batch 吞吐量真正转化为训练时长的缩短。

与Adam对比

如果 Adam 是只发统一登山鞋的教练，那 LAMB 就是会按每个人脚码配鞋的向导。Adam 的“自适应”体现在参数维度：每个参数根据自身历史梯度调整这次更新的数值大小。但它并不会问一句“你这个层整体有多大”。对于一个 MobileNet 这样只有几层的模型，问题不大；可对于一个几十层的 Transformer，层间参数尺度差异会被极大的 batch 放大。

LAMB 与 Adam 有三个关键区别。第一，Adam 使用全局学习率乘以每个参数的归一化梯度，LAMB 则额外计算每一层的信任比率。第二，Adam 的更新范数可能远小于或远大于参数范数，而 LAMB 会让更新范数与参数范数保持一个相对稳定的比例。第三，在 batch size 较小时，LAMB 往往和 Adam 表现接近，优势不明显；但当 batch 增大到 64K 甚至 128K 时，Adam 需要极其精细的学习率搜索，LAMB 却依然稳如磐石。

这并不意味着 LAMB 永远更好。小 batch 下 Adam 已经足够稳定，LAMB 多出的范数计算和每层缩放反而增加了一点开销。真正的定位是“大批量训练的钥匙”，用来开启数据并行、模型并行和超长预训练带来的规模红利。

工程调优点

最好的登山队长，也要有补给计划、天气窗口和应急预案。LAMB 不是魔法，工程上还要注意几个细节：

第一，学习率热身（warm-up）几乎是必需品。训练最初几百步，梯度估计还不稳，先用一个较小的学习率起步，再逐步升高到峰值，可以有效避免早期爆炸。第二，信任比率需要裁剪。如果不加限制，某层参数范数偶尔接近零时，信任比率会异常大，导致参数猛烈跳动。常见做法是把 \(r\) 限制在 \([0.1, 10]\) 左右。第三，默认超参数可以直接沿用 Adam 的习惯：\(\beta_1=0.9\)、\(\beta_2=0.999\)、\(\varepsilon=10^{-6}\) 或 \(10^{-8}\)，在此基础上再对峰值学习率做一次小范围搜索。第四，永远要监控损失曲线和逐层更新范数。如果某层更新范数突然异常，及时减小裁剪上限或调整热身步数。

把这些工程细节看成给登山队配的安全绳：LAMB 解决“步长分层”的核心问题，而热身和裁剪解决“极端天气”下的意外风险。两者配合，才敢真正开足马力冲上超大批量这座高峰。

AI 应用场景

Python 演示

import numpy as np  class LAMB:    def __init__(self, params, lr=0.01, betas=(0.9, 0.999), eps=1e-6, max_trust=10.0):        self.params = params        self.lr = lr        self.beta1, self.beta2 = betas        self.eps = eps        self.max_trust = max_trust        self.ms = [np.zeros_like(p) for p in params]        self.vs = [np.zeros_like(p) for p in params]        self.t = 0     def step(self, grads):        self.t += 1        for i in range(len(self.params)):            p = self.params[i]            g = grads[i]            self.ms[i] = self.beta1 * self.ms[i] + (1 - self.beta1) * g            self.vs[i] = self.beta2 * self.vs[i] + (1 - self.beta2) * g * g            m_hat = self.ms[i] / (1 - self.beta1 ** self.t)            v_hat = self.vs[i] / (1 - self.beta2 ** self.t)            update = m_hat / (np.sqrt(v_hat) + self.eps)            p_norm = np.linalg.norm(p)            u_norm = np.linalg.norm(update)            trust = 1.0 if u_norm == 0 else p_norm / u_norm            trust = min(trust, self.max_trust)            p -= self.lr * trust * update  def make_data(n=128):    x = np.random.randn(n, 4)    w = np.array([1.5, -2.0, 3.0, -0.5])    y = x.dot(w) + 0.1 * np.random.randn(n)    return x, y  def train():    x, y = make_data()    hidden = np.random.randn(4, 3) * 0.2    out = np.random.randn(3, 1) * 0.2    params = [hidden, out]    opt = LAMB(params, lr=0.1)    for epoch in range(50):        h = np.maximum(x.dot(hidden), 0)        pred = h.dot(out)        loss = np.mean((pred - y[:, None]) ** 2)        d = 2 * (pred - y[:, None]) / len(x)        g_out = h.T.dot(d)        g_hidden = (d.dot(out.T) * (x.dot(hidden) > 0))        g_hidden = x.T.dot(g_hidden)        opt.step([g_hidden, g_out])        if epoch % 10 == 0:            print('epoch', epoch, 'loss', round(float(loss), 4))  if __name__ == '__main__':    train()

1. self.ms[i] = self.beta1 * self.ms[i] + (1 - self.beta1) * g维护一阶矩，记录梯度方向的平滑趋势。
2. self.vs[i] = self.beta2 * self.vs[i] + (1 - self.beta2) * g * g维护二阶矩，记录梯度幅度的平滑估计。
3. update = m_hat / (np.sqrt(v_hat) + self.eps)先算出标准的 Adam 更新方向。
4. trust = 1.0 if u_norm == 0 else p_norm / u_norm计算信任比率，即参数范数对更新范数的比值。
5. p -= self.lr * trust * update用学习率、信任比率和更新方向一起修该层参数。

常见误区

面试题