Rennala NSGD: Asynchronous Stochastic Optimization Beyond Euclidean Geometry¶

会议: NeurIPS 2026
arXiv: 2610.38773
领域: 优化/理论
关键词: 非欧几何优化、半异步并行SGD、双对偶空间正则性、轻尾噪声模型、运行时间复杂度

一句话总结¶

本文提出了针对一般赋范空间平滑非凸随机优化的半异步并行算法 Rennala NSGD，借助对偶空间的 $\kappa$-正则性框架在原生对偶范数轻尾噪声下建立了高概率预言复杂度界，并在异构计算延迟下给出了极小化极值形式的真实时间（wall-clock time）收敛保障。

研究背景与动机¶

随着深度学习规模的扩大，现代矩阵自感知优化器（如采用矩阵预条件的 Shampoo，以及采用正交化矩阵更新的 Muon 和 Dion）在 Transformer 大语言模型训练中展现出显著超越 AdamW 的经验收敛效率。从几何视角来看，这类更新本质上对应于特定非欧范数（如算子谱范数或各类 Schatten 范数）下的最速下降法（steepest descent），能够有效尊重注意力与前馈层权重矩阵固有的矩阵代数结构。

然而，非欧范数下的随机优化理论在分布式与并行计算场景下仍极其匮乏。现有针对非欧最速下降分析的核心理论痛点在于对随机梯度噪声（oracle noise）的刻画方式：绝大多数理论工作通常假定欧氏或 Frobenius 范数下的有界方差/矩条件，随后通过有限维空间中的范数等价性常数将分析“硬转换”至目标非欧对偶范数。这种转换通常不可避免地引入严重依赖模型参数维度甚至矩阵秩的高阶多项式惩罚因子，彻底掩盖了非欧几何更新在特定噪声结构下应有的理论优势。

与此同时，在分布式异构并行系统的实际工程中，不同计算节点（workers）的硬件能力与计算延迟往往存在显著差异。传统的同步小批量方法会受制于最慢的“掉队者（straggler）”，而经典的完全异步 SGD（ASGD）虽然摆脱了同步屏障，但其引入的过期梯度（stale gradients）会破坏无偏估计并导致次优的真实时间复杂度。本文的目标正是在统一的非欧几何与计算异构视角下解决上述核心矛盾。核心 idea：直接在更新范数所诱导的对偶范数下施加轻尾噪声条件，利用对偶空间的 $\kappa$-正则性建立紧凑的小批量集中不等式，进而构建仅接收当前迭代点最新梯度、丢弃过期梯度的半异步 Rennala NSGD 算法，消除欧氏转换引入的维度惩罚并取得匹配计算异构特性的最优时间复杂度。

方法详解¶

整体框架¶

算法在有限维赋范实向量空间 $(E, \|\cdot\|)$ 上优化无约束期望损失函数 $f(x) := \mathbb{E}_{\xi \sim \mathcal{D}}[f(x; \xi)]$，目标是找到满足 $\|\nabla f(x)\|_*^2 \le \varepsilon$ 的高概率 $\varepsilon$-驻点。整个优化流包含两个核心阶段：首先是基于尖锐算子（sharp operator）的非欧最速下降更新；其次是中心服务端与 $N$ 个具有固定异构计算时间（$\tau_1 \le \tau_2 \le \dots \le \tau_N$）的并行计算节点之间的半异步新鲜小批量收集协议。

