深度学习中的优化器之Adafactor算法原理与自适应参数缩放机制

好的，我们将详细讲解Adafactor算法。这是一个用于训练大规模模型的优化器，尤其在资源受限场景（如内存）下表现优异，是Adam和AdaGrad等自适应优化器的一种高效替代方案。

题目描述

Adafactor算法是一种自适应学习率优化算法，由谷歌在2018年提出。它的核心目标是：在保持与Adam等优化器相当性能的同时，显著减少参数更新所需的存储开销（特别是动量和二阶矩估计所占的内存）。这对于动辄拥有数十亿甚至万亿参数的现代大模型训练至关重要。我们需深入理解其如何通过移除动量显式存储、分解二阶矩矩阵以及对学习率进行缩放来实现这一目标。

解题过程（原理与机制详解）

让我们从优化器需要存储什么开始，循序渐进地推导Adafactor的每个设计步骤。

第一步：回顾基准算法——Adam的内存瓶颈

首先，我们回顾Adam优化器的更新规则。对于第 t 步，模型参数为 θ_t，梯度为 g_t。

1. 动量（一阶矩）更新: m_t = β1 * m_{t-1} + (1 - β1) * g_t
2. 二阶矩更新: v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 (这里的平方是按元素计算)。
3. 偏差校正: m̂_t = m_t / (1 - β1^t), v̂_t = v_t / (1 - β2^t)。
4. 参数更新: θ_t = θ_{t-1} - α * m̂_t / (sqrt(v̂_t) + ϵ)。

内存瓶颈分析：

• 为了计算 v_t 和 m_t，Adam需要为模型中的每一个可训练参数都存储一个对应的动量和二阶矩值。
• 因此，Adam的额外存储开销是参数数量的2倍。对于一个拥有10亿参数的模型，这额外的20亿个浮点数会占用约8GB（假设float32）的显存，这是一个巨大的负担。

Adafactor的目标就是大幅减少这部分开销。

第二步：Adafactor的核心设计——分解二阶矩矩阵

Adafactor的第一个关键创新是针对二阶矩 v_t 的。

1. 移除非时变维度并重新定义二阶矩：

• 在Adam中，v_t 是一个与参数 θ 形状完全相同的张量。
• Adafactor观察到，对于多维参数（如权重矩阵 W ∈ R^{m×n}），其梯度 G_t 也是一个 m×n 矩阵。Adam会为每个元素 (i, j) 单独计算一个二阶矩 v_t(i,j)。
• Adafactor提出，我们可以按行和按列分别聚合二阶矩信息，以近似完整的二阶矩矩阵。具体来说：
  • 行均值（Row RMS）: 计算每行梯度的平方的均值。这得到一个长度为 m 的向量 R_t，其中 R_t[i] = RMS(G_t[i, :])^2。它反映了第 i 行参数的梯度变化幅度。
  • 列均值（Col RMS）: 计算每列梯度的平方的均值。这得到一个长度为 n 的向量 C_t，其中 C_t[j] = RMS(G_t[:, j])^2。它反映了第 j 列参数的梯度变化幅度。

2. 用外积近似完整二阶矩：

• 完整的二阶矩矩阵 V_t (Adam中的 v_t 按矩阵排布) 被近似为向量 R_t 和 C_t 的外积，再乘以一个缩放因子：
  V_t ≈ (R_t * C_t^T) / (sum(C_t)) 或更常用的一种简化形式：V_t ≈ R_t * C_t^T (在后续更新规则中体现)。
• 内存节省：原本需要存储 m * n 个值的 V_t，现在只需要存储 m + n 个值的 R_t 和 C_t。当 m 和 n 都很大时（例如，transformer中的大权重矩阵），这带来了数量级的内存节省。

第三步：更新规则推导与最终形式

基于上述分解思想，Adafactor的更新规则如下（对于二维权重矩阵 W）：

