深度学习中的优化器之SM3（平方和最小化内存优化）算法原理与内存高效自适应机制
字数 3302 2025-12-17 18:02:26

深度学习中的优化器之SM3（平方和最小化内存优化）算法原理与内存高效自适应机制

1. 题目描述

SM3（Square-Minimized Minimal Memory Adaptive Optimizer）是一种专门为大规模模型训练设计的自适应优化算法。其核心目标是在保持类似Adam优化器自适应学习率优势的同时，大幅减少内存占用。尤其在处理包含海量参数的现代大语言模型（LLM）或推荐系统模型时，SM3通过一种巧妙的“平方和最小化”策略，将每个参数的梯度二阶矩估计的内存消耗从 \(O(n)\) 降低到接近 \(O(\sqrt{n})\) 的级别，而不显著牺牲收敛性能。

2. 问题背景与动机

在深度学习训练中，自适应优化器（如Adam、AdaGrad）通常需要为每个可训练参数维护一个或多个“状态”（如梯度的指数移动平均）。对于一个包含 \(n\) 个参数的模型，这需要额外的 \(O(n)\) 内存开销。当 \(n\) 极大（例如数十亿）时，这部分内存会成为训练瓶颈。SM3试图解决这个根本矛盾：如何在不显著增加内存的前提下，实现对每个参数的个性化自适应学习率？

3. 核心思想：因式分解与共享估计

SM3的核心洞察是：模型参数通常以高维张量（如矩阵、向量）形式组织。一个 \(d_1 \times d_2\) 的权重矩阵有 \(n = d_1 \times d_2\) 个参数。SM3不再为每个标量参数单独维护一个二阶矩估计，而是维护两组更小的向量：

  • 一组与行维度 \(d_1\) 相关的估计量 \(\mathbf{v}_r \in \mathbb{R}^{d_1}\)
  • 一组与列维度 \(d_2\) 相关的估计量 \(\mathbf{v}_c \in \mathbb{R}^{d_2}\)

然后，通过这两组向量的外积近似，来得到原始参数矩阵中每个位置的二阶矩估计。这好比用两个“小盒子”里的信息，组合起来描述一个“大表格”中每个格子的情况，从而极大节约了存储空间。

4. 算法原理与逐步推导

我们以一个二维权重矩阵 \(\mathbf{W} \in \mathbb{R}^{d_1 \times d_2}\) 为例，说明SM3的工作流程。设其梯度为 \(\mathbf{G}_t \in \mathbb{R}^{d_1 \times d_2}\)。

步骤1：定义分解的估计量

  • 维护行向量 \(\mathbf{v}^{(r)}_t \in \mathbb{R}^{d_1}\)，其初始值为零向量。
  • 维护列向量 \(\mathbf{v}^{(c)}_t \in \mathbb{R}^{d_2}\)，其初始值为零向量。
    这两个向量是SM3算法额外需要存储的状态，总大小为 \(O(d_1 + d_2)\)，远小于 \(O(d_1 \times d_2)\)。

步骤2：更新估计量（平方和最小化）
在每一步 \(t\)，算法收到当前梯度矩阵 \(\mathbf{G}_t\)。SM3的核心操作是更新 \(\mathbf{v}^{(r)}_t\) 和 \(\mathbf{v}^{(c)}_t\)，使得它们的外积的每个元素尽可能接近“理想的”二阶矩估计（即梯度的平方）。
具体更新规则为：

\[\begin{aligned} \mathbf{v}^{(r)}_{t+1}[i] &= \max\left( \mathbf{v}^{(r)}_t[i],\ \max_{j} \mathbf{G}_t[i, j]^2 \right) \\ \mathbf{v}^{(c)}_{t+1}[j] &= \max\left( \mathbf{v}^{(c)}_t[j],\ \max_{i} \mathbf{G}_t[i, j]^2 \right) \end{aligned} \]

这里的“\(\max\)”操作是关键。它确保 \(\mathbf{v}^{(r)}\) 的第 \(i\) 个分量记录了历史上第 \(i\) 行所有梯度平方的最大值，而 \(\mathbf{v}^{(c)}\) 的第 \(j\) 个分量记录了历史上第 \(j\) 列所有梯度平方的最大值。

步骤3：构造参数的自适应学习率
对于矩阵 \(\mathbf{W}\) 中位置 \((i, j)\) 的参数 \(w_{ij}\)，其对应的“自适应学习率分母项”（类似于Adam中的 \(\hat{v}_t\)）构造为：

\[u_{t+1}[i, j] = \max\left( \mathbf{v}^{(r)}_{t+1}[i],\ \mathbf{v}^{(c)}_{t+1}[j] \right) \]

这个构造保证了 \(u_{t+1}[i, j] \ge \mathbf{G}_t[i, j]^2\)，即它总是大于等于当前梯度平方。这起到了与AdaGrad/Adam中累积平方梯度类似的作用：梯度大的方向，学习率会被减小。

步骤4：参数更新
最终的参数更新公式与AdaGrad/Adam类似：

