技术摘要：LoRA-Muon：低秩流形上的谱陡峭下降法

1. 问题陈述

低秩自适应（LoRA）因其在减少训练 FLOPs 和内存需求方面的效率，被广泛应用于大型深度学习模型的微调。然而，当 AdamW 等标准优化器独立应用于 LoRA 因子（$A$ 和 $B$）时，与全秩训练相比，会引入显著的调优困难。这些问题包括：

• 对初始化的敏感性： 性能高度依赖于初始化的选择。
• 糟糕的超参数迁移性： 最优学习率无法在不同的秩（rank）、宽度、深度或因子缩放约定之间良好迁移。
• 发散的动力学： 即使从相同的组合权重（$W = AB^T$）开始，如果因子规模发生漂移，训练动力学也会发生分歧。
• 次优性能： 原生的低秩预训练在没有额外稳定化的情况下，往往难以达到全秩基准的水平。

核心挑战在于设计一种针对低秩流形的优化器，使其能够保持全秩优化器（特别是 Muon 和 Shampoo 系列）的几何特性，同时确保最优学习率能够在秩、宽度、深度和因子重缩放之间进行迁移。

2. 方法论

2.1 几何基础

本文将优化器设计重新定义为一个几何问题：确定优化器在何种范数下执行最陡峭下降。

• Muon 优化器： 在环境矩阵空间 $\mathbb{R}^{m \times n}$ 中，在谱范数（$\|\cdot\|_{2\to2}$）下执行最陡峭下降。其更新方向由矩阵符号函数 $-\text{msign}(G)$ 给出。
• 低秩流形： LoRA 将权重限制在流形 $\mathcal{M}_r = \{W \in \mathbb{R}^{m \times n} : \text{rank}(W) = r\}$ 上。作者提出直接在这一流形的切空间 $T_W \mathcal{M}_r$ 上求解最陡峭下降问题，而不是将 $A$ 和 $B$ 作为独立的坐标进行优化。

2.2 LoRA-Muon 的推导

作者通过在低秩切空间下、在谱范数约束下求解线性最小化问题，推导出了 LoRA-Muon。

1. 切空间分解： 对 $W$ 的更新分解为 $\Delta W = \Delta A B^T + A \Delta B^T$。
2. 置信区域拆分： 为了使问题变得可处理，置信区域预算 $\eta$ 被平均分配给两个切分分量。这基于一个经验观察，即这两个分量在训练过程中会迅速对齐，使得三角不等式几乎取等号。
3. 投影器公式化： 利用酉不变性，问题被简化为在投影梯度上寻找线性最小化算子（LMOs）。理想的更新为：
   $$\Delta W = \frac{\eta}{2} \text{LMO}(\nabla W f \cdot P_B) + \frac{\eta}{2} \text{LMO}(P_A \cdot \nabla W f)$$
   其中 $P_A$ 和 $P_B$ 是向 $A$ 和 $B$ 列空间的投影。
4. 谱特化： 对于谱范数，LMO 为 $-\text{msign}(\cdot)$。由此产生的因子更新（公式 14–15）涉及通过当前的 Gram 矩阵（$S_A = A^TA, S_B = B^TB$）对因子梯度进行白化，并应用谱符号函数：
   $$\Delta A^* = -\frac{\eta}{2} \text{msign}(\nabla_A f \cdot S_B^{-1/2}) S_B^{-1/2}$$
   $$\Delta B^* = -\frac{\eta}{2} \text{msign}(\nabla_B f \cdot S_A^{-1/2}) S_A^{-1/2}$$

2.3 分离权重衰减（Split Weight Decay）

标准的因子级权重衰减会导致全秩与低秩更新动力学之间的不匹配。作者推导了一种分离权重衰减规则，其中衰减因子 $s = \sqrt{1 - \lambda \eta}$ 以非对称方式应用：
$$A_{t+1} = s A_t + \frac{1}{s} \Delta A^*_t, \quad B_{t+1} = s B_t + \frac{1}{s} \Delta B^*_t$$
这确保了组合权重的更新与预期的全秩动力学相匹配，从而保持了超参数的可迁移性。

