Rennala NSGD: Asynchronous Stochastic Optimization Beyond Euclidean Geometry

AI summary

Overview

• Research area: Optimization theory for machine learning — nonconvex stochastic gradient methods in general normed spaces, and distributed/asynchronous parallel optimization.
• Technical level: Advanced. The paper is a theory contribution built on $\kappa$-regular normed spaces, martingale concentration inequalities, and worst-case complexity analysis; there are no experiments in the content provided.
• Scope (one sentence): The paper derives high-probability oracle-complexity bounds for minibatch SGD under a smoothness/noise model stated in the dual of an arbitrary update norm, and uses them to prove wall-clock time guarantees for a new semi-asynchronous distributed method, Rennala NSGD, when workers have fixed but different computation speeds.

What This Paper Is About

Modern deep-learning optimizers such as Muon and Shampoo are "geometry-aware": their updates correspond to steepest descent in non-Euclidean norms (for example the spectral norm on weight matrices) rather than in the usual Euclidean norm. The theory of stochastic optimization in such geometries is underdeveloped, especially in distributed settings, and most existing analyses bound the gradient noise in a Euclidean or Frobenius norm and then convert to the target norm, which can inflate the bound by dimension-dependent factors. This paper asks whether stating the noise assumption directly in the dual norm induced by the update norm gives sharper bounds, and combines that answer with a time-aware analysis of asynchronous training on workers with heterogeneous speeds.

Key Contributions

1. Minibatch non-Euclidean SGD under a dual-norm light-tail oracle. The authors prove high-probability oracle-complexity bounds for minibatch SGD that is L-smooth and lower bounded with respect to an arbitrary norm, with unbiased noise satisfying a light-tail condition stated in the induced dual norm. The bound is $\tilde{\mathcal{O}}(L\Delta/\varepsilon + \kappa L\Delta\sigma^2/\varepsilon^2)$ stochastic gradient evaluations to reach $\min_{0\le k\le T-1}|\nabla f(x^k)|_*^2 \le \varepsilon$ with probability at least $1-\delta$, where $\kappa$ is the regularity constant of the dual space, $\Delta = f(x^0) - f^{\inf}$, and $\sigma$ is the dual-norm noise scale. The technical core is a new minibatch concentration inequality for fresh i.i.d. samples in regular dual spaces.

2. Sharp regularity bounds and a comparison with norm-equivalence conversions. The paper gives matching upper and lower bounds on the optimal $\kappa$-regularity parameter for $\ell_p^d$ and Schatten-$p$ spaces — exact for $1 \le p < 2$ and up to absolute constants for $2 \le p \le \infty$. It then quantifies when the native-norm analysis beats the "Euclidean noise plus norm conversion" route, for $\ell_p^d$ and Schatten-$p$ dual geometries.

3. Wall-clock guarantees for semi-asynchronous non-Euclidean SGD (Rennala NSGD). The authors introduce Rennala NSGD, a semi-asynchronous method that collects fresh minibatches at the current iterate and discards stale gradients, and prove a time-complexity bound under fixed heterogeneous worker times $\tau_i$.

4. Consistency with the Euclidean special case. When $\kappa = 1$, both the oracle bound and the wall-clock bound match the parameter dependence of the standard minibatch SGD rate and of classical Rennala SGD (Tyurin and Richtárik, 2023) up to logarithmic factors.

Main Findings

• Oracle rate with explicit geometry: The total number of stochastic-gradient evaluations satisfies $N_* = \tilde{\mathcal{O}}\left(\frac{L\Delta}{\varepsilon} + \frac{\kappa L\Delta\sigma^2}{\varepsilon^2}\right)$ with probability at least $1-\delta$, where $\tilde{\mathcal{O}}$ hides logarithmic factors in $T$ and $1/\delta$. The proof uses the fixed stepsize $\gamma_k \equiv \frac{1}{2L}$, $T = \lceil 128L\Delta/(15\varepsilon)\rceil$ iterations, and minibatch size $B = \max{\lceil 16\sigma^2(\sqrt{2\kappa} + \sqrt{6\log(T/\delta)})^2/\varepsilon\rceil, 1}$.