在单次迭代中，各节点并行计算随机梯度。服务端并不等待所有节点完成，而是按到达先后严格收集 $B$ 个在当前迭代点 $x^k$ 处评估的“新鲜梯度”（带时间戳 $k' = k$），立即将先前迭代滞后计算出的过期梯度（$k' < k$）丢弃；一旦集齐 $B$ 个新鲜样本，服务端计算对偶空间平均梯度 $g^k$，应用尖锐算子 $(g^k)^\sharp \in \partial \phi^*(g^k)$ 进行非欧更新得到 $x^{k+1}$，并率先将新参数返回给触发该更新的节点，其余节点则在完成当前计算后接收最新参数。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：初始点 x0、步长 γk=1/(2L)、目标精度 ε、批大小 B"] --> B["阶段 1：半异步新鲜梯度收集<br/>服务端仅接纳带当前迭代戳 k'=k 的梯度，过滤并丢弃滞后过期梯度"]
    B --> C["阶段 2：对偶空间正则集中与对偶梯度平均<br/>在正则对偶空间内对 B 个独立轻尾噪声梯度进行无偏平均得到 gk"]
    C --> D["阶段 3：尖锐算子映射与非欧最速下降更新<br/>求解 LMO 支撑点生成 (gk)♯ 并执行 xk+1 = xk - γk(gk)♯"]
    D --> E["输出：高概率满足 ||∇f(xk)||*² ≤ ε 的优化参数解"]

关键设计¶

1. 基于尖锐算子与 LMO 的非欧最速下降更新：消除欧氏结构绑定

在一般的赋范空间中，目标函数的梯度 $\nabla f(x)$ 和随机梯度估计 $g^k$ 是对偶空间 $E^*$ 中的线性泛函，而模型参数增量必须位于原始参数空间 $E$ 中。为了在不依赖欧氏内积结构的前提下定义最速下降，作者定义凸势函数 $\phi(x) := \frac{1}{2}\|x\|^2$ 及其 Fenchel 共轭 $\phi^*(s) := \frac{1}{2}\|s\|_*^2$。尖锐算子（sharp selection）定义为一个 Borel 可测映射 $s \mapsto s^\sharp \in \partial \phi^*(s)$，严格满足几何相容性条件： $$\|s^\sharp\|^2 = \|s\|_*^2, \quad \langle s, s^\sharp \rangle = \|s\|_*^2$$ 在非欧小批量更新中，参数更新公式为 $x^{k+1} = x^k - \gamma_k (g^k)^\sharp$。该更新等价于在以 $x^k$ 为中心、半径为 $t_k = \gamma_k \|g^k\|_*$ 的原范数闭球上调用线性最小化预言机（LMO）：$x^{k+1} \in \operatorname{argmin}_{\|x - x^k\| \le t_k} \langle g^k, x \rangle$。该构造统一了多种典型算法：当 $\|\cdot\|$ 为 $\ell_2$ 范数时，$(g^k)^\sharp = g^k$，退化为标准 SGD；当 $\|\cdot\|$ 为 $\ell_\infty$ 范数时，对偶范数为 $\ell_1$，$(g^k)^\sharp = \|g^k\|_1 \operatorname{sign}(g^k)$，产生类似于 signSGD 的带符号更新；当 $\|\cdot\|$ 为矩阵谱范数 $\|\cdot\|_{\mathrm{op}}$ 时，对偶范数为核范数 $\|\cdot\|_{\mathrm{nuc}}$，标准选择给出矩阵极分解正交投影 $(G^k)^\sharp = \|G^k\|_{\mathrm{nuc}} \operatorname{orth}(G^k)$，精确刻画了 Muon 和 unregularized Shampoo 的几何更新本质。

2. 对偶空间的 $\kappa$-正则性与原生对偶轻尾集中：打破多项式维度爆炸

