1 Introduction↩︎

Neural networks have evolved from multilayer perceptrons (MLP) trained by backpropagation [1] into a dominant paradigm of modern AI. Key milestones include mixtures of experts (MoE) [2], convolutional neural networks (CNN) for vision [3], recurrent neural networks (RNN) for sequential data [4], [5], and deep belief networks (DBN) that helped revive interest in deep architectures [6]. Since the breakthrough of AlexNet on ImageNet [7], deep networks have transformed vision, speech, language and generative AI, with residual networks (ResNet) enabling much deeper models [8] and transformers enabling scalable sequence modeling [9]. This progress has relied not only on data and hardware, but on nonconvex stochastic optimization methods, from SGD and momentum [10]–[13] to adaptive methods such as AdaGrad, RMSProp, Adam, and AdamW [14]–[17]. Yet, these methods are designed around Euclidean and coordinatewise parameter geometries and standard stochastic gradient noises, rather than layerwise, heterogeneous geometry and heavy-tailed stochastic gradient noises observed in training modern neural networks [18]–[20].

Norm-based matrix optimizers, such as Muon and Scion, have addressed this mismatch by scaling updates according to layerwise matrix geometry. The guiding principle is simple: for a neural-network layer, the update scale should reflect how the layer transforms its inputs, rather than only the Euclidean length of all entries in its weight matrix. Muon implements this idea through matrix-sign updates [21]–[23], while Scion provides a broader framework based on linear minimization oracles over input-output matrix norm balls [24]. Thus, the use of general norms is not merely a theoretical abstraction but is already present in modern scale-invariant optimizer design. These methods are also connected to architecture-aware scaling, modular norm viewpoints, and hyperparameter transfer principles [19], [25]–[30]. Compared with coordinatewise rescaling used in Adam-type methods, scale-invariant matrix optimizers make update magnitudes more comparable across heterogeneous layers and have become increasingly appealing in large-scale neural network training [31], [32]. This leads to a natural question: despite their practical adoption, what guarantees can we prove for scale-invariant methods under general norm geometry and heavy-tailed stochastic gradients?

Recent works have begun to clarify why normalized, sign-based, and matrix-sign optimizers can be effective when stochastic gradients are heavy-tailed. For standard vector methods, a common strategy is to reduce sensitivity to unusually large gradients through clipping or normalization [33]–[36]. Complementary results give optimal unclipped guarantees in the Euclidean setting [37], while adaptive methods may require additional robustification under heavy-tailed noise [38]. More recent work also provides heavy-tailed guarantees for sign-based optimizers, including Muon and its variants, under optimizer-specific noise and smoothness models [39]. However, these results do not fully explain general-norm scale-invariant methods such as Scion. In particular, Euclidean or Frobenius normalization corresponds to one special geometry, while scale-invariant matrix optimizers use non-Euclidean layerwise geometries whose dual norms can introduce dimension-dependent estimation errors. Existing analyses thus leave open the sharp dimension dependence of scale-invariant methods under heavy-tailed noise, and whether higher-order smoothness can accelerate their convergence.

In this paper, we study the scale-invariant methods through nonconvex smooth stochastic optimization over \(\mathbb{R}^{m\times n}\) with general norms under \(p^\textrm{th}\)-moment heavy-tailed noise. The general-norm formulation reflects the geometry already used by scale-invariant neural-network optimizers, while the heavy-tailed noise model captures the non-Gaussian stochasticity often observed in training. Our analysis has three parts. For the lower bound, we combine the framework of [40] and the hard instance of [37] by embedding independent hard signals across rows, showing that the dimension dependence is unavoidable. For the matching upper bound, we modify the Scion method [24] using batched momentum. The key step is to control the momentum error in the dual norm before applying the linear minimization oracle. For acceleration, we adapt the implicit gradient transport mechanism [41] to Scion and show that the transported point cancels the first-order drift in the momentum recursion, leaving only a Hessian-Lipschitz curvature error. The analysis is nontrivial because scale-invariant updates are nonlinear, discard gradient magnitudes, and must be controlled in non-Euclidean dual norms under heavy-tailed noise.

1.0.0.1 Contributions.

We focus on two central questions: (i) what dimension dependence is unavoidable for scale-invariant methods beyond Frobenius geometry, and (ii) whether higher-order smoothness accelerates such methods under heavy-tailed noise. Our contributions can be summarized as follows.

1. We establish the sharp dimension dependence for scale-invariant stochastic optimization under heavy-tailed noise. In the spectral-normed space \((\mathbb{R}^{m\times n},\|\cdot\|_{\mathrm{op}})\), we prove that when \(\frac{\max\{m,n\}}{(\min\{m,n\})^2}\) is large enough, any scale-invariant first-order method requires \(\Omega(\min\{m, n\}\epsilon^{-\frac{3p-2}{p-1}})\) oracles to find an \(\epsilon\)-stationary point. We then show that a batched Scion method achieves the matching upper bound of \(O(\min\{m, n\}\epsilon^{-\frac{3p-2}{p-1}})\).

2. We propose a transported Scion method that leverages higher-order smoothness and prove an improved bound of \(O(\min\{m, n\}\epsilon^{-\frac{5p-3}{2p-2}})\) when the norm is spectral and the Hessian is Lipschitz. We incorporate practical heuristics into our transported method and use it to train CNNs and transformers, showing the flexibility and compatibility of the transportation technique in neural network training.

1.0.0.2 Related work.

Our work is mostly related to the literature on neural network optimization methods and optimization under heavy-tailed noise. Due to space limitations, we defer our comments on other relevant topics to Appendix 6. Earlier matrix optimizers exploit layerwise matrix structure through spectral or Kronecker-factored preconditioning [21], [42]–[52]. A recent line treats scale invariance itself as the optimization primitive: Muon applies the matrix sign of layerwise gradients or momentum, typically via Newton-Schulz iterations, while Scion casts related updates as linear minimization oracles over operator-norm balls [22], [24]. This viewpoint has led to a growing family of related matrix optimizers [31], [53]–[67] and faster matrix-sign routines [68], [69].

Empirical studies have found heavy-tailed stochastic gradient noise in training neural networks and language models, motivating updates that reduce sensitivity to raw gradient magnitudes [20], [33], [70]–[72]. [37] showed that normalized SGD with momentum attains the optimal heavy-tailed nonconvex rate without clipping, but when applied to matrix-valued parameters this corresponds to Frobenius normalization. More recently, [39] established heavy-tailed guarantees for sign-based optimizers under optimizer-specific noise and smoothness models: coordinatewise for vector signs and matrix-weighted for matrix signs. Furthermore, [73] proved heavy-tailed convergence for Muon under Frobenius noise and smoothness conditions. In contrast, our work gives a norm-agnostic analysis of scale-invariant stochastic LMO methods under smoothness and a general \(p^\textrm{th}\) moment noise condition in the dual norm. The resulting factor \(\tau(\|\cdot\|_\star,m,n,p)\) can be dimension dependent, and we show that this dependence is unavoidable for input-output matrix norm geometry. Our dimension-dependent lower bound does not contradict their work since our assumptions are weaker.

2 Preliminaries and Technical Background↩︎

We provide an overview of scale-invariant Scion methods and the specific norm geometry that has contributed to their popularity in neural network optimization, along with formal definitions of the function classes and heavy-tailed noise models considered in this paper.

2.1 Scale-invariant methods and norm geometry↩︎

Throughout this paper, we equip \(\mathbb{R}^{m\times n}\) with a general norm \(\|\cdot\|\) and denote \(\|\cdot\|_\star\) as the dual norm. We assume that \(F^\star := \inf_{X \in \mathbb{R}^{m\times n}} F(X) > -\infty\) and consider \[\min_{X \in \mathbb{R}^{m\times n}} F(X).\] The key component of scale-invariant methods is the linear minimization oracle (\(\operatorname{LMO}\)) over the unit norm ball in the following form of \[\operatorname{lmo}(S) \in \mathop{\rm argmin}_{\|X\|\le 1}\langle S,X\rangle .\] By definition, we have \(\|\operatorname{lmo}(S)\| \leq 1\) and \(\langle S,\operatorname{lmo}(S)\rangle=-\|S\|_\star\). Thus, this oracle is scale invariant: \(\operatorname{lmo}(\alpha S)=\operatorname{lmo}(S)\) for all \(\alpha>0\), up to the choice of minimizer. An update \(X^+=X+\eta\operatorname{lmo}(S)\) therefore fixes the step length in the chosen geometry and uses the direction of \(\operatorname{lmo}(S)\).

In practice, Scion methods choose the norm ball layerwise using input-output matrix norms for neural network optimization [24]. In particular, for a weight matrix \(W \in \mathbb{R}^{d_{\mathrm{out}}\times d_{\mathrm{in}}}\) and vector norms \(\|\cdot\|_\alpha,\|\cdot\|_\beta\), we define \[\|W\|_{\alpha \to \beta} := \sup_{\|z\|_\alpha\le 1}\|Wz\|_\beta .\] If the layer input is bounded in \(\|\cdot\|_\alpha\), then bounding \(\|W\|_{\alpha\to\beta}\) directly controls the output in \(\|\cdot\|_\beta\). This input-output interpretation explains why general matrix norms arise in scale-invariant neural network optimizers. Frobenius normalization views \(W\) as a vector in Euclidean space. However, it does not measure the layer map \(z \mapsto Wz\) through an input-output matrix norm of the form \(\|\cdot\|_{\alpha \to \beta}\).

For hidden layers, a common choice is the RMS-to-RMS norm, where \(\|z\|_{\mathrm{RMS}}=\frac{1}{\sqrt{d}}\|z\|_2\) for \(z \in \mathbb{R}^d\). The corresponding matrix norm is \[\|W\|_{\mathrm{RMS}\to\mathrm{RMS}} = \sqrt{\tfrac{d_{\mathrm{in}}}{d_{\mathrm{out}}}}\, \|W\|_{\mathrm{op}}.\] If \(S=U\Sigma V^\top\) is SVD, we have \[\operatorname{lmo}_{\mathrm{RMS}\to\mathrm{RMS}}(S) = -\sqrt{\tfrac{d_{\mathrm{out}}}{d_{\mathrm{in}}}}UV^\top,\] which is the scaled matrix-sign update used by the Muon optimizer and its variants [22], [23], [31], [53]. Other norms also recover normalized updates: the \(1 \to \mathrm{RMS}\) norm gives column-normalized updates, the \(\mathrm{RMS} \to \infty\) norm gives row-normalized updates, and the \(1 \to \infty\) norm gives sign updates.

For a network parameter \(\Theta=(W_1,b_1,\ldots,W_L,b_L)\), the layerwise geometry can be summarized by the block norm as follows, \[\|\Theta\| := \max_{\ell\in[L]} \;\tfrac{1}{\rho_\ell}\max\{\|W_\ell\|_{\alpha_\ell \to \beta_\ell}, \|b_\ell\|_{\beta_\ell}\},\] where \(\rho_\ell>0\) is a layerwise radius. The LMO over this product ball decomposes across layers, so each layer receives a radius-scaled normalized update in its own input-output geometry. For example, in the RMS-to-RMS case, the update has spectral norm of order \(\sqrt{d_{\mathrm{out}}/d_{\mathrm{in}}}\), matching the scaling used for width-stable feature learning and hyperparameter transfer. Our theory keeps the norm abstract but focuses on one weight matrix for simplicity.

2.2 Function class and heavy-tailed noise model↩︎

We present definitions for generalized smooth functions and heavy-tailed noise model.

Assumption 1 recovers \(L_0\)-smoothness when \(L_1=0\), while allowing the local smoothness scale to grow with the dual gradient norm \(\|\nabla F(X)\|_\star\). This lets our analysis cover objectives with relaxed or unbounded smoothness while still retaining the descent inequality needed for normalized updates.

Assumption 2 reduces to the classical finite \(p^\textrm{th}\) moment heavy-tailed noise model when \(\sigma_1=0\), and the additional \(\sigma_1\) term allows the noise scale to grow with the local gradient. In particular, for \(p<2\) we do not assume bounded variance or bounded stochastic gradients. Given a query point \(X\) and batch size \(B\), we write

where \(\xi^1,\ldots,\xi^B\) are i.i.d. samples. Thus, when an algorithm chooses \(X_t\) from the past and then draws a mini-batch, Assumption 2 implies the conditional unbiasedness, conditional independence, and conditional \(p^\textrm{th}\) moment bounds used in the analysis.

The key difference between our analysis and the analysis of [37] is that the momentum error is controlled in a general dual norm. To highlight this norm-dependent effect, we define \[\label{def:factor} \tau(\|\cdot\|_\star,m,n,p):=\sup_{\{Z_t\}_{t=1}^T}\tfrac{\mathbb{E}\|\sum_{t=1}^T Z_t\|_\star}{ \mathbb{E}(\sum_{t=1}^T \|Z_t\|_\star^p)^{1/p}},\tag{1}\] where the supremum is over all \(T \in \mathbb{N}\) and all integrable \(\mathbb{R}^{m\times n}\)-valued martingale difference sequences \(\{Z_t\}_{t=1}^T\) with respect to their natural filtrations. The ratio is \(0\) when the denominator is \(0\). The following lemma records the dimension dependence of this martingale factor.

Lemma 1 identifies the theoretical difference between Frobenius normalization and general norms used by scale-invariant LMO methods. For Frobenius normalization, the factor \(\tau\) is dimension-free, matching the setting in which a matrix parameter is treated as a vector [37]. For other matrix norms, the estimator error is measured in the corresponding dual norm, and this can introduce dimension dependence. For example, spectral-norm LMO updates require controlling stochastic error in nuclear norm, where \(\tau(\|\cdot\|_{\rm nuc},m,n,p)=\Theta(\min\{m,n\}^{1-1/p})\). Thus, the dimension dependence studied in this paper is a consequence of combining heavy-tailed noise with general input-output matrix norm geometries used by neural network optimizers. The proof of Lemma 1 is deferred to Appendix 7.

Assumption 3 is implied by the standard Hessian Lipschitzness condition in Euclidean geometry. This additional higher-order smoothness assumption allows us to design accelerated algorithms.

3 Main Results↩︎

We establish a dimension-dependent lower bound for stochastic scale-invariant first-order methods under spectral-norm geometry and heavy-tailed noise. We then show that a batched Scion method achieves the matching upper bound. Finally, we introduce a transported Scion method and prove an improved rate under Hessian Lipschitzness. The lower bound is stated for the spectral-norm geometry, where stationarity is measured in the nuclear norm. The upper bounds are stated for a general norm \(\|\cdot\|\) on \(\mathbb{R}^{m\times n}\) with dual norm \(\|\cdot\|_\star\). Throughout this section, when stating upper bounds, we write \(\tau_\star := \tau(\|\cdot\|_\star,m,n,p)\) where \(\tau(\cdot)\) is the martingale factor. For spectral-norm LMO updates, we have \(\|\cdot\|=\|\cdot\|_{\mathrm{op}}\) and \(\|\cdot\|_\star=\|\cdot\|_{\mathrm{nuc}}\). In this case, we have \(\tau_\star=\Theta(\min\{m,n\}^{1-1/p})\). All proofs are deferred to Appendix 7.

3.1 Dimension-dependent lower bound↩︎

