OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality

技术摘要："OptMuon：具有零噪声最优性的随机优化闭环正交化动量法"

1. 问题陈述

本文研究了矩阵值非凸目标函数的随机一阶优化问题：
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \mathbb{E}_{\xi \sim \mathcal{D}}[f(X; \xi)]$$
目标是在 $T$ 次随机梯度调用内，找到一个随机迭代值 $X_{out}$，使得期望 Frobenius 范数梯度 $\mathbb{E}[\|\nabla f(X_{out})\|_F]$ 最小化。

虽然 Muon 式优化器通过将动量矩阵转化为近似正交方向（通过极分解）展示了经验稳定性，但现有方法通常将这些方向与开环或恒定幅度的规则相结合。这些规则通常依赖于预设的平滑常数 ($L$) 或方差水平 ($\sigma^2$)，这使得它们对超参数调优非常敏感，且无法根据观察到的优化轨迹显式校准更新幅度。相反，像 STORM 和 AdaGrad-Norm 这样的自适应方法提供了“闭环”标量自适应和噪声自适应，但它们通常作用于原始梯度而非正交化后的动量方向。

核心挑战在于：能否将闭环标量自适应与 Muon 式正交化动量成功结合，从而在保持几何稳定性的同时，实现噪声自适应和零噪声最优性，而无需手动重新调优。

2. 方法论：OptMuon

作者提出了 OptMuon 框架，该框架将更新方向与更新幅度解耦。

• 方向（正交化）： 更新方向是动量矩阵 $M_t$ 的极因子，记作 $\text{Orth}(M_t) = U W^\top$（其中 $M_t = U \Sigma W^\top$ 是其 SVD 分解）。这确保了方向经过几何归一化，并将方向与奇异值缩放分离。
• 幅度（闭环自适应）： 步长由基于观测轨迹（梯度和动量历史）得出的标量系数 $\gamma_t$ 决定，而非固定的问题常数。
  • 系数调度： 幅度控制的核心是一个运行最大值修正的 AdaGrad-Norm 调度。它在分母中累积随机梯度能量以抑制噪声较大的阶段，但在分子中包含了一个项 ($\max g_i^2$)，以防止系数因孤立的梯度尖峰而过度坍缩。
  • 无 Lipschitz 常数要求： 该算法不需要了解平滑常数 $L$、方差 $\sigma^2$ 或有界梯度常数 $\hat{G}$ 来进行参数选择，尽管这些常数出现在理论公式中。

两个变体：

1. OptMuon-A（平均平滑度）： 专为平均平滑假设设计。它在每次迭代中使用一个随机梯度，并采用标准动量累积 ($M_t = G_t + (1-\alpha_t)M_{t-1}$)。其调度的指数为 $q=1/2$。
2. OptMuon-I（个体平滑度）： 专为个体平滑假设设计。它采用了类 STORM 的递归动量估计器 ($M_t = (1-\alpha_t)(M_{t-1} - \nabla f(X_{t-1}; \xi_t))$)，利用在同一样本上的两次梯度评估来降低方差。其指数为 $q=2/3$。

两个变体的参数更新方式均为：
$$X_{t+1} = X_t - \theta \gamma_t \|M_t\|_F \text{Orth}(M_t)$$

3. 核心贡献

论文声称了以下具体贡献：

• 闭环动量正交化： 一个清晰地将矩阵几何（通过极因子）与幅度校准（通过自归一化标量调度）分离的框架。
• 运行最大值修正调度： 一种新颖的系数设计，它既保留了 AdaGrad-Norm 的自归一化抑制特性，又减轻了由单个梯度尖峰导致的系数坍缩风险。
• 噪声自适应收敛率： 理论保证显示，随着方差 $\sigma$ 的减小，收敛界限中的随机项会自动改善。
• 零噪声最优性： 证明了当 $\sigma = 0$ 时，界限会自动降低到近乎最优的确定性速率 $\tilde{O}(T^{-1/2})$，且无需任何超参数重调。

4. 理论结果

在满足下界性、无偏随机梯度且方差有界、平滑性（平均或个体）以及几乎处处有界随机梯度等假设下，本文证明了：

• OptMuon-A（平均平滑度）： 实现收敛率为：
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/2}T^{-1/4}\right)$$
• OptMuon-I（个体平滑度）： 实现更精确的随机速率：
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/3}T^{-1/3}\right)$$

结果的意义：

• 零噪声恢复： 在零噪声情形下 ($\sigma = 0$)，两个界限均简化为 $\tilde{O}(T^{-1/2})$，达到了最优的确定性一阶速率。这与以往的 Muon 分析形成对比，后者往往产生 $\tilde{O}(T^{-1/4})$ 或需要固定步长。
• 噪声自适应性： 随机项 ($\sigma^{1/2}T^{-1/4}$ 和 $\sigma^{1/3}T^{-1/3}$) 是自适应的；随着噪声水平降低，它们会自动收紧，这与使用固定步长的算法不同。
• 对比： 论文指出，在低噪声环境下，OptMuon-A 改进了近期方差缩减型 Muon 变体（如 Chang et al., 2025）所建立的 $\tilde{O}(T^{-1/4})$ 速率，同时保持了零噪声最优性。

5. 重要性与局限性

重要性：
本文确立了闭环标量自适应可以与 Muon 式正交化动量相结合，且不会牺牲几何稳定性。它将“零噪声自适应性”（此前发现于标量自适应步长规则中）的概念扩展到了极因子动量更新中。这使得优化过程既稳健，又能自动调整以适应噪声水平，并在无需人工干预的情况下恢复最优的确定性性能。

局限性（如作者所述）：

• 精确正交化： 理论假设精确计算极因子。虽然 Newton-Schulz 迭代可以近似实现，但目前的分析并未量化不精确正交化所引入的误差。
• 有界梯度： 分析依赖于几乎处处有界的随机梯度条件（假设 4）。本文并未声称拥有无界梯度的理论。
• 经验验证： 本文侧重于理论收敛保证，未包含大规模经验评估或向复合/约束目标的扩展，这些被视为未来的研究方向。