Constrained Stochastic Spectral Preconditioning Converges for Nonconvex Objectives¶

会议: NeurIPS 2026
arXiv: 2605.11850
领域: 优化/理论
关键词: 随机谱预条件, 近端梯度法, 非凸约束优化, 重尾噪声, 矩阵符号函数

一句话总结¶

针对深度学习中 Muon 与 Scion 等谱梯度优化器缺乏通用约束与正则化支持的理论瓶颈，本文提出了一套结合非线性谱对偶预条件与各向异性近端算子的随机近端预条件梯度算法框架，首次证明了在重尾噪声和非凸约束下向驻点的自适应收敛性，并严格刻画了 Polar Express 等多项式迭代的实际收敛行为。

研究背景与动机¶

近年来以 Muon、Scion 及谱下降法为代表的随机谱梯度优化器在深度神经网络（尤其是大语言模型和 Transformer 架构）的训练中展现出远超标准一阶动量法的收敛效率与稳定性。这类优化器区别于传统随机梯度算法的核心机制在于：它们不对含噪梯度向量直接施加欧氏更新，而是通过奇异值分解或其多项式迭代逼近矩阵符号函数（$\mathrm{sign}(D) = U V^\top$），从几何层面上对更新步幅进行自适应谱归一化。在直观上，这种操作极大地抑制了梯度爆炸并能自适应于更宽泛的非各向同性平滑性，因而与梯度裁剪和符号梯度法共同构成了自适应谱归一化优化器族。

然而，现有针对此类谱优化器的理论分析几乎无一例外局限于单一光滑目标的无约束极小化问题。在实际深度学习训练中，权重约束与结构化显式正则化扮演着不可或缺的角色——例如为保证网络 Lipschitz 连续性而施加的谱范数球约束、正交性要求的 Stiefel 流形约束，以及权重范数控制等。当试图将约束引入谱梯度方法时，现有文献要么仅能处理极特殊的凸紧集指示函数（借助 Frank-Wolfe / 线性极小化预言机 LMO 机制，如 Scion），要么依赖欧氏近端算子；但后一种简单结合会导致算法的不动点与原约束优化问题的真实一阶驻点完全脱节，破坏优化的一致性。

该挑战的核心矛盾在于：矩阵符号函数等非线性谱预条件操作是在对偶空间中重塑梯度的几何结构，而传统的欧氏近端步（Euclidean Prox）无法在非线性变换后保持对偶测度与次微分锥结构的一致性，导致常规近端项退化；加之深度模型训练中随机梯度的重尾噪声（Heavy-tailed Noise）特性与非线性变换引入的随机偏差交织，使得非凸约束下的收敛性分析极其困难。本文的切入角度是借助凸分析中的广义各向异性近端映射（Anisotropic Proximal Mapping）与 Bregman 散度对偶测度，将前向谱预条件更新与后向各向异性近端投影有机解耦。核心 idea：建立基于参考函数凸共轭的前向非线性谱预条件与各向异性后向近端算子两步迭代架构，构建几何匹配的对偶 Bregman 间隙测度，首次在无需预知光滑常数与噪声阶数的条件下，证明了带一般非凸约束/正则项及重尾噪声的随机谱近端梯度法的次线性收敛性，并揭示了多项式矩阵符号近似本质上等价于非线性预条件系统。

方法详解¶

整体框架¶

本文考虑非凸复合优化目标函数 $\min_{x \in E} F(x) := f(x) + g(x)$，其中 $f \in \mathcal{C}^1(E)$ 为可能非凸但满足 Lipschitz 光滑（或各向异性光滑）的损失函数，而 $g: E \to \overline{\mathbb{R}}$ 为可能非光滑且非凸的正则项或闭集约束指示函数 $\delta_C$（如 Stiefel 流形、谱球或谱球面）。

算法在每一步迭代 $k$ 包含两个基本阶段：首先在对偶空间中利用强凸参考函数 $\phi$ 的凸共轭梯度 $\nabla \phi^*$ 对（随机）动量搜索方向 $d^k$ 进行非线性谱预条件变换，得到原始空间中的前向试验点 $y^k$；若该试验点脱离了约束可行域或破坏了正则化结构，则执行后向各向异性近端步（Backward Anisotropic Proximal Step），通过求解由 episcaling 尺度化参考函数 $(\gamma_k \star \phi)$ 诱导的非欧近端子问题投影回可行域，生成新迭代点 $x^{k+1}$。