We show that the dimension dependence induced by spectral geometry is unavoidable. Let \(\Delta,L>0\). Define \(\mathcal{F}_{\mathrm{op}}(m,n,\Delta,L)\) as the class of differentiable functions \(F:\mathbb{R}^{m\times n} \to \mathbb{R}\) satisfying \[F(0)-\inf_{X \in \mathbb{R}^{m\times n}} F(X) \leq \Delta, \quad \|\nabla F(X)-\nabla F(Y)\|_{\mathrm{nuc}} \leq L\|X-Y\|_{\mathrm{op}} \textrm{ for all } X,Y \in \mathbb{R}^{m\times n}.\] Since we focus on the matrix normed space with spectral norm, the corresponding dual stationarity measure is defined under the nuclear norm.

We use the stochastic first-order oracle model [40]. Each oracle \(\mathsf{O}\) consists of a distribution \(P_\xi\) on a measurable space \(\Xi\) and a mapping \(\mathsf{O}_F(X,\xi)=(F(X),g(X,\xi))\) such that, for every \(F \in \mathcal{F}_{\mathrm{op}}(m,n,\Delta,L)\), we have \[\mathbb{E}[g(X,\xi) \mid X] = \nabla F(X), \quad \mathbb{E}[\|g(X,\xi)-\nabla F(X)\|_{\mathrm{nuc}}^p \mid X] \leq \sigma_0^p.\] We denote the set of all such oracles by \(\mathcal{O}_p(\sigma_0)\).

We use the randomized algorithm model [40]. Let \(\mathcal{A}_{\rm rand}\) denote the class of randomized first-order algorithms using the oracle \(\mathsf{O}_F\). At round \(t\), an algorithm \(A[\mathsf{O}_F] \in \mathcal{A}_{\rm rand}\) chooses a query point \(X_{A[\mathsf{O}_F]}^{(t)}\) that is measurable with respect to its internal random seed and all previous oracle observations. The oracle then draws \(\xi^{(t)} \sim P_\xi\) independently and returns \(F(X_{A[\mathsf{O}_F]}^{(t)})\) and \(g(X_{A[\mathsf{O}_F]}^{(t)},\xi^{(t)})\).

Let \(\mathcal{P}(\mathcal{F}_{\mathrm{op}}(m,n,\Delta,L))\) be the set of probability measures over \(\mathcal{F}_{\mathrm{op}}(m,n,\Delta,L)\). For any tolerance \(\epsilon>0\) and \(p \in (1,2]\), we define the worst-case number of oracle rounds needed to output an expected \(\epsilon\)-stationary point by \[\mathfrak{m}^{\mathrm{rand}}_{\epsilon,p}(m,n,\Delta,L,\sigma_0):=\sup_{\mathsf{O} \in \mathcal{O}_p(\sigma_0)} \sup_{P_F \in \mathcal{P}(\mathcal{F}_{\mathrm{op}}(m,n,\Delta,L))} \inf_{A \in \mathcal{A}_{\mathrm{rand}}} \inf\left\{N \mid \mathbb{E}_{F\sim P_F, A[\mathsf{O}_F]} \|\nabla F(X_{A[\mathsf{O}_F]}^{(N)})\|_{\mathrm{nuc}} \le\epsilon \right\}.\] The expectation is taken over the random problem instance, the internal randomness of the algorithm, and the stochastic oracle.

Theorem 1 extends the Euclidean heavy-tailed lower bound to spectral-norm matrix optimization. When \(\min\{m,n\}=1\), the spectral norm and the nuclear norm reduce to the vector \(\ell_2\)-norm. Our lower bound matches the existing bounded variance [40] and heavy-tailed lower bounds [37]. The key to our new results is the multiplicative factor \(\min\{m,n\}\), which comes from the interaction between spectral-norm updates and nuclear-norm stochastic gradient estimation. [39] establishes a convergence rate in the same setting that avoids explicit dimension dependence by a stronger assumption. [73] provides a dimension-independent convergence rate under Frobenius-norm-based smoothness and noise assumptions. In contrast, Theorem 1 demonstrates that under only standard smoothness and heavy-tailed noise assumptions with gradient norm measured by the dual norm, dimension dependence is unavoidable. In the lower-bound construction, independent hard instances are embedded across matrix rows, and an additional random row index controls the information leaked by each oracle response. Since standard \(L\)-smoothness implies Assumption 1 with \(L_0=L\) and \(L_1=0\), the lower bound also applies to the generalized smoothness setting used in our upper bound analysis.

3.2 Smooth and nonconvex problems↩︎

We move to the upper bound. Algorithm 1 is a batched momentum variant of unconstrained Scion. The LMO fixes the update scale in the primal norm, while the momentum estimator is controlled in the dual norm. This dual-norm control is precisely where the martingale factor \(\tau_\star\) enters.

Theorem 2 matches the lower bound of Theorem 1 for spectral-norm geometry, up to constants. When \(\min\{m,n\}=1\), it recovers the existing heavy-tailed upper bound [37]. Under spectral-norm updates, we have \(\tau_\star^{p/(p-1)} = \Theta(\min\{m,n\})\), so the dimension dependence is the cost of estimating heavy-tailed gradient noise in the nuclear norm. Our lower bound requires the larger matrix dimension to be sufficiently large relative to the tolerance and noise level. Therefore, in regimes not covered by this condition, such as some square-matrix regimes common in neural network layers, the lower bound does not rule out sharper dimension dependence.

The preceding theorem uses parameters depending on \(p\) which is usually unknown in practice. The next theorem gives a parameter choice that does not require knowing \(p\) yet leads to a worse bound.

Theorem 3 trades optimal tuning for robustness to an unknown tail index. When \(T=\Omega(L_1^4)\), the chosen step size becomes independent of \(L_1\), and this yields the bound with worse dependence on \(\min\{m,n\}\) and \(\epsilon\) under spectral-norm updates. The additional factor \(\min\{m,n\}\) compared with Theorem 2 comes from using \(B=1\). Without a batch size tuned to \(\tau_\star\) and \(p\), the stochastic dual-norm error cannot be calibrated directly to remove one power of the spectral-geometry dimension factor.

3.3 Highly smooth and nonconvex problems↩︎

To accelerate Algorithm 1, we introduce an additional higher-order smoothness condition (Assumption 3). Under this assumption, we propose Algorithm [algorithm:transported-uSCG]. Inspired by the gradient transportation technique [41], this algorithm utilizes an auxiliary sequence \(Y_t\) for gradient evaluation while maintaining the primary sequence \(X_t\) for gradient-descent-style updates. By evaluating the gradient at \(Y_t\), certain first-order drift terms are replaced by more controllable Hessian curvature error terms.

The convergence rate improves the \(\epsilon\)-dependence by a factor of \(\Theta(\epsilon^{\frac{1}{2}})\) over Theorem 2, albeit at the cost of the extra smoothness assumption and the dependence of \(L_0\) being replaced by \(\sqrt{L_2}\).

To achieve acceleration over Algorithm 1 in the unknown \(p\) setting, we choose parameters independent of \(p\) in Theorem 4 and prove the following theorem.

The convergence rate improves the \(\epsilon\)-dependence by a factor of \(\Theta(\epsilon^{\frac{p}{4p-4}})\) over Theorem 3. Moreover, the \(\sigma_0\)-dependence improves by \(\sigma_0^{\frac{p}{4p-4}}\) and the dimension dependence improves by \(\min\{m,n\}^{1/4}\) over Theorem 3. It is also worth remarking that the choice of \(\eta\) is independent of \(L_2\).

4 Experiment↩︎

We evaluate Algorithms 1 and [algorithm:transported-uSCG] for pretraining deep neural networks, including CNNs and transformer-based large language models (LLMs). For the LLM experiments, we train nanochat models [74] of different sizes on the NVIDIA ClimbMix dataset [75] and report training loss, validation loss, and downstream performance using the CORE metric [76]. For the CNN experiments, we train the 2M-parameter CIFARNET model on CIFAR-10 and CIFAR-100. All experiments are implemented in Python 3.12 and PyTorch 2.9.1. The LLM experiments use 10 NVIDIA A40 GPUs, and the CNN experiments use a single NVIDIA A40 GPU. Each GPU has 46 GB of memory and runs Ubuntu 22.04.05 LTS. Additional details on datasets, architectures, benchmarks, and hyperparameter searches are provided in Appendix 8.

4.0.0.1 Implementation details.

Using \(\|\cdot\|_{\mathrm{op}}\) in the LMO of Algorithm 1 recovers the Muon update rule [22], [24] with heavy-ball momentum. The subroutine based on Newton-Schulz iteration or PolarExpress algorithm [68] provides fast and accurate LMO approximations; we use PolarExpress whenever applicable. We also use the normalization techniques of [53] to improve the LMO approximation and balance neuronwise norms. Except for the baseline Muon optimizer, these normalization techniques are used throughout our experiments. Algorithm [algorithm:transported-uSCG] evaluates gradients at the auxiliary sequence \(Y_t\) while updating the model weights along the sequence \(X_t\), analogous to Nesterov momentum. The NAdam optimizer [77] showed how to incorporate this idea into neural network optimizers through the scheme \[\bar{G}_t = \tfrac{1}{B}\sum_{i=1}^B G(X_t,\xi_t^i), \quad m_{t+1} = \beta m_{t} + (1-\beta)\bar G_t, \quad X_{t+1} = X_t+\eta \operatorname{lmo}\left((1-\beta)m_{t+1} + \beta\bar{G}_t\right),\] where \((1-\beta)m_{t+1} + \beta\bar{G}_t\) is the Nesterov momentum passed to the LMO. This practical scheme is used in several recent works [22], [31], [74]. In our experiments, we evaluate both standard Nesterov momentum and the exact update rule from Algorithm [algorithm:transported-uSCG]. For the latter, we use the auxiliary update \(Y_{t}=X_t+\alpha(X_t-X_{t-1})\), and tune \(\alpha\) as a hyperparameter instead of fixing it to \(\beta/(1-\beta)\).

4.0.0.2 LLM Experiment.

We pretrain nanochat models with depths 12, 16, and 24 on the NVIDIA ClimbMix dataset using a data-to-scalable-parameter ratio of 8. These models contain 287M, 539M, and 1.39B parameters and are trained on 882M, 1.88B, and 5.84B tokens, respectively. We compare five optimizers: AdamW [17], Muon [22], NorMuon [53], NorMuon with Nesterov momentum, and NorMuon with transportation, corresponding to Algorithm [algorithm:transported-uSCG]. Learning rates and momentum factors are selected by grid search, as described in Appendix 8. Table 1 reports the LLM results. Consistent with prior observations, NorMuon outperforms Muon, and Muon outperforms AdamW. This confirms the effectiveness of normalization and scale-invariant optimization under spectral geometry. Moreover, NorMuon with transportation improves over NorMuon, demonstrating the practical benefit of the transportation technique.

4.0.0.3 CNN Experiment.

We train CIFARNET [78], [79] on CIFAR-10 and CIFAR-100 [80]. In addition to the five optimizers used in the LLM experiments, we include SGD with momentum as a CNN baseline. Details on the setup, hyperparameter selection, and grid searches are given in Appendix 8. Table 2 summarizes the CNN results. The Muon-family methods substantially outperform both AdamW and SGD with momentum on both datasets. The transportation variant achieves the best CIFAR-10 test loss and the best CIFAR-100 test accuracy, while maintaining CIFAR-10 test accuracy comparable to the best heavy-ball NorMuon result.

5 Conclusion↩︎

We studied scale-invariant optimization methods under general norm geometry and heavy-tailed stochastic noise. The motivation is that general-norm scale-invariant methods are already used in modern matrix optimizers and are closely connected to hyperparameter transfer, but their theoretical guarantees remain incomplete beyond Frobenius geometry. We showed that, under heavy-tailed noise, the choice of geometry affects the difficulty of stochastic optimization through the dual-norm martingale factor, which can introduce unavoidable dimension dependence. We established a dimension-dependent lower bound, proved that a batched Scion method achieves the matching upper bound, and proposed a transported Scion method with an improved rate under Hessian Lipschitzness. Experiments on CNNs and transformer models show that transported techniques are compatible with practical training pipelines. Future directions include developing principled momentum and transportation techniques for other matrix optimizers [42], [44], [50] and applying these ideas to larger-scale LLM pretraining with schedule-free parameter tuning [81].

Acknowledgments↩︎