2.4 规范不变性（Gauge Invariance）

LoRA-Muon 的一个关键属性是规范不变性。组合权重 $W$ 上的更新对于任意可逆变换 $(A, B) \to (AR, BR^{-T})$ 保持不变。这是因为更新仅取决于投影器 $P_A$ 和 $P_B$，而投影器本身是规范不变的。

3. 核心贡献

1. LoRA-Muon 的推导： 一类新型优化器，通过在低秩流形上、在酉不变范数（特化为谱范数）下求解最陡峭下降问题而得出。
2. 分离权重衰减规则： 一种特定的更新规则，使低秩训练的动力学与全秩训练对齐，从而实现超参数迁移。
3. 规范不变性的证明： 提供了理论证明，表明 LoRA-Muon 的更新在全规范群作用下是不变的，确保优化器取决于诱导的权重更新，而非任意的因子坐标。
4. 与 Spectron 的比较： 证明了 Spectron 不是规范不变的，并且对任意因子缩放具有敏感性，这与 LoRA-Muon 不同。
5. 与 LoRA-RITE 的关系： 证明了 LoRA-RITE 简化的 QR 坐标核心在代数上与 LoRA-Muon 谱更新是完全一致的，尽管 LoRA-Muon 避免了 QR 分解和二阶矩存储的计算开销。
6. 实验验证： 证据表明 LoRA-Muon 可以作为全秩 Muon 学习率搜索的有效代理。

4. 实验结果

实验是在使用 TinyShakespeare 数据集的字符级语言建模上进行的，采用了计算量匹配的密集（dense）和低秩基准。

• 学习率迁移： 密集型 Muon 的最优学习率（$lr = 0.1$）在经过秩（2, 4, 8, 16, 32）、宽度、深度和标量因子缩放的扫略后，完美地迁移到了 LoRA-Muon 上。
  • 在 秩为 2 时，Lo-RA-Muon 恢复了密集型模型测试出的最佳学习率。
  • 在 秩为 32 时，在种子平均扫略中，LoRA-Muon 实现的平均验证损失（$1.776 \pm 0.002$）低于密集型基准（$1.789 \pm 0.002$）。
• 规范重平衡（Gauge Rebalancing）： 当因子以较大的标量不平衡（规模 9）进行初始化时，应用重平衡步骤对 LoRA-Muon 的影响微乎其微（损失变化 $\approx -5.1 \times 10^{-5}$），证实了其规范不变性。相比之下，Spectron 的性能发生了显著变化（$2.0 \times 10^{-2}$），证实了其对因子缩放的敏感性。
• 效率： 一个秩为 2 的代理模型仅使用密集模型约 4.3% 的可训练参数，同时在相同的计算约束下支持约 9.2 倍更大的优化器步长预算。

5. 重要性与主张

本文声称 LoRA-Muon 为 LoRA 的调优难题提供了一个有理论依据的解决方案。通过将 LoRA 视为低秩流形上的优化问题而非因子化坐标系，作者实现了：

• 鲁棒性： 优化器对于任意因子分解的选择具有不变性（规范对称性）。
• 可迁移性： 最优学习率（以及通过缩放法则实现的其它超参数）可以在秩、宽度、深度和因子规模之间无缝迁移。
• 效率： 它在不进行昂贵的 QR 分解或二阶矩存储的情况下计算谱更新，使其对加速器友好。

作者将 LoRA-Muon 定位为全秩 Muon 和 Shampoo 系列优化器的“良好低秩代理”。他们强调，低秩参数化不应决定优化器的几何结构；相反，更新应该由低秩矩阵流形本身的几何结构来定义。

局限性： 作者指出，目前的经验性主张是在 TinyShakespeare 规模的模型上建立的，尚未在大型预训练或下游微调基准上得到验证。该研究侧重于学习率迁移，而非完整的超参数扫略。