1. 计算梯度平方的指数移动平均（EMA）：

   • 对于行方向：R_t = β2 * R_{t-1} + (1 - β2) * (mean(G_t^2, axis=1)) (对列求平均)。
   • 对于列方向：C_t = β2 * C_{t-1} + (1 - β2) * (mean(G_t^2, axis=0)) (对行求平均)。
   • 这里 mean 操作符是Adafactor的另一个关键，它通过聚合进一步节省了计算和存储。在一些实现中，它直接使用 RMS（均方根）值。

2. 计算参数更新量：

   • 近似二阶矩：V̂_t[i, j] = R_t[i] * C_t[j] (这是 R_t 和 C_t 外积的简化表达)。
   • 更新量：ΔW_t = -α_t * G_t / (sqrt(V̂_t) + ϵ)。
   • 注意：这里的学习率 α_t 是随时间衰减的，见下一步。同时，分母的 sqrt(V̂_t) 是按元素计算的。

3. 移除动量并引入学习率缩放：

   • Adafactor的默认配置完全移除了显式动量（β1=0）。这是为了进一步节省内存（不再存储 m_t）。模型训练的“惯性”主要依靠二阶矩的EMA来隐式提供。
   • 为了补偿没有动量的影响，并确保训练稳定，Adafactor引入了一个基于二阶矩幅度的学习率裁剪（或称为缩放）。
   • 学习率计算：α_t = max(α_{base}, α_{relative} * RMS(ΔW_{t-1}))。
     1. α_{base} 是一个很小的绝对下限（如 1e-30），保证学习率不为零。
     2. α_{relative} 是相对学习率，是一个重要的超参数（例如 0.01）。
     3. RMS(ΔW_{t-1}) 是上一步参数更新量的均方根。这意味着学习率会根据最近更新的幅度进行自适应缩放：如果更新幅度大，学习率会相应增大；反之则减小。这在一定程度上模拟了动量带来的稳定效果。

4. 一维参数的处理：

   • 对于偏置（bias）等一维参数，无法进行行列分解。Adafactor的处理方式更简单：直接存储该一维向量的每个元素的二阶矩EMA，或者退化使用类似AdaGrad的策略，但通过周期性重置（clipping）来防止累积值过大。

第四步：算法总结与优势分析

最终，Adafactor算法的伪代码简化如下：

初始化：对于每个权重矩阵W，初始化行向量R=0，列向量C=0。设置超参数β2, ϵ, α_relative, α_base。
对于每个训练步t：
    1. 计算当前梯度 G_t。
    2. 更新二阶矩的行/列估计：
        R_t = β2 * R_{t-1} + (1 - β2) * mean(G_t^2, axis=1)
        C_t = β2 * C_{t-1} + (1 - β2) * mean(G_t^2, axis=0)
    3. 近似完整的二阶矩：V̂_t = R_t (外积) C_t。
    4. 计算未缩放的更新：Δ_unscaled = G_t / (sqrt(V̂_t) + ϵ)。
    5. 计算学习率：α_t = max(α_base, α_relative * RMS(Δ_{t-1}))。
    6. 执行参数更新：W_t = W_{t-1} - α_t * Δ_unscaled。
    7. 记录本次更新量 Δ_t = α_t * Δ_unscaled 用于下一步计算RMS。

核心优势：

1. 内存高效：主要内存开销从 O(2n)（Adam）降为 O(n + m + p + ...)，其中 n, m, p 是各参数矩阵的维度之和，远小于元素总数 n。对于大矩阵，节省效果显著。
2. 性能相当：在翻译、语言建模等任务上，实验表明Adafactor在达到与Adam相近精度的情况下，能节省大量显存。
3. 无动量设计：简化了算法，并通过对更新量进行RMS缩放来维持稳定性。

应用场景：
Adafactor特别适合训练超大规模模型，例如Transformer-based的语言模型，当GPU/TPU内存成为主要瓶颈时。它也是许多现代大模型训练库（如T5）中的默认或推荐优化器之一。

通过以上步骤，我们详细剖析了Adafactor如何通过二阶矩矩阵分解和移除显式动量，在保证优化性能的同时，实现内存占用的革命性降低。