技术摘要：“无追踪子空间：低秩训练中的不可识别性与优化器状态”

问题陈述

像 GaLore (Zhao et al., 2024) 这样具有内存效率的优化器，通过将梯度投影到一个定期重新计算的低秩子空间（秩为 $r$）并在该子空间内运行 Adam，从而减少大型语言模型的内存占用。这种方法的底层几何前提是：前 $r$ 个梯度的子空间是一个有意义的、缓慢漂移的对象，可以在刷新间隔（$T$ 步）内被追踪。如果这一前提成立，连续估计的子空间应该是相似的，并且在刷新时携带优化器状态仅需要进行微小的修正。

本文对这一前提提出了质疑，认为除了一个小的可复现核心之外，前 $r$ 个子空间并不是一个稳定的对象，而是统计学上不可识别的。因此，每次刷新时子空间表现出的“旋转”主要是由估计器噪声而非稳定几何结构的运动所主导的。

研究方法

作者采用了一个严谨的测量框架来量化子空间稳定性以及不可识别性对优化器性能的影响：

1. 子空间旋转测量： 本文引入了一种符号校正的弦距离度量（$\|\Delta U\|_F$）来衡量每次刷新时旧基底（$U_{old}$）与新基底（$U_{new}$）之间的旋转。该度量范围从 0（完全相同）到 $\sqrt{2r}$（正交）。
2. 同步步数控制（“无时间”测试）： 为了区分真实的随时间漂移与估计器噪声，作者将单个小批量（minibatch）拆分为两个不相交的部分，并位于相同的训练步数。他们计算每个部分的前 $r$ 个子空间，并测量其差异。如果子空间是梯度的稳定属性，那么同一步数的估计值应该是一致的。如果它们的差异与跨越 $T$ 步的估计值差异相当，则说明旋转是噪声。
3. 谱分析： 作者分析了梯度的奇异值谱，以确定是否存在一个在秩 $r$ 处存在的谱间隙（spectral gap），从而证明一个稳定的子空间截断是合理的。
4. 平均实验： 本文测试了通过对 $N$ 步进行平均（或对基底进行随时间平均）是否能通过打开一个谱间隙来“挽救”可识别性。
5. 补救措施评估： 作者测试了三种减轻子空间重绘成本的策略：
   • 基底平均（Basis Averaging）： 使用协方差矩阵的指数移动平均（EMA）来稳定基底。
   • 状态传输（State Transport）： 通过旋转矩阵 $R = U_{new}^\top U_{old}$ 显式地传输优化器动量（一阶和二阶矩）（如 LDAdam 中所做）。
   • 缩短记忆（Memory Shortening）： 降低 Adam 的衰减参数 $\beta_2$，以缩短二阶矩的记忆，防止其整合来自之前子空间的陈旧统计数据。

实验涵盖了四个模型家族（Pythia, GPT-2, Qwen2.5, Llama-3），参数量从 70M 到 6.9B 不等，以及一个视觉 Transformer (Vi-B)。

核心发现

1. 子空间是不可识别的

• 最大旋转： 在每次刷新时，测得的旋转 $\|\Delta U\|_F$ 达到了几何最大值 $\sqrt{2r}$ 的 96.7–99.6%。新基底与旧基底几乎正交。
• 噪声主导： 同步步数控制显示，来自同一批次两个不相交部分的估计值之间的差异为 $0.725\sqrt{2r}$，这与跨时间旋转的 $0.742\sqrt{2r}$ 几乎相同。这种“旋转”不是运动，而是估计器噪声。
• 可复现核心 ($k^\star$)： 只有极小比例的方向是可复现的。对于 $r=128$，只有 $k^\star \approx 39$ 个方向在不同的小批量之间保持一致。剩余的 $r - k^\star$ 个方向实际上是从无间隙的谱尾部中随机抽取的。

2. 谱成因与平均法的失效

• 无间隙谱： 梯度谱由一个小的占优峰值（几个到约 20 个方向）以及一个平滑的、近乎各向同性的尾部组成，在秩 $r$ 处没有间隙。在 $r=128$ 时，相邻的奇异值比率约为 $1.005$。
• 信号而非噪声： 深层的谱尾部并非采样噪声，而是具有幂律谱的信号。在 $N$ 倍平均下，该尾部仅按 $N^{-1/4}$ 缩小，而非纯噪声预期的 $N^{-1/2}$。因此，任何有限的平均预算都无法打开间隙或稳定子空间。
• 平均法失效： 将基底在 10,000 步内进行平均，仅略微降低了每次刷新的旋转（从 97% 降至 84% 的最大值），且对最终困惑度（perplexity）没有明显的测量影响（跨度仅为 0.36）。

3. 为什么 GaLore 仍然有效

尽管缺乏稳定的子空间，GaлоRE 之所以有效，是因为投影的主要作用是能量捕获，而非追踪。由于梯度尾部接近各向同性，任何包含主导峰值的随机 $r$ 维框架都能捕获大致恒定比例的梯度能量（$\gamma_g \approx 0.67–0.73$）。只要能量比例得以保留，特定方向的选择对下降方向的影响微乎其微。

4. 有效的补救措施

重绘的代价是携带的优化器状态的陈旧性，特别是二阶矩（方差）。

• 状态传输： 通过 $R$ 传输一阶矩，并通过 $R$ 的平方项传输二阶矩是最优的。盲目携带二阶矩在理论上是不优的，其损失因子约为 $(r - k^\star)/2$。在受控测试中，在默认 $\beta_2=0.999$ 时进行状态传输达到了 18.73 的困惑度，优于未传输的 GaLore (19.28)。
• 缩短记忆： 将 $\beta_2$ 从 0.999 降低到 0.99 有助于刷新优化器，因为它防止了整合来自先前（现已失效）子空间的陈旧统计数据。这使困惑度提高了 1.8 到 26 个点，具体取决于优化器和方案。
• 基底平均的无效性： 如前所述，通过平均来稳定基底是行不通的，因为子空间在根本上是不可识别的。

意义与主张

本文声称提供了对低秩梯度优化器背后几何前提的首次直接测量，并发现该前提在通常使用的秩下是不成立的。

• 机制澄清： 本文阐明了 GaLore 之所以有效，不是因为它追踪了一个稳定的结构，而是因为它在一个实际上是重新绘制的框架中捕获了恒定比例的梯度能量。
• 诊断工具： 作者提出了一个实用的探测方法：通过拆分批次来测量可复现秩 $k^\star$。如果选定的秩 $r > k^\star$，则多余的方向即为噪声。
• 工程处方： 本文为低秩训练提供了两种可操作的修复方案：
  1. 通过旋转进行状态传输（如 LDAdam 中所示），而不是盲目携带。
  2. 将 $\beta_2$ 降低至 0.99，以减少二阶矩的记忆，从而减轻由于周期性重绘导致的陈旧性问题。
• 适用范围： 这些发现专门适用于显式计算并刷新梯度子空间的优化器（GaLore 系列）。它们不适用于像 LoRA 这样不计算梯度子空间的参数化方法。

作者总结道，“对象”（稳定的低秩子空间）在小规模核心之外并不存在，试图追踪它是在追逐幻影。前进的方向在于承认重绘的存在，并据此调整优化器状态。