• Optimality in the Euclidean case: The bound is optimal in $\varepsilon$ and $\sigma$. Jin (2026, Theorem 4.1 and the proof of Proposition 8.9) gives a lower bound of $\Omega(L\Delta/\varepsilon + L\Delta\sigma^2/\varepsilon^2)$ oracle calls for general unbiased first-order oracles with uniformly bounded noise, so with $\kappa = 1$ the new rate matches it up to logarithmic factors. For a general update norm the two are not literally comparable, because $\sigma$ is then measured in $|\cdot|_*$ rather than $|\cdot|_2$; what the bound adds in that case is the single geometry factor $\kappa$.

• $\kappa$ is never vacuous but can be small: Every nonzero finite-dimensional dual space is $\kappa$-regular for some $\kappa \le \dim(E)$, and this worst-case linear dependence on dimension cannot be improved in general (Proposition 11). For $\ell_p^d$ and Schatten-$p$ norms the regularity parameter can be much smaller.

• Gains from tuning the auxiliary norm: For dual $\ell_p^d$ spaces with $2 \le p \le \infty$, tuning the auxiliary norm gives $\kappa = \Theta(\min{p, \ln(ed)})$, versus $\kappa = d^{1-2/p}$ for the Euclidean choice. At $p = \infty$ this turns a linear dimension dependence into a logarithmic one; at $p = 2$ both choices give $\kappa = 1$; for $1 \le p < 2$ a scaled Euclidean proxy is already optimal. These gains require no algorithmic change — the auxiliary norm is used only in the analysis.

• Where native geometry helps and where it does not (Table 1): The table reports statistical oracle-cost terms, omitting the factor $\nu^2$ and the common leading term $\tilde{\mathcal{O}}(L\Delta/\varepsilon)$, for $d \ge 1$. For the dual space $(\mathbb{R}^d, |\cdot|_1)$, the Euclidean/Frobenius conversion gives $\tilde{\mathcal{O}}(d^2L\Delta/\varepsilon^2)$ while the geometry-native bounds give $\tilde{\mathcal{O}}(d^3L\Delta/\varepsilon^2)$ for both Rademacher and Gaussian noise. For $(\mathbb{R}^d, |\cdot|2)$ all three give $\tilde{\mathcal{O}}(dL\Delta/\varepsilon^2)$. For $(\mathbb{R}^d, |\cdot|\infty)$ — corresponding to an $\ell_1$ update norm — the Euclidean conversion gives $\tilde{\mathcal{O}}(dL\Delta/\varepsilon^2)$, the native Rademacher bound gives $\tilde{\mathcal{O}}(\ln(ed)L\Delta/\varepsilon^2)$, and the native Gaussian bound gives $\tilde{\mathcal{O}}((\ln(ed))^2L\Delta/\varepsilon^2)$. In words: under dense Rademacher perturbations the statistical coefficient can grow logarithmically rather than linearly with dimension.

• Wall-clock bound with heterogeneous workers: With worker computation times $\tau_i$ satisfying $0 < \tau_1 \le \tau_2 \le \cdots \le \tau_N$, Rennala NSGD achieves wall-clock time $\tilde{\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)$, where $S$ indexes only the fastest workers considered in the bound. With $\kappa = 1$ this matches the parameter dependence of classical Rennala SGD up to logarithmic factors.

• Minibatch size scaling: Theorem 1 uses $B = \tilde{\mathcal{O}}(1 + \kappa\sigma^2/\varepsilon)$ with $B \ge 1$, the same dependence on $\sigma^2/\varepsilon$ as the Euclidean analyses of Tyurin and Richtárik (2023) and Tyurin (2024) up to logarithmic factors. For exact gradients, $B = 1$ and $N_* = \mathcal{O}(L\Delta/\varepsilon)$, which explains the noise-independent leading term.

• No empirical evaluation is reported: The content provided contains theory, tables of asymptotic orders, and appendices; it does not report experiments, benchmark datasets, or measured speedups.

Methodology in Plain English

The authors set up a plain minimization problem over a finite-dimensional space $E$ — which could be vectors in $\mathbb{R}^d$ or matrices in $\mathbb{R}^{m\times n}$ — and analyze an update that takes a step in the direction of steepest descent for a chosen norm rather than the Euclidean norm. Different norm choices recover familiar methods: the $\ell_2$ norm gives ordinary SGD, the $\ell_\infty$ norm gives a sign-type update related to signSGD, the $\ell_1$ norm gives a greedy coordinate-type step, and the spectral norm gives an orthogonalized matrix update of the Muon/Shampoo family.