We sincerely appreciate Buzz High Performance Computing (https://www.buzzhpc.ai, info@buzzhpc.ai) for providing computational resources and support for this work.

6 Further Related Works↩︎

We make some comments on other topics, including more discussions on neural network optimization methods and optimization under heavy-tailed noise, the theoretical understanding of matrix optimizers, and hyperparameter transfer. For an overview of neural network optimization methods, we refer to the recent monograph [82].

6.0.0.1 More discussions on neural network optimization methods.

Beyond matrix optimizers discussed in the main text, many neural network optimizers operate through vector updates. The classical baseline is SGD with momentum, whose practical relevance in deep learning is tied to initialization and momentum tuning [10]–[13]. Coordinatewise adaptive methods instead maintain diagonal statistics of past gradients: AdaGrad and RMSProp accumulate squared-gradient information, Adam combines first- and second-moment exponential averages, AdamW decouples weight decay, Adafactor reduces optimizer memory through factored second-moment estimates, and more recent variants such as Adan and MARS modify the momentum or variance-reduction component [14]–[17], [51], [83], [84]. Theoretical analyses of vector optimizers are extensive but optimizer-specific. Indeed, SGD-type methods have nonconvex upper and lower bounds under standard smoothness [40], [85]. Adaptive methods have been analyzed under relaxed, coordinatewise, or anisotropic smoothness assumptions and through sign-magnitude interpretations of updates [86]–[93]. Layerwise-scaled methods such as LARS and LAMB control update scale, which is useful in large-batch regimes [94], [95]. Finally, recent works have framed optimizer design through explicit norm control: modular-norm and operator-norm perspectives motivate scale-invariant layerwise updates and hyperparameter transfer, Scion framework unifies normalized, sign-based, and spectral updates as norm-ball steps, and modern constrained optimization views connect many existing methods to implicit or explicit norm constraints [23]–[26], [96]–[98].

6.0.0.2 More discussions on optimization under heavy-tailed noise.

Most existing works focus on vector optimizers and separate into clipping, normalization, and coordinatewise sign updates. Clipping-based stochastic methods established the \(O(\epsilon^{-\frac{3p-2}{p-1}})\) rate for smooth nonconvex objectives and later extended this to high-probability guarantees, nonsmooth and convex problems, and variational inequalities [33], [34], [99]–[103], while matching lower-bound refinements further clarified the dependence on the initial gap, smoothness, and noise scale [33], [37]. A second line shows that clipping is not the only robustification mechanism: momentum-based normalized SGD attains \(O(\epsilon^{-\frac{2p}{p-1}})\) rate when \(p\) is unknown [35]–[37]. Coordinatewise sign methods provide another non-Euclidean route, beginning with SignSGD and its majority-vote and error-feedback variants [104]–[108]. In particular, the Lion optimizer combines sign updates with two momentum parameters and decoupled weight decay, with analyses via constrained dynamics and stochastic Frank-Wolfe interpretations [109]–[111]. Heavy-tailed modeling also connects this optimization literature to robust learning, online learning, and bandits [112]–[120], and to empirical studies of neural-network and language-model gradient statistics, Zipfian imbalance, and sign-like adaptivity [20], [70]–[72], [121], [122].

6.0.0.3 Theoretical understanding of matrix optimizers.

Existing theory has developed along two basic directions: convergence analysis and mechanistic interpretation. On the former side, early analyses of Muon study an idealized polar-step version, replacing the finite Newton–Schulz (NS) orthogonalization by an exact matrix sign [123]–[125]. [79] is particularly relevant since it narrows this gap: it proves nonconvex convergence for practical Muon with a finite number of NS steps, shows that the NS error only introduces a multiplicative factor relative to the exact-polar rate, and proves that this factor approaches one doubly exponentially in the number of NS iterations and improves with the polynomial degree. Its comparison with SGD with momentum identifies a rank-dependence advantage for matrix orthogonalization under a nuclear-norm optimality criterion. Closely related works analyze Scion or Muon-like methods as norm-constrained linear minimization steps [24], [111], but they do not provide the sharp dimension dependence of scale-invariant stochastic optimization under heavy-tailed noise. A second line of work aims to understand why matrix-sign updates are useful. In particular, the normalized steepest descent view identifies Muon as a spectral-geometry analogue of sign or normalized gradient descent [23], [54], [126], [127]. Other works focus on implicit bias, simplicity bias, nonsmooth analysis, local quadratic models, nonconvex matrix factorization problems, or preconditioning interpretations [57], [63], [128]–[135].

6.0.0.4 Hyperparameter transfer.

Hyperparameter transfer seeks scaling rules under which hyperparameters tuned on small proxy models, such as initialization, learning rates, residual scales, and regularization, remain near-optimal as width, depth, or compute grows. The maximal-update parametrization \(\mu\)P gives the canonical width rule: scale initialization and layerwise learning rates so all layers maintain stable, nontrivial feature learning in the infinite-width limit, enabling zero-shot transfer [28], [29]. A finite-width complement is the spectral condition, which preserves layerwise input-output geometry by requiring weight updates to have \(\Theta(\sqrt{d_{\mathrm{out}}/d_{\mathrm{in}}})\) scale in operator norm [30]. Transfer also extends to depth, where residual branches and learning rates must be co-scaled to sustain feature learning [136]. Finally, learning rate transfer alone is insufficient for compute-optimal training: regularization can depend on model size, with inverse-width scaling often improving transfer [137], [138]. Thus, optimizer comparisons should transfer initialization, learning rate, depth scaling, and regularization rules jointly.

7 Missing Proofs↩︎

7.1 Martingale property↩︎

Proof of Lemma 1. We claim that for any norm \(\|\cdot\|\) on \(\mathbb{R}^{m\times n}\) and any \(p\in(1,2]\), we have \(1\leq \tau(\|\cdot\|,m,n,p)\leq 2\sqrt{2mn}\).

For the lower bound, take a nonzero matrix \(A\in\mathbb{R}^{m\times n}\) and let \(Z_1=\epsilon A\), where \(\epsilon\) is a fair sign. Letting \(T=1\) gives \(\tau(\|\cdot\|,m,n,p)\ge1\).

For the upper bound, by John’s theorem, there exists a Euclidean norm \(|\cdot|\) on \(\mathbb{R}^{m\times n}\) such that \(|A|\le \|A\|\le \sqrt{mn}\,|A|\) for every \(A\in\mathbb{R}^{m\times n}\). Hence, for any \(\mathbb{R}^{m\times n}\)-valued martingale difference sequence \((Z_t)_{t=1}^T\), it follows that \[\begin{align} \mathbb{E}\biggl\|\sum_{t=1}^T Z_t\biggr\| \le \sqrt{mn}\,\mathbb{E}\biggl|\sum_{t=1}^T Z_t\biggr| \le 2\sqrt{2mn}\,\mathbb{E}\biggl(\sum_{t=1}^T |Z_t|^p\biggr)^{\frac{1}{p}} \le 2\sqrt{2mn}\,\mathbb{E}\biggl(\sum_{t=1}^T \|Z_t\|^p\biggr)^{\frac{1}{p}}, \end{align}\] where the middle inequality is Lemma 4.3 of [37].

We claim that for any \(p\in(1,2]\), there exists a constant \(C_p\) such that \({\min\{m,n\}}^{1-1/p}\le\tau(\|\cdot\|_{\mathrm{nuc}},m,n,p)\le C_p\min\{m,n\}^{1-1/p}\).

Let \(r=\min\{m,n\}\). For the lower bound, let \(Z_k=\epsilon_k E_{kk}\) for \(k=1,\dots,r\), where \((\epsilon_k)\) are independent fair signs. Then \(\|\sum_{k=1}^r Z_k\|_{\mathrm{nuc}}=r\) and \((\sum_{k=1}^r \|Z_k\|_{\mathrm{nuc}}^p)^{\frac{1}{p}}=r^{\frac{1}{p}}\) almost surely, so \(\tau(\|\cdot\|_{\mathrm{nuc}},m,n,p) \ge r^{1-1/p}\).

In what follows, we prove the upper bound. Let \(\|\cdot\|_{S_p}\) denote the Schatten \(p\)-norm. By Hölder’s inequality for the singular values, for every \(A\in\mathbb{R}^{m\times n}\), we have \[\label{eq:nuclear-sp-comparison} \|A\|_{\mathrm{nuc}}\le r^{1-1/p}\|A\|_{S_p}\leq r^{1-1/p} \|A\|_{\mathrm{nuc}}.\tag{2}\]

Next, by [139], the Schatten class \(S_p^{m,n}\) is \(p\)-uniformly smooth for every \(1<p\le 2\). By [140], this implies that there exists a constant \(C_p>0\), depending only on \(p\), such that every \(S_p^{m,n}\)-valued martingale \(f=(f_n)_{n\ge 0}\) satisfies \[\mathbb{E}[\sup_{n\ge 0}\|f_n\|_{S_p}]\le C_p\mathbb{E} [\Bigl(\sum_{n\ge 0}\|df_n\|_{S_p}^p\Bigr)^{\frac{1}{p}}].\]

Now let \((Z_k)_{k=1}^N\) be any matrix-valued martingale difference sequence, and define the martingale \(f_n:=\sum_{k=1}^n Z_k\) with the convention \(f_0=0\). Then \(df_n=Z_n\) for \(1\le n\le N\), and hence \[\mathbb{E}\Bigl\|\sum_{k=1}^N Z_k\Bigr\|_{S_p} = \mathbb{E}\|f_N\|_{S_p} \le \mathbb{E}[\sup_{n\ge 0}\|f_n\|_{S_p}] \le C_p\, \mathbb{E}\Bigl(\sum_{k=1}^N \|Z_k\|_{S_p}^p\Bigr)^{\frac{1}{p}}.\] Using 2 on both sides, we obtain \[\begin{align} \mathbb{E}\Bigl\|\sum_{k=1}^N Z_k\Bigr\|_{\mathrm{nuc}} \le r^{1-1/p}\, \mathbb{E}\Bigl\|\sum_{k=1}^N Z_k\Bigr\|_{S_p} &\le C_p\,r^{1-1/p}\, \mathbb{E}\Bigl(\sum_{k=1}^N \|Z_k\|_{S_p}^p\Bigr)^{\frac{1}{p}} \\ &\le C_p\,r^{1-1/p}\, \mathbb{E}\Bigl(\sum_{k=1}^N \|Z_k\|_{\mathrm{nuc}}^p\Bigr)^{\frac{1}{p}}. \end{align}\] This proves the claim. \(\square\)

7.2 Dimension-dependent lower bound↩︎

We prove Theorem 1 in full and restate it with the logarithmic dimension condition made explicit after stating the key lemmas. Throughout this section, \(\mathrm{St}(d_1,d_2):=\{U\in\mathbb{R}^{d_1\times d_2}:U^\top U=I_{d_2}\}\) denotes the Stiefel manifold. For \(x\in\mathbb{R}^d\), let \(\operatorname{prog}_{\alpha}(x):=\sup(\{i\in[d]:|x_i|>\alpha\}\cup\{0\})\) and \(\operatorname{supp}(x):=\{i\in[d]:x_i\ne0\}\).

The lower bound is obtained from a distribution over hard objectives and stochastic gradient oracles. The construction starts from the following zero-chain instance of [40].

We embed independent copies of the chain in disjoint row–column blocks. Fix \(m\le n\) and \(n_{\rm blk}\ge1\) with \(mn_{\rm blk}\le n\). Choose disjoint column sets \(C_1,\ldots,C_m\subseteq[n]\) with \(|C_i|=n_{\rm blk}\), and let \(c_{i,1}<\cdots<c_{i,n_{\rm blk}}\) enumerate \(C_i\). The unfold map and its right-inverse are \[[\Pi(X)]_{i,k}:=X_{i,c_{i,k}}, \qquad [\Pi^*(y)]_{i',j}:=\begin{cases} y_{i,k} & (i',j)=(i,c_{i,k}),\\ 0 & \text{else,}\end{cases}\] mapping \(\mathbb{R}^{m\times n}\leftrightarrow\mathbb{R}^{m\times n_{\rm blk}}\). The image of \(\Pi^*\) is the linear subspace \(\mathcal{S}^*:=\{Y\in\mathbb{R}^{m\times n}:Y_{i',j}=0\text{ for }(i',j)\notin\bigcup_{i\in[m]}\{i\}\times C_i\}\), and for smooth \(h:\mathbb{R}^{m\times n_{\rm blk}}\to\mathbb{R}\) we have \(\nabla(h\circ\Pi)(X)=\Pi^*(\nabla h(\Pi(X)))\in\mathcal{S}^*\).

On this subspace we average the row-wise chains. For \(W\in\mathbb{R}^{m\times T}\), define \[\label{eq:unscaled-tensor-chain} H_T(W):=\frac{1}{m}\sum_{i=1}^m \phi_T(W_{i,:}),\tag{3}\] and let \(I\sim\operatorname{Unif}([m])\), \(Z\sim\operatorname{Bernoulli}(q)\) be independent, \(q\in(0,1]\). The stochastic gradient oracle \(\bar g_T(W,I,Z)\in\mathbb{R}^{m\times T}\) is given by \[\label{eq:unscaled-prob-oracle} [\bar g_T(W,I,Z)]_{i,j} :=\tfrac{1}{m}[\nabla\phi_T(W_{i,:})]_j\bigl[1+\mathbf{1}\{j>\operatorname{prog}_{1/4}(W_{i,:})\}(m\mathbf{1}\{I=i\}Z/q-1)\bigr].\tag{4}\]

We next hide the active coordinates by independent random rotations in each block. For \(U=(U_1,\ldots,U_m)\in\mathrm{St}(n_{\rm blk},T)^m\) and \(X\in\mathbb{R}^{m\times n}\), define \([\Phi_U(X)]_{i,:}:=U_i^\top X_{i,C_i}\). The rotated oracle is \[\label{eq:bounded-rotated-oracle} \widetilde{H}_{T,U}(X):=H_T(\Phi_U(X)),\quad [\widetilde{g}_{T,U}(X,I,Z)]_{i,C_i}:=U_i\bigl[\bar g_T(\Phi_U(X),I,Z)\bigr]_{i,:},\tag{5}\] with \([\widetilde{g}_{T,U}(X,I,Z)]_{i,j}=0\) for \(j\notin C_i\). It follows that \(\nabla\widetilde{H}_{T,U}(X)\in\mathcal{S}^*\) with \([\nabla\widetilde{H}_{T,U}(X)]_{i,C_i}=U_i[\nabla H_T(\Phi_U(X))]_{i,:}\). This is the matrix analogue of the rotated oracle construction in [40].

The proof now separates into three ingredients. First, bounded queries cannot reveal enough hidden coordinates under the rotated oracle. Second, a soft projection removes the bounded-query assumption with the compressed objectives and oracles. Third, the compressed objectives lie in the desired family \(\mathcal{F}_{\rm{op}}(m,n,\Delta,L)\).

We now compress arbitrary queries into a bounded block before applying the rotated oracle. For constants \(C_\rho\) and \(\eta\) to be fixed below, set \(R=C_\rho\sqrt T\) and define the blockwise soft projection \[[\rho_R^{\rm blk}(X)]_{i,j}:=\begin{cases}[\rho_R(X_{i,C_i})]_k & j=c_{i,k}\in C_i,\\ 0 & j\notin C_i,\end{cases} \qquad \rho_R(z):=\tfrac{z}{\sqrt{1+\|z\|_2^2/R^2}},\] equivalently \(\rho_R^{\rm blk}=\Pi^*\circ\rho_R^{\otimes m}\circ\Pi\). Setting \(h_{T,U}(z):=\phi_T(U^\top\rho_R(z))+\tfrac{\eta}{2}\|z\|_2^2\), the compressed objective and its stochastic gradient oracle are \[\begin{align} \tag{6} \widehat F_{T,U}(X) &:=H_T(\Phi_U(\rho_R^{\rm blk}(X)))+\tfrac{\eta}{2m}\sum_{i=1}^m\|X_{i,C_i}\|_2^2=\frac{1}{m}\sum_{i=1}^m h_{T,U_i}(X_{i,C_i}),\\ \tag{7} [\widehat g_{T,U}(X,I,Z)]_{i,C_i} &:=J_R(X_{i,C_i})^\top U_i\bigl[\bar g_T(\Phi_U(\rho_R^{\rm blk}(X)),I,Z)\bigr]_{i,:}+\tfrac{\eta}{m}X_{i,C_i}, \end{align}\] with \([\widehat g_{T,U}(X,I,Z)]_{i,j}=0\) for \(j\notin C_i\), where \(J_R(z):=\nabla\rho_R(z)\). By construction, we have \(\widehat g_{T,U}(X,I,Z)\in\mathcal{S}^*\), \(\nabla\widehat F_{T,U}(X)\in\mathcal{S}^*\), and \(\widehat F_{T,U}\) does not depend on \(X_{i,j}\) for \(j\notin C_i\).

The next lemma transfers the bounded hardness to this compressed instance.

With these constants fixed, the final ingredient records the deterministic properties needed to embed the construction in the desired function class \(\mathcal{F}_{\rm{op}}(m,n,\Delta,L)\).

7.2.0.1 Theorem 1 (restated).

For any \(p\in(1,2]\), there exist constants \(c_p,c_p',C_{\rm dim}>0\) depending only on \(p\) such that for any \(m,n\ge1\) and \(\Delta,L,\sigma_0>0\), if \(0<\epsilon\le c_p'\min\{\sqrt{\Delta L},\sigma_0\}\) and \[\max\{m,n\}\ge C_{\rm dim}\min\{m,n\}^2 \tfrac{(\Delta L)^2}{\epsilon^4}(\sigma_0/\epsilon)^{p/(p-1)}\log\bigl(C_{\rm dim}\min\{m,n\}^2 \tfrac{(\Delta L)^2}{\epsilon^4}(\sigma_0/\epsilon)^{p/(p-1)}\bigr),\] then \[\mathfrak m^{\rm rand}_{\epsilon,p}(m,n,\Delta,L,\sigma_0)\ge c_p\min\{m,n\}\Delta L\sigma_0^{p/(p-1)}\epsilon^{-(3p-2)/(p-1)}.\]

The transposition invariance of \(\|\cdot\|_{\rm op}\) and \(\|\cdot\|_{\rm nuc}\) makes \(\mathfrak m^{\rm rand}_{\epsilon,p}\) invariant under \((m,n)\mapsto(n,m)\), so we assume \(m\le n\). Let \(\kappa,\Delta_0,\ell_1,\varsigma,C_\rho\) be the constants in Proposition 6 and Lemma 4, set \(A_2:=\kappa^2/(16\Delta_0\ell_1)\), and choose \(c_p'>0\) small enough that \[\label{eq:cp-prime-checks} c_p'\le 1,\quad c_p'\le \kappa/(4\varsigma),\quad A_2/(c_p')^2\ge 8\log 4+8.\tag{8}\] Set \[\label{eq:scaling-lambda-T-q} \lambda:=\tfrac{4\ell_1\epsilon}{\kappa L},\quad T:=\bigl\lfloor\tfrac{\Delta}{\Delta_0(L\lambda^2/\ell_1)}\bigr\rfloor,\quad q:=\bigl(\tfrac{\varsigma L\lambda}{\ell_1\sigma_0}\bigr)^{p/(p-1)},\tag{9}\] which gives the identities \(\Delta/[\Delta_0(L\lambda^2/\ell_1)]=A_2\Delta L/\epsilon^2\) and \(q=(4\varsigma\epsilon/(\kappa\sigma_0))^{p/(p-1)}\). The assumption \(\epsilon\le c_p'\min\{\sqrt{\Delta L},\sigma_0\}\) and 8 imply \(0<q\le 1\), \(A_2\Delta L/\epsilon^2\ge 8\log 4+8\), and hence \[\label{eq:T-lower-check} T\ge \tfrac{A_2}{2}\tfrac{\Delta L}{\epsilon^2}\ge 4\log 4+4.\tag{10}\] Set \(\delta=1/2\). Then \(N_{1/2}:=\lfloor(mT-2\log 4)/(4q)\rfloor\) satisfies \(N_{1/2}\ge 1\) and \(N_{1/2}+1>mT/(8q)\).

For a constant \(C_1\) depending only on the construction constants, the block size \(n_{\rm blk}:=\lceil T+N_{1/2}+C_{\rm rot}C_\rho^2 TN_{1/2}\log(2C_{\rm rot}mN_{1/2}T)\rceil\) satisfies \(mn_{\rm blk}\le C_1m^2T^2/q\cdot\log(C_1m^2T^2/q)\). By 9 there is \(C_2\) (depending only on \(p\) and construction constants) with \(T^2/q\le C_2(\Delta L)^2\epsilon^{-4}(\sigma_0/\epsilon)^{p/(p-1)}\). Taking \(C_{\rm dim}\ge \max\{C_1C_2,e\}\) ensures \(n\ge mn_{\rm blk}\), and we fix disjoint column blocks \(C_1,\ldots,C_m\subseteq[n]\) of size \(n_{\rm blk}\). The dimension condition in Lemma 4 holds with \(\delta=1/2\).

Let \(\mu\) be the product Haar measure on \(\mathrm{St}(n_{\rm blk},T)^m\). For \(U\sim\mu\), define \[\label{eq:main-def-F-g} F^*_{T,U}(X):=\tfrac{L\lambda^2}{\ell_1}\widehat F_{T,U}(X/\lambda),\quad g^*_{T,U}(X,I,Z):=\tfrac{L\lambda}{\ell_1}\widehat g_{T,U}(X/\lambda,I,Z),\tag{11}\] and let \(P_F\) be the push-forward of \(\mu\) by \(U\mapsto F^*_{T,U}\), \(P_{I,Z}\) the law of \((I,Z)\). By Proposition 6(4) and 11 , the map \(U\mapsto F^*_{T,U}\) is injective, so for \(F\) in the hard family \(\mathcal{F}_{\rm hard}:=\{F^*_{T,U}:U\in\mathrm{St}(n_{\rm blk},T)^m\}\) there is a unique inverse \(U(F)\). We define the oracle \[\label{eq:main-oracle-scheme} \mathsf O_F(X,I,Z):=\begin{cases} (F(X),g^*_{T,U(F)}(X,I,Z)), & F\in\mathcal{F}_{\rm hard},\\ (F(X),\nabla F(X)), & F\notin\mathcal{F}_{\rm hard}, \end{cases}\tag{12}\] so \(\mathsf O_{F^*_{T,U}}(X,I,Z)=(F^*_{T,U}(X),g^*_{T,U}(X,I,Z))\). For \(F\notin\mathcal{F}_{\rm hard}\) admissibility is immediate, so it suffices to verify admissibility on \(\mathcal{F}_{\rm hard}\).

Proposition 6(1) and 9 give \(F^*_{T,U}(0)-\inf_X F^*_{T,U}(X)\le (L\lambda^2/\ell_1)\Delta_0 T\le\Delta\). Proposition 6(2) yields \(\|\nabla F^*_{T,U}(X)-\nabla F^*_{T,U}(Y)\|_{\rm nuc}\le L\|X-Y\|_{\rm op}\), so \(P_F\in\mathcal{P}(\mathcal{F}_{\rm op}(m,n,\Delta,L))\). For the noise moment, Proposition 6(3) and 9 give \[\mathbb{E}_{I,Z}\|g^*_{T,U}(X,I,Z)-\nabla F^*_{T,U}(X)\|_{\rm nuc}^p \le (L\lambda/\ell_1)^p\varsigma^p q^{1-p}=\sigma_0^p.\] Thus \(\mathsf O\in\mathcal{O}_p(\sigma_0)\).

Fix \(\mathsf A\in\mathcal{A}_{\rm rand}\). Rescaling queries by \(\lambda^{-1}\) turns the trajectory of \(\mathsf A\) on \(\mathsf O\) into that of a randomized algorithm \(\mathsf B\) on the compressed oracle \(\widehat{\mathsf O}\). Lemma 4 applied to \(\mathsf B\) with \(\delta=1/2\) gives, with probability at least \(1/2\), \[\|\nabla\widehat F_{T,U}(X_{\mathsf A[\mathsf O]}^{(t)}/\lambda)\|_{\rm nuc}\ge\kappa\quad\forall t\le N_{1/2},\] and multiplication by \(L\lambda/\ell_1=4\epsilon/\kappa\) gives \(\|\nabla F^*_{T,U}(X_{\mathsf A[\mathsf O]}^{(t)})\|_{\rm nuc}\ge 4\epsilon\) on the same event. Hence \[\mathbb{E}\|\nabla F^*_{T,U}(X_{\mathsf A[\mathsf O]}^{(t)})\|_{\rm nuc}\ge 2\epsilon>\epsilon\] for every \(t\le N_{1/2}\), and consequently \[\inf_{\mathsf A}\inf\bigl\{N\in\mathbb{N}:\mathbb{E}_{F\sim P_F,\mathsf A,\mathsf{O}_F}\|\nabla F(X_{\mathsf A[\mathsf O_F]}^{(N)})\|_{\rm nuc}\le\epsilon\bigr\}\ge N_{1/2}+1.\] Combining with \(N_{1/2}+1>mT/(8q)\ge \tfrac{m}{8}\cdot\tfrac{A_2}{2}\tfrac{\Delta L}{\epsilon^2}(\kappa\sigma_0/(4\varsigma\epsilon))^{p/(p-1)}\) yields the claimed lower bound \(c_p m\Delta L\sigma_0^{p/(p-1)}\epsilon^{-(3p-2)/(p-1)}\). \(\Box\)

7.3 Proof of Lemma 3↩︎

To begin, we introduce some notation. Let \(\theta\) be the algorithm’s internal seed and write \(X^{(t)}\in\mathbb{R}^{m\times n}\) for the round-\(t\) query, set \(\xi^{(t)}:=(I^{(t)},Z^{(t)})\), and define \(\Gamma_{t,i}:=\max\bigl(\{j\le T:\exists s\le t,\,[\bar g_T(\Phi_U(X^{(s)}),\xi^{(s)})]_{i,j}\ne 0\}\cup\{0\}\bigr)\) with \(\Gamma_{0,i}=0\). Let \[\mathcal{E}_i:=\sigma(\theta,\{U_a:a\ne i\},\{\xi^{(t)}:t\le N_\delta\}),\quad \mathsf S_{j,i}^{(s)}:=\operatorname{span}\bigl(\{U_ie_\ell:\ell\le j\}\cup\{X_{i,C_i}^{(a)}:a\le s\}\bigr),\] and let \(P_j^{(s)}\) be the orthogonal projection onto \((\mathsf S_{j,i}^{(s)})^\perp\), \(\Pi_j^{(s)}:=I-P_j^{(s)}\). Define \[\mathcal{U}_{j,i}^{(s)}:=\mathcal{E}_i\vee\sigma\bigl(\{U_ie_\ell:\ell\le j\},\{U_i^\top X_{i,C_i}^{(a)}:a\le s\}\bigr),\] \[\mathcal{T}_{j,i}^{(s-1)}:=\{j>\Gamma_{s-1,i}\},\quad \mathcal{G}_{j,i}^{(s)}:=\bigl\{\|\Pi_{j-1}^{(s)}P_{j-1}^{(s-1)}U_ie_j\|_2<\tfrac{1}{4R\sqrt{N_\delta}}\bigr\}.\] Here, \(\mathcal{T}_{j,i}^{(s-1)}\) is the event that coordinate \(j\) is still unrevealed, while \(\mathcal{G}_{j,i}^{(s)}\) controls how much the new query leaks about the residual part of \(U_ie_j\).

On \(\mathcal{T}_{j,i}^{(s-1)}\), every previous row-\(i\) response is supported on \([j-1]\) (by definition of \(\Gamma_{s-1,i}\)), so \[\label{eq:block-r-response-reconstruction} [\widetilde{g}_{T,U}(X^{(a)},\xi^{(a)})]_{i,C_i}=\sum_{\ell=1}^{j-1}[\bar g_T(\Phi_U(X^{(a)}),\xi^{(a)})]_{i,\ell}U_ie_\ell\quad(a<s).\tag{13}\] Function values factor as \(\widetilde{H}_{T,U}(X^{(a)})=m^{-1}\phi_T(U_i^\top X_{i,C_i}^{(a)})+m^{-1}\sum_{i'\ne i}\phi_T(U_{i'}^\top X_{i',C_{i'}}^{(a)})\). Since the first query depends only on \(\theta\), induction over rounds shows the queries \(\{X_{i,C_i}^{(a)}\}_{a\leq s}\) are \(\mathcal{U}_{j-1,i}^{(s-1)}\)-measurable on \(\mathcal{T}_{j,i}^{(s-1)}\), and hence so are \(\Gamma_{s-1,i}\), \(\mathcal{T}_{j,i}^{(s-1)}\), \(P_{j-1}^{(s-1)}\), and \(\Pi_{j-1}^{(s)}\).

