Papers

Rennala NSGD:超越欧氏几何的异步随机优化算法

Rennala NSGD: Asynchronous Stochastic Optimization Beyond Euclidean Geometry

Why we picked it

这篇论文把 Muon、Shampoo 这类优化器背后的非欧氏几何做了严格理论化,还给出半异步分布式版本的时间复杂度分析,做优化理论研究的话很值得细看。

arXiv 论文提出 Rennala NSGD,一个面向任意范数空间(非欧氏几何)的半异步分布式随机梯度下降算法。理论部分基于 κ-正则赋范空间框架,推导出高概率 oracle 复杂度界,复杂度依赖通过对偶空间的正则常数 κ 显式表达。在欧氏情形(κ=1)下,该界在相差对数因子内恢复标准 minibatch SGD 的速率,并将某些非欧氏情形下的维度依赖从线性降到对数级。论文还在固定且异构的 worker 计算时间下给出了 wall-clock 时间复杂度保证,呼应 Muon、Shampoo 等矩阵感知优化器的实践经验。

Original · arXiv cs.LG

Rennala NSGD: Asynchronous Stochastic Optimization Beyond Euclidean Geometry

The empirical success of matrix-aware optimizers, such as Muon (Jordan et al., 2024) and Shampoo (arXiv:1802.09568), highlights the potential of non-Euclidean geometry for deep learning, including large language model (LLM) training. Despite this progress, the theoretical understanding of stochastic optimization in general normed spaces remains limited, particularly in distributed settings. We analyze minibatch stochastic gradient descent (SGD) for smooth nonconvex optimization with respect to an arbitrary norm on a finite-dimensional space. Using the framework of $κ$-regular normed spaces (arXiv:0809.0813), we derive high-probability oracle complexity bounds under a light-tail condition stated directly in the dual norm induced by the update norm. The bounds make the dependence on geometry explicit through the regularity constant $κ$ of the dual space. In the Euclidean case, our bound recovers the standard minibatch SGD rate up to logarithmic factors. For some non-Euclidean geometries and noise models, it is sharper than bounds obtained by converting a Euclidean noise assumption, reducing a linear dimension dependence to a logarithmic one. Building on this framework, we introduce Rennala NSGD, a semi-asynchronous distributed algorithm, and establish wall-clock time complexity guarantees under fixed, heterogeneous worker computation times. In the Euclidean case ($κ=1$), our bound matches the parameter dependence of classical Rennala SGD up to logarithmic factors.