针对现有研究将欧氏噪声硬转为目标范数导致的维数爆炸痛点，作者将分析锚定于 Juditsky-Nemirovski 的 $\kappa$-正则空间框架。若对偶空间 $(E^*, \|\cdot\|_*)$ 存在一个辅助光滑范数 $\|\cdot\|_{*, +}$ 满足 $\|\cdot\|_*^2 \le \|\cdot\|_{*, +}^2 \le \frac{\kappa}{\kappa_+} \|\cdot\|_*^2$，则称该空间具有正则性常数 $\kappa$。作者直接在诱导对偶范数下提出轻尾假设： $$\mathbb{E}_{\xi}\left[\exp\left(\frac{\|\nabla f(x; \xi) - \nabla f(x)\|_*^2}{\sigma^2}\right)\right] \le \exp(1)$$ 基于此，作者证明了核心引理（Lemma 1）：对 $B$ 个独立同分布的轻尾噪声向量，其经验平均 $\bar{\eta} = \frac{1}{B}\sum_{j=1}^B \eta^j$ 满足非欧集中不等式： $$\mathbf{Prob}\left\{\|\bar{\eta}\|_* \ge (\sqrt{2\kappa} + \sqrt{2}\lambda)\frac{\sigma}{\sqrt{B}}\right\} \le \exp(-\lambda^2 / 3)$$ 进而通过在所有 $T = \lceil \frac{128 L \Delta}{15 \varepsilon} \rceil$ 步迭代中施加联合界，证明只需设定小批量尺寸 $B = \widetilde{\mathcal{O}}(1 + \kappa \sigma^2 / \varepsilon)$，即可保证整个优化过程以至少 $1-\delta$ 的概率单调下降，产生全域预言复杂度 $N_* = \widetilde{\mathcal{O}}\left(\frac{L\Delta}{\varepsilon} + \frac{\kappa L \Delta \sigma^2}{\varepsilon^2}\right)$。在稠密 Rademacher 噪声及 $\ell_1$ 原范数更新下，正则参数 $\kappa$ 仅为 $\mathcal{O}(\ln(ed))$，而传统的欧氏转换界包含 $\mathcal{O}(d)$ 因子，实现了从多项式维度依赖到对数维度的指数级压缩。

3. 半异步新鲜梯度收集协议：消除梯度过期偏差与异构计算掉队瓶颈