\[\mathbf{W}_{t+1}[i, j] = \mathbf{W}_t[i, j] - \frac{\eta}{\sqrt{u_{t+1}[i, j] + \epsilon}} \cdot \mathbf{G}_t[i, j] \]

其中 \(\eta\) 是全局学习率，\(\epsilon\) 是一个小的常数（如 \(10^{-8}\)）用于数值稳定性。

5. 内存节省分析

  • 传统方法（如Adam）：需要一个与 \(\mathbf{W}\) 同形的矩阵 \(\mathbf{V} \in \mathbb{R}^{d_1 \times d_2}\) 来存储二阶矩估计。内存为 \(O(d_1 \times d_2)\)。
  • SM3方法：只需要两个向量 \(\mathbf{v}^{(r)} \in \mathbb{R}^{d_1}\) 和 \(\mathbf{v}^{(c)} \in \mathbb{R}^{d_2}\)。内存为 \(O(d_1 + d_2)\)。
  • 节省倍数：近似为 \(O\left(\frac{d_1 \times d_2}{d_1 + d_2}\right)\)。当 \(d_1 = d_2 = d\) 时，节省倍数为 \(O(d/2)\)。例如，对于一个 \(10000 \times 10000\) 的矩阵，Adam需要存储1亿个状态值，而SM3只需存储2万个，节省了5000倍内存。

6. 算法特性与直观理解

  1. 最大操作的意义：通过取行和列上的历史梯度平方最大值，SM3为每个参数构建了一个“足够大”的 \(u_{t+1}[i, j]\)，确保学习率不会过大，避免了训练不稳定。这是一种保守但安全的策略。
  2. 信息共享：同一行（或同一列）的不同参数会共享部分二阶矩估计信息（\(\mathbf{v}^{(r)}\) 或 \(\mathbf{v}^{(c)}\)）。这是内存节省的本质，也意味着SM3的自适应性是“粗粒度”的，但实践表明这对最终模型性能影响很小。
  3. 扩展到高维张量：对于维度为 \(d_1 \times d_2 \times ... \times d_k\) 的张量，SM3会为每个维度维护一个估计向量，内存开销为 \(O(\sum_i d_i)\)，仍然远小于 \(O(\prod_i d_i)\)。

7. 总结

SM3优化器通过因式分解和共享估计的巧妙设计，在自适应学习率和内存效率之间取得了卓越的平衡。它特别适用于训练参数量极大的模型，是推动大模型发展的重要基础设施算法之一。其核心思想“用低维向量的组合来模拟高维张量的状态”也为后续的其他高效优化器设计提供了灵感。