Rather than assuming the gradient noise is small in a Euclidean sense and then translating, they assume the noise has "light tails" measured directly in the dual of the update norm — a single exponential-moment inequality controlled by a scale $\sigma$. To make concentration work in an arbitrary norm, they borrow the $\kappa$-regularity machinery of Juditsky and Nemirovski (2008): the dual space is equipped with a smooth auxiliary norm comparable to the real one, with the mismatch captured by the constant $\kappa$. This yields a martingale-style concentration lemma for minibatch averages directly in the target norm.

They then combine that lemma with a standard descent argument to choose a batch size large enough that the gradient-estimation error is controlled simultaneously across all iterations, producing a high-probability stationarity guarantee and an oracle count. Finally, they embed the same minibatch step in an asynchronous architecture: a server keeps the current iterate, broadcasts it, applies the next batch of gradients that arrives, and discards anything computed at an old iterate. Because stale work is thrown away, the analysis can charge each step to the workers that actually contributed, which is what produces the harmonic-mean-style dependence on worker speeds instead of the slowest-worker bottleneck.

Why This Matters

Impact on research. The paper gives a template for analyzing stochastic nonconvex optimization directly in whatever geometry an optimizer uses, rather than analyzing it in Euclidean space and converting. It shows precisely when this changes the answer: helpful for some geometries and noise models (logarithmic instead of linear dimension dependence), unchanged in others, and in at least one case (dual $\ell_1$, i.e. $\ell_\infty$ update) worse than conversion. It also connects two literatures that had stayed separate — non-Euclidean/matrix-aware optimizer theory and time-optimal asynchronous parallel SGD — and it supplies matching regularity bounds for $\ell_p$ and Schatten-$p$ spaces that others can reuse.

Real-world applications (as motivated by the paper):

• Large language model training with matrix-aware optimizers such as Muon, Shampoo, and Dion, whose updates correspond to spectral-norm or related Schatten-norm steepest descent on weight matrices in Transformer attention and feed-forward layers.
• Federated learning, which the paper explicitly cites as the motivating setting, where the finite-sum objective over $N$ workers and per-device computation times $\tau_i$ are the natural model.
• Distributed training on heterogeneous clusters, where slow workers ("stragglers") make synchronous minibatch SGD wait for $\tau_{\max}$ while an asynchronous schedule can be governed by a harmonic-mean-like quantity.
• Matrix-parameterized models generally, where respecting the matrix structure of parameters through a chosen norm is the whole point of the geometry-aware update.

Industry relevance. Any organization training large models across many uneven devices or machines has a direct stake in wall-clock, not just iteration-count, guarantees. The paper's selling point for practitioners is that the geometry factor enters the time bound as a single multiplicative constant $\kappa$ that can be computed per norm, so system designers can reason about whether a matrix-aware optimizer's geometry is cheap or expensive in a given distributed regime. The absence of experiments means these are still theoretical predictions rather than measured speedups.

Future Directions

• Empirical validation. The paper reports no experiments. Testing whether the predicted wall-clock gains of Rennala NSGD materialize with spectral-norm/Muon-style updates on real Transformer training, on hardware with genuinely heterogeneous worker speeds, is the obvious next step.
• Statistical heterogeneity. The analysis deliberately assumes all workers share the same sample loss and distribution $\mathcal{D}$ in order to isolate device heterogeneity. Extending the guarantees to differing $\mathcal{D}_i$ (the standard federated setting) remains open.
• Closing the gaps in the geometry comparison. The native-norm bound is worse than the Euclidean conversion for the dual $\ell_1$ case (Table 1), and the regularity theory is worked out for the $\ell_p$ and Schatten-$p$ families. Whether sharper native analyses exist for other geometries, or for the sign-type $\ell_\infty$ update specifically, is unresolved.
• Tightness of the batch-size choice. As the authors note, the lower bound of Jin (2026) constrains the total cost $BT$ but not its split between iteration count and batch size, so it does not establish that their particular $B = \tilde{\mathcal{O}}(1 + \kappa\sigma^2/\varepsilon)$ is necessary.

Target Audience

Optimization theorists and graduate students working on nonconvex stochastic methods, non-Euclidean and matrix-aware optimizers, or distributed and asynchronous SGD will get the most from this paper. It is also relevant to machine-learning systems researchers interested in wall-clock rather than iteration complexity, and to practitioners working on Muon/Shampoo-style training or federated learning who want to understand the theoretical assumptions under which geometry-aware, asynchronous updates provably help. Readers without a background in normed spaces, duality, and high-probability convergence analysis will find the main text demanding.

Authors’ abstract

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.