在分布式多节点计算中，完全异步方案会导致梯度在旧参数 $x^{k-\delta_k}$ 处被计算，造成严重的无偏性丢失与方差漂移。Rennala NSGD 继承并扩展了 Rennala SGD 的半异步思想：服务端为每次参数分发打上时间戳 $k$。当节点完成计算并返回梯度与时间戳 $(g, k')$ 时，服务端仅在 $k' = k$ 时将其计入当前批次累加器，若 $k' < k$ 则果断丢弃。更进一步，作者改进了服务端响应顺序：当收到完成当前批次最后一份梯度的消息时，服务端先执行参数更新 $x^{k+1} = x^k - \gamma_k (g^k)^\sharp$，再向该节点回复带最新戳 $k+1$ 的参数，使其零等待直接进入下一轮有效计算。对于其余慢速节点，它们在此期间被丢弃的滞后计算至多浪费自身一个周期的计算时间 $\tau_i$。该机制在数学上严格保证了每个被接纳的梯度在条件历史下均为当前迭代点 $x^k$ 处的无偏独立采样，从而将复杂的异步并行优化直接投影到串行高概率收敛的理论轨迹上。

实验关键数据¶

原论文是一篇严谨的理论优化工作，主要通过系统性的构造性理论推导、渐近复杂度严谨对比（Table 1 及附录 J）以及与已知下界的匹配性分析来验证理论优越性。

主实验：不同空间下统计预言机开销阶数对比¶

下表提取自论文 Table 1，展示了在维度为 $d$ 的空间中，不同更新几何所诱导的对偶空间 $(E^*, \|\cdot\|_*)$ 在两种典型稠密噪声模型（Rademacher 扰动与标准高斯扰动）下，原生对偶范数分析与欧氏/Frobenius 转换分析的统计预言项阶数对比（省略通用系数 $\nu^2$ 及确定性收敛首项 $\widetilde{\mathcal{O}}(L\Delta/\varepsilon)$）：

对偶空间 $(E^, \\|\cdot\\|_)$	对应原空间更新几何	欧氏/Frobenius转换分析	原生对偶范数分析（Rademacher噪声）	原生对偶范数分析（高斯噪声）
$(\mathbb{R}^d, \\|\cdot\\|_1)$	$\ell_\infty$ 范数（signSGD型）	$\widetilde{\mathcal{O}}(d^2 L\Delta / \varepsilon^2)$	$\widetilde{\mathcal{O}}(d^3 L\Delta / \varepsilon^2)$	$\widetilde{\mathcal{O}}(d^3 L\Delta / \varepsilon^2)$
$(\mathbb{R}^d, \\|\cdot\\|_2)$	$\ell_2$ 范数（标准 Euclidean SGD）	$\widetilde{\mathcal{O}}(d L\Delta / \varepsilon^2)$	$\widetilde{\mathcal{O}}(d L\Delta / \varepsilon^2)$	$\widetilde{\mathcal{O}}(d L\Delta / \varepsilon^2)$
$(\mathbb{R}^d, \\|\cdot\\|_\infty)$	$\ell_1$ 范数（贪心坐标下降型）	$\widetilde{\mathcal{O}}(d L\Delta / \varepsilon^2)$	$\widetilde{\mathcal{O}}(\ln(ed) L\Delta / \varepsilon^2)$	$\widetilde{\mathcal{O}}((\ln(ed))^2 L\Delta / \varepsilon^2)$
Schatten-$S_\infty$ 空间	矩阵核范数更新（Schatten-$S_1$）	$\widetilde{\mathcal{O}}(r L\Delta / \varepsilon^2)$	$\widetilde{\mathcal{O}}(\ln(er) L\Delta / \varepsilon^2)$	$\widetilde{\mathcal{O}}((\ln(er))^2 L\Delta / \varepsilon^2)$

消融与理论属性对比：计算时间复杂度与正则常数界¶

下表对比了在异构并行计算架构下，不同分布式优化范式在收敛保证、梯度过期容忍性以及真实墙上时间（wall-clock time）复杂度上的理论差异，并总结了 $\ell_p$ / Schatten-$S_p$ 空间的正则参数 $\kappa$ 边界：

算法机制 / 范式配置	梯度有效性保障	支持更新几何	理论时间复杂度上界 $T_{\mathrm{total}}$	关键瓶颈 / 优势说明
同步 Minibatch SGD	严格无偏（等待全量 $N$ 节点）	欧氏几何 ($\kappa=1$)	$\mathcal{O}\left(\tau_{\max} \left(\frac{L\Delta}{\varepsilon} + \frac{\sigma^2 L\Delta}{N \varepsilon^2}\right)\right)$	受制于最慢节点 $\tau_{\max}=\max_i \tau_i$，存在严重掉队木桶效应
传统异步 ASGD	存在过期梯度偏差 $\delta_k$	主要是欧氏几何	依赖于调和平均但包含延迟项及方差常数，非最优	引入偏差，且非欧更新下缺乏直接的高概率时间界
Rennala NSGD (Ours)	仅采新鲜梯度，丢弃过期残差	任意有限维赋范空间	$\widetilde{\mathcal{O}}\left(\min_{S \in [N]} \left(\frac{1}{S}\sum_{i=1}^S \frac{1}{\tau_i}\right)^{-1} \left(\frac{L\Delta}{\varepsilon} + \frac{\kappa L\Delta \sigma^2}{S \varepsilon^2}\right)\right)$	自适应平衡最快 $S$ 个节点的调和平均计算耗时与方差缩减，达极小化最优
$\ell_p$ 空间正则性常数 ($1 \le p < 2$)	采用加权欧氏辅助范数	$\ell_p$ / Schatten-$S_p$	最优常数 $\kappa = d^{\frac{2}{p}-1}$（矩阵为 $r^{\frac{2}{p}-1}$）	存在多项式维度放大，原生分析在 $p=1$ 时劣于欧氏转换
$\ell_p$ 空间正则性常数 ($2 \le p \le \infty$)	采用最优调优 $\ell_\rho$ 辅助范数	$\ell_p$ / Schatten-$S_p$	最优常数 $\kappa = \Theta(\min\{p, \ln(ed)\})$	对数级维数缩放，在 $p=\infty$ 时实现极大幅度的方差收敛增益

关键发现¶

几何原生分析并非无条件占优：在对偶 $\ell_\infty$ 空间（对应 $\ell_1$ 范数更新或矩阵谱更新）下，原生分析将多项式维数惩罚 $d$ 降至 $\ln(ed)$，带来极显著提升；然而在对偶 $\ell_1$ 空间下，由于 $\kappa = d$，原生统计系数反倒达到 $\widetilde{\mathcal{O}}(d^3)$，劣于欧氏转换的 $\widetilde{\mathcal{O}}(d^2)$。这表明选择哪种噪声假设必须紧密结合实际问题的几何与扰动结构。
异构加速比的最优权衡：Rennala NSGD 的时间界展现出对前 $S$ 个最快工作节点的自适应权衡能力。当网络包含极端慢节点时，公式中的 $\min_{S \in [N]}$ 会自动将慢节点从计算队列中“逻辑剔除”，使系统运行时间仅由最快子集的调和平均速度主导，完全免疫掉队者阻塞。
欧氏情形下的理论闭环：当退化为欧氏范数（$\kappa = 1$）时，该算法的高概率复杂度在对数因子意义下完全恢复了 Tyurin & Richtárik (2023) 针对经典 Rennala SGD 在零占优算法类中所证明的极小化极大最优时间界。

亮点与洞察¶

直接在诱导对偶范数下度量随机噪声：避免了传统分析依赖范数等价性不等式导致的维数惩罚，首次将非欧矩阵优化算法（如 Muon/Shampoo 理论内核）的分析建立在原生几何度量上。
巧用 $\kappa$-正则空间与辅助光滑范数：将不可微的非欧范数映射到光滑辅助范数上进行鞅差序列集中不等式推导，既保证了高概率大偏差界的严格性，又使最终常数仅由对偶空间正则常数 $\kappa$ 控制。
更新优于回复（Update-before-reply）的微协议改进：在服务端接收到满足批量的最后一个梯度时，先执行尖锐更新再向该节点回复新参数，消除了慢节点等待时间，维持了并行更新下的无偏性。

局限与展望¶

计算模型未考虑尖锐算子本身的耗时：对于矩阵谱范数更新，求尖锐算子需要对大矩阵进行紧凑 SVD 或极分解，服务端计算开销不可忽略；论文将通信及服务端处理简化为瞬时完成，未纳入时间复杂度核算。
统计同质性假设与联邦异质性脱节：算法要求所有节点共享相同的数据采样分布 $\mathcal{D}$；若存在各节点数据异质性（non-IID），优先采纳快节点的策略将导致严重的模型漂移与抽样偏差。
高概率界依赖较强的轻尾假设：相比于常见的有界方差弱假设，指数矩轻尾条件更为苛刻，重尾分布下的非欧分布式优化仍有待进一步拓展。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次建立了一般非欧空间下半异步分布式 SGD 的高概率预言与真实时间复杂度理论，给出了紧凑的正则常数推导。
实验充分度: ⭐⭐⭐☆☆ 作为纯理论工作，包含了完备的引理体系与多维几何复杂度阶数对比，但缺乏在大规模分布式集群训练 Transformer 上的经验实测。
写作质量: ⭐⭐⭐⭐⭐ 数学论证环环相扣，从对偶几何、尖锐算子到集中不等式与异构排队调度的逻辑链条极其清晰严谨。
价值: ⭐⭐⭐⭐ 为 Muon、Dion 等现代非欧架构的异构分布式训练提供了坚实的理论基石，具有重要的方法学指导意义。

对偶空间 \((E^, \\|\cdot\\|_)\)	对应原空间更新几何	欧氏/Frobenius转换分析	原生对偶范数分析（Rademacher噪声）	原生对偶范数分析（高斯噪声）
\((\mathbb{R}^d, \\|\cdot\\|_1)\)	\(\ell_\infty\) 范数（signSGD型）	\(\widetilde{\mathcal{O}}(d^2 L\Delta / \varepsilon^2)\)	\(\widetilde{\mathcal{O}}(d^3 L\Delta / \varepsilon^2)\)	\(\widetilde{\mathcal{O}}(d^3 L\Delta / \varepsilon^2)\)
\((\mathbb{R}^d, \\|\cdot\\|_2)\)	\(\ell_2\) 范数（标准 Euclidean SGD）	\(\widetilde{\mathcal{O}}(d L\Delta / \varepsilon^2)\)	\(\widetilde{\mathcal{O}}(d L\Delta / \varepsilon^2)\)	\(\widetilde{\mathcal{O}}(d L\Delta / \varepsilon^2)\)
\((\mathbb{R}^d, \\|\cdot\\|_\infty)\)	\(\ell_1\) 范数（贪心坐标下降型）	\(\widetilde{\mathcal{O}}(d L\Delta / \varepsilon^2)\)	\(\widetilde{\mathcal{O}}(\ln(ed) L\Delta / \varepsilon^2)\)	\(\widetilde{\mathcal{O}}((\ln(ed))^2 L\Delta / \varepsilon^2)\)
Schatten-\(S_\infty\) 空间	矩阵核范数更新（Schatten-\(S_1\)）	\(\widetilde{\mathcal{O}}(r L\Delta / \varepsilon^2)\)	\(\widetilde{\mathcal{O}}(\ln(er) L\Delta / \varepsilon^2)\)	\(\widetilde{\mathcal{O}}((\ln(er))^2 L\Delta / \varepsilon^2)\)

算法机制 / 范式配置	梯度有效性保障	支持更新几何	理论时间复杂度上界 \(T_{\mathrm{total}}\)	关键瓶颈 / 优势说明
同步 Minibatch SGD	严格无偏（等待全量 \(N\) 节点）	欧氏几何 (\(\kappa=1\))	\(\mathcal{O}\left(\tau_{\max} \left(\frac{L\Delta}{\varepsilon} + \frac{\sigma^2 L\Delta}{N \varepsilon^2}\right)\right)\)	受制于最慢节点 \(\tau_{\max}=\max_i \tau_i\)，存在严重掉队木桶效应
传统异步 ASGD	存在过期梯度偏差 \(\delta_k\)	主要是欧氏几何	依赖于调和平均但包含延迟项及方差常数，非最优	引入偏差，且非欧更新下缺乏直接的高概率时间界
Rennala NSGD (Ours)	仅采新鲜梯度，丢弃过期残差	任意有限维赋范空间	\(\widetilde{\mathcal{O}}\left(\min_{S \in [N]} \left(\frac{1}{S}\sum_{i=1}^S \frac{1}{\tau_i}\right)^{-1} \left(\frac{L\Delta}{\varepsilon} + \frac{\kappa L\Delta \sigma^2}{S \varepsilon^2}\right)\right)\)	自适应平衡最快 \(S\) 个节点的调和平均计算耗时与方差缩减，达极小化最优
\(\ell_p\) 空间正则性常数 (\(1 \le p < 2\))	采用加权欧氏辅助范数	\(\ell_p\) / Schatten-\(S_p\)	最优常数 \(\kappa = d^{\frac{2}{p}-1}\)（矩阵为 \(r^{\frac{2}{p}-1}\)）	存在多项式维度放大，原生分析在 \(p=1\) 时劣于欧氏转换
\(\ell_p\) 空间正则性常数 (\(2 \le p \le \infty\))	采用最优调优 \(\ell_\rho\) 辅助范数	\(\ell_p\) / Schatten-\(S_p\)	最优常数 \(\kappa = \Theta(\min\{p, \ln(ed)\})\)	对数级维数缩放，在 \(p=\infty\) 时实现极大幅度的方差收敛增益