深度学习中的优化器之SM3（平方和最小化内存优化）算法原理与内存高效自适应机制 1. 题目描述 SM3（Square-Minimized Minimal Memory Adaptive Optimizer）是一种专门为 大规模模型训练 设计的自适应优化算法。其核心目标是在保持类似Adam优化器自适应学习率优势的同时， 大幅减少内存占用 。尤其在处理包含 海量参数 的现代大语言模型（LLM）或推荐系统模型时，SM3通过一种巧妙的“平方和最小化”策略，将每个参数的梯度二阶矩估计的内存消耗从 \(O(n)\) 降低到接近 \(O(\sqrt{n})\) 的级别，而不显著牺牲收敛性能。 2. 问题背景与动机 在深度学习训练中，自适应优化器（如Adam、AdaGrad）通常需要为每个可训练参数维护一个或多个“状态”（如梯度的指数移动平均）。对于一个包含 \(n\) 个参数的模型，这需要额外的 \(O(n)\) 内存开销。当 \(n\) 极大（例如数十亿）时，这部分内存会成为训练瓶颈。SM3试图解决这个根本矛盾： 如何在不显著增加内存的前提下，实现对每个参数的个性化自适应学习率？ 3. 核心思想：因式分解与共享估计 SM3的核心洞察是：模型参数通常以 高维张量 （如矩阵、向量）形式组织。一个 \(d_ 1 \times d_ 2\) 的权重矩阵有 \(n = d_ 1 \times d_ 2\) 个参数。SM3不再为每个标量参数单独维护一个二阶矩估计，而是维护两组更小的向量： 一组与 行维度 \(d_ 1\) 相关的估计量 \(\mathbf{v}_ r \in \mathbb{R}^{d_ 1}\) 一组与 列维度 \(d_ 2\) 相关的估计量 \(\mathbf{v}_ c \in \mathbb{R}^{d_ 2}\) 然后，通过这两组向量的 外积近似 ，来得到原始参数矩阵中每个位置的二阶矩估计。这好比用两个“小盒子”里的信息，组合起来描述一个“大表格”中每个格子的情况，从而极大节约了存储空间。 4. 算法原理与逐步推导 我们以一个二维权重矩阵 \(\mathbf{W} \in \mathbb{R}^{d_ 1 \times d_ 2}\) 为例，说明SM3的工作流程。设其梯度为 \(\mathbf{G}_ t \in \mathbb{R}^{d_ 1 \times d_ 2}\)。 步骤1：定义分解的估计量 维护行向量 \(\mathbf{v}^{(r)}_ t \in \mathbb{R}^{d_ 1}\)，其初始值为零向量。 维护列向量 \(\mathbf{v}^{(c)}_ t \in \mathbb{R}^{d_ 2}\)，其初始值为零向量。 这两个向量是SM3算法额外需要存储的状态，总大小为 \(O(d_ 1 + d_ 2)\)，远小于 \(O(d_ 1 \times d_ 2)\)。 步骤2：更新估计量（平方和最小化） 在每一步 \(t\)，算法收到当前梯度矩阵 \(\mathbf{G}_ t\)。SM3的核心操作是更新 \(\mathbf{v}^{(r)}_ t\) 和 \(\mathbf{v}^{(c)} t\)，使得它们的 外积的每个元素 尽可能接近“理想的”二阶矩估计（即梯度的平方）。 具体更新规则为： \[ \begin{aligned} \mathbf{v}^{(r)} {t+1}[ i] &= \max\left( \mathbf{v}^{(r)} t[ i],\ \max {j} \mathbf{G} t[ i, j ]^2 \right) \\ \mathbf{v}^{(c)} {t+1}[ j] &= \max\left( \mathbf{v}^{(c)} t[ j],\ \max {i} \mathbf{G}_ t[ i, j ]^2 \right) \end{aligned} \] 这里的“\(\max\)”操作是关键。它确保 \(\mathbf{v}^{(r)}\) 的第 \(i\) 个分量记录了历史上第 \(i\) 行所有梯度平方的最大值，而 \(\mathbf{v}^{(c)}\) 的第 \(j\) 个分量记录了历史上第 \(j\) 列所有梯度平方的最大值。 步骤3：构造参数的自适应学习率 对于矩阵 \(\mathbf{W}\) 中位置 \((i, j)\) 的参数 \(w_ {ij}\)，其对应的“自适应学习率分母项”（类似于Adam中的 \(\hat{v} t\)）构造为： \[ u {t+1}[ i, j] = \max\left( \mathbf{v}^{(r)} {t+1}[ i],\ \mathbf{v}^{(c)} {t+1}[ j ] \right) \] 这个构造保证了 \(u_ {t+1}[ i, j] \ge \mathbf{G}_ t[ i, j ]^2\)，即它总是大于等于当前梯度平方。这起到了与AdaGrad/Adam中累积平方梯度类似的作用：梯度大的方向，学习率会被减小。 步骤4：参数更新 最终的参数更新公式与AdaGrad/Adam类似： \[ \mathbf{W}_ {t+1}[ i, j] = \mathbf{W} t[ i, j] - \frac{\eta}{\sqrt{u {t+1}[ i, j] + \epsilon}} \cdot \mathbf{G}_ t[ i, j ] \] 其中 \(\eta\) 是全局学习率，\(\epsilon\) 是一个小的常数（如 \(10^{-8}\)）用于数值稳定性。 5. 内存节省分析 传统方法 （如Adam）：需要一个与 \(\mathbf{W}\) 同形的矩阵 \(\mathbf{V} \in \mathbb{R}^{d_ 1 \times d_ 2}\) 来存储二阶矩估计。内存为 \(O(d_ 1 \times d_ 2)\)。 SM3方法 ：只需要两个向量 \(\mathbf{v}^{(r)} \in \mathbb{R}^{d_ 1}\) 和 \(\mathbf{v}^{(c)} \in \mathbb{R}^{d_ 2}\)。内存为 \(O(d_ 1 + d_ 2)\)。 节省倍数：近似为 \(O\left(\frac{d_ 1 \times d_ 2}{d_ 1 + d_ 2}\right)\)。当 \(d_ 1 = d_ 2 = d\) 时，节省倍数为 \(O(d/2)\)。例如，对于一个 \(10000 \times 10000\) 的矩阵，Adam需要存储1亿个状态值，而SM3只需存储2万个，节省了5000倍内存。 6. 算法特性与直观理解 最大操作的意义 ：通过取行和列上的历史梯度平方最大值，SM3为每个参数构建了一个“足够大”的 \(u_ {t+1}[ i, j ]\)，确保学习率不会过大，避免了训练不稳定。这是一种保守但安全的策略。 信息共享 ：同一行（或同一列）的不同参数会共享部分二阶矩估计信息（\(\mathbf{v}^{(r)}\) 或 \(\mathbf{v}^{(c)}\)）。这是内存节省的本质，也意味着SM3的自适应性是“粗粒度”的，但实践表明这对最终模型性能影响很小。 扩展到高维张量 ：对于维度为 \(d_ 1 \times d_ 2 \times ... \times d_ k\) 的张量，SM3会为每个维度维护一个估计向量，内存开销为 \(O(\sum_ i d_ i)\)，仍然远小于 \(O(\prod_ i d_ i)\)。 7. 总结 SM3优化器通过 因式分解和共享估计 的巧妙设计，在 自适应学习率 和 内存效率 之间取得了卓越的平衡。它特别适用于训练参数量极大的模型，是推动大模型发展的重要基础设施算法之一。其核心思想“用低维向量的组合来模拟高维张量的状态”也为后续的其他高效优化器设计提供了灵感。