%%{init: {'theme': 'base', 'themeVariables': {'primaryColor': '#EEF2FF', 'primaryBorderColor': '#6366F1', 'lineColor': '#4F46E5', 'fontSize': '13px'}, 'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：随机梯度估计与当前模型参数"] --> B["阶段 1：自适应动量估计器<br/>Polyak 动量或 STORM 方差缩减估计器更新方向 d^k"]
    B --> C["阶段 2：对偶空间谱预条件前向步<br/>利用共轭梯度计算前向更新点 y^k = x^k - γ_k ∇ϕ*(d^k)"]
    C --> D["阶段 3：各向异性近端后向步<br/>通过奇异值分解将非凸约束/正则项转化为一维近端映射并投影"]
    D --> E["输出：收敛至临界驻点的新参数 x^{k+1}"]

关键设计¶

1. 非线性共轭谱预条件与矩阵符号函数统一：消除前向更新的尺度敏感性 传统动量 SGD 与非线性自适应优化器之间的割裂，本质上源于更新方向映射函数的不同。本文引入定义在欧氏空间或实正交不变矩阵空间上的下半连续、$\mu$-强凸偶参考函数 $\phi$，其凸共轭函数 $\phi^*$ 处处可微且 $\nabla \phi^*$ 严格单调、奇对称。在前向步骤中，算法计算： $$y^k = x^k - \gamma_k \nabla \phi^*(d^k)$$ 对于矩阵参数 $D^k \in \mathbb{R}^{m \times n}$（设奇异值分解为 $D^k = U \mathrm{Diag}(\sigma) V^\top$），通过构建谱各向异性参考函数 $\Phi_{\mathrm{aniso}}(X) = \sum_{i=1}^{\min(m,n)} h(\sigma_i(X))$，其共轭梯度呈现为奇异值维度的逐元非线性缩放： $$\nabla \Phi_{\mathrm{aniso}}^*(D^k) = U \, \mathrm{Diag}\left({h^*}'(\sigma_1), \dots, {h^*}'(\sigma_q)\right) V^\top$$ 当选取近似障碍函数 $h(t) = -\varepsilon (\ln(1 - |t|) + |t|)$ 时，${h^*}'(\sigma_i) = \frac{\sigma_i}{\varepsilon + \sigma_i}$。当平滑参数 $\varepsilon \to 0^+$ 时，前向步精确恢复为 Muon 中采用的理想化矩阵符号更新 $U V^\top$；而当 $\varepsilon > 0$ 时，该构造天然消除了零奇异值处符号函数的不连续性，为多项式逼近提供了更真实的数学映射。

2. 广义各向异性近端映射：保持非凸约束不动点与一阶驻点的一致性 若直接将欧氏投影算子 $P_C(y^k)$ 作用于经谱预条件处理的前向点，所得不动点系统仅满足 $-\nabla \phi^*(\nabla f(x^*)) \in \partial g(x^*)$，因非线性映射与次微分加法的不兼容性，该不动点通常与原问题驻点 $0 \in \nabla f(x^*) + \partial g(x^*)$ 并不等价。为此，本文采用由 episcaling $(\gamma_k \star \phi)(z) = \gamma_k \phi(\gamma_k^{-1} z)$ 诱导的各向异性近端步： $$x^{k+1} \in \arg\min_{x \in E} \left\{ g(x) + (\gamma_k \star \phi)(x - y^k) \right\}$$ 根据最优性一阶条件，该极小化问题解满足 $-\nabla \phi\left(\frac{x^{k+1} - y^k}{\gamma_k}\right) \in \partial g(x^{k+1})$。由于 $\nabla \phi$ 与 $\nabla \phi^*$ 互为逆映射，代入 $y^k$ 定义即得 $x^{k+1} - x^k = -\gamma_k \nabla \phi^*(d^k) - \gamma_k \nabla \phi^*(\widetilde{\nabla} g(x^{k+1}))$。当且仅当算法收敛达到不动点 $x^{k+1} = x^k$ 时，严格单调性保证了 $-\nabla f(x^*) = \widetilde{\nabla} g(x^*)$，严格保持了一阶变分几何的一致性。

对于实正交不变的正则项或约束集合（如 Stiefel 流形 $\{X \mid X^\top X = r^2 I\}$、谱球及谱球面），本文通过绝对对称函数与弱优序（Weak Majorization）工具证明了矩阵近端问题的奇异值解耦定理：矩阵近端步可精确转化为奇异值向量上的一维各向异性近端标量求解，从而避免了高维耦合非凸优化问题，给出了显式或易于二分求解的高效闭式实现。

3. 对偶 Bregman 散度驻点度量与重尾噪声自适应收敛分析 在非凸随机优化与非线性几何下，传统的梯度范数无法直接度量复合约束的下降幅度。本文设计了与算法内在几何匹配的对偶 Bregman 间隙函数： $$\mathrm{gap}(x^k) := D_{\phi^*}(\nabla f(x^k), -\widetilde{\nabla} g(x^k)) = \phi^*(\nabla f(x^k)) - \phi^*(-\widetilde{\nabla} g(x^k)) - \langle \nabla \phi^*(-\widetilde{\nabla} g(x^k)), \nabla f(x^k) + \widetilde{\nabla} g(x^k) \rangle$$ 由于 $\phi^*$ 的严格凸性，$\mathrm{gap}(x^k) \ge 0$ 且该度量为零当且仅当系统处于一阶驻点。针对神经网络梯度中广泛存在的 $p$ 阶中心矩有界重尾噪声（$\mathbb{E}[\|\nabla f(x,\xi) - \nabla f(x)\|^p] \le \sigma^p, p \in (1, 2]$），本文借助动量估计器序列在鞅差差分序列（Martingale Difference Sequence）下的广义 Pinelis 概率不等式，完全摆脱了对全梯度光滑常数的先验依赖。对于 Polyak 动量（参数 $\alpha = (K+1)^{-1/2}, \gamma = \bar{\gamma}(K+1)^{-3/4}$），证明了自适应收敛速率： $$\frac{1}{K+1}\sum_{k=0}^K \mathbb{E}[\mathrm{gap}(x^{k+1})] \le \mathcal{O}\left((K+1)^{-\min\left\{\frac{1}{4}, \frac{p-1}{2p}\right\}}\right)$$ 在标准有限方差假设（$p=2$）及逐样本光滑条件下，进一步结合 STORM（STochastic Recursive Momentum）方差缩减估计器，将收敛速率大幅提升至 $\widetilde{\mathcal{O}}((K+1)^{-1/3})$。

4. Polar Express 多项式近似的非线性预条件本质刻画 实际中 Muon 依赖 Newton-Schulz 或 Polar Express 等奇次多项式迭代来逼近矩阵符号。以往理论强行将其视作理想矩阵符号并引入有界截断误差分析，忽视了多项式在原点附近天然趋于 0 而非 $\pm 1$ 的本质特征。本文证明，Polar Express 结合前置 Frobenius 归一化步的实际计算流，实质上等价于一类基于超几何函数导出的高阶广义非线性预条件子： $$\nabla \phi^*(y)_i = \frac{y_i}{(\varepsilon^\kappa + |y_i|^\kappa)^{1/\kappa}}, \quad \kappa=4$$ 在输入归一化方向 $d^k / (\|d^k\| + \epsilon)$ 作用下，本文克服了传统下降引理在非单调多项式逼近下的失效，建立了直接针对实际工程实现算法的非渐进收敛界：$\frac{1}{K+1}\sum_{k=0}^K \mathbb{E}[\|\nabla f(x^k)\|] \le \mathcal{O}((K+1)^{-\min\{1/4, \frac{p-1}{2p}\}})$，弥合了工业级底层算子实现与收敛理论之间的鸿沟。

损失函数 / 训练策略¶

算法优化策略针对随机环境提供了两种动量配置方案： - 方案 A（Polyak 动量，抗重尾噪声）：动量估计器为 $d^k = \alpha \nabla f(x^k, \xi^k) + (1-\alpha) d^{k-1}$，取静态超参数 $\alpha = (K+1)^{-1/2}$，步长 $\gamma = \bar{\gamma}(K+1)^{-3/4}$，能够自适应耐受任意 $p \in (1, 2]$ 的重尾扰动。 - 方案 B（STORM 方差缩减动量，高精度收敛）：动量迭代式为 $d^k = (1-\alpha_k) d^{k-1} + \alpha_k \nabla f(x^k, \xi^k) + (1-\alpha_k)(\nabla f(x^k, \xi^k) - \nabla f(x^{k-1}, \xi^k))$，动态参数设为 $\alpha_k = (k+1)^{-2/3}, \gamma_k = (k+1)^{-2/3}$，实现无偏差梯度累积。

实验关键数据¶

主实验¶

论文在三个代表性深度学习基准任务上验证了所提约束谱预条件优化框架的有效性与收敛行为：模块化算术（DivMod97 Grokking 泛化任务）、CIFAR-10 卷积网络分类，以及 1.24 亿参数的 NanoGPT 预训练。

下表汇总了 CIFAR-10 图像分类与 124M NanoGPT 在 10,200 次迭代下的验证困惑度/精度主实验对比（对照无约束 Scion 与施加谱球、Frobenius 球约束的近端谱方法）：

数据集 / 任务	评估指标	本文近端谱方法 (Ours + 约束)	无约束基线 (uScion / AdamW)	性能增益 / 相对提升
CIFAR-10 (CNN)	测试准确率 (%)	78.4% (奇异值截断谱球) / 78.2% (Frobenius 球近端)	75.8% (uScion 基准)	+2.6% / +2.4%
NanoGPT 124M (10.2k 步)	验证集交叉熵 Loss / 困惑度	3.27 (Frobenius 球约束谱预条件)	3.32 (uScion + Polar Express)	验证损失下降 -0.05
DivMod97 (Transformer)	达到 95% 泛化精度所需 Epochs (中位数)	180 (谱球约束 Scion)	360 (AdamW 无约束) / 240 (uScion)	泛化加速 2.0× (vs AdamW) / 1.33× (vs uScion)

消融实验¶

下表给出了在 DivMod97 算法推理泛化任务中，针对谱层施加不同约束策略及预条件机制对泛化（Grokking）速度与训练稳定性的消融对比：

配置	到达 95% 验证精度 Epochs (中位数)	最大稳定初始学习率	机制说明
Full model (谱层奇异值裁剪约束 + 谱预条件)	180	0.00036	完整模型：谱范数球约束强力限制权重半径，大幅拓宽稳定学习率区间
替换为 Frobenius 球约束 (Frobenius-ball)	205	0.00032	采用一维二次多项式根式映射替代精确 SVD 截断，兼顾速度与控制
无约束 Scion (uScion, 无后向投影)	240	0.00026	缺少后向近端投影，权重范数持续发散导致 Grokking 延迟
标准 AdamW (强权重衰减 baseline)	360	0.00018	依赖经验性权重衰减，缺乏对谱几何边界的主动约束

关键发现¶

权重视谱范数控制加速泛化（Grokking）：在模块化算术任务中，对 Transformer 线性注意力层施加严格的谱范数边界，使得模型进入算法泛化阶段所需的训练 Epochs 减少达 50%，同时使优化器能够承受 2 倍以上的初始学习率而不会发生数值发散。
近端后向步在保持计算效率的同时保证稳定性：实验表明，采用各向异性近端步投影至 Frobenius 球或谱球，能够完全复现甚至超越精细调优权重衰减的效果，且实验曲线证明利用一维标量多项式求根求解的 Frobenius 近端步在卷积神经网络上几乎不增加总体墙钟训练时间。

亮点与洞察¶

重塑 Muon / Scion 的理论基石：揭示了工程中广泛采用的 Polar Express 多项式迭代并非“带误差的符号函数”，而严格对应于超几何势函数诱导的非线性共轭预条件算子。这一视角将现代谱优化器的经验性技巧彻底纳入了变分分析与各向异性凸分析的严密理论版图。
解耦的前向-后向两阶段设计：将复杂的谱规范化（前向对偶变换）与流形约束/范数截断（后向近端步）清晰解耦，并证明了正交不变矩阵约束可降维为奇异值标量方程，为后续针对大模型开发自适应正交化优化器提供了通用的即插即用工具箱。
对重尾噪声的天然鲁棒性：通过利用参考函数共轭梯度的有界性（$\|\nabla \phi^*\| \le D$），从理论上天然锁定了每次迭代的移动步幅上界，使得算法无须引入繁琐的人工梯度裁剪阈值超参即可对非高斯、重尾分布的训练噪声免疫。

局限与展望¶

精确 SVD 计算开销：尽管理论提供了针对任意实正交不变约束的一维降维求解方案，但在极大规模 Transformer 参数矩阵上频繁执行精确奇异值分解的计算与显存开销依然显著；未来需进一步探索带误差界的随机化 SVD 或低秩近似算子的近端相容性。
多层解耦与自适应步长调度：目前理论假定所有网络层共享统一的各向异性势函数缩放，未针对网络深浅层梯度方差的不均衡性构建逐层自适应半径的理论闭环。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次建立统一的约束随机谱预条件近端框架，并从理论上澄清了 Polar Express 的非线性本质。
实验充分度: ⭐⭐⭐⭐☆ 覆盖了 Grokking 算法任务、CIFAR-10 及 124M NanoGPT，兼顾理论验证与预训练场景。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，凸分析与变分分析工具运用精湛，理论脉络清晰。
价值: ⭐⭐⭐⭐⭐ 为 Muon / Scion 等前沿大模型优化器的理论演进与约束扩展奠定了坚实的理论基石。