We extract key linear algebra components from [40] and adapt them to our setting.

Nesting gives \(\Pi^{(i')}=\Pi^{(i'-1)}+P^{(i'-1)}\Pi^{(i')}P^{(i'-1)}\). Iterating, we get \[\Pi^{(i)}=\Pi^{(0)}+\sum_{i'\le i}P^{(i'-1)}\Pi^{(i')}P^{(i'-1)}.\] The summands lie in pairwise orthogonal subspaces \(S^{(i')}\cap(S^{(i'-1)})^\perp\), and \(\Pi^{(0)}u=0\) since \(u\in(S^{(0)})^\perp\). Hence, \(\|\Pi^{(i)}u\|_2^2=\sum_{i'\le i}\|P^{(i'-1)}\Pi^{(i')}P^{(i'-1)}u\|_2^2\le\sum_{i'\le i}\|\Pi^{(i')}P^{(i'-1)}u\|_2^2<i/(16R^2N)\). For \(x\in S^{(i)}\), \(|\langle u,x\rangle|=|\langle \Pi^{(i)}u,x\rangle|\le\|\Pi^{(i)}u\|_2\|x\|_2\). \(\Box\)

We follow the techniques from [40] to show the conditional uniformity in our block-wise setting.

We write \(U_{i,<j}:=[U_i e_1,\ldots,U_i e_{j-1}]\), \(U_{i,\ge j}:=[U_i e_j,\ldots,U_i e_T]\), \(\mathcal{U}:=\mathcal{U}_{j-1,i}^{(s-1)}\), and \(\mathcal{T}:=\mathcal{T}_{j,i}^{(s-1)}\). Let \(p_{\rm rot}\) denote the Haar law of \(U_i\) on \(\mathrm{St}(n_{\rm blk},T)\), with marginal density \(p_{{\rm rot},<j}\) for the first \(j-1\) columns and conditional probability \(p_{{\rm rot},\ge j\mid <j}\) for the last \(T-j+1\) columns given the first \(j-1\), and hence \[\label{eq:haar-disintegration} p_{\rm rot}(U_i) = p_{{\rm rot},<j}(U_{i,<j})\, p_{{\rm rot},\ge j\mid<j}(U_{i,\ge j}\mid U_{i,<j}).\tag{14}\] We define \(Y^{(<s)}:=\bigl\{U_i^\top X^{(a)}_{i,C_i}:1\le a<s\bigr\}\). The factor denoted below by \(\widetilde{p}^{(<s)}(Y^{(<s)},\mathcal{T}\mid U_i,\mathcal{E}_i)\) is a Dirac delta of \((Y^{(<s)},\mathbf{1}_{\mathcal{T}})\) given \((U_i,\mathcal{E}_i)\). Once \(U_i\) and \(\mathcal{E}_i\) are fixed, \((Y^{(<s)},\mathbf{1}_{\mathcal{T}})\) is deterministic. Moreover, let \(\widetilde{p}^{(<s)}_{<j}(Y^{(<s)},\mathcal{T}\mid U_{i,<j},\mathcal{E}_i)\) denote the conditional probability density of \(Y^{(<s)},\mathcal{T}\) given \(U_{i,<j}\) and \(\mathcal{E}_i\). Bayes’ rule gives \[\label{eq:block-bayes-density} p_{\ge j}(U_{i,\ge j}\mid \mathcal{U},\mathcal{T}) = \frac{ \widetilde{p}^{(<s)}(Y^{(<s)},\mathcal{T}\mid U_i,\mathcal{E}_i)\, p_{{\rm rot},\ge j\mid<j}(U_{i,\ge j}\mid U_{i,<j}) }{ \widetilde{p}^{(<s)}_{<j}(Y^{(<s)},\mathcal{T}\mid U_{i,<j},\mathcal{E}_i) },\tag{15}\] where the prior density \(p_{{\rm rot},<j}(U_{i,<j})\) appears in both numerator and denominator and has cancelled.

Let \(W\) be any orthogonal map of \(\mathbb{R}^{n_{\rm blk}}\) that fixes \(\mathsf S_{j-1,i}^{(s-1)}\) pointwise, i.e., \[\label{eq:block-W-fixes-S} W^\top W=I,\;Wz=z,\;\forall z\in \{U_i e_\ell:1\le\ell\le j-1\}\cup\{X^{(a)}_{i,C_i}:1\le a\le s-1\}.\tag{16}\] We set \(U_i':=WU_i\), leaving all other \(\{U_{i'}\}_{i'\not=i}\), the oracle seeds \(\xi^{(<s)}\), and the algorithm seed \(\theta\) unchanged. Let \(X^{\prime(a)}\) be the matrix queries produced in this alternative execution. We first prove, by induction over rounds, that the full previous matrix queries agree, given by \[\label{eq:block-induction-claim} X^{\prime(a)}=X^{(a)},\quad \forall a<s.\tag{17}\] If \(s=1\), there are no previous rounds to compare. Otherwise, the first query depends only on \(\theta\), so the claim holds for \(a=1\). We assume it holds through round \(a-1\), where \(2\le a\le s-1\). Since the row-\(i\) blocks \(X^{(a')}_{i,C_i}\) with \(a'<a\) belong to \(\mathsf S_{j-1,i}^{(s-1)}\) and \(W\) fixes \(\mathsf S_{j-1,i}^{(s-1)}\), we have \[(U_i')^\top X^{\prime(a')}_{i,C_i} =(WU_i)^\top X^{(a')}_{i,C_i} =U_i^\top W^\top X^{(a')}_{i,C_i} =U_i^\top X^{(a')}_{i,C_i}, \quad \forall a'<a.\] Matrices \(\{U_{i'}\}_{i'\not=i}\) are unchanged, and by the induction hypothesis the corresponding query rows are also unchanged. Hence, we have \[\Phi_{U'}(X^{\prime(a')})=\Phi_U(X^{(a')}), \quad \forall a'<a.\] Therefore, the unrotated function values and unrotated stochastic gradient responses agree for all previous queries. The rotated function values also agree. The rotated stochastic gradients for rows \(i'\ne i\) agree because the unrotated stochastic gradients agree. For row \(i\), if the unrotated stochastic gradient is \(c\in\mathbb{R}^T\), then on \(\mathcal{T}=\{j>\Gamma_{s-1,i}\}\) every previous row-\(i\) stochastic gradient is supported on \([j-1]\). Thus \(U_i c\in\operatorname{span}\{U_i e_\ell:1\le\ell\le j-1\}\subseteq\mathsf S_{j-1,i}^{(s-1)}\), and \(U_i'c=WU_i c=U_i c\). So the full oracle responses through round \(a-1\) are identical in the two executions. Since the algorithm is a measurable function of \(\theta\) and the previous responses, it produces the same round-\(a\) matrix queries. This closes the induction and proves 17 .

It follows from 17 and \(W|_{\mathsf S_{j-1,i}^{(s-1)}}=I\) that \(Y^{(<s)}\) is unchanged when \(U_i\) is replaced by \(WU_i\). The same argument also shows that the event \(\mathcal{T}=\{j>\Gamma_{s-1,i}\}\) is unchanged. Hence, we have \[\label{eq:block-consistency-invariance} \widetilde{p}^{(<s)}(Y^{(<s)},\mathcal{T}\mid WU_i,\mathcal{E}_i) = \widetilde{p}^{(<s)}(Y^{(<s)},\mathcal{T}\mid U_i,\mathcal{E}_i).\tag{18}\] Moreover, because \(W\) fixes \(U_{i,<j}\), we have \[\label{eq:block-haar-left-invariance} p_{{\rm rot},\ge j\mid<j}(W U_{i,\ge j}\mid U_{i,<j}) = p_{{\rm rot},\ge j\mid<j}(U_{i,\ge j}\mid U_{i,<j}).\tag{19}\] The denominator in 15 depends on the conditioning only through \(U_{i,<j}\) and \(\mathcal{E}_i\), so it is also unchanged. Combining 15 , 18 , and 19 gives \[\label{eq:block-tail-density-invariance} p_{\ge j}(W U_{i,\ge j}\mid \mathcal{U},\mathcal{T}) = p_{\ge j}(U_{i,\ge j}\mid \mathcal{U},\mathcal{T}).\tag{20}\]

Conditional on \(\mathcal{U}\) and \(\mathcal{T}\), let \(p_j(U_ie_j\mid\mathcal{U},\mathcal{T})\) denote the marginal density of \(U_ie_j\), obtained from the joint conditional density of \(U_{i,\ge j}\) by integrating out \(U_{i,\ge j+1}\): \[p_j(U_ie_j\mid\mathcal{U},\mathcal{T}) =\int p_{\ge j}\bigl(U_{i,\geq j}\bigm|\mathcal{U},\mathcal{T}\bigr) \,d(U_{i,\geq j+1}).\] Let \(W\) be any orthogonal map of \(\mathbb{R}^{n_{\rm blk}}\) fixing \(\mathsf S_{j-1,i}^{(s-1)}\) pointwise. Applying 20 to the matrix \(WU_i=[U_{i,<j},WU_ie_j,WU_ie_{j+1},\ldots,WU_ie_T]\) and recalling that \(W\) fixes \(U_{i,<j}\), we have \[p_{\ge j}\bigl([WU_ie_j,WU_ie_{j+1},\ldots,WU_ie_T]\bigm|\mathcal{U},\mathcal{T}\bigr) =p_{\ge j}\bigl([U_ie_j,U_ie_{j+1},\ldots,U_ie_T]\bigm|\mathcal{U},\mathcal{T}\bigr).\] The substitution \(d(U_{i,\geq j+1})\mapsto(WU_ie_{j+1},\ldots,WU_ie_{T})\) is an orthogonal transformation of \(\mathbb{R}^{n_{\rm blk}}\times\cdots\times\mathbb{R}^{n_{\rm blk}}\), so it preserves the integration measure \(d(U_{i,\geq j+1})\). Changing variables in the integral on the right-hand side gives \[\begin{align} p_j(WU_ie_j\mid\mathcal{U},\mathcal{T}) &=\int p_{\ge j}\bigl([WU_ie_j,u_{j+1}',\ldots,u_T']\bigm|\mathcal{U},\mathcal{T}\bigr) \,d(u_{j+1}',\ldots,u_T')\\ &=\int p_{\ge j}\bigl([WU_ie_j,WU_ie_{j+1},\ldots,WU_ie_T]\bigm|\mathcal{U},\mathcal{T}\bigr) \,d(U_{i,\geq j+1})\\ &=\int p_{\ge j}\bigl(U_{i,\geq j}\bigm|\mathcal{U},\mathcal{T}\bigr) \,d(U_{i,\geq j+1}) =p_j(U_ie_j\mid\mathcal{U},\mathcal{T}), \end{align}\] where the second equality is the change of variables \(u_\ell'=WU_ie_\ell\) for \(\ell\ge j+1\), and the third equality is 20 applied to the integrand.

Hence the conditional law of \(U_ie_j\) given \(\mathcal{U}\) and \(\mathcal{T}\) is invariant under \(u\mapsto Wu\) for every orthogonal map \(W\) of \(\mathbb{R}^{n_{\rm blk}}\) fixing \(\mathsf S_{j-1,i}^{(s-1)}\) pointwise. Since any orthogonal map \(Q\) of \((\mathsf S_{j-1,i}^{(s-1)})^\perp\) extends to such a \(W\) by acting as the identity on \(\mathsf S_{j-1,i}^{(s-1)}\), and since \(W\) fixes \(\mathsf S_{j-1,i}^{(s-1)}\) and acts as \(Q\) on \((\mathsf S_{j-1,i}^{(s-1)})^\perp\), we have \(P_{j-1}^{(s-1)}(WU_ie_j)=W(P_{j-1}^{(s-1)}U_ie_j)=Q(P_{j-1}^{(s-1)}U_ie_j)\). Thus, the conditional law of \(P_{j-1}^{(s-1)}U_ie_j\) is invariant under every orthogonal transformation of \((\mathsf S_{j-1,i}^{(s-1)})^\perp\). \(\Box\) The following lemma establishes an upper bound for the information leaked in one step.

If \(n_{\rm blk}-s-j\le 0\), then the right-hand side is at least \(2\). Hence, we assume \(n_{\rm blk}-s-j>0\). Write \(P:=P_{j-1}^{(s-1)}\) and \(\Pi:=\Pi_{j-1}^{(s)}\). Since \(\mathsf S_{j-1,i}^{(s-1)}\) is spanned by at most \((j-1)+(s-1)\) vectors, \(d_0:=\dim(\operatorname{range}P)\ge n_{\rm blk}-s-j\). The operator \(\Pi|_{\operatorname{range}P}\) is the orthogonal projection from \(\operatorname{range}P\) onto \(\mathsf S_{j-1,i}^{(s)}\cap\operatorname{range}P\), which has dimension \(d_1\le 1\).

By Lemma 6, the conditional law of \(PU_ie_j\) on \(\mathcal{T}_{j,i}^{(s-1)}\) given \(\mathcal{U}_{j-1,i}^{(s-1)}\) is invariant under every orthogonal transformation of \(\operatorname{range}P\). Letting \(V\sim\operatorname{Unif}(\mathbb{S}(\operatorname{range}P))\) and choosing a \(\mathcal{U}_{j-1,i}^{(s-1)}\)-measurable orthonormal basis \((e'_1,\ldots,e'_{d_0})\) of \(\operatorname{range}P\) whose first \(d_1\) vectors span \(\Pi(\operatorname{range}P)\), we have \(\|\Pi PU_i e_j\|_2^2\leq \sum_{\ell=1}^{d_1}\langle e'_\ell,V\rangle^2\) since \(\|PU_ie_j\|_2\leq 1\). Thus, we have \[\mathbb{P}\bigl((\mathcal{G}_{j,i}^{(s)})^c \cap \mathcal{T}_{j,i}^{(s-1)}\mid\mathcal{U}_{j-1,i}^{(s-1)}\bigr) \le \mathbb{P}\!\bigl(\textstyle\sum_{\ell=1}^{d_1}v_\ell^2\ge \tfrac{1}{16R^2N_\delta}\bigr) \le \mathbb{P}\!\bigl(v_1^2\ge\tfrac{1}{16R^2N_\delta}\bigr)\le 2e^{-d_0/(64R^2N_\delta)},\] where \(v\sim\operatorname{Unif}(\mathbb{S}^{d_0-1})\), the second inequality uses \(d_1\le 1\), and the third is the spherical tail bound \(\mathbb{P}(v_1^2\ge\alpha)\le 2e^{-\alpha d_0/4}\) [141]. This completes the proof. \(\Box\)

We next connect leakage control to the zero-chain progress recursion.

Unbiasedness follows from \(\mathbb{E}[m\mathbf{1}\{I=i\}Z/q]=1\). For the moment bound, we set \(a_i:=m^{-1}\nabla\phi_T(W_{i,:})\odot\mathbf{1}_{>\operatorname{prog}_{1/4}(W_{i,:})}\), so \(\|a_i\|_2\le G_{\rm ch}/m\) by Lemma 2(3), and \(\bar g_T-\nabla H_T\) has \(i\)-th row \(a_i(m\mathbf{1}\{I=i\}Z/q-1)\). On \(\{Z=0\}\) the row-norm sum is at most \(G_{\rm ch}\); on \(\{Z=1,I=i\}\) it is at most \(G_{\rm ch}+G_{\rm ch}/q\le 2G_{\rm ch}/q\). The \(p\)-th moment is \(\le G_{\rm ch}^p+q(2G_{\rm ch}/q)^p\le(3G_{\rm ch})^p q^{1-p}\). For the progress recursion, Lemma 2(4) gives \(\operatorname{supp}\nabla\phi_T(W_{i,:}^{(t)})\subseteq[\Gamma_{t-1,i}+1]\) on the event, and coordinates beyond \(\operatorname{prog}_{1/4}(W_{i,:}^{(t)})\) are nonzero only when \(I^{(t)}=i\), \(Z^{(t)}=1\), so \(\Gamma_{t,i}-\Gamma_{t-1,i}\le\mathbf{1}\{I^{(t)}=i\}Z^{(t)}\); summing gives the claim. \(\Box\)

Proof of Lemma 3. Condition on \(\theta\). Define \(V_{t,i}:=\{\operatorname{prog}_{1/4}(U_i^\top X_{i,C_i}^{(t)})\le\Gamma_{t-1,i}\}\) and \(\mathcal{A}_i:=\bigcap_{s\le N_\delta,j\le T}(\mathcal{G}_{j,i}^{(s)}\cup\{j\le\Gamma_{s-1,i}\})\). We claim \(\mathcal{A}_i\subseteq\bigcap_{s\le N_\delta}V_{s,i}\). If \(j>\Gamma_{s-1,i}\) on \(\mathcal{A}_i\), then \(\mathcal{G}_{j,i}^{(a)}\) holds for all \(a\le s\), and Lemma 5 gives \(|\langle U_ie_j,X_{i,C_i}^{(s)}\rangle|<1/4\), so \(j\) does not contribute to \(\operatorname{prog}_{1/4}(U_i^\top X_{i,C_i}^{(s)})\). A union bound with Lemma 7 gives \(\mathbb{P}(\mathcal{A}_i^c)\le 2N_\delta T e^{-(n_{\rm blk}-N_\delta-T)/(64R^2N_\delta)}\), and a further union over \(i\in[m]\) yields \[\label{eq:rotation-failure-bound} \mathbb{P}\bigl(\bigl[\textstyle\bigcap_{t,i}V_{t,i}\bigr]^c\bigr)\le 2mN_\delta T e^{-(n_{\rm blk}-N_\delta-T)/(64R^2N_\delta)}\le\delta/2,\tag{21}\] provided \(C_{\rm rot}\) is taken large enough.

Let \(\Gamma_t:=\sum_i\Gamma_{t,i}\). On \(\bigcap_{s\le t,i}V_{s,i}\), Lemma 8 gives \(\Gamma_t-\Gamma_{t-1}\le Z^{(t)}\), so by the Chernoff bound \(\mathbb{P}(S\ge a)\le e^{2\mathbb{E} S-a}\) for a Bernoulli sum \(S\) with \(a\ge 2\mathbb{E} S\), taking \(a=mT/2\) and \(\mathbb{E} S=qN_\delta\le(mT-2\log(2/\delta))/4\) (by ?? ), we have \[\label{eq:progress-chernoff} \mathbb{P}\bigl(\Gamma_{N_\delta}\ge mT/2,\,\textstyle\bigcap_{t,i}V_{t,i}\bigr)\le e^{2qN_\delta-mT/2}\le\delta/2.\tag{22}\] Combining 21 and 22 , with probability at least \(1-\delta\), \(V_{t,i}\) holds for all \(t\le N_\delta\), \(i\in[m]\), and \(\Gamma_t<mT/2\). Since fewer than \(m/2\) rows can have \(\Gamma_{t-1,i}=T\), at least \(m/2\) rows have \(\Gamma_{t-1,i}<T\), and for each such \(i\), \(V_{t,i}\) gives \(\operatorname{prog}_{1/4}(U_i^\top X_{i,C_i}^{(t)})<T\). The bound is uniform in \(\theta\), hence unconditional. \(\square\)

7.4 Proof of Lemma 4↩︎

The reduction to the bounded case uses the soft projection together with two facts: a norm identity on \(\mathcal{S}^*\) and a gradient lower bound for one soft-projected block.

Nonzero rows of \(Y\in\mathcal{S}^*\) have disjoint column supports by construction, hence are pairwise orthogonal. Then \(YY^\top\) is diagonal with \((YY^\top)_{ii}=\|Y_{i,:}\|_2^2\), so the nonzero singular values of \(Y\) are \(\{\|Y_{i,:}\|_2:Y_{i,:}\ne 0\}\), giving the nuclear-norm identity. The operator bound follows from \(\|X_{i,C_i}\|_2\le\|X_{i,:}\|_2=\|X^\top e_i\|_2\le\|X\|_{\rm op}\). \(\Box\)

Let \(s(z):=\sqrt{1+\|z\|_2^2/R^2}\); then \(\|\rho_R(z)\|_2=\|z\|_2/s(z)\le R\). The Jacobian \(J_R(z)=s^{-1}I-s^{-3}zz^\top/R^2\) has eigenvalues \(s^{-1}\) (on \(z^\perp\)) and \(s^{-3}\) (on \(\operatorname{span}(z)\)), so \(\|J_R\|_{\rm op}\le 1\), and the Lipschitz bound for \(\rho_R\) follows by the mean-value theorem. For the Jacobian Lipschitz bound it suffices to set \(R=1\). Writing \(a(x):=(1+\|x\|_2^2)^{-1/2}\), we have \(J_1(x)=aI-a^3xx^\top\), and a direct computation yields \(\nabla J_1(x)[h]=-a^3\langle x,h\rangle I+3a^5\langle x,h\rangle xx^\top-a^3(hx^\top+xh^\top)\). For unit \(h\) and \(r=\|x\|_2\), \(\|\nabla J_1(x)[h]\|_{\rm op}\le 3a^3r+3a^5r^3\le 6\) since \(a^3r,a^5r^3\le 1\). Rescaling completes the proof. \(\Box\)

For \(h^0_{T,U}(z):=\phi_T(U^\top\rho_R(z))\), the chain rule gives \(\nabla h^0_{T,U}(z)=J_R(z)^\top U\nabla\phi_T(U^\top\rho_R(z))\). Using \(\|J_R\|_{\rm op}\le 1\), \(U^\top U=I\), Lemma 2, and Lemma 10, \[\|\nabla h^0_{T,U}(z)-\nabla h^0_{T,U}(z')\|_2\le \ell_{\rm ch}\|z-z'\|_2+ (6/R)\|z-z'\|_2\cdot G_{\rm ch}\sqrt T\le(\ell_{\rm ch}+6G_{\rm ch}/C_\rho)\|z-z'\|_2.\] The quadratic adds \(\eta\) in the smoothness parameter, so \(\ell_1:=\ell_{\rm ch}+6G_{\rm ch}/C_\rho+\eta\) works.

Choose \(\eta\in(0,1)\) with \(\eta\sqrt 5/2\le 1/16\), and \(C_\rho\ge 1\) with \(G_{\rm ch}/(2C_\rho)\le 1/16\) and \(\eta C_\rho/2-G_{\rm ch}\ge 1\). Let \(y:=\rho_R(z)\), \(v:=\nabla\phi_T(U^\top y)\), \(s:=\sqrt{1+\|z\|_2^2/R^2}\) (so \(z=sy\)), \(j:=\operatorname{prog}_1(U^\top y)+1\in[T]\), and \(u:=Ue_j\). Lemma 2(5) gives \(|\langle u,Uv\rangle|=|v_j|\ge 1\) and \(|\langle u,y\rangle|=|(U^\top y)_j|\le 1\). With \(\nabla h_{T,U}(z)=J_R(z)^\top Uv+\eta z\), we consider two cases.

If \(\|z\|_2\le R/2\), then \(s^{-1}\ge 2/\sqrt 5\) and \(\|y\|_2\le R/2\), giving \[|\langle u,\nabla h_{T,U}(z)\rangle|\ge s^{-1}|v_j|-s^{-1}\|y\|_2\|v\|_2/R^2-\eta s|\langle u,y\rangle| \ge \tfrac{2}{\sqrt 5}-\tfrac{G_{\rm ch}}{2C_\rho}-\tfrac{\eta\sqrt 5}{2}\ge \tfrac{2}{\sqrt 5}-\tfrac{1}{8}.\] If \(\|z\|_2>R/2\), then \(\|\nabla h_{T,U}(z)\|_2\ge \eta\|z\|_2-\|J_R^\top Uv\|_2>\eta R/2-G_{\rm ch}\sqrt T=(\eta C_\rho/2-G_{\rm ch})\sqrt T\ge 1\). Thus \(\kappa_0:=\min\{1,2/\sqrt 5-1/8\}\) suffices. \(\Box\)

Proof of Lemma 4. Condition on the seed of \(\mathsf A\). Define a deterministic simulator \(\mathsf B\) for the bounded rotated oracle: when \(\mathsf A\) proposes \(X^{(t)}\), \(\mathsf B\) submits \(Y^{(t)}:=\rho_R^{\rm blk}(X^{(t)})\), with \(\|Y_{i,C_i}^{(t)}\|_2\le R\) for each \(i,t\). Given the returned \((\widetilde{H}_{T,U}(Y^{(t)}),\widetilde{g}_{T,U}(Y^{(t)},\xi^{(t)}))\), \(\mathsf B\) reconstructs \[\begin{align} \widehat F_{T,U}(X^{(t)}) &= \widetilde{H}_{T,U}(Y^{(t)})+\tfrac{\eta}{2m}\textstyle\sum_i\|X_{i,C_i}^{(t)}\|_2^2,\\ [\widehat g_{T,U}(X^{(t)},\xi^{(t)})]_{i,C_i}&= J_R(X_{i,C_i}^{(t)})^\top[\widetilde{g}_{T,U}(Y^{(t)},\xi^{(t)})]_{i,C_i}+\tfrac{\eta}{m}X_{i,C_i}^{(t)} \end{align}\] using \(\Phi_U(\rho_R^{\rm blk}(X))=\Phi_U(Y)\). By induction on \(t\), the simulated trajectory of \(\mathsf A\) is distributed exactly as that of \(\mathsf A\) on \(\widehat{\mathsf O}\), with \(Y^{(t)}=\rho_R^{\rm blk}(X_{\mathsf A[\widehat{\mathsf O}]}^{(t)})\).

Since \(\mathsf B\) is \(R\)-bounded with \(R=C_\rho\sqrt T\), dimension assumption ?? matches that of Lemma 3, giving an event of probability \(\ge 1-\delta\) on which \(|\mathcal{I}_t|\ge m/2\) for every \(t\le N_\delta\), where \(\mathcal{I}_t:=\{i:\operatorname{prog}_{1/4}(U_i^\top\rho_R(X_{i,C_i}^{(t)}))<T\}\). For each \(i\in\mathcal{I}_t\), Lemma 11 gives \(\|\nabla h_{T,U_i}(X_{i,C_i}^{(t)})\|_2\ge\kappa_0\). Since \(\nabla\widehat F_{T,U}(X^{(t)})\in\mathcal{S}^*\) has \(i\)-th row \(m^{-1}\nabla h_{T,U_i}(X_{i,C_i}^{(t)})\), Lemma 9 yields \[\|\nabla\widehat F_{T,U}(X^{(t)})\|_{\rm nuc}=\textstyle\sum_i\|[\nabla\widehat F_{T,U}(X^{(t)})]_{i,:}\|_2\ge m^{-1}\sum_{i\in\mathcal{I}_t}\|\nabla h_{T,U_i}(X_{i,C_i}^{(t)})\|_2\ge(m/2)\kappa_0/m=:\kappa.\] Removing the conditioning on the algorithm seed yields ?? . \(\square\)

7.5 Proof of Proposition 6↩︎

We verify the four items in order.

Since \(\rho_R(0)=0\), \(\widehat F_{T,U}(0)=\phi_T(0)\), and \(\widehat F_{T,U}(X)\ge\inf_u\phi_T(u)\). By Lemma 2(1), \(\widehat F_{T,U}(0)-\inf_X\widehat F_{T,U}(X)\le\Delta_{\rm ch}T\), so \(\Delta_0:=\Delta_{\rm ch}\) suffices.

Since \(\widehat F_{T,U}=m^{-1}\sum_i h_{T,U_i}(X_{i,C_i})\) and \(\nabla\widehat F_{T,U}\in\mathcal{S}^*\), Lemma 9 and Lemma 11 give \[\begin{align} \|\nabla\widehat F_{T,U}(X)-\nabla\widehat F_{T,U}(Y)\|_{\rm nuc}=m^{-1}\textstyle\sum_i\|\nabla h_{T,U_i}(X_{i,C_i})-\nabla h_{T,U_i}(Y_{i,C_i})\|_2 &\le m^{-1}\ell_1\sum_i\|X_{i,C_i}-Y_{i,C_i}\|_2 \\ &\le\ell_1\|X-Y\|_{\rm op}. \end{align}\]

Let \(W:=\Phi_U(\rho_R^{\rm blk}(X))\). By Lemma 8, \(\mathbb{E}_{I,Z}\bar g_T(W,I,Z)=\nabla H_T(W)\); the chain rule on \(h_{T,U_i}\) then gives \(\mathbb{E}_{I,Z}[\widehat g_{T,U}(X,I,Z)]_{i,C_i}=J_R(X_{i,C_i})^\top U_i[\nabla H_T(W)]_{i,:}+(\eta/m)X_{i,C_i}=[\nabla\widehat F_{T,U}(X)]_{i,C_i}\). The deterministic quadratic cancels in the noise, so \(\widehat g_{T,U}-\nabla\widehat F_{T,U}\in\mathcal{S}^*\). Lemma 9, \(\|J_R\|_{\rm op}\le 1\), and \(U_i^\top U_i=I_T\) give \[\|\widehat g_{T,U}-\nabla\widehat F_{T,U}\|_{\rm nuc}=\textstyle\sum_i\|[\widehat g_{T,U}-\nabla\widehat F_{T,U}]_{i,:}\|_2\le\sum_i\|[\bar g_T(W,I,Z)-\nabla H_T(W)]_{i,:}\|_2,\] and raising to the \(p\)-th power and taking expectation yields the moment bound via Lemma 8.

We first show that the single-block map \(U\mapsto h_{T,U}\) is injective on \(\mathrm{St}(n_{\rm blk},T)\). Suppose \(U,V\in\mathrm{St}(n_{\rm blk},T)\) satisfy \(h_{T,U}(z)-h_{T,V}(z)=\text{const}\) on \(\mathbb{R}^{n_{\rm blk}}\). Since \(\rho_R\) is a diffeomorphism from \(\mathbb{R}^{n_{\rm blk}}\) onto the open ball \(B_R:=\{y:\|y\|_2<R\}\) and the quadratic terms in \(h_{T,U},h_{T,V}\) coincide, \(\phi_T(U^\top y)-\phi_T(V^\top y)=\text{const}\) on \(B_R\). Differentiating, \[\label{eq:hidden-ident-grad} U\nabla\phi_T(U^\top y)=V\nabla\phi_T(V^\top y),\quad y\in B_R.\tag{23}\] We recover the columns inductively. At \(y=0\), Lemma 2(4)–(5) gives \(\nabla\phi_T(0)=a_1 e_1\) with \(a_1\ne 0\), so 23 yields \(Ue_1=Ve_1\). Assume \(Ue_\ell=Ve_\ell\) for \(\ell<j\), \(2\le j\le T\). Choose \(a_j>1\) with \(a_j\sqrt{j-1}<R\) (possible since \(R=C_\rho\sqrt T\ge \sqrt T>\sqrt{j-1}\)), and set \(u^{(j)}:=a_j\sum_{\ell<j}e_\ell\), \(y^{(j)}:=Uu^{(j)}=Vu^{(j)}\in B_R\). Then \(U^\top y^{(j)}=V^\top y^{(j)}=u^{(j)}\), \(\operatorname{prog}_1(u^{(j)})=j-1<T\), and \(\operatorname{prog}_{1/2}(u^{(j)})=j-1\), so Lemma 2(4)–(5) gives \(\operatorname{supp}\nabla\phi_T(u^{(j)})\subseteq[j]\) and \(|[\nabla\phi_T(u^{(j)})]_j|>1\). Applying 23 at \(y=y^{(j)}\) and using the induction hypothesis to cancel columns \(1,\ldots,j-1\) gives \([\nabla\phi_T(u^{(j)})]_j(Ue_j-Ve_j)=0\), hence \(Ue_j=Ve_j\). By induction, \(U=V\).

For the multi-block case, suppose \(\widehat F_{T,U}=\widehat F_{T,V}\) for \(U,V\in\mathrm{St}(n_{\rm blk},T)^m\). Fix a row \(i\) and set all blocks other than \(i\) to zero: from 6 , \(h_{T,U_i}(z)-h_{T,V_i}(z)=\sum_{r\ne i}(h_{T,V_r}(0)-h_{T,U_r}(0))\) is independent of \(z\), so the single-block claim gives \(U_i=V_i\). Hence \(U=V\), proving identifiability.

7.6 Smooth and nonconvex problems↩︎

By the fundamental theorem of calculus, we have \[\begin{align} F(Y)-F(X)-\langle\nabla F(X),Y-X\rangle &= \int_0^1 \langle \nabla F(X+t(Y-X))-\nabla F(X),\,Y-X\rangle\,dt \\ &\le \int_0^1 \|\nabla F(X+t(Y-X))-\nabla F(X)\|_\star\,\|Y-X\|\,dt \\ &\le \int_0^1 (L_0+L_1\|\nabla F(X)\|_\star)\,t\,\|Y-X\|^2\,dt \\ &= \frac{L_0+L_1\|\nabla F(X)\|_\star}{2}\|Y-X\|^2. \end{align}\] This finishes the proof. \(\Box\)

Since \(\sum_{s=0}^{t-1} a_s\bar\zeta_s =\sum_{s=0}^{t-1}\sum_{i=1}^B \frac{a_s}{B}\zeta_s^i\), we first check that the lexicographically ordered array \(Z_{s,i}:=\frac{a_s}{B}\zeta_s^i\), \(s=0,\dots,t-1\), \(i\in[B]\), is a martingale difference sequence with respect to its natural filtration. Let \(\mathcal{N}_{s,i}:=\sigma(Z_{r,j}:(r,j)<_{\rm lex}(s,i))\). Since \(\mathcal{N}_{s,i}\subseteq \mathcal{H}_s\vee\sigma(\zeta_s^1,\dots,\zeta_s^{i-1})\), conditional independence and \(\mathbb{E}[\zeta_s^i\mid\mathcal{H}_s]=0\) imply \[\mathbb{E}[Z_{s,i}\mid\mathcal{N}_{s,i}] = \mathbb{E}\!\left[ \mathbb{E}[Z_{s,i}\mid \mathcal{H}_s\vee\sigma(\zeta_s^1,\dots,\zeta_s^{i-1})] \mid \mathcal{N}_{s,i}\right] =0.\] Therefore, by the definition of \(\tau_\star\) in Eq. 1 , we have \[\label{eq:E-sum-as-zetas} \mathbb{E}\Bigl\|\sum_{s=0}^{t-1} a_s\bar\zeta_s\Bigr\|_\star = \mathbb{E}\Bigl\|\sum_{s=0}^{t-1}\sum_{i=1}^B \frac{a_s}{B}\zeta_s^i\Bigr\|_\star \le \frac{\tau_\star}{B} \mathbb{E}\left[ \left( \sum_{s=0}^{t-1}\sum_{i=1}^B a_s^p\|\zeta_s^i\|_\star^p \right)^{\frac{1}{p}} \right].\tag{24}\] It remains to bound the last expectation. For any \(k\) and any nonnegative \(\mathcal{H}_k\)-measurable random variable \(U\), we have \[\begin{align} \mathbb{E}\left[ \left. \left(U+a_k^p\sum_{i=1}^B\|\zeta_k^i\|_\star^p\right)^{\frac{1}{p}} \right|\mathcal{H}_k \right] &\le \left( U+a_k^p\sum_{i=1}^B \mathbb{E}[\|\zeta_k^i\|_\star^p\mid\mathcal{H}_k] \right)^{\frac{1}{p}} \\ &\le \left(U+B a_k^p\sigma_0^p+B a_k^p\sigma_1^p\Gamma_k^p\right)^{\frac{1}{p}}\\ &\le \left(U+B a_k^p\sigma_0^p\right)^{\frac{1}{p}} +B^{\frac{1}{p}}a_k\sigma_1\Gamma_k. \end{align}\] Here the first inequality follows from the concavity of \(x\mapsto x^{1/p}\) and Jensen’s inequality, the second uses the conditional moment assumption, and the third uses \((a+b)^{1/p}\le a^{1/p}+b^{1/p}\) for \(a,b\ge 0\). Applying this inequality successively for \(k=t-1,t-2,\dots,0\) gives \[\mathbb{E}\left[ \left( \sum_{s=0}^{t-1}\sum_{i=1}^B a_s^p\|\zeta_s^i\|_\star^p \right)^{\frac{1}{p}} \right] \le B^{\frac{1}{p}}\sigma_0 \Bigl(\sum_{s=0}^{t-1} a_s^p\Bigr)^{\frac{1}{p}} + B^{\frac{1}{p}}\sigma_1 \sum_{s=0}^{t-1} a_s\,\mathbb{E}\Gamma_s.\] Substituting the above into Eq. 24 yields the desired result. \(\Box\)

Let \(\zeta_s^i:=G(X_s,\xi_s^i)-\nabla F(X_s)\) and \(\bar\zeta_s:=\bar G_s-\nabla F(X_s)\). By Assumption 2 and the fresh i.i.d. mini-batches at each query point \(X_s\), Lemma 12 applies with \(\Gamma_s=\|\nabla F(X_s)\|_\star\). Also, \(m_1-\nabla F(X_0)=\bar\zeta_0\). For \(t\ge 1\), \[m_{t+1}-\nabla F(X_t) = \beta(m_t-\nabla F(X_{t-1})) +\beta\bigl(\nabla F(X_{t-1})-\nabla F(X_t)\bigr) +(1-\beta)\bar\zeta_t.\] Iterating, we get \[m_{t+1}-\nabla F(X_t) = \beta^t\bar\zeta_0 +\sum_{s=0}^{t-1} \beta^{\,t-s}\bigl(\nabla F(X_s)-\nabla F(X_{s+1})\bigr) +\sum_{s=1}^{t} (1-\beta)\beta^{\,t-s}\bar\zeta_s,\] where the sums are empty when \(t=0\). Using Assumption 1, we have \(\|\nabla F(X_s)-\nabla F(X_{s+1})\|_\star \le \eta(L_0+L_1\|\nabla F(X_s)\|_\star)\). Therefore, \[\begin{align} \label{eq:E-norm-mt-gt} \mathbb{E}\|m_{t+1}-\nabla F(X_t)\|_\star &\le \beta^t\,\mathbb{E}\|\bar\zeta_0\|_\star +\sum_{s=0}^{t-1} \beta^{\,t-s}\eta\bigl(L_0+L_1\mathbb{E}\|\nabla F(X_s)\|_\star\bigr) \notag\\ &\quad+ \mathbb{E}\Bigl\|\sum_{s=1}^{t} (1-\beta)\beta^{\,t-s}\bar\zeta_s\Bigr\|_\star. \end{align}\tag{25}\] It remains to control the batched noise terms. Taking \(t=1\) and \(a_0=1\) in Lemma 12 gives \[\mathbb{E}\|\bar\zeta_0\|_\star \le \frac{\tau_\star}{B^{\frac{p-1}{p}}} \bigl(\sigma_0+\sigma_1\|\nabla F(X_0)\|_\star\bigr).\] For \(t\ge 1\), applying Lemma 12 after shifting the index set to the batches \(1,\dots,t\) with weights \(a_{s-1}=(1-\beta)\beta^{\,t-s}\) gives \[\begin{align} &\mathbb{E}\Bigl\|\sum_{s=1}^{t} (1-\beta)\beta^{\,t-s}\bar\zeta_s\Bigr\|_\star\\ &\le \frac{\tau_\star}{B^{\frac{p-1}{p}}} \left[ \sigma_0(1-\beta)\left(\sum_{s=1}^{t}\beta^{p(t-s)}\right)^{\frac{1}{p}} +\sigma_1\sum_{s=1}^{t} (1-\beta)\beta^{\,t-s}\mathbb{E}\|\nabla F(X_s)\|_\star \right] \\ &\le \frac{\tau_\star}{B^{\frac{p-1}{p}}} \left[ \sigma_0(1-\beta)^{\frac{p-1}{p}} +\sigma_1\sum_{s=1}^{t} (1-\beta)\beta^{\,t-s}\mathbb{E}\|\nabla F(X_s)\|_\star \right]. \end{align}\] Combining the above two inequalities with Eq. 25 yields the desired result. \(\Box\)

Proof of Theorem 2. Let \(\mathcal{A}_T:=\mathbb{E}\|\nabla F(\widetilde{X}_T)\|_\star\), \(r=\frac{p-1}{p}\), \(S_0:=\sigma_0+\sigma_1\|\nabla F(X_0)\|_\star\). Since \(\langle m_{t+1},\operatorname{lmo}(m_{t+1})\rangle=-\|m_{t+1}\|_\star\) and \(\|\operatorname{lmo}(m_{t+1})\|\le1\), we have \[\langle \nabla F(X_t),\operatorname{lmo}(m_{t+1})\rangle \leq -\|\nabla F(X_t)\|_\star+2\|m_{t+1}-\nabla F(X_t)\|_\star .\] Using Proposition 7 and \(\eta\le1/L_1\), we obtain \[F(X_{t+1}) \le F(X_t)-\eta\|\nabla F(X_t)\|_\star +2\eta\|m_{t+1}-\nabla F(X_t)\|_\star +\frac{\eta^2}{2}\bigl(L_0+L_1\|\nabla F(X_t)\|_\star\bigr).\] Summing over \(t=0,\ldots,T-1\) gives \[\label{eq:avg-descent-lemma} \left(1-\frac{L_1\eta}{2}\right)\mathcal{A}_T \le \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{2} +\frac{2}{T}\sum_{t=0}^{T-1}\mathbb{E}\|m_{t+1}-\nabla F(X_t)\|_\star .\tag{26}\]

Summing Lemma 13 over \(t=0,\ldots,T-1\) and using \[\sum_{t=0}^{T-1}\beta^t\le\frac{1}{\alpha}, \qquad \sum_{t=0}^{T-1}\sum_{s=0}^{t-1}\beta^{t-s}\le\frac{T}{\alpha}, \qquad \sum_{t=0}^{T-1}\sum_{s=0}^{t-1} \beta^{t-s}\mathbb{E}\|\nabla F(X_s)\|_\star \le \frac{T}{\alpha}\mathcal{A}_T,\] and \[\sum_{t=0}^{T-1}\sum_{s=1}^{t} \alpha\beta^{t-s}\mathbb{E}\|\nabla F(X_s)\|_\star \le T\mathcal{A}_T,\] we obtain \[\begin{align} \frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\|m_{t+1}-\nabla F(X_t)\|_\star &\le \frac{\tau_\star S_0}{B^r\alpha T} +\frac{\tau_\star\sigma_0\alpha^r}{B^r} +\frac{L_0\eta}{\alpha} +\left(\frac{L_1\eta}{\alpha}+\frac{\tau_\star\sigma_1}{B^r}\right)\mathcal{A}_T . \end{align}\] Combining this with Eq. 26 yields \[\label{eq:unknown-p-key} \begin{align} \left(1-\frac{L_1\eta}{2}-\frac{2L_1\eta}{\alpha}-\frac{2\tau_\star\sigma_1}{B^r}\right)\mathcal{A}_T &\le \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{2} +\frac{2\tau_\star S_0}{B^r\alpha T} +\frac{2\tau_\star\sigma_0\alpha^r}{B^r} +\frac{2L_0\eta}{\alpha}. \end{align}\tag{27}\] By the choices of \(B\) and \(\eta\), \(2\tau_\star\sigma_1/B^r\le1/8\), \(2L_1\eta/\alpha\le1/4\), and \(L_1\eta/2\le1/16\). Hence the left coefficient in Eq. 27 is at least \(9/16\), and therefore \[\mathcal{A}_T \le \frac{40}{9}\left[ \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{\alpha} +\frac{\tau_\star S_0}{B^r\alpha T} +\frac{\tau_\star\sigma_0\alpha^r}{B^r} \right].\] Since \(\frac{\Delta_0}{\eta T}\le\frac{8\Delta_0L_1}{\alpha T}+\sqrt{\frac{\Delta_0L_0}{\alpha T}}\) and \(\frac{L_0\eta}{\alpha}\le\sqrt{\frac{\Delta_0L_0}{\alpha T}}\), we have \[\label{eq:sharp-buscg-reduced} \mathcal{A}_T \le \frac{320}{9}\left[ \frac{A_0}{\alpha T} + \sqrt{\frac{\Delta_0L_0}{\alpha T}} + \frac{\tau_\star\sigma_0\alpha^r}{B^r} \right].\tag{28}\]

Let \(U_\tau =\tfrac{A_0^{\frac{p}{2p-1}}B^{\frac{p-1}{2p-1}}}{(\tau_\star\sigma_0T)^{\frac{p}{2p-1}}}\) and \(V_\tau=\tfrac{(\Delta_0L_0)^{\frac{p}{3p-2}}B^{\frac{2p-2}{3p-2}}}{(\tau_\star\sigma_0)^{\frac{2p}{3p-2}}T^{\frac{p}{3p-2}}}\). If \(\alpha=1\), then either \(U_\tau\ge1\) or \(V_\tau\ge1\). If \(U_\tau\geq 1\), then \(\tau_\star\sigma_0/B^r\le A_0/T\). If \(V_\tau\geq 1\), then \(\tau_\star\sigma_0/B^r\le\sqrt{\Delta_0L_0/T}\). This implies the desired result. It remains to consider \(\alpha<1\). Then \(\alpha\ge U_\tau\), \(\alpha\geq V_\tau\), and \(\alpha^r\le U_\tau^r+V_\tau^r\). Hence, we have \[\frac{A_0}{\alpha T} \le \frac{A_0^{\frac{r}{1+r}}(\tau_\star\sigma_0)^{\frac{1}{1+r}}}{(BT)^{\frac{r}{1+r}}}, \quad \sqrt{\frac{\Delta_0L_0}{\alpha T}} \le \frac{(\Delta_0L_0)^{\frac{r}{1+2r}}(\tau_\star\sigma_0)^{\frac{1}{1+2r}}}{(BT)^{\frac{r}{1+2r}}},\] and \[\frac{\tau_\star\sigma_0\alpha^r}{B^r} \le \frac{A_0^{\frac{r}{1+r}}(\tau_\star\sigma_0)^{\frac{1}{1+r}}}{(BT)^{\frac{r}{1+r}}} + \frac{(\Delta_0L_0)^{\frac{r}{1+2r}}(\tau_\star\sigma_0)^{\frac{1}{1+2r}}}{(BT)^{\frac{r}{1+2r}}}.\] Substituting these bounds into Eq. 28 proves the desired result. The oracle complexity follows by setting \(N=BT\) and solving the leading term \(\frac{(\Delta_0L_0)^{\frac{r}{1+2r}}(\tau_\star\sigma_0)^{\frac{1}{1+2r}}}{N^{\frac{r}{1+2r}}}\le\epsilon\), which gives \[N=O(\Delta_0L_0(\tau_\star\sigma_0)^{1/r}\epsilon^{-(1+2r)/r}),\] equivalently the stated bound. \(\square\)

Proof of Theorem 3. Letting \(\sigma_1=0\) in Eq. 27 , we obtain \[\begin{align} \left( 1-\frac{L_1\eta}{2} -\frac{2L_1\eta}{\alpha} \right)\mathcal{A}_T &\le \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{2} +\frac{2\tau_\star\sigma_0}{\alpha T} +{2\tau_\star\sigma_0\alpha^{\frac{p-1}{p}}} +\frac{2L_0\eta}{\alpha}. \end{align}\] This implies \[\label{eq:unknown-p-master} \mathcal{A}_T \le \frac{16}{11} \left[ \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{2} +\frac{2\tau_\star\sigma_0}{\alpha T} + 2\tau_\star\sigma_0\alpha^{\frac{p-1}{p}} + \frac{2L_0\eta}{\alpha} \right],\tag{29}\] which yields the desired result. The complexity bound follows immediately. \(\square\)

7.7 Highly smooth and nonconvex problems↩︎

Proof of Theorem 4. Let \(r:=\frac{p-1}{p}\), \(\alpha:=1-\beta\), \(\tau_\star:=\tau(\|\cdot\|_\star,m,n,p)\), \(S_0:=\sigma_0+\sigma_1\|\nabla F(X_0)\|_\star\), and \(\Delta_0:=F(X_0)-F^\star\). Define \(\mathcal{A}_T:=\mathbb{E}\|\nabla F(\widetilde{X}_T)\|_\star\). Since Algorithm [algorithm:transported-uSCG] uses \(X_{t+1}=X_t+\eta\operatorname{lmo}(m_{t+1})\), the proof of Eq. 26 gives \[\label{eq:avg-descent-lemma-2} \left(1-\frac{L_1\eta}{2}\right)\mathcal{A}_T \le \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{2} +\frac{2}{T}\sum_{t=0}^{T-1} \mathbb{E}\|m_{t+1}-\nabla F(X_t)\|_\star.\tag{30}\] We next bound the fresh estimator error. Let \(g_t:=\nabla F(X_t)\), \(\bar\zeta_0:=\bar G_0-g_0\), and, for \(t\ge1\), \(\bar\zeta_t:=\bar G_t-\nabla F(Y_t)\). Set \(\zeta_0^i:=G(X_0,\xi_0^i)-\nabla F(X_0)\) and, for \(s\ge1\), \(\zeta_s^i:=G(Y_s,\xi_s^i)-\nabla F(Y_s)\). Then Lemma 12 applies with \(\Gamma_0=\|\nabla F(X_0)\|_\star\) and \(\Gamma_s=\|\nabla F(Y_s)\|_\star\) for \(s\ge1\), by Assumption 2 and the fresh i.i.d. mini-batches at the query points \(X_0,Y_1,\dots,Y_{T-1}\). For \(t\ge1\), \[\begin{align} m_{t+1}-g_t &= \beta (m_t-g_{t-1}) +\beta(g_{t-1}-g_t) +\alpha\bigl(\nabla F(Y_t)-g_t\bigr) +\alpha\bar\zeta_t. \end{align}\] Let \(Z(a,b):=\nabla F(a)-\nabla F(b)-\nabla^2F(b)[a-b]\). Since \(Y_t-X_t=\frac{\beta}{\alpha}(X_t-X_{t-1})\), the linear Hessian terms in \(\beta(g_{t-1}-g_t)+\alpha(\nabla F(Y_t)-g_t)\) cancel, and therefore \[m_{t+1}-g_t = \beta (m_t-g_{t-1}) +\beta Z(X_{t-1},X_t) +\alpha Z(Y_t,X_t) +\alpha\bar\zeta_t.\] Unrolling the recursion gives, for \(t=0,\dots,T-1\), \[\label{eq:transported-generalized-unroll} \begin{align} m_{t+1}-g_t &= \beta^t\bar\zeta_0 +\sum_{s=1}^t \beta^{\,t-s} \bigl[\beta Z(X_{s-1},X_s)+\alpha Z(Y_s,X_s)\bigr] +\sum_{s=1}^t \alpha\beta^{\,t-s}\bar\zeta_s. \end{align}\tag{31}\]

We bound the three terms on the right-hand side. First, Lemma 12 gives \[\mathbb{E}\|\bar\zeta_0\|_\star \le \frac{\tau_\star}{B^r} \bigl(\sigma_0+\sigma_1\|\nabla F(X_0)\|_\star\bigr).\]

Second, for the curvature term, Assumption 3 and \(\eta\le\alpha/(8L_1)\) give \(\|Z(X_{s-1},X_s)\|_\star\le L_2\eta^2\) and \(\|Z(Y_s,X_s)\|_\star\le L_2\eta^2/\alpha^2\). Hence \[\|\beta Z(X_{s-1},X_s)+\alpha Z(Y_s,X_s)\|_\star \le \beta L_2\eta^2+\alpha L_2\frac{\eta^2}{\alpha^2} \le \frac{2L_2\eta^2}{\alpha}.\] Thus \[\sum_{s=1}^t \beta^{\,t-s} \mathbb{E}\|\beta Z(X_{s-1},X_s)+\alpha Z(Y_s,X_s)\|_\star \le \frac{2L_2\eta^2}{\alpha^2}.\]

Third, Lemma 12 yields \[\mathbb{E}\Bigl\|\sum_{s=1}^t \alpha\beta^{t-s}\bar\zeta_s\Bigr\|_\star \le \frac{\tau_\star}{B^r} \left[ \sigma_0\alpha^r + \sigma_1\sum_{s=1}^t \alpha\beta^{\,t-s}\,\mathbb{E}\|\nabla F(Y_s)\|_\star \right].\]

We next compare \(\nabla F(Y_s)\) to \(\nabla F(X_s)\). Since \(\|Y_s-X_s\|\le \eta/\alpha\le 1/L_1\), Assumption 1 gives \[\|\nabla F(Y_s)-\nabla F(X_s)\|_\star \le \bigl(L_0+L_1\|\nabla F(X_s)\|_\star\bigr)\|Y_s-X_s\| \le \frac{L_0\eta}{\alpha} + \frac{L_1\eta}{\alpha}\|\nabla F(X_s)\|_\star.\] It follows that \[\|\nabla F(Y_s)\|_\star \le \Bigl(1+\frac{L_1\eta}{\alpha}\Bigr)\|\nabla F(X_s)\|_\star + \frac{L_0\eta}{\alpha}.\] Combining the bounds for the three terms in Eq. 31 and averaging over \(t=0,\dots,T-1\), we get \[\label{eq:transported-generalized-error-avg} \begin{align} \frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\|m_{t+1}-\nabla F(X_t)\|_\star &\le \frac{\tau_\star S_0}{B^r\alpha T} +\frac{\tau_\star\sigma_0\alpha^r}{B^r} +\frac{\tau_\star\sigma_1L_0\eta}{B^r\alpha} +\frac{2L_2\eta^2}{\alpha^2} +\frac{\tau_\star\sigma_1}{B^r} \Bigl(1+\frac{L_1\eta}{\alpha}\Bigr)\mathcal{A}_T. \end{align}\tag{32}\] Here we used \(\sum_{t=0}^{T-1}\sum_{s=1}^t \alpha\beta^{t-s}\mathbb{E}\|\nabla F(X_s)\|_\star\le T\mathcal{A}_T\).

Substituting 32 into 30 yields \[\begin{align} &\left( 1-\frac{L_1\eta}{2} -\frac{2\tau_\star\sigma_1}{B^r}\Bigl(1+\frac{L_1\eta}{\alpha}\Bigr) \right)\mathcal{A}_T \\ &\le \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{2} +\frac{2\tau_\star S_0}{B^r\alpha T} +\frac{2\tau_\star\sigma_0\alpha^r}{B^r} +\frac{2\tau_\star\sigma_1L_0\eta}{B^r\alpha} +\frac{4L_2\eta^2}{\alpha^2}. \end{align}\]

By \(\eta\le \alpha/(8L_1)\) and \(B\ge(8\tau_\star\sigma_1)^{1/r}\), the left-hand coefficient is at least \(1-\frac{1}{16}-\frac{1}{4}(1+\frac{1}{8})=\frac{21}{32}\). Therefore, we have \[\label{eq:theorem-transported-uSCG-key} \begin{align} \mathcal{A}_T &\le \frac{32}{21} \Biggl[ \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{2} +\frac{2\tau_\star S_0}{B^r\alpha T} +\frac{2\tau_\star\sigma_0\alpha^r}{B^r} +\frac{2\tau_\star\sigma_1L_0\eta}{B^r\alpha} +\frac{4L_2\eta^2}{\alpha^2} \Biggr]. \end{align}\tag{33}\] This implies the stated bound.

It remains to prove the sample complexity bound. For sufficiently small \(\epsilon>0\), choose the parameters in the theorem and \(T=20\lceil \frac{\Delta_0}{\eta\epsilon}\rceil\). The choice of \(B\) then satisfies \(B^r\ge 8\tau_\star\sigma_1\), and the left-hand coefficient above is at least \(7/8\). Hence Eq. 33 holds with the prefactor \(8/7\) in place of \(32/21\). Moreover, \(\alpha\le1\) and \(\eta\le\alpha/(8L_1)\). By construction, \[\frac{\Delta_0}{\eta T}\le\frac{\epsilon}{20},\qquad \frac{2\tau_\star\sigma_0\alpha^r}{B^r}\leq \frac{2\epsilon}{3}, \qquad \frac{4L_2\eta^2}{\alpha^2}\leq 0.01\,\epsilon, \qquad \frac{2\tau_\star\sigma_1L_0\eta}{B^r\alpha} \le 0.1\,\epsilon.\] Also, since \(B^r=O(\epsilon^{-1/2})\), we have \(\alpha=O(\epsilon^{1/(2r)})\), and hence \(L_0\eta=o(\epsilon)\). Similarly, \(\frac{2\tau_\star S_0}{B^r\alpha T}=o(\epsilon)\). Therefore \(\mathcal{A}_T\le\epsilon\) for all sufficiently small \(\epsilon\). It remains to count oracle calls. Letting \(N=BT\), we have \[N\le 20B+20\frac{B\Delta_0}{\eta\epsilon} = 20B+400\frac{B\Delta_0\sqrt{L_2}}{\alpha\epsilon^{3/2}} .\] Using \(\alpha=B\left(\epsilon/(3\tau_\star\sigma_0)\right)^{1/r}\), we get \(B/\alpha=3^{1/r}(\tau_\star\sigma_0/\epsilon)^{1/r}\). Absorbing the lower-order terms, we obtain \[N = O\left( 3^{1/r} \Delta_0\sqrt{L_2} (\tau_\star\sigma_0)^{1/r} \epsilon^{-3/2-1/r} \right).\] Since \(r=\frac{p-1}{p}\), this proves the desired sample complexity bound. \(\square\)

Proof of Theorem 5. Let \(r=\frac{p-1}{p}\) and \(\alpha=1-\beta\). With \(\sigma_1=0\), the key inequality preceding Eq. 33 and \(\eta\leq\frac{\alpha}{8L_1}\) imply \[\mathcal{A}_T \le 2 \left[ \frac{\Delta_0}{\eta T} +\frac{L_0\eta}{2} +\frac{2\tau_\star\sigma_0}{B^r\alpha T} +\frac{2\tau_\star\sigma_0\alpha^r}{B^r} +\frac{4L_2\eta^2}{\alpha^2} \right].\] By the definition of \(\eta\), \(\frac{1}{\eta}\le T^{5/7}+8L_1T^{4/7}\), so we have \(\frac{\Delta_0}{\eta T}\le\frac{\Delta_0}{T^{2/7}}+\frac{8\Delta_0L_1}{T^{3/7}}\). Moreover, \[\frac{L_0\eta}{2}\le \frac{L_0}{2T^{5/7}}, \qquad \frac{4L_2\eta^2}{\alpha^2}\le \frac{4L_2}{T^{2/7}},\qquad \frac{2\tau_\star\sigma_0}{B^r\alpha T} = \frac{2\tau_\star\sigma_0}{B^rT^{3/7}}, \qquad \frac{2\tau_\star\sigma_0\alpha^r}{B^r} = \frac{2\tau_\star\sigma_0}{B^rT^{4r/7}}.\] Since \(r\le 1/2\), we have \(T^{-3/7}\le T^{-4r/7}\) for \(T\ge1\). It follows that \(\frac{2\tau_\star\sigma_0}{B^rT^{3/7}} + \frac{2\tau_\star\sigma_0}{B^rT^{4r/7}} \le \frac{4\tau_\star\sigma_0}{B^rT^{4r/7}}.\) Combining the above gives the desired inequality and the sample complexity bound follows. \(\square\)

8 Additional Experiments↩︎

8.0.0.1 LLM experimental setup.

All LLM experiments use nanochat [74] trained on the NVIDIA ClimbMix dataset [75], with the number of training tokens specified in Table 3 and the number of split validation tokens being 42M for all models. The model backbone follows the nanochat GPT implementation and includes rotary position embeddings [142], RMSNorm [143], QK normalization [144], and local sliding-window attention [145] implemented with FlashAttention2 [146]. The model also uses untied token embeddings and language-model head, ReLU-squared MLPs, and value embeddings [147]. We use a vocabulary size of 32,768 and the standard nanochat tokenizer [74]. Table 3 lists the model configurations. The required training-token count is computed from the number of scalable parameters, defined as the sum of transformer matrix parameters and language-model head parameters.

We conducted our experiments using three nodes, each equipped with ten A40 GPUs. One node exhibited slightly slower execution times compared to the two nearly identical nodes. For the 287M and 539M models, we performed multiple runs across all nodes and reported the average runtime for Algorithm [algorithm:transported-uSCG] as 19.8 and 81.9 minutes, respectively; in comparison, other Muon-family methods ranged from 19.2–19.5 and 80.2–80.7 minutes. For the 1.39B model, the runtime on a faster node was 10.46 hours, compared to 10.26 hours for Muon with heavy ball momentum and Muon with Nesterov momentum.

8.0.0.2 LLM layerwise learning-rate multipliers.

Only transformer matrix parameters switch between AdamW and the Muon-family optimizers; all other parameters are trained with AdamW. Let \(s_B=\sqrt{B/2^{19}}\) denote the automatic batch-size multiplier, and let \(s_d=(d_{\mathrm{model}}/768)^{-1/2}\) denote the width multiplier applied to AdamW. Table 4 reports the learning rates and weight-decay settings for each parameter group, and Table 5 lists the batch-size and width multipliers for each model size. The learning-rate schedule uses a 40-step linear warmup, followed by a constant phase, and then a linear warmdown over the final 65% of training steps to a final multiplier of 0.05 [148]. For Muon, weight decay is cosine-decayed to zero. Its initial value is scaled as \(\lambda = \lambda_{\mathrm{ref}}\sqrt{B/B_{\mathrm{ref}}}\,(D_{\mathrm{ref}}/D)\), where \(\lambda_{\mathrm{ref}}=0.28\) and \(B_{\mathrm{ref}}=2^{19}\).

8.0.0.3 LLM hyperparameter.

We select LLM hyperparameters by grid search using validation loss. For AdamW parameters, we run a single-seed learning-rate sweep on 287M and 539M models. We use \(\beta\)s \((0.9,0.999)\) and test base learning rates \(\{0.0006, 0.0008, 0.0010, 0.0012, 0.0015, 0.0018, 0.0022\}\), before applying the automatic multipliers in Table 5. The best common learning rate is 0.0012, which we use for the 287M and 539M AdamW rows in Table 1. For the 1.39B model, we use an AdamW matrix learning rate of 0.0010, which gives a lower validation loss than 0.0012. For Muon without variance reduction and for NorMuon with heavy-ball momentum, Nesterov momentum, and transportation, we sweep learning rates on the 287M and 539M models to identify the best setting and evaluate the stability of Muon-family methods. We use heavy-ball momentum 0.90 and Nesterov momentum 0.95. The validation losses are reported in Tables 6 and 7. These methods are robust across learning rates, with the best base learning rate typically between 0.02 and 0.03. After selecting the best learning rate, we sweep the heavy-ball and Nesterov momentum factors on the 287M model over \(\{0.90,0.91,0.92,0.93,0.94,0.95\}\). On the 539M model, we sweep heavy-ball momentum over \(\{0.88,0.89,0.90,0.91,0.92,0.93\}\), and Nesterov momentum over \(\{0.91,0.92,0.93,0.94,0.95,0.96\}\). For transported NorMuon, we use the best heavy-ball momentum factors and set \(\alpha\) to 0.005, 0.0075, and 0.010 for the 287M, 539M, and 1.39B models, respectively. The 1.39B Muon-family settings are extrapolated from the 539M hyperparameter search rather than tuned with a separate sweep. The hyperparameters are listed in Table 8.

8.0.0.4 CNN experimental setup.

For the CNN experiments, we use the CIFARNET architecture [78], [79]. We use CIFAR-10 and CIFAR-100 [80] with the standard training and validation split. CIFARNET consists of a frozen \(2\times 2\) whitening convolution with a trainable bias, followed by three groups of convolutional layers with widths \((64,256,256)\) and a linear classifier. Each group applies a \(3\times 3\) convolution, max pooling, BatchNorm, GELU, another \(3\times 3\) convolution, BatchNorm, and a final GELU. Only the main \(3\times 3\) convolutional filters switch between AdamW, SGD with momentum, and the Muon-family optimizers. In all runs, the whitening bias, BatchNorm biases, and linear head are trained with SGD with Nesterov momentum. Before applying the Muon-family LMO, we flatten the convolutional filters into matrices of shape \(d_{\mathrm{out}}\times(d_{\mathrm{in}}k^2)\), where \(k\) is the convolution kernel size and \(d_{\mathrm{out}}\) and \(d_{\mathrm{in}}\) are the numbers of output and input channels. We train for 50 epochs with batch size 512, label smoothing 0.2, and global gradient clipping at 1.0. The learning-rate schedule uses a 5% linear warmup followed by cosine decay to zero. Data augmentation consists of random horizontal flips and reflection-padded translations of up to two pixels.

8.0.0.5 CNN hyperparameter.

For CIFARNET, we select hyperparameters by single-seed grid search using validation accuracy, and then report final results over five random seeds. For AdamW, we sweep the learning rate for the convolutional-filter over \(\{0.0003,0.001,0.003,0.01\}\). For SGD with momentum and the Muon-family methods, we sweep the learning rate for the convolutional-filter over \(\{0.01,0.02,0.03,0.0430316,0.05,0.075,0.10,0.15,0.20,0.24\}\), and we sweep the momentum factor over \(\{0.6,0.7,0.8,0.9,0.95\}\). The value 0.0430316 is the default learning rate from [78]. The selected hyperparameters are listed in Table 9.

8.0.0.6 Additional results.

For the LLM experiments, our results are consistent with prior empirical findings that Nesterov momentum outperforms heavy-ball momentum [22], [31], [74]. In Table 1, NorMuon with Nesterov momentum achieves better CORE metrics, which measure downstream performance. Scheduled Nesterov momentum is currently state of the art in nanochat [74], but it requires momentum warmup and warmdown. A full comparison among scheduled Nesterov momentum, heavy-ball momentum, and transportation would require tuning three momentum factors together with the warmup duration, which we leave for future work. The scheduled Nesterov configuration in nanochat linearly warms the Muon momentum from 0.85 to 0.97 over the first 400 steps, keeps it at 0.97, and then linearly decays it to 0.90 during the learning-rate warmdown phase, which spans the final 65% of training. We report these results in Table 10.

References↩︎