1 Introduction↩︎

Optimization algorithms are fundamental to the training of deep neural networks (DNNs). Since the introduction of stochastic gradient descent (SGD) [1], numerous optimizers have been proposed to accelerate convergence. Among these, adaptive gradient algorithms such as Adam [2] and its subsequent variant AdamW [3] have emerged as the de facto standard in modern deep learning, valued for their rapid convergence and robust performance across a wide range of tasks. A common characteristic of these widely used first-order optimizers is that they treat the weight parameters of neural networks, which are inherently matrices or higher-order tensors, as high-dimensional vectors. However, this vector-based perspective, while effective, disregards the underlying geometric and algebraic structure within the parameter matrices.

The recently proposed Muon optimizer [4] introduces a distinct paradigm that departs from the conventional vector-based viewpoint. The core idea of Muon is to preserve and leverage the intrinsic matrix structure of the network parameters. Instead of using the gradient vector, Muon computes its search direction by orthogonalizing the gradient momentum matrix. Specifically, for a given momentum matrix \(C_t \in \mathbb{R}^{m \times n}\), the search direction \(O_t \in \mathbb{R}^{m \times n}\) is the orthogonal matrix that is closest in Frobenius norm, found by solving \[\begin{align} \label{eq:ot} O_t := \underset{O \in \left\{O \in \mathbb{R}^{m\times n}:O^\top O=I_n \right\}}{\operatorname{argmin}} \| O-C_t\|_{\rm{F}}, \end{align}\tag{1}\] where \(I_n\) denotes a \(n \times n\) identity matrix. As established by classic matrix theory, if the singular value decomposition (SVD) of \(C_t\) is \(U_tS_tV_t^\top\), the optimal search direction is simply \(O_t = U_tV_t^\top\), where \(U_t \in \mathbb{R}^{m\times r}, S_t \in \mathbb{R}^{r\times r}, V_t \in \mathbb{R}^{n\times r}\), and \(r>0\) is the rank of \(C_t\). This process of gradient orthogonalization is aimed at finding a search direction that is independent of the gradient’s magnitude, potentially leading to more stable and effective training dynamics. While SVD provides an exact analytical solution, its high computational cost renders it impractical for the large matrices found in modern DNNs. As detailed by [5], the key computational innovation of Muon is the use of the Newton-Schulz iteration [6]–[9], a classic and remarkably efficient numerical method, to approximate this orthogonalization. This iterative algorithm enables Muon to compute the search direction without performing an explicit SVD, making it a computationally feasible optimizer for large-scale applications. This elegant fusion of a novel optimization perspective with a powerful numerical technique positions Muon as a promising, theoretically grounded alternative to existing optimizers.

Several studies have reported the strong empirical performance of Muon. [4] showed that Muon outperforms Shampoo [10] and SOAP [11], both on a per-step and wall-clock basis. [12] showed that by dividing the parameters and updating them with SOAP and Muon, optimization performance equivalent to or better than SOAP can be achieved while substantially reducing memory usage. [13] demonstrated that Muon is effective for training large-scale LLMs and suggested it has the potential to replace AdamW as the standard optimizer. [14] showed that Muon expands AdamW’s Pareto frontier on the compute-time plane, enlarging the practitioner’s flexibility in resource allocation. However, the theoretical understanding of Muon’s convergence behavior remains underdeveloped, and a formal justification for its strong performance over AdamW is still lacking. This work aims to bridge this gap by providing a rigorous convergence analysis of Muon.

Furthermore, since the batch size is a critical hyperparameter for managing computational costs in large-scale training, we also consider Muon’s critical batch size. This is defined as the batch size that minimizes the computational cost of training. In other words, the critical batch size is the point at which further increases in batch size yield diminishing returns in hardware throughput (i.e., the number of samples processed per unit time). By understanding and utilizing the critical batch size, one can maximize GPU utilization, thereby shortening training times and reducing overall computing costs. Therefore, to theoretically understand Muon and maximize its performance, analyzing its critical batch size is essential. Following previous studies, we aim to clarify the critical batch size of Muon.

Our main contributions are as follows:

• We present a convergence analysis for four variants of Muon—with and without Nesterov momentum and with and without weight decay (Theorems 1 and 2). We showed that the upper bound on the average expected gradient norm is as follows: \[\begin{align} \frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\| \nabla f(W_t) \|_* \right] = \begin{cases} \displaystyle \mathcal{O}\left( \frac{1}{\eta T} + \sqrt{\frac{(1-\beta)r_1}{b}} + \bar{r}\eta \right) & \text{without weight decay,} \\[10pt] \displaystyle \mathcal{O}\left( \frac{1}{\eta T} + \sqrt{\frac{(1-\beta)r_1}{b}} + \hat{r}\eta \right) & \text{with weight decay,} \end{cases} \end{align}\] where \(W_t \in \mathbb{R}^{m \times n}\), \(\eta > 0\) is a learning rate, \(\beta \in [0,1)\) is momentum, \(b\) is the batch size, \(r_1 \leq n\) is maximum rank of \(C_t - \nabla f(W_t)\), \(r_2 \leq n\) is maximum rank of \(O_t\), \(r_3 \leq n\) is maximum rank of \(W_t\), \(\bar{r}:=\max\{r_1, r_2\}\), and \(\hat{r} := \max\{r_1, r_2, r_3 \}\). For detailed definitions, see Section 2 and Theorems 1 and 2. The variant combining both Nesterov momentum and weight decay is of particular interest as it mirrors common practical settings. Our analysis thus offers direct insights into Muon’s real-world behavior.

• We prove that incorporating weight decay ensures almost-sure boundedness of both the parameter and gradient norms (Propositions 1 and 2) without assuming bounded gradients, and present experimental results to support this finding. We also show that for Muon to converge, the learning rate \(\eta\) and weight decay coefficient \(\lambda\) must satisfy \(\eta \leq \frac{1}{\lambda}\), a condition supported by our experimental results (see Figure 2).

• We derive a lower bound on the critical batch size for the four Muon variants (Proposition 6). For example, for Muon with Nesterov momentum and weight decay we show \[\begin{align} b_{\text{Muon}}^\star > \dfrac{9(1-\beta)(1+\sqrt{2}\beta)^2r_1\sigma^2}{(1-\lambda)^2\epsilon^2}, \end{align}\] where \(\beta \in (0,1]\) is momentum, \(\lambda>0\) is the weight decay coefficient, \(\sigma^2 > 0\) is the variance of the stochastic gradient, \(r_1\) is the maximum rank of \(C_t - \nabla f(W_t)\), and \(\epsilon > 0\) is the stopping threshold. Because the formula contains problem-dependent quantities (\(\sigma^2\), \(r_1\), \(\epsilon\)) that are not directly controlled by the practitioner, it does not predict the absolute value of the critical batch size. Instead, it identifies the qualitative scaling law: how the hyperparameters \(\beta\) and \(\lambda\) govern \(b^\star\), and how the ratio between variants depends only on \(\beta\) and \(\lambda\) (the unknown quantities cancel). Our experiments validate these hyperparameter-dependent predictions (see Figures 4–5).

2 Preliminaries↩︎

2.1 Notations and Definition↩︎

Let \(\mathbb{N}\) be the set of nonnegative integers. For \(p \in \mathbb{N} \setminus \{0\}\), define \([p] := \{1,2,\ldots,p\}\). Let \(\mathbb{R}^d\) be a \(d\)-dimensional Euclidean space. We use lowercase letters for scalars (e.g., \(a \in \mathbb{R})\), bold lowercase letters for vectors (e.g., \(\boldsymbol{a} \in \mathbb{R}^d\)), and uppercase letters for matrices (e.g., \(A \in \mathbb{R}^{m \times n}\)). \(\boldsymbol{a}^\top \in \mathbb{R}^{1\times d}\) and \(A^\top \in \mathbb{R}^{n \times m}\) denote the transposes of \(\boldsymbol{a} \in \mathbb{R}^d\) and \(A \in \mathbb{R}^{m\times n}\), respectively. For a square matrix \(A = (a_{ij}) \in \mathbb{R}^{n \times n}\), the trace is defined as \(\operatorname{tr}(A) := \sum_{i=1}^{n} a_{ii}\). For all vectors \(\boldsymbol{x}, \boldsymbol{y} \in \mathbb{R}^d\), the Euclidean inner product is defined as \(\langle \boldsymbol{x}, \boldsymbol{y} \rangle_2 := \boldsymbol{x}^\top \boldsymbol{y}\) and the Euclidean norm is defined as \(\| \boldsymbol{x} \|_2 := \sqrt{\langle \boldsymbol{x}, \boldsymbol{x} \rangle}_2\). For all matrices \(A, B \in \mathbb{R}^{m \times n}\), the Frobenius inner product is defined as \(\langle A, B \rangle_{\rm{F}} := \operatorname{tr}(A^\top B)\), the Frobenius norm defined as \(\| A \|_{\rm{F}} := \sqrt{\langle A, A \rangle_{\rm{F}}}\), the operator norm defined as \(\| A \|_{\rm{op}} := \max_{\| \boldsymbol{x} \|_2 \leq 1} \| A\boldsymbol{x} \|_2\), and the nuclear norm defined as \(\| A \|_{*} := \sum_{i=1}^{\min\{m, n\}} s_i(A)\), where \(s_i(A)\) denotes the \(i\)-th singular value of \(A\). The model is parameterized by a matrix \(W \in \mathbb{R}^{m \times n}\) \((m \geq n)\), which is optimized by minimizing the empirical loss function \(f(W) := \frac{1}{N} \sum_{i \in [N]} f_i(W)\), where \(N \in \mathbb{R}\) is the number of training samples and \(f_i(W)\) denotes the loss associated with the \(i\)-th training sample \(\boldsymbol{z}_i\) \((i \in [N])\). We define \(W^\star := \operatorname{argmin}_{W\in\mathbb{R}^{m \times n}} f(W)\). Let \(\xi\) be a random variable that does not depend on \(W \in \mathbb{R}^{m \times n}\), and let \(\mathbb{E}_{\xi}[X]\) denote the expectation with respect to \(\xi\) of a random variable \(X\). \(\xi_{t,i}\) is a random variable generated from the \(i\)-th sampling at time \(t\), and \(\xi_{t,i}\) and \(\xi_{t,j}\) are independent (\(i\neq j\)). \(\boldsymbol{\xi}_t := (\xi_{t,1}, \xi_{t,2}, \ldots, \xi_{t,b})^\top\) is independent of sequence \((W_k)_{k=0}^{t} \subset \mathbb{R}^{m \times n}\) generated by Muon (Algorithm 1), where \(b\) \((\leq N)\) is the batch size. The independence of \(\boldsymbol{\xi}_0, \boldsymbol{\xi}_1, \ldots\) allows us to define the total expectation \(\mathbb{E}\) as \(\mathbb{E}=\mathbb{E}_{\boldsymbol{\xi}_0} \mathbb{E}_{\boldsymbol{\xi}_1} \cdots \mathbb{E}_{\boldsymbol{\xi}_t}\). Let \(\mathsf{G}_{\xi}(W)\) be the stochastic gradient of \(f(\cdot)\) at \(W \in \mathbb{R}^{m \times n}\). The mini-batch \(\mathcal{S}_t\) consists of \(b\) samples at time \(t\), and the mini-batch stochastic gradient of \(f(W_t)\) for \(\mathcal{S}_t\) is defined as \(\nabla f_{\mathcal{S}_t} (W_t) := \frac{1}{b} \sum_{i\in [b]} \mathsf{G}_{\xi_{t,i}}(W_t) = \frac{1}{b} \sum_{i \in \mathcal{S}_t} \nabla f_i(W_t).\)

r0.5

Algorithm 1 presents the most common variant of Muon, which incorporates Nesterov momentum and weight decay. Our implementation of Nesterov momentum and decoupled weight decay [3] follows the original formulations of [4]. Specifically, the Nesterov step \(C_t = \beta M_t + (1-\beta) \nabla f_{\mathcal{S}_t}(W_t)\) is the “accelerated gradient” form of Nesterov momentum, in which the gradient is blended with the momentum buffer rather than computed at a look-ahead point, following the reparametrization of [15] that is standard in modern deep learning frameworks, and adopted by [16] for Adam. It is equivalent to the classical Nesterov form under a change of variables. Muon optimizers with both Nesterov momentum and weight decay are often used in practice [4], [14]. \(\text{NewtonSchulz5}(\cdot)\) receives \(C_t\), sets \(X_0 := C_t / \| C_t \|_{\rm{F}}\), performs the following iterations from \(k=0\) to \(4\), and returns \(X_5\). \[\begin{align} \label{eq:ns} X_{k+1} := a X_k + b(X_kX_k^\top)X_k + c(X_kX_k^\top)^2X_k, \end{align}\tag{2}\] where \(a=3.4445, b=-4.7750, c=2.0315\). The sequence \((X_k)_{k\in\mathbb{N}}\) converges to \(O_t\) defined as Eq. (1 ) [6], [7]. Our main theoretical analysis assumes that \(O_t := \text{NewtonSchulz5}(C_t)\) satisfies Eq. (1 ). Therefore, there may be a gap between the Muon we consider theoretically and its practical implementation. To quantitatively assess the impact of the approximation error introduced by the Newton–Schulz iteration on the convergence rate of Muon, we provide an analysis of Muon using the approximate solution obtained via the Newton–Schulz iteration in Section 3.3.

2.2 Assumptions↩︎

We make the following standard assumptions:

3 Analysis of Muon’s convergence↩︎

3.1 Muon without weight decay↩︎

We now present a convergence analysis of Muon (Algorithm 1) without weight decay. The proofs of Theorem 1(i) and (ii) are in Appendix 9.

Theorem 1 shows that Muon, both with and without Nesterov momentum, converges to a neighborhood of a stationary point, and achieves same upper bounds on convergence. Note that some terms in the upper bound are slightly smaller when Nesterov momentum is used. The term \(\Delta\) arises from the initialization \(\boldsymbol{m}_{-1} := \boldsymbol{0}\) used in Algorithm 1, and its effect diminishes as the number of steps \(T\) increases. From a theoretical perspective, initializing \(\boldsymbol{m}_{0} := \boldsymbol{0}\) or \(\boldsymbol{m}_0 := \nabla f_{\mathcal{S}_0}(W_0)\) would yield a more refined upper bound for \(\Delta\). However, note that the Muon implementation [4] adopts \(\boldsymbol{m}_{-1} := \boldsymbol{0}\).

3.1.0.1 Interpretation of the convergence rate.

Theorem 1 suggests that setting the learning rate small and the momentum and batch size large results in a tighter bounds. For example, using \(\eta = T^{-\frac{3}{4}}\) and \(1-\beta = T^{-\frac{1}{2}}\) yields \(\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\left[\| \nabla f(W_t) \|_{*}\right] = \mathcal{O}\left( T^{-\frac{1}{4}} \right)\). This rate matches the standard rate for nonconvex SGD-type methods, although the convergence criteria are different. For example, under exactly the same assumptions as ours, [17] show that SGD with momentum using a learning rate \(\eta=T^{-\frac{1}{2}}\) achieves \(\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\| \nabla f(\boldsymbol{\theta}_t) \|_{2} \right] = \mathcal{O}\left( T^{-\frac{1}{4}} \right).\) Note that in the matrix context, the Euclidean norm of the vectorized parameters corresponds to the Frobenius norm. Since \(\| \nabla f(W_t) \|_{\rm{F}} \leq \| \nabla f(W_t) \|_{*}\), a convergence guarantee in terms of the nuclear norm is stronger than that in terms of the Frobenius norm. Therefore, compared to conventional methods such as SGD with momentum, which effectively ignore the matrix structure through vectorization, Muon respects the matrix structure and is shown to converge under a stricter analytical criterion given by the nuclear norm. This suggests that Muon provides a more structure aware and potentially more informative convergence guarantee in practice.

3.1.0.2 Interpretation of the role of rank \(r_1\).

The distinguishing feature of Muon’s bound lies in the variance term \(\sqrt{(1-\beta)r_1/b}\), which depends on the rank \(r_1 = \sup_t \mathrm{rank}(C_t - \nabla f(W_t))\) rather than the full dimension \(mn\). In practical deep networks, the effective rank of gradient matrices is often much smaller than the full dimension (i.e., \(r_1 \ll n\)), particularly in over-parameterized regimes where gradients exhibit low-rank structure. Recent empirical evidence supports this: [18] showed that a rank-fraction of \(1/4\) (i.e., \(r_1 \approx n/4\)) suffices to match full-rank Muon performance in LLMs up to 3B parameters, and that larger models tolerate even lower rank fractions. This suggests that Muon’s bound may be tighter than bounds for vector-based optimizers whose variance terms scale with the full parameter count. To make this comparison concrete, Appendix 16 measures \(r_1\) directly on a controlled FashionMNIST MLP and finds \(r_1 \leq 329\) for the \(2048 \times 784\) layer, compared to \(d = mn \approx 1.6\times 10^{6}\). Since the SGD-type variance term scales with \(\sqrt{d/b}\) while the Muon bound scales with \(\sqrt{r_1/b}\), this translates to a variance-term reduction factor of \(\sqrt{d/r_1} \gtrsim 70\times\) at fixed \(b\); equivalently, matching Muon’s variance contribution would require SGD to use a batch roughly \(d/r_1 \gtrsim 4.9 \times 10^{3}\) times larger. We emphasize, however, that this is a property of the upper bound, not a proof of strict superiority; the actual convergence behavior depends on problem-specific constants that are difficult to compare across optimizers. The empirical advantage of Muon over AdamW observed in our experiments (Figures 3–6) is consistent with this interpretation but may also arise from other mechanisms (e.g., implicit regularization from gradient orthogonalization) not captured by the worst-case bound.

3.2 Muon with weight decay↩︎

The following proposition establishes a key result for Muon (Algorithm 1) with weight decay. The proofs of Propositions 1 and 2 are in Appendix 10.

Proposition 1 establishes that when \(\eta\leq\frac{1}{\lambda}\), weight decay ensures the parameter norm remains almost surely bounded. Furthermore, the upper bound decreases monotonically with \(t\), converging to \(\frac{\sqrt{r_2}}{\lambda}\) as \(t \to \infty\). The bound is minimized uniformly across all \(t\) when \(\eta=\frac{1}{\lambda}\). Proposition 2 extends the result of Proposition 1 to the full gradient norm, which is likewise almost surely bounded. This bound decreases monotonically with \(t\), converging to \(\frac{L\sqrt{r_2}}{\lambda} + L\| W^\star \|_{\rm{F}}\) as \(t \to \infty\). From these results, Corollary 1 establishes an almost surely bound of Muon. In both cases, the upper bounds are minimized when \(\eta=\frac{1}{\lambda}\). Under He initialization, the expected operator norm \(\|W_0\|_{\mathrm{op}}\) is approximately \(\sqrt{2}\) for a square matrix, so the condition \(\eta \le 1/\lambda\) is easily satisfied for practical weight decay values (\(\lambda \le 0.2\) in all our experiments). This condition is a consequence of the proof technique (ensuring the iterate norms remain bounded); we conjecture it is not tight.

While these results suggest that a larger weight decay \(\lambda\) yields a tighter bound, the condition \(\eta\leq\frac{1}{\lambda}\) necessitates a smaller learning rate \(\eta\). These desirable properties stem from the fact that Muon’s search direction is inherently bounded. A key advantage of this feature is that our analysis does not rely on the common-and often restrictive-assumption of bounded gradients. We emphasize that the advantage of weight decay demonstrated by Propositions 1 and 2 is not a tighter asymptotic convergence rate: the rate \(\mathcal{O}(1/(\eta T) + \sqrt{(1{-}\beta)r_1/b} + \hat{r}\eta)\) is the same with or without weight decay (compare Theorems 1 and 2), and the explicit \(1/(1-\lambda)\) prefactor appearing in Theorem 2 in fact slightly loosens the upper bound. The practical benefit of weight decay is instead the almost-sure boundedness of \(\|W_t\|_{\rm{F}}\) and \(\|\nabla f(W_t)\|_{\rm{F}}\) established above, which removes the need for the common (and often restrictive) bounded-gradient assumption and provides implicit regularization during training.

The following is a convergence analysis of Muon (Algorithm 1) with weight decay (the proofs of Theorems 2(i) and (ii) are in Appendix 10).

Similar conclusions follow from Theorem 2, which again demonstrate a modest advantage from incorporating Nesterov momentum. These results build on Propositions 1 and 2 and therefore inherit the assumption that \(\eta \leq \frac{1}{\lambda}\). In other words, for Muon with weight decay to attain the stated convergence rate, it must satisfy \(\eta \leq \frac{1}{\lambda}\). Practically speaking, since the weight decay coefficient \(\lambda\) is typically less than \(1\), this assumption is realistic and does not materially constrain the choice of learning rate. Furthermore, Theorem 2 requires that the weight decay coefficient satisfy \(\lambda < \min \left\{ \frac{1}{\| W_0 \|_{\rm{op}}}, 1\right\}\), which recommends using a small \(\lambda\). In practice, \(\lambda\) is often used with values much smaller than 1, and this assumption can also be considered realistic. For practical reference, under He initialization [19] the operator norm \(\|W_0\|_{\mathrm{op}}\) is \(\mathcal{O}(1)\) for each layer: we measured \(\|W_0\|_{\mathrm{op}} \approx 0.76\) for \(3{\times}3\) convolution layers and \(\approx 1.41\) for \(1{\times}1\) shortcuts in ResNet-18, giving \(1/\max_\ell \|W_0^{(\ell)}\|_{\mathrm{op}} \approx 0.71\). Hence the condition \(\lambda < 1/\|W_0\|_{\mathrm{op}}\) is satisfied for all weight decay values used in our experiments (\(\lambda \leq 0.5\)). Compared to Theorem 1, combining Muon and weight decay multiplies several terms in the upper bound by \(\frac{1}{1-\lambda}\), which slightly loosens the rate bound rather than tightening it. When \(\lambda\) is sufficiently small we have \(\frac{1}{1-\lambda} \approx 1\), so the bound does not grow excessively large even with weight decay. Consistent with the discussion after Propositions 1–2, the faster empirical convergence of the weight-decay variant observed in Figure 3 should therefore be attributed to the almost-sure boundedness of \(\|W_t\|_{\rm{F}}\) and \(\|\nabla f(W_t)\|_{\rm{F}}\) (i.e., implicit regularization), not to a smaller leading constant in the convergence rate. As in Theorem 1, these bounds represent convergence to a neighborhood of a stationary point, suggesting that setting the learning rate small and the momentum and batch size large results in a smaller upper bound.

3.3 Muon with Newton-Schulz↩︎

In practice, the search direction \(O_t\) in Muon is approximated via the Newton-Schulz iteration, whereas in the previous sections we have assumed that \(O_t\) is computed exactly via SVD. In this section, we show that the approximation error of \(O_t\) introduced by the Newton-Schulz iteration has a limited effect on the convergence of Muon. Let \(O_t\) be the exact orthogonalization matrix obtained via SVD, and let \(\tilde{O}_t^{(k)}\) be the approximate matrix after \(k\) iterations of the Newton-Schulz iteration (2 ). According to [20], for all \(k \in \mathbb{N}\), there exists \(\delta_{(k)}>0\), such that \(\| \tilde{O}_t^{(k)}-O_t\|_{\rm{op}} \leq \delta_{(k)} \leq c_0^{3^k}\), where \(c_0 < 1\). Using this fact, the following proposition characterizes the effect of the approximation error introduced by the Newton-Schulz iteration on the convergence of Muon with Nesterov momentum and weight decay. We present the analysis for this variant, as it is the most involved, and omit the results for the other three variants. Note that similar results can be derived for the other variants as well. The proof of Proposition 3 is in Appendix 11.

Comparing Proposition 3 with Theorem 2(ii), we observe that the Newton-Schulz approximation introduces an additional \(\delta_{(k)}^2\)-dependent term, while the existing coefficients are also mildly affected through the factor \(1-\lambda-\delta_{(k)}\). Since \(\delta_{(k)}\) decays exponentially with \(k\), it is small for the fixed number of Newton-Schulz iterations used in practice, and its impact is therefore limited. In fact, according to [21], the experimental performance of the SVD-based Muon is equivalent to that of the Newton-Schulz-based version, with the main difference being the higher computational cost of the SVD procedure. Our own experiments comparing \(k{=}3\) and \(k{=}5\) Newton-Schulz iterations across four network widths (Appendix, Table ¿tbl:tab:ns95focus35?) also show that while \(k{=}5\) reduces the approximation error, the test accuracy gain is at most \(+0.15\) percentage points, confirming that the gap is negligible in practice.

4 Analysis of Muon’s critical batch size↩︎

Our theoretical analysis characterizes the critical batch size as a function of the gradient variance \(\sigma^2\) and optimization hyperparameters. While explicitly modeling the dependence of \(\sigma^2\) on model width or depth is beyond the scope of this single-matrix analysis, our results establish the fundamental relationship \(b^\star \propto \sigma^2\). This suggests that scaling behaviors observed in larger models are mediated through changes in their gradient noise properties. Indeed, recent large-scale empirical studies report that Muon remains efficient at increasingly large batch sizes for large language models [13]; our theory supports this observation, predicting that if larger models entail distinct gradient variance characteristics, the critical batch size will shift accordingly. With this motivation, we now formalize the notion of the critical batch size used in our analysis.

We next introduce the concept of the critical batch size, defined as the batch size that minimizes computational complexity. This complexity is measured in terms of the stochastic first-order oracle (SFO) complexity, which is the total number of stochastic gradient computations. Since the optimizer computes \(b\) stochastic gradients per step, an optimizer that runs for \(T\) steps with batch size \(b\) incurs a total of \(Tb\) SFO complexity. Empirically, for batch sizes up to a certain threshold \(b^\star\) (the critical batch size), the number of training steps \(T\) required to train a DNN scales inversely with \(b\) [22]–[24]. Beyond \(b^\star\), increasing the batch size yields diminishing returns in reducing \(T\). The critical batch size is therefore the batch size that minimizes SFO complexity \(Tb\). Prior work has shown that \(b^\star\) depends on both the optimizer [25] and dataset size [26] and has established a theoretical framework for proving its existence and estimating its lower bound [27], [28]. Concurrently, [29] derived a non-Euclidean gradient noise scale tailored to spectral descent methods (including Muon) and proposed an adaptive batch-size schedule based on this noise scale. Their analysis characterizes the critical batch size of the spectral gradient step itself but does not consider the effects of Nesterov momentum or weight decay. In contrast, our work adopts the SFO-minimization framework and aims to clarify the relationship between the Muon hyperparameters \(\beta\) and \(\lambda\) and its critical batch sizes.

4.1 Relationship between batch size and number of steps needed for training↩︎

Suppose Assumptions 1 and 2 hold. Then, by Theorems 1 and 2, the following inequality holds: \[\begin{align} \frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\left[ \| \nabla f(W_t) \|_{\rm{F}} \right] \leq \frac{X}{T} + \frac{Y}{\sqrt{b}} + Z, \end{align}\] where \(X,Y,Z>0\) are nonnegative constants. Let \(\epsilon>0\) be an arbitrarily fixed threshold. \[\begin{align} \label{eq:new} \exists T, \exists b \;\colon \frac{X}{T} + \frac{Y}{\sqrt{b}} + Z < \epsilon, \end{align}\tag{3}\] The relationship between \(b\) and the number of steps \(T_b\) satisfying Eq. (3 ) is as follows:

Proposition 4 implies that the required number of steps \(T_b\) decreases as \(\sqrt{b}\) increases. In fact, the reduction in the required number of steps achieved by increasing the batch size has a ceiling. The next section will show that a critical batch size exists that minimizes computational complexity.

4.2 Existence of a critical batch size↩︎

The critical batch size minimizes the computational complexity for training. Here, we use SFO complexity as a measure of computational complexity. Since the stochastic gradient is computed \(b\) times per step, SFO complexity is defined as

\[\begin{align} \label{eq:Tbb} T_bb = \frac{Xb\sqrt{b}}{(\epsilon-Z)\sqrt{b}-Y}. \end{align}\tag{4}\]

The following theorem guarantees the existence of critical batch sizes that are global minimizers of \(T_bb\) defined by Eq. (4 ).

In Proposition 5, we obtained a lower bound for the critical batch size using \(\frac{3Y}{2(\epsilon-Z)} > \frac{3Y}{2\epsilon}\). We provide additional context regarding the tightness of this inequality. In any variant we are considering, the term \(Z\) is expressed using the parameter dimension \(n\) and the learning rate \(\eta\). For example, in the case of Muon with Nesterov momentum and weight decay, \(Z=\frac{4L\beta\sqrt{r_1r_2}\eta}{(1-\beta)(1-\lambda)} + \frac{(r_2+\lambda^2r_3)L\eta}{1-\lambda} = \mathcal{O}\left( \hat{r}\eta\right)\). Therefore, by choosing a sufficiently small learning rate \(\eta\), the term \(Z\) also becomes sufficiently small, and this inequality can be considered tight. On the basis of Theorems 1 and 2 and Proposition 5, we derive the following proposition, which gives \(b_{\text{Muon}}^\star\).

To connect Table 1 with the proxy \(X/T + Y/\sqrt{b} + Z < \epsilon\), the correspondence for Muon with Nesterov and weight decay (Theorem 2(ii)) is \(Y = \frac{2(1+\sqrt{2}\beta)\sqrt{1-\beta}}{1-\lambda}\sqrt{r_1\sigma^2}\) and \(Z = \frac{4L\beta\sqrt{r_1r_2}\eta}{(1-\beta)(1-\lambda)} + \frac{(r_2+\lambda^2r_3)L\eta}{1-\lambda}\). The critical batch size \(b^\star > 9Y^2/(4\epsilon^2)\) then yields the formula in Table 1. The results in Table 1 indicate that, for a given accuracy \(\epsilon\), combining Nesterov momentum significantly increases the critical batch size lower bound for Muon. In contrast, the increase in the critical batch size lower bound achieved by combining weight decay is modest when using sufficiently small values of \(\lambda\). Furthermore, in all cases, \(Y^2\) becomes smaller as \(\beta \to 1\) (e.g., \(Y^2=(1-\beta)(1+\sqrt{2}\beta)^2 \to 0\)), suggesting that the larger the momentum \(\beta\), the smaller the critical batch size \(b_{\text{Muon}}^\star\). We stress that the formulas in Table 1 are lower bounds derived from the convergence upper bound and contain the problem-dependent quantities \(\sigma^2\), \(\epsilon\), and \(r_1\). Accordingly, they should be read as qualitative scaling laws that identify which hyperparameters govern \(b^\star\) and how, rather than as numerically predictive estimates. When comparing two Muon variants on the same workload, the unknown quantities cancel and the ratio depends only on \(\beta\) and \(\lambda\) (see Section 5). Although Table 1 is instantiated at \(\lambda = 0.0625\) for readability, the \(1/(1-\lambda)^2\) factor produces a substantially larger shift for larger \(\lambda\): the predicted CBS ratio between the Nesterov+weight-decay and Nesterov-only variants scales as \(1/(1-\lambda)^2\), yielding \(\approx 1.56\) at \(\lambda = 0.2\) and \(\approx 4.0\) at \(\lambda = 0.5\). The widened-range ResNet-18/CIFAR-10sweep in Appendix 13.4 (Table ¿tbl:tab:variant95cbs95cifar?; 66 runs, \(b \in \{32, \ldots, 4096\}\)) confirms this quantitatively: the empirical CBS shifts from \(\approx 1024\) for the Nesterov-only variant to \(\approx 4096\) at \(\lambda = 0.5\), i.e., a ratio of \(4096/1024 = 4.0\) that matches the predicted \(1/(1-0.5)^2 = 4.0\) within the power-of-two grid. In the next section, we provide numerical experiments focusing on this relationship between the critical batch sizes and hyperparameters \(\beta\) and \(\lambda\) (See Figures 4–7).

4.2.0.1 Novelty of the CBS analysis for Muon.

We emphasize that, to our knowledge, the CBS derivation for Muon is among the first to address three key aspects compared to existing CBS analyses for SGD-type optimizers [27], [28]. First, our CBS lower bound formula suggests a rank-dependent scaling \(b^\star \propto r_1\), where \(r_1\) is the rank of the momentum error matrixand does not occur in SGD analysis [28]. This structural difference may stem from Muon’s matrix-recognized update rule. Second, the momentum parameter \(\beta\) appears explicitly in the CBS lower bound formula (e.g., \(b^\star \propto (1-\beta)(1+\sqrt{2}\beta)^2\) for the Nesterov variant), suggesting how momentum tuning may affect the optimal batch size—a relationship not emphasized in prior CBS results. Third, under our current analysis, we provide a two-level validation framework that combines a controlled full-Muon toy experiment designed to match the theoretical assumptions (with proxy fit \(R^2 \approx 0.914\)) with practical-scale hybrid experiments across vision and language workloads, offering a comprehensive validation approach for a matrix-aware optimizer.

5 Numerical Experiments↩︎

We evaluate Muon on three workloads: (i) ResNet-18 on CIFAR-10, (ii) VGG-16 on CIFAR-100, and (iii) Llama3.1 (320M) on the C4 corpus. We first analyze convergence and critical batch size on CIFAR-10 with ResNet-18, then report language–modeling results on C4. Results for VGG-16 on CIFAR-100 are summarized in Appendix 15, where we observe the same qualitative trends.

The experiments are organized into three complementary tiers. First, a controlled full-Muon setting (Teacher-Student task, Appendix 13) that exactly matches the theoretical assumptions serves as a direct, quantitative test of the convergence proxy and CBS lower bound formulas. Second, practical-scale vision experiments (ResNet/CIFAR, VGG/CIFAR-100) test whether qualitative predictions (CBS ordering among variants, \(\beta\)–CBS monotonicity, and weight-decay effects) transfer to realistic architectures under the standard hybrid optimizer. Third, a language-modeling workload (Llama/C4) probes whether the key trends generalize to a substantially different domain, serving as a practical relevance check rather than a strict theory validation.

5.0.0.1 Experimental Setup

For CIFAR-10 and CIFAR-100 experiments, unless otherwise stated, we tuned the learning rate by grid search at a base batch size of 512 and applied square-root scaling for Muon and AdamW; for Momentum SGD we tried both square-root and linear scaling. Each configuration was run five times with different seeds and we report mean and standard deviation. For C4 dataset ¹ we trained Llama3.1 (320M) with sequence length 2048 and batch sizes from 64 to 8192. SFO complexity is measured as steps \(\times\) batch size. Further details of the experimental protocol are in Appendix 12. Details of the tuning protocol and compute fairness considerations are provided in Appendix 12.

In the vision workloads, we follow common practice and use a hybrid optimizer (Muon on matrix-shaped parameter blocks and AdamW on the remaining parameters). This reflects the standard deployment of Muon: since the original release [4], Muon has been applied exclusively to matrix-shaped (2D) parameters while a separate optimizer (typically AdamW) handles biases, normalization layers, and embedding vectors. We therefore adopt the same hybrid configuration in our practical-scale experiments to ensure our results are directly relevant to practitioners.

5.0.0.2 Two-level validation strategy: full Muon vs.hybrid Muon.

Our theory analyzes a single matrix parameter updated entirely by Muon, whereas practical implementations use a hybrid optimizer (Muon on matrix-shaped blocks, AdamW on the rest). To bridge this gap, we employ a two-level strategy: (i) a controlled full-Muon setting (Appendix 13) that exactly matches the theoretical assumptions, confirming the convergence proxy \(X/T + Y/\sqrt{b} + Z\) (\(R^2 \approx 0.914\)) and CBS scaling \(b^\star \propto 1/\varepsilon^2\); and (ii) practical-scale hybrid experiments (ResNet/CIFAR, VGG/CIFAR, Llama/C4) showing that qualitative predictions (CBS ordering, \(\beta\)–CBS monotonicity) transfer to the hybrid regime. This transfer is expected because the Muon component dominates, gradient noise scales are highly correlated across layer types [30], and the CBS predictions depend only on hyperparameter ratios, not problem-specific constants.

5.0.0.3 Theory-aligned stopping proxy: gradient-norm threshold.

Our theory upper-bounds the average expected full-gradient norm. To better align experiments with this proxy, we additionally report a gradient-norm–based stopping metric. For each run, we track the Frobenius norm of the (mini-batch) gradient, \(g_t := \|\nabla f_{\mathcal{S}_t}(W_t)\|_{\rm F}\), and define a smoothed estimate \(\tilde{g}_t\) via an exponential moving average (EMA) over steps.² Given a target threshold \(\epsilon\), we define the stopping time \(T_\varepsilon(b)\) as the first step such that \(\tilde{g}_t \le \epsilon\), and report both steps \(T_\varepsilon(b)\) and SFO complexity \(b\cdot T_\varepsilon(b)\). We emphasize that our original loss/accuracy targets remain useful for practitioner-facing comparisons, while the gradient-norm criterion is introduced specifically to validate the theory-aligned proxy.

5.0.0.4 Rationale for workload-specific stopping targets.

Our experiments use different stopping targets across workloads: \(90\%\) test accuracy and \(95\%\) training accuracy for ResNet-18/CIFAR-10(Figure 4), the training loss for Llama3.1 (320M)/C4(Figure 6), and the gradient-norm threshold \(\varepsilon\) for the controlled full-Muon toy (Appendix 13). These choices are not arbitrary but follow a single rule: for each workload, we choose the most commonly reported metric in that domain, selecting a threshold within the regime where all compared optimizers reach the target so that SFO is well-defined. For vision classification, test accuracy is the de facto benchmark metric and \(90\%\) on CIFAR-10 sits in the well-trained regime for ResNet-18(near but below the \(\approx 93\%\) ceiling), so every optimizer in our sweep reaches it; training accuracy at \(95\%\) additionally removes any generalization-gap confound. For language modeling, test accuracy is not standard; we therefore use training loss, which is the reported metric in the Muon-LLM literature [4], [13], [14]. For the controlled toy, the theoretical quantity itself is the full-gradient norm, so the \(\varepsilon\)-threshold criterion directly matches the theorem statement. To verify that our qualitative conclusions are not an artifact of any particular target, Appendix 13.5 (Table ¿tbl:tab:stopping95sensitivity?) sweeps 13 different stopping criteria spanning test accuracy, training accuracy, training loss, and EMA gradient norm on a common workload, and shows that the SFO-minimizing batch size consistently falls in the range \(32\)–\(256\), with tighter thresholds shifting the CBS to larger values as \(b^\star \propto 1/\varepsilon^2\) predicts. This robustness supports the workload-specific choices adopted in the main text.

5.0.0.5 Convergence Analysis

We empirically validated the stability condition from Proposition [prop:3952]. Figure 2 shows final gradient norm and training loss for ResNet-18 on CIFAR-10 across learning rates \(\eta\) at fixed weight decay \(\lambda{=}0.0625\). The vertical dashed line marks the threshold \(\eta{=}1/\lambda{=}16.0\). The lowest gradient norm occurs near this threshold; for larger \(\eta\) training becomes unstable. The same behavior holds for other values of \(\lambda\) (Appendix 15). This result validates a key practical implication of Propositions 1 and 2: the condition \(\eta \le 1/\lambda\) is not merely a theoretical convenience but a sharp stability boundary. Practitioners tuning Muon with weight decay should therefore set the learning rate at or below \(1/\lambda\).

We next compared the four Muon variants with AdamW and Momentum SGD. Figure 3 shows that Muon with Nesterov momentum and weight decay attains the fastest decrease in both loss and gradient norm. From Theorems 1 and 2, the four variants of Muon exhibit almost identical convergence rates, making it difficult to clearly compare their convergence speeds based on the theory alone. On the other hand, Muon with weight decay consistently demonstrates better performance than its counterpart without weight decay, which can be explained by Propositions 1 and 2. In particular, weight decay guarantees the boundedness and monotonicity of both the parameter norm and the gradient norm, improving stability and thereby leading to better convergence behavior in practice.

5.0.0.6 Critical Batch Size

We measured the number of steps and SFO needed to reach 90% test accuracy and 95% training accuracy. Figure 4 shows that Muon scales better with batch size than the baselines. SFO is the lowest for Muon over the entire range, and Nesterov shifts the SFO-minimizing batch size to larger values. Momentum SGD requires more steps within the same schedule but eventually reaches comparable accuracy (Appendix 15).

Table 1 predicts that, among the four Muon variants, adding Nesterov momentum increases the CBS coefficient more than adding weight decay alone. In Figure 4 (Right), the SFO-minimizing batch sizes of the four Muon variants are close to each other, with slight differences that are directionally consistent with the theoretical ordering. Across all tested batch sizes, Muon achieves fewer steps and lower SFO than AdamW. We note that the SFO curve is relatively flat near its minimum for all Muon variants, so several neighboring batch sizes yield near-optimal efficiency. In this flat regime, the small inter-variant CBS differences predicted by Table 1 are well within the plateau width, which is why the four curves largely overlap. A companion FashionMNISTsweep at batch sizes down to \(b{=}4\) (Table ¿tbl:tab:sfo95plateau? in Appendix 13.3) quantifies this: at the optimal batch size the four Muon variants agree to within \(4.6\%\), and across the plateau they lie within a \(1.58\times\) band; the controlled full-Muon experiment, which eliminates the hybrid-optimizer confound, recovers a statistically significant ordering consistent with the theory. To demonstrate that the CBS differences become consequential when hyperparameters are varied more widely, we conducted an additional experiment on the same ResNet-18/CIFAR-10setup sweeping Nesterov on/off, \(\beta \in \{0.90, 0.95\}\), and \(\lambda \in \{0, 0.2, 0.5\}\) across batch sizes up to \(4096\) (66 runs; details in Appendix 13.4). In this setting the CBS shifts from \(\approx 128\) (w/o Nesterov) to \(\ge 4096\) (w/ Nesterov, \(\lambda{=}0.5\)), and the weight-decay CBS ratio between \(\lambda{=}0.5\) and \(\lambda{=}0\) equals \(4.0\), matching the predicted \(1/(1{-}0.5)^2\) exactly.

5.0.0.7 Role of rank \(r_1\) and unknown quantities in the CBS lower bound formula.

The CBS lower bound formulas in Table 1 contain several problem-dependent quantities that are not directly controlled by the practitioner: the gradient variance \(\sigma^2\), the target precision \(\epsilon\), and the maximum rank \(r_1 := \sup_t \mathrm{rank}(C_t - \nabla f(W_t))\), where \(r_1 \le n\). Because the primary goal of our analysis is to clarify how the hyperparameters \(\beta\) and \(\lambda\) govern the critical batch size, the quantities \(\sigma^2\), \(\epsilon\), and \(r_1\) play the role of task-dependent constants that cancel when comparing different Muon variants on the same workload. For instance, the ratio of CBS with Nesterov to CBS without Nesterov is \((1+\sqrt{2}\beta)^2 / 2\), which depends only on \(\beta\). Our experiments in Figures 4–7 are designed to test precisely these hyperparameter-dependent predictions and do not require knowledge of \(r_1\) or \(\sigma^2\). We note that recent work by [18] provides indirect evidence that \(r_1 \ll n\) in practice: their Dion optimizer achieves near-full-rank Muon performance with only \(r_1 \approx n/4\) in LLMs up to 3B parameters, suggesting that the low-rank structure assumed in our bound is realistic. We furthermore directly measure the rank and time-profile of \(C_t - \nabla f(W_t)\) on a controlled FashionMNIST MLP in Appendix 16: the hard rank saturates within \(\approx 500\) steps at \(\leq 0.42\min(m,n)\) and does not exhibit unbounded growth with \(T\), and the stable rank stays \(\approx 1.15\)–\(1.25\), confirming that the worst-case \(r_1 \leq \min(m,n)\) used in the proof is not attained in practice. The same appendix shows that the \(\sqrt{(1-\beta)r_1/b}\) variance term saturates quickly in \(b\) while the accompanying \(\mathcal{O}(\eta L\hat{r})\) floor can be suppressed by a diminishing step-size schedule, so large batches are not required for convergence. A detailed decomposition of how \(r_1\), \(\sigma^2\), and \(\varepsilon\) individually affect the CBS is provided in Appendix 14.

5.0.0.8 Tightness of the CBS lower bound.

Proposition 5 provides a lower bound on the critical batch size via \(b_{\rm Muon}^\star > 9Y^2/(4\epsilon^2)\), where the gap arises from bounding \(\epsilon - Z\) by \(\epsilon\). As noted in Section 4.2, \(Z = \mathcal{O}(\hat{r}\eta)\) and becomes negligible for sufficiently small learning rates. In our controlled toy experiment (Appendix 13), the fitted value of \(Z\) is close to zero (\(\hat{Z} \approx -0.017\), effectively negligible), and the simplified model without \(Z\) still achieves \(R^2 \approx 0.853\), empirically supporting that the lower bound is reasonably tight in practice.

5.0.0.9 Width-varying CBS: empirical evidence for rank scaling.

To probe the rank dependence, we varied the ResNet-18 width multiplier \(w \in \{0.125, \ldots, 3.0\}\) on CIFAR-10. The effective rank increases monotonically with width, and the SFO complexity at the optimal batch size scales as \(\mathrm{SFO} \propto r_1^{-0.99}\), consistent with \(b^\star \propto r_1\). Direct measurement of the gradient variance \(\sigma^2\) shows that \(r_1 \cdot \sigma^2\) is approximately constant across widths, explaining the flat CBS–rank scaling. Full results, including CBS estimates, gradient variance decomposition, and predicted vs.empirical CBS, are reported in Appendix 14.

5.0.0.10 Critical batch size under the gradient-norm stopping proxy (controlled full Muon).

As the first level of our validation strategy, we run a controlled experiment where Muon is applied to all parameters (full Muon) on a Teacher-Student Tanh Regression task (Appendix 13). This setting exactly matches the theoretical assumptions: a single matrix parameter updated entirely by Muon, with no hybrid optimizer confound. We define the empirical SFO as \(b \cdot T_\varepsilon(b)\) where \(T_\varepsilon(b)\) is the first time the EMA-smoothed gradient norm drops below \(\epsilon\). The EMA window (smoothing coefficient \(\alpha{=}0.01\), corresponding to an effective window of \(\approx 100\) steps) was chosen to be long enough to suppress stochastic fluctuations while remaining shorter than the convergence timescale; varying \(\alpha\) by a factor of two did not materially change the identified CBS. We observe a clear U-shaped SFO—batch curve and a well-defined minimizer \(b^\star\) (e.g., \(b^\star{=}32\) for \(\epsilon{=}0.08\)), providing evidence that a critical batch size exists when using the theory-aligned gradient-norm criterion.

5.0.0.11 Quantitative validation of the \(X/T + Y/\sqrt{b} + Z\) proxy and predicted \(b^\star\).

We further quantify how well the theoretical decomposition matches observations by fitting \[\bar g(T,b) \;\approx\; \frac{X}{T} + \frac{Y}{\sqrt{b}} + Z\] to the measured average gradient norms collected across multiple \((T,b)\) pairs in the controlled full-Muon MLP setting. A linear regression in the features \((1/T,\,1/\sqrt{b},\,1)\) achieves a high goodness-of-fit (e.g., \(R^2{=}0.914\)), supporting that the proxy captures the dominant scaling with \(T\) and \(b\). From the fitted coefficients, we obtain \(\hat{Y}\) and \(\hat{Z}\) and compute a predicted critical batch size via \[b^\star_{\rm pred} \;=\; \frac{9\hat{Y}^2}{4(\epsilon-\hat{Z})^2},\] which matches the empirically SFO-minimizing batch size in the regime \(\epsilon>\hat{Z}\). Moreover, when \(\epsilon\) is not too close to the floor \(\hat{Z}\), we observe the predicted scaling \(b^\star \propto 1/(\epsilon-\hat{Z})^2\), consistent with the theoretical prediction from Proposition 5.

5.0.0.12 Additional CBS validation.

We validated several additional predictions of the CBS theory. (i) Varying the target accuracy across four levels confirms \(b^\star \propto 1/\varepsilon^2\) (log-log exponent \(1.02\), \(R^2{=}0.984\)). (ii) A three-optimizer comparison (Muon, AdamW, Shampoo) shows that Muon achieves the lowest SFO at every batch size (Appendix 15.1). (iii) Wall-clock measurements on a single A100 GPU confirm that Muon’s per-step cost (\(\approx 16\) ms) is moderate compared to AdamW (\(\approx 8.5\) ms) and lower than Shampoo (\(\approx 22\) ms; Appendix 15.1). (iv) To verify that the CBS conclusions are robust to the choice of stopping criterion, we analyzed FashionMNIST (16 seeds) training logs under stopping criteria spanning four metric types: test accuracy, training accuracy, training loss, and EMA-smoothed gradient norm (Appendix 13.5, Table ¿tbl:tab:stopping95sensitivity?). Across all criteria, the SFO-minimizing batch size falls in the range 32–256. The EMA gradient norm sweep yields the most direct evidence for the theoretical prediction: the CBS shifts monotonically as \(32 \to 64 \to 128 \to 256\) (for \(\varepsilon{=}0.60, 0.55, 0.40, 0.32\)), supporting \(b^\star \propto 1/\varepsilon^2\). (v) A controlled comparison of exact SVD vs.NS(\(k{=}3\)) vs.NS(\(k{=}5\)) orthogonalization in the full-Muon setting (10 seeds) confirms that all three methods converge to comparable gradient norms despite differing approximation errors, with NS error \(<0.04\) across all batch sizes (Appendix 13.2, Table ¿tbl:tab:ns95controlled?). (vi) An additional CBS experiment on ResNet-18/CIFAR-10(66 runs, batch sizes \(32\)–\(4096\), 1–3 seeds) sweeping Nesterov on/off, \(\beta \in \{0.90, 0.95\}\), and \(\lambda \in \{0, 0.2, 0.5\}\) demonstrates clear CBS separation across configurations (Appendix 13.4, Table ¿tbl:tab:variant95cbs95cifar?). The CBS ordering, w/o Nesterov (\(b^\star{\approx}128\)) \(<\) w/ Nesterov (\(b^\star{\approx}1024\)) \(<\) w/ Nesterov, \(\lambda{=}0.2\) (\(b^\star{\approx}2048\)) \(<\) w/ Nesterov, \(\lambda{=}0.5\) (\(b^\star{\ge}4096\)), matches the direction predicted by Table 1, and the weight-decay CBS ratio between \(\lambda{=}0.5\) and \(\lambda{=}0\) equals \(4.0\), agreeing with the predicted \(1/(1{-}0.5)^2\). This directly substantiates the boundedness/stability interpretation of the \(1/(1{-}\lambda)^2\) factor predicted by the theory. In summary, the CBS range remains stable across stopping metrics (accuracy, loss, gradient norm), and tighter thresholds consistently shift the CBS to larger values. Both observations are consistent with the theoretical prediction \(b^\star \propto 1/\varepsilon^2\).

5.0.0.13 Effect of \(\beta\) on Muon’s Critical Batch Size

Theory in Section 4.2 predicts that the critical batch size decreases as \(\beta\!\to\!1\). Figure 5 confirms this trend for ResNet-18/CIFAR-10, regardless of weight decay or Nesterov. Quantitatively, the CBS lower bound formulas in Table 1 contain the factor \((1-\beta)\) (or \((1-\beta)(1+\sqrt{2}\beta)^2\) with Nesterov), both of which vanish as \(\beta \to 1\), driving the critical batch size toward zero. In Figure 5, the SFO-minimizing batch size shifts leftward as \(\beta\) increases from \(0.7\) to \(0.999\), consistent with this monotone relationship. Practically, this suggests that high-momentum configurations can be trained efficiently at relatively small batch sizes, reducing the hardware resources required. A quantitative regression of normalized CBS \(b^\star/(r_1\sigma^2)\) on the full-Muon Teacher-Student task confirms this monotone relationship (\(R^2{=}0.90\); see Appendix 13.7).

5.0.0.14 Language-Model Workload: C4 on Llama3.1 (320M)

We include a language-modeling workload as a practical relevance check: the goal is to verify that the main qualitative trends, namely Muon’s SFO advantage over AdamW and the \(\beta\)–CBS relationship, generalize beyond vision, rather than to perform a strict quantitative validation of every theoretical constant. We train Llama3.1 (320M)(18 layers, hidden dimension 768, 12 attention heads; approximately 320M parameters) on the C4 corpus with batch sizes ranging from 32 to 4096 (details in Appendix 12). Muon is applied to all matrix-shaped Transformer parameters (approximately 127M parameters), while the remaining parameters (embeddings, layer norms) are updated by AdamW. We define the stopping criterion as the number of steps to reach a target training loss, and report SFO complexity \(= \mathrm{steps} \times \mathrm{batch\;size}\). Figure 6 reports steps to reach the target loss and SFO versus batch size. Muon generally requires fewer steps and achieves lower SFO than AdamW, particularly at small to moderate batch sizes. At the largest batch sizes, the Muon variants show increasing SFO, narrowing the gap with AdamW. The inter-variant loss differences among the four Muon configurations are small (typically \(< 0.05\)), consistent with Theorems 1 and 2 predicting the same asymptotic rate for all variants; accordingly, the lack of consistent Nesterov/weight-decay gains in this setting is expected rather than a negative finding. Using the gradient-norm stopping criterion (which directly corresponds to the quantity bounded by the theorems), the SFO-minimizing batch size for Muon with Nesterov at \(\beta{=}0.95\) shifts from \(32\) at loose thresholds (\(\|\nabla f\| \le 0.40\)) to \(128\) at tighter thresholds (\(\|\nabla f\| \le 0.25\)), with a clear U-shaped SFO curve spanning GBS \(= 16\) to \(8192\) at every threshold (Table ¿tbl:tab:llm95cbs95target95sensitivity? in Appendix 17).

5.0.0.15 Momentum Sweeps on C4

We varied \(\beta\) on C4 to examine the critical batch size. Figure 7 shows that a moderate value, \(\beta\approx0.95\), gives the best loss and SFO. As \(\beta\) decreases, the critical batch size increases; as \(\beta\) increases toward 1, the critical batch size decreases, but extreme values are suboptimal. These observations are consistent with Section 4.2 and mirror the vision results. The fact that the same \(\beta\)–CBS relationship holds across both vision (Figure 5) and language modeling (Figure 7) workloads supports the generality of the theoretical prediction, even though the two settings differ substantially in model architecture, loss landscape, and the ratio of Muon-updated to AdamW-updated parameters. Quantitatively, using the gradient-norm threshold \(\|\nabla f\| \le 0.20\), the SFO at \(\beta{=}0.95\) is \(1.30 \times 10^5\) (CBS \({=} 32\)), compared to \(4.02 \times 10^5\) at \(\beta{=}0.80\) (CBS \({=} 128\)) and \(1.37 \times 10^6\) at \(\beta{=}0.70\) (SFO nearly flat across batch sizes, indicating that the CBS exceeds the tested range). This confirms that \(\beta{=}0.95\) offers the best efficiency. Values \(\beta \ge 0.999\) fail to converge at most tested batch sizes.

6 Related Works↩︎

Several studies have investigated the theoretical properties and convergence behavior of Muon. [6] connected Muon to momentum applied in the steepest descent direction under a spectral norm constraint. [31] provided a pioneering analysis assuming Frobenius norm Lipschitz smoothness. [32] studied Muon in the context of optimization methods that use linear minimization oracles for norm balls, and established a convergence rate. Connections to other optimizers have also been explored. Shah et al. showed that Shampoo and SOAP reduce to Muon under simplifying assumptions [14]. [33] proposed and analyzed a stochastic non-Euclidean trust-region method that includes Muon as a special case. Similarly, [34] proposed an adaptive structured gradient optimization algorithm that matches Muon in its momentum-free variant. Other studies have explored specific properties and extensions of Muon. [35] proposed and analyzed a zeroth-order version . Chen et al. demonstrated Muon’s compatibility with the Lion-\(\mathcal{K}\) algorithm [36] and showed that Muon with weight decay implicitly solves an optimization problem with a spectral norm constraint [37]. [21] presented a comprehensive analysis of Muon’s convergence rate in comparison to gradient descent. [38] introduced PolarGrad, a unifying framework for matrix-aware preconditioned methods including Muon, and established convergence rates. The core concepts of gradient orthogonalization and dualization, which are central to Muon, were introduced in the foundational works by [39] and [40]. On the practical side, [29] derived a non-Euclidean gradient noise scale for spectral descent (including Muon) and used it to build an adaptive batch-size schedule, providing a complementary, empirically driven perspective on the critical batch size of Muon; however, their analysis does not account for the effects of momentum \(\beta\) or weight decay \(\lambda\). The convergence rates of major related works are summarized in Table 2. Our contribution is, to our knowledge, the first study to analyze the Muon gradient norm using both Nesterov momentum and weight decay.

The framework for assumptions and proofs presented by [21] is the most similar to our own. The crucial difference between our study and all previous studies, including this one, is that we also consider most common variant of Muon, which incorporates Nesterov momentum and weight decay. Technically, the technique of transforming \(\langle C_t,O_t\rangle_{\rm{F}}\) using the dual norm and inverse triangle inequality follows the proof of [32] (see proof of Theorem 3 in Appendix 9).

7 Conclusion↩︎

Through a comprehensive theoretical analysis of the Muon optimizer, we established convergence guarantees for four practical configurations (with and without Nesterov momentum and with and without weight decay). Our primary theoretical contribution is demonstrating the crucial role of weight decay. We proved that it enforces a strict decrease in parameter and gradient norms, a clear advantage over the standard Muon configuration. This theoretical insight, along with the necessary condition relating the learning rate and weight decay coefficient, was empirically validated by our experimental results. Additionally, we derived the critical batch size for Muon, revealing its dependence on fundamental hyperparameters such as momentum and weight decay. Collectively, our findings provide both a deeper theoretical understanding of Muon and actionable guidance for practitioners aiming to leverage this promising optimizer in large-scale settings.

Finally, we note three limitations. First, our single-matrix analysis does not directly model the layer-wise heterogeneity of gradient statistics in deep networks. However, [30] show that gradient noise scales of different layers are highly correlated, suggesting that our theory captures the fundamental functional dependencies of the critical batch size on hyperparameters. Second, in practice Muon is used as a hybrid optimizer (Muon on matrix-shaped parameters, AdamW on the rest), whereas our theory assumes full Muon. We address this gap through a two-level validation strategy: a controlled full-Muon experiment that exactly matches the theoretical assumptions, and practical-scale hybrid experiments that confirm the qualitative predictions transfer to the hybrid regime (Section 5). Third, while Muon consistently outperforms Shampoo in our experiments (Appendix 15.1), deriving analogous convergence bounds and CBS lower bound formulas for Shampoo and SOAP [11] remains an open problem.

8 Tools for Proof of All Theorems↩︎

The results presented in this section are not new and are given simply for reference and completeness.

The following lemma was established by [41]. In their setting, the algorithm without momentum corresponds to the case in which \(\beta_1=1\), whereas in ours it corresponds to \(\beta_1=0\). As a result, the statements may appear slightly different.

9 Proof of Theorems for Muon without weight decay↩︎

9.1 Proof of Theorem 1(i)↩︎

9.2 Proof of Theorem 1(ii)↩︎

10 Proof of Theorems for Muon with weight decay↩︎

10.1 Proof of Proposition 1↩︎

10.2 Proof of Proposition 2↩︎

10.3 Proof of Theorem 2(i)↩︎

10.4 Proof of Theorem 2(ii)↩︎

11 Proof of Proposition 3 for Muon with Newton-Schulz↩︎

12 Experimental Details↩︎

This section details the experimental setup for evaluating the optimizers. Our code will be available at https://anonymous.4open.science/r/critical_batchsize_muon-0D38/README.md.

12.0.0.1 Workloads and General Setup (Vision)

All CIFAR experiments used ResNet-18 on CIFAR-10 and VGG-16 on CIFAR-100, trained from scratch. Training was performed for a fixed number of samples corresponding to 100 epochs at a batch size of 512. For any batch size \(B\), the number of epochs was scaled as \(E_B = 100 \times (512 / B)\) to keep the number of seen samples constant. Each configuration was repeated with five random seeds, and we report mean and standard deviation of test accuracy.

12.0.0.2 Workload and System Setup (Language Modeling: C4/ Llama3.1 (320M))

For the language–modeling workload we trained Llama3.1 (320M) on the C4 dataset with sequence length 2048. We used the torchtitan codebase with PyTorch FSDP (full sharding) and activation checkpointing (mode=full). Gradient clipping with max norm 1.0 was applied. Due to sharding, Muon requires an additional all-gather of gradients before the optimizer step, which introduces one extra collective communication per step. We therefore report SFO-based metrics (steps \(\times\) batch size) in addition to loss to separate algorithmic efficiency from system overhead.

Unless otherwise noted, the training budget was fixed to \(3.2\)B tokens. At a base global batch size (GBS) of 128 this corresponds to 12,208 steps; for other GBS values the number of steps scales inversely to keep the token budget constant. We ran on up to 8 H100 GPUs with per-GPU micro-batch size 8 and no tensor, pipeline, or context parallelism (TP=PP=CP=1).

12.0.0.3 Learning-Rate Schedules

Vision workloads used a grid search at base batch size 512 and square-root scaling to other batch sizes for Muon and AdamW. Momentum SGD was tested with both square-root and linear scaling. For Llama3.1 (320M), we used linear warmup followed by cosine decay. The base configuration uses 2000 warmup steps; in sweeps we also parameterized warmup as 10% of total steps to maintain a comparable schedule across different GBS.

12.0.0.4 Hyperparameter Tuning Protocol (Vision)

We performed an extensive grid search over base learning rates and weight decay. For \(B \neq 512\), learning rates were scaled from the base value using \(\eta_B=\eta_{512}\sqrt{B/512}\) for AdamW and Muon; for Momentum SGD we report both \(\sqrt{B/512}\) and \({B/512}\) scaling. The shared search space is summarized in Table 3, and optimizer-specific spaces in Table 4.

12.0.0.5 Optimizer Configuration (Vision)

For the hybrid Muon optimizer, we applied the Muon update to all convolutional layers in ResNet-18 (excluding the input stem). Biases, batch-norm parameters, and the final linear layer were updated by AdamW. For standard AdamW and Momentum SGD, the respective optimizer was applied to all parameters.

12.0.0.6 Search Space (Language Modeling)

For C4/ Llama3.1 (320M), we performed sweeps using grid search. GBS took values in \(\{64,256,1024,4096\}\) when resources permitted. We tuned Muon and the AdamW baseline per GBS without automatic LR scaling.

12.0.0.7 Tuning fairness.

To ensure a fair comparison across optimizers, we applied the same tuning protocol to all methods: at the base batch size, we performed a grid search over learning rates; at other batch sizes, we applied standard scaling rules (square-root scaling for Muon and AdamW, square-root and linear for Momentum SGD). We selected the configuration that minimizes SFO complexity to the target accuracy (or target loss for LLM experiments). This ensures that all optimizers are given equal opportunity to find their best operating point, and that differences in SFO complexity reflect genuine algorithmic advantages rather than tuning asymmetry.

12.0.0.8 Compute fairness and the choice of SFO complexity.

We use stochastic first-order oracle (SFO) complexity (steps \(\times\) batch size) as the primary efficiency metric because it is hardware-independent: it counts the total number of individual gradient evaluations regardless of the parallelism, memory bandwidth, or accelerator type used. This makes SFO comparisons reproducible across different hardware setups. We acknowledge that wall-clock time is ultimately the metric of practical interest; however, wall-clock comparisons conflate algorithmic efficiency with implementation-specific factors (e.g., the extra all-gather communication required by Muon under FSDP sharding). By reporting SFO complexity, we isolate the algorithmic contribution. Note that [4] demonstrated that Muon achieves favorable wall-clock performance relative to both Shampoo and SOAP, and [14] showed that Muon expands AdamW’s Pareto frontier on the compute-time plane.

13 Additional Diagnostics: Gradient-norm proxy and full-Muon toy setting↩︎

13.0.0.1 Toy task and full-Muon configuration.

As described in Section 5, we use a controlled full-Muon experiment to provide an exact test of the theoretical predictions, free from the hybrid optimizer confound. We conduct experiments on a controlled synthetic task where Muon is applied to all parameters. We employ a Teacher-Student Tanh Regression problem. A fixed teacher matrix \(W^\star \in \mathbb{R}^{m \times n}\) and a student matrix \(W \in \mathbb{R}^{m \times n}\) are initialized with entries drawn from \(\mathcal{N}(0, 1/n)\). At each step \(t\), we sample inputs \(x_t \sim \mathcal{N}(0, I_n)\) and generate labels \(y_t = \tanh(x_t (W^\star)^\top) + \xi_t\), where \(\xi_t \sim \mathcal{N}(0, \sigma^2 I)\) represents label noise. The student minimizes the loss \(\mathcal{L}(W) = \frac{1}{2} \| \tanh(x_t W^\top) - y_t \|^2\). We set the dimensions to \(m=256, n=128\), and the noise level to \(\sigma=0.1\). Unlike the large-scale experiments, here Muon (Algorithm 1) is applied to the entire matrix \(W\) with learning rate \(\eta=0.05\), momentum \(\beta=0.95\) (Nesterov), and 5 Newton-Schulz iterations.

13.0.0.2 Measured quantities and stopping rule.

To rigorously validate the theoretical convergence bound \(\mathbb{E}[\| \nabla \mathcal{L} \|] \lesssim \frac{X}{T} + \frac{Y}{\sqrt{b}} + Z\), we track the cumulative mean of the gradient norm, denoted as \(\bar{g}_t = \frac{1}{t} \sum_{i=1}^t \| \nabla \mathcal{L}(W_i) \|_F\). We define the convergence step \(T_\varepsilon(b)\) as the first step \(t\) where \(\bar{g}_t \le \varepsilon\). The stochastic first-order oracle (SFO) complexity is computed as \(\text{SFO} = b \cdot T_\varepsilon(b)\). We perform a grid search over batch sizes \(b \in \{4, 8, \dots, 512\}\), extended to \(b \in \{4, 8, \dots, 2048\}\) in a follow-up sweep.

13.0.0.3 Results and Validity of Approximations.

We confirm three key findings in this controlled setting: (1) Existence of Critical Batch Size: As shown in Figure 8 (Top Left), the SFO complexity exhibits a clear convex shape with respect to \(b\), identifying an empirical optimum \(b^\star\). (2) Model Fit: The observed data fits the theoretical proxy \(\frac{X}{T} + \frac{Y}{\sqrt{b}} + Z\) with high accuracy (\(R^2 \approx 0.914\)), validating our convergence analysis. (3) Justification for \(Z \approx 0\): To address the theoretical concern regarding the non-vanishing term \(Z\), we compared the full fit against a simplified model \(\frac{X}{T} + \frac{Y}{\sqrt{b}}\) (forcing \(Z=0\)). Figure 8 (Bottom Right) demonstrates that the simplified model still predicts the observed gradient norms with reasonable accuracy (\(R^2 \approx 0.853\)), supporting the practical validity of the \(Z \approx 0\) approximation. This empirically justifies our derivation of the critical batch size \(b^\star \propto 1/\varepsilon^2\) by treating \(Z\) as negligible during the primary optimization phase.

13.0.0.4 Interpretation of negative \(\hat{Z}\).

The fitted intercept \(\hat{Z}\) is negative in all four variants (ranging from \(-1.64\) to \(-2.56\)), whereas the theoretical term \(Z = \mathcal{O}(L\eta \hat{r})\) is non-negative. This discrepancy does not undermine the proxy validation. The regression is performed over a finite \((T,b)\) grid, and \(\hat{Z}\) represents the best linear-in-\((1/T, 1/\sqrt{b})\) approximation to the gradient norm surface within that grid. Since the true convergence bound is an upper bound, the actual gradient norm can lie well below the bound; a negative fitted intercept simply reflects this slack. Moreover, the primary purpose of the proxy fit is to validate the functional form (how the gradient norm depends on \(T\) and \(b\)), not the absolute constants. The high \(R^2 \approx 0.914\) confirms that the \(X/T + Y/\sqrt{b}\) terms capture the dominant variation, which is the theoretically meaningful finding.

13.0.0.5 Predicted CBS exceeds the sweep range.

The CBS values predicted from the fitted proxy (\(\hat{b}^\star = 9\hat{Y}^2/(4\varepsilon^2)\), accounting for \(\hat{Z} < 0\)) range from approximately 1500 to 2000 for \(\varepsilon \in [0.05, 0.20]\), which exceeds the maximum batch size in our sweep (\(b = 2048\)). Consequently, no empirical CBS is observed: the cumulative mean gradient norm does not fall below any \(\varepsilon\) threshold within \(T{=}10{,}000\) steps, even at \(b{=}2048\). This is consistent with the theory, as the predicted CBS is simply larger than the experimental sweep range. Importantly, the primary purpose of the toy experiment is to validate the functional form of the convergence bound (the \(X/T + Y/\sqrt{b} + Z\) decomposition), not to observe the CBS transition itself. The high proxy fit quality (\(R^2 \approx 0.93\)) confirms that the theory accurately describes the gradient norm surface, and the CBS formula follows algebraically from this validated functional form. We note, however, that when the predicted CBS exceeds the experimental sweep range, the \(R^2{=}0.914\) proxy fit confirms the functional form but the CBS point estimate extrapolates beyond the observed data, which reduces confidence in the specific CBS value relative to settings where the transition is directly observed.

13.0.0.6 Newton-Schulz follow-up and remaining experimental gap.

We also ran a seeded follow-up study around the practical Newton-Schulz choice \(k \in \{3,5\}\) on ResNet-18/CIFAR-10 across widths \(\{0.25,0.5,1.0,2.0\}\), with three seeds per configuration. The result is that \(k{=}3\) remains the best global default: at widths \(0.25\) and \(0.5\), moving from \(k{=}3\) to \(5\) does not produce meaningful test-accuracy gains; at width \(2.0\), the approximation error decreases but the test accuracy is essentially unchanged; and width \(1.0\) is the clearest case where \(k{=}5\) is justified. This supports our choice to use five iterations as a robust practical default while indicating that fewer iterations are often sufficient in vision workloads. Table ¿tbl:tab:ns95focus35? reports the aggregated test accuracy and Newton-Schulz approximation error (mean \(\pm\) std over seeds). The delta analysis shows that the accuracy gain from \(k{=}5\) is at most \(0.15\)% (at width \(1.0\)), while the NS error reduction is substantial only at widths \(1.0\) and \(2.0\). This indicates that for widths where the NS approximation is already tight (\(w \le 0.5\)), additional iterations do not translate into accuracy improvements.

13.0.0.7 Rank tracking: SFO comparison across optimizers.

We ran a rank-tracking experiment on ResNet-18/CIFAR-10 across batch sizes \(\{8, 32, 128, 512, 2048\}\) for Muon, AdamW, and Momentum SGD. Table 8 reports the best SFO complexity to reach 80% test accuracy for each (optimizer, batch size) pair. The heuristic CBS (the batch size with the smallest SFO) is \(b^\star{=}32\) for Muon and \(b^\star{=}8\) for AdamW, confirming that Muon’s CBS exceeds AdamW’s. We also measured the early-training effective rank at larger batch sizes (\(b \ge 128\)): Muon exhibits higher effective rank than AdamW and Momentum SGD across all tested batch sizes (e.g., \(\mathrm{erank} \approx 148\) for Muon vs.\(\approx 103\) for AdamW at \(b{=}128\)), and the mean Newton-Schulz approximation error decreases with batch size for all optimizers, reaching \(\approx 0.30\)–\(0.49\) at \(b{=}2048\).

13.1 Proxy Calibration: Full-Gradient Norm vs.Mini-Batch Surrogates↩︎

To directly address the discrepancy between the theoretical quantity (the full-batch gradient norm \(\|\nabla f(W_t)\|_{\rm F}\)) and the experimental proxy (the EMA-smoothed mini-batch gradient norm), we compute both quantities at every training step in the controlled full-Muon Teacher-Student setting described above.

Table ¿tbl:tab:proxy95calibration? reports the Pearson correlation between the full-batch gradient norm and two stochastic surrogates (the raw mini-batch gradient norm and the EMA-smoothed proxy) across batch sizes and two noise levels (\(\sigma{=}0.1\) and \(\sigma{=}0.2\)), aggregated over 6 independent seeds each.

The high correlation (\(>0.97\) in all cases) confirms that mini-batch gradient norms are a well-calibrated surrogate for the full gradient norm. We note that the absolute CBS point estimate can differ between the full-gradient metric and the stochastic proxy, because the stochastic metrics include mini-batch noise that the full-gradient metric does not. However, the high correlation indicates that both metrics track the same underlying optimization dynamics, and the qualitative CBS behavior (U-shaped SFO curve, variant ordering direction) is preserved.

13.2 Newton-Schulz Approximation Gap: Controlled Comparison↩︎

To complement the practical NS(\(k\)) comparison on ResNet-18(Table ¿tbl:tab:ns95focus35?), we compare exact SVD orthogonalization against NS(\(k{=}3\)) and NS(\(k{=}5\)) in the controlled full-Muon Teacher-Student setting (\(m{=}512\), \(n{=}256\), \(\sigma{=}0.1\)). Table ¿tbl:tab:ns95controlled? reports the final gradient norm and mean orthogonalization error across batch sizes, aggregated over 10 seeds.

All three methods converge successfully; the final gradient norms differ due to slightly different optimization trajectories induced by the approximation, but all reach the same order of magnitude. The NS approximation error is stable across batch sizes and small enough (\(< 0.04\)) to be negligible relative to the gradient scale.

13.3 Four-Variant CBS Ordering in the Controlled Full-Muon Setting↩︎

13.3.0.1 SFO plateau and inter-variant agreement on FashionMNIST.

The ResNet-18/CIFAR-10runs underlying Figure 4 start at batch size \(b{=}8\), which already sits at (or above) the CBS for every Muon variant and therefore provides limited resolution for comparing variants at and below the CBS. To obtain a more controlled microscope for variant-level differences near the plateau, we ran a compact FashionMNISTsweep on a 2-layer MLP (hidden dim \(1024\), target test accuracy \(0.84\)), which allows us to sweep batch sizes down to \(b{=}4\) within a short time budget and tune learning rates densely. For each (variant, \(b\)) cell we report the best SFO (lowest median across three seeds) over a learning-rate grid (\(\{0.02, 0.05, 0.1\}\) for Muon variants and \(\{3\!\times\!10^{-4}, 10^{-3}, 3\!\times\!10^{-3}\}\) for AdamW), matching the “best-over-LR” protocol used for CIFAR-10.

Table ¿tbl:tab:sfo95plateau? shows the resulting SFO grid. Three observations stand out: (i) at the optimal batch size \(b{=}8\), the four Muon variants achieve best SFO in the range \(21.6\)–\(22.6\mathrm{K}\), a maximum spread of \(4.6\%\); (ii) across the plateau (\(b\in\{8,16,32,64\}\)) every Muon variant’s best SFO lies in a \(19.2\)–\(30.4\mathrm{K}\) band (within a factor of \(1.58\times\)), so the predicted inter-variant CBS differences (Table 1) are again smaller than the plateau width; (iii) unlike ResNet-18/CIFAR-10, AdamW is comparable to Muon on this simpler MLP/FashionMNISTtask: the AdamW-to-best-Muon SFO ratio ranges from \(0.96\times\) to \(1.75\times\) across batch sizes. We flag (iii) explicitly: it illustrates that the Muon-vs-AdamW gap documented on ResNet-18/CIFAR-10is task-dependent and narrows on simpler MLP benchmarks. The CBS mechanism analysed in this paper concerns inter-variant closeness (observations (i) and (ii)) and does not rely on any minimum Muon-vs-AdamW gap, so point (iii) does not affect the conclusion that practical CBS differences among the four Muon variants are within the plateau width.

We emphasise two caveats: the \(b{=}4\) column shows a wider inter-variant spread (driven by Muon/Muon+WD at lr \(\in\{0.02, 0.05\}\) failing to reach the target within the step budget, giving a pessimistic best-LR estimate on a narrow grid), and the empirical CBS location varies across variants (\(b^\star \in \{8, 16, 64\}\)). Neither affects the main quantitative claim in the plateau window \(b\in\{8,16,32,64\}\).

13.4 Variant CBS Separation on ResNet-18/CIFAR-10↩︎

The four Muon variants in Figure 4 share similar CBS values because the hyperparameter differences (Nesterov on/off, small \(\lambda\)) produce only modest CBS shifts relative to the plateau width. To test whether the CBS formulas predict consequential differences when hyperparameters are varied more widely, we conducted a dedicated experiment on the same ResNet-18/CIFAR-10setup with five configurations: Muon without Nesterov (\(\beta{=}0.95\), \(\lambda{=}0\)); Muon with Nesterov at \(\beta{=}0.95\) and \(\beta{=}0.90\) (both \(\lambda{=}0\)); and Muon with Nesterov (\(\beta{=}0.95\)) combined with weight decay \(\lambda \in \{0.2, 0.5\}\). We swept batch sizes \(b \in \{32, 64, 128, 256, 512, 1024, 2048, 4096\}\) using the same LR scaling \(\eta = 0.005 \cdot \sqrt{b/512}\) and epoch scaling \(E_b = \lceil 100 \cdot 512/b \rceil\) (66 runs, 1–3 seeds per cell).

Table ¿tbl:tab:variant95cbs95cifar? reports the SFO complexity (steps \(\times\) batch size) to reach 95% training accuracy. A dash (—) indicates the target was not reached within the allocated epochs.

CBS ordering matches theory. The empirical CBS ordering, w/o Nesterov (\(\approx 128\)) \(<\) w/ Nesterov (\(\approx 1024\)) \(<\) w/ Nesterov+WD\(_{0.2}\) (\(\approx 2048\)) \(<\) w/ Nesterov+WD\(_{0.5}\) (\(\ge 4096\)), is consistent with the theoretical lower bounds in Table 1: Nesterov shifts the CBS rightward, and weight decay provides an additional shift proportional to \(1/(1{-}\lambda)^2\).

Weight-decay CBS ratio matches quantitatively. The ratio of CBS between the \(\lambda{=}0.5\) and \(\lambda{=}0\) Nesterov configurations equals \(4096/1024 = 4.0\), which agrees with the predicted factor \(1/(1{-}0.5)^2 = 4.0\). The corresponding ratio for \(\lambda{=}0.2\) is \(2048/1024 = 2.0\), consistent with the predicted \(1/(1{-}0.2)^2 \approx 1.56\) within the power-of-two grid resolution. The Nesterov-to-non-Nesterov ratio (\(1024/128 = 8.0\)) exceeds the lower-bound prediction of \(\approx 2.75\); this is expected because the formulas in Table 1 provide a lower bound, not an equality.

13.5 Stopping-Target Sensitivity: CBS Robustness Across Metrics↩︎

To verify that the CBS conclusion is robust to the choice of stopping criterion, we analyzed FashionMNIST training logs (2-layer MLP, hidden dimension 1024, full Muon, 16 seeds) under stopping criteria spanning four metric types: test accuracy, train accuracy, train loss, and EMA-smoothed gradient norm. Table ¿tbl:tab:stopping95sensitivity? summarizes the SFO-minimizing batch size for each criterion.

Two key observations emerge: (i) Across all four metric types, the CBS falls consistently in the range \(32\)–\(256\), confirming that the qualitative CBS conclusion is robust to the choice of stopping criterion. (ii) Within each metric type, tighter thresholds systematically shift the CBS to larger values, supporting the theoretical prediction \(b^\star \propto 1/\varepsilon^2\). The EMA gradient norm sweep provides the most striking evidence: the CBS shifts monotonically as \(32 \to 64 \to 128 \to 256\) as \(\varepsilon\) decreases from \(0.60\) to \(0.32\).

13.6 Full-Muon Weight Decay Sweep↩︎

To isolate the effect of weight decay \(\lambda\) on the CBS in a controlled setting free from the hybrid-optimizer confound, we swept \(\lambda \in \{0, 0.01, 0.05, 0.1, 0.2\}\) on the Teacher-Student task using the full-Muon configuration above (\(\beta{=}0.95\), Nesterov, 5 Newton-Schulz iterations). Table ¿tbl:tab:fullmuon95wd95sweep? reports the predicted CBS (from gradient-norm proxy fit), the fit quality \(R^2\), and the estimated gradient variance \(\hat{\sigma}^2\).

The CBS increases from 438 (\(\lambda{=}0\)) to 646 (\(\lambda{=}0.1\)), a \(1.48\times\) increase. The theoretical factor \(1/(1{-}\lambda)^2\) predicts only a \(1.23\times\) increase at \(\lambda{=}0.1\). The discrepancy is explained by the concurrent increase in \(\hat{\sigma}^2\): from 1,308 to 2,201 (\(1.68\times\)), so the product \((1/(1{-}\lambda)^2) \cdot (\sigma^2(\lambda)/\sigma^2(0)) \approx 1.23 \times 1.68 = 2.07\times\), which is consistent with the observed CBS shift of \(1.48\times\) (the CBS lower bound formula involves \(\sigma^2\) under a square root). At \(\lambda{=}0.2\), the CBS decreases slightly to 587, possibly reflecting overly aggressive regularization that alters the optimization landscape.

13.7 \(\beta\)-CBS Quantitative Analysis↩︎

Figure 9 shows the regression of normalized CBS \(b^\star/(r_1\sigma^2)\) against the theoretical momentum factor \(\mathrm{MF}(\beta) := (1{-}\beta)(1{+}\sqrt{2}\beta)^2\) for the full-Muon Teacher-Student task. Here \(b^\star\) is the predicted CBS from the proxy fit (Section 13), and \(r_1\sigma^2\) is the empirically estimated gradient noise level for each \(\beta\). Normalizing by \(r_1\sigma^2\) is essential because the noise amplitude varies by more than \(10\times\) across \(\beta\) values; without normalization, the raw CBS shows essentially no correlation with \(\mathrm{MF}\) (\(R^2 < 0.09\)), since the variation in \(r_1\sigma^2\) masks the momentum-factor dependence. After normalization, the power-law fit \(b^\star/(r_1\sigma^2) \propto \mathrm{MF}^{0.634}\) achieves \(R^2{=}0.897\) (rounded to \(0.90\) in the main text). The sub-linear exponent (\(0.634 < 1.0\)) is expected because Proposition 6 provides a lower bound: the empirical CBS need not scale as steeply as the bound predicts. The full momentum factor \(\mathrm{MF}\) provides a slightly better fit (\(R^2{=}0.897\)) than \((1{-}\beta)\) alone (\(R^2{=}0.864\)), supporting the Nesterov-aware functional form in Table 1. We caution that the regression has only \(n{=}6\) data points, which limits its statistical power. The result is suggestive of the predicted monotone relationship, but should not be interpreted as definitive evidence for a precise exponent value.

13.8 Gradient Variance Sensitivity↩︎

The gradient variance \(\sigma^2\) is a key component of the CBS lower bound formula. To understand its sensitivity to training hyperparameters, we measured \(\sigma^2\) across learning rates \(\eta \in \{0.005, 0.01, 0.02\}\) and batch sizes \(b \in \{128, 512, 2048\}\) on ResNet-18/CIFAR-10(width \(= 1.0\), Muon with Nesterov). Table ¿tbl:tab:grad95var95sensitivity? reports the measured \(\sigma^2\) and effective rank.

Two observations are noteworthy. First, \(\sigma^2\) varies by \(2\)–\(3\) orders of magnitude across batch sizes at a given learning rate (e.g., \(0.00113\) to \(9.31\) at \(\eta{=}0.005\)), confirming that batch size has a dominant effect on the gradient noise level. Second, higher learning rates yield lower \(\sigma^2\), reflecting faster convergence to a region of lower gradient variance. By contrast, the effective rank is relatively stable across all conditions (\(94\)–\(134\)), consistent with the rank being a structural property of the network rather than a training-dynamics quantity. These observations support treating \(r_1\) as approximately constant and \(\sigma^2\) as the primary problem-dependent variable in the CBS lower bound formula.

14 Width-Varying CBS Experiment↩︎

This section provides the detailed results for the width-varying CBS experiment described in Section 5. We trained ResNet-18 on CIFAR-10 with width multipliers \(w \in \{0.25, 0.5, 1.0, 2.0\}\), batch sizes \(\{8, 16, 32, 64, 128, 256, 512, 1024, 2048\}\), and two optimizers (Muon and AdamW).

14.0.0.1 SFO complexity across widths.

Table 9 shows the CBS and minimum SFO for each (optimizer, width) pair, estimated from the training accuracy target (\(\ge 95\%\)). At every width, Muon achieves a lower minimum SFO than AdamW, confirming its efficiency advantage across model scales. Notably, wider networks yield lower SFO for both optimizers, reflecting the fact that increased model capacity accelerates convergence to the training target.

14.0.0.2 SFO sensitivity near the critical batch size.

The CBS point estimates in Table 9 are constrained to a power-of-2 grid, which limits the resolution of the CBS estimate. To assess the uncertainty, Table ¿tbl:tab:sfo95sensitivity? reports the SFO complexity at the CBS and at the two adjacent batch sizes (CBS\(/2\) and CBS\({\times}2\)) for each width. At most widths, the relative SFO increase when moving to a neighboring batch size is less than \(5\%\), confirming that the SFO curve is extremely flat near the minimum. This flatness explains why the empirical CBS exponent with respect to rank (\(-0.17\); 95% bootstrap CI: \([-0.45, 0.00]\), \(R^2{=}0.20\)) is noisy and statistically indistinguishable from zero. By contrast, the minimum SFO itself scales cleanly as \(\mathrm{SFO}^\star \propto r_1^{-0.89}\) (95% CI: \([-1.17, -0.38]\), \(R^2{=}0.89\)), in excellent agreement with the theoretical prediction \(\mathrm{SFO}^\star \propto 1/r_1\) that follows from \(b^\star \propto r_1\).

14.0.0.3 Effective rank vs.width.

Table 10 reports the early-training gradient effective rank and momentum-error effective rank for Muon at batch size 512 (representative) across all eight widths. Both measures increase monotonically with width, and the ratio of momentum-error rank to gradient rank remains within \(1.00\)–\(1.05\), confirming that changing the width multiplier effectively controls the rank of the gradient matrices. Together with the CBS data in Table 9, these results provide directional evidence for the rank-dependent CBS scaling predicted by the theory: the effective rank spans an order of magnitude across widths, and Muon’s CBS is consistently larger than AdamW’s at every width.

14.0.0.4 Momentum error rank vs.gradient rank.

To validate that the gradient effective rank used throughout our experiments is a faithful proxy for the theoretical quantity \(r_1 = \sup_t \mathrm{rank}(C_t - \nabla f(W_t))\), we tracked the effective rank of the momentum error matrix \(M_{t-1} - \nabla f_{\mathcal{S}_t}(W_t)\) during training (whose rank equals that of \(C_t - \nabla f_{\mathcal{S}_t}(W_t)\) up to a scalar factor). Table 11 shows that the momentum error effective rank and the gradient effective rank are nearly identical across all eight widths (ratio within \(1.00\)–\(1.05\)). This confirms that the proxy gap between the measured and theoretical rank quantities is negligible, supporting the use of gradient rank as a practical stand-in for the CBS lower bound formula’s rank term. We note that both the gradient effective rank and the momentum-error effective rank are measured using stochastic gradients at batch size \(b{=}512\), rather than the full gradient \(\nabla f(W_t)\) appearing in the theoretical quantity \(r_1\). At \(b{=}512\) (comprising approximately \(1\%\) of the CIFAR-10 training set), the stochastic gradient is a high-fidelity estimate of the full gradient, and the stability of rank measurements across batch sizes (Table ¿tbl:tab:grad95var95sensitivity?) confirms that the stochastic rank is a reliable proxy for the full-gradient rank.

14.0.0.5 Gradient variance \(\sigma^2\) across widths.

To complete the CBS lower bound formula decomposition, we measured the per-sample gradient variance \(\sigma^2\) across all eight widths. Table 12 reports the gradient standard deviation \(\sigma\), variance \(\sigma^2\), and the full-batch gradient norm \(\|\nabla f(W)\|\) at each width. Gradient variance decreases sharply with width: \(\sigma^2 \approx 0.53\) at \(w{=}0.125\) versus \(\sigma^2 \approx 0.03\) at \(w{=}3.0\), a \(17.5\times\) decrease. A log-log regression yields \(\sigma^2 \propto r_1^{-1.17}\) (\(R^2{=}0.881\)), indicating that gradient variance scales approximately inversely with rank.

14.0.0.6 \(r_1 \cdot \sigma^2\) product analysis.

The CBS lower bound formula (Table 1) predicts \(b^\star \propto r_1 \cdot \sigma^2 / \varepsilon^2\). While \(r_1\) increases \(12\times\) from the narrowest to the widest model, Table 13 shows that the product \(r_1 \cdot \sigma^2\) varies by only \(\approx 3.4\times\) (range 3.66–12.28). The U-shaped profile of \(r_1 \cdot \sigma^2\)—high at narrow widths (due to large \(\sigma^2\)), low in the middle range, and moderately high at wide widths (due to large \(r_1\))—explains why the empirical CBS point estimates appear flat across widths. This result resolves the apparent tension between the theoretical prediction \(b^\star \propto r_1\) and the empirical CBS scaling \(\propto r_1^{-0.03}\): when \(\sigma^2 \propto r_1^{-1}\) (as observed), the product \(r_1 \cdot \sigma^2\) is approximately constant, yielding a flat CBS across widths.

15 Additional Results (Vision)↩︎

This section provides supplementary results to support the analysis in the main text.

15.0.0.1 Convergence (Vision)

We conducted experiments under various weight decay configurations to examine whether the upper bound of the learning rate is determined by the weight decay (\(\lambda = 0.125\)), similar to the analysis presented in Figure 2 of the main text (\(\lambda = 0.0625\)). As shown in Figure 10, the results are consistent with the earlier findings, confirming that the gradient norm begins to increase once the learning rate exceeds \(1/\lambda\).

Figure 12 compares convergence rates for ResNet-18 on CIFAR-10 across three batch sizes. This complements Figure 3 by providing a more detailed view of the effect of batch size.

15.0.0.2 Extended-budget convergence at \(b{=}2048\).

To verify that all optimizer variants reach full convergence at large batch sizes, we re-ran the \(b{=}2048\) experiment with \(100\) epochs (\(4\times\) the standard schedule), using 5 learning-rate settings per optimizer (30 runs total). Table ¿tbl:tab:sklk95bs204895ext? confirms that all Muon variants reach train loss \(< 0.015\) and EMA gradient norms \(\le 0.21\). The ordering Muon (\(\sim\)​92% test accuracy) \(>\) AdamW (\(\sim\)​87.5%) \(>\) Momentum SGD (\(\sim\)​87%) is preserved, consistent with Figure 3.

15.0.0.3 Critical Batch Size (Vision)

In the main text, we defined the critical batch size as the smallest batch size that reaches the test target accuracy in the fewest steps. For comparison, we also report results using the training target accuracy in Figure 11. While the overall trends remain similar, the gap between AdamW and Muon becomes significantly larger in this setting.

In the main text, Figure 4 presents results exclusively for AdamW and Muon. For a broader comparison, results for Momentum SGD are included and shown in Figure 13.

15.0.0.4 Additional Vision Workload

Due to space limits, the main text focuses on ResNet-18/CIFAR-10. Figure 14 reports VGG-16/CIFAR-100, which shows the same qualitative behavior.

15.0.0.5 VGG16/CIFAR-100: CBS generalization across architectures.

To verify that the CBS concept generalizes beyond ResNet-18/CIFAR-10, we trained VGG16 on CIFAR-100 with Muon and AdamW across batch sizes from 8 to 2048. Table 14 reports the SFO complexity (test) for each optimizer-batch size pair. Muon exhibits a clear U-shaped SFO curve with CBS \(\approx 64\) and best SFO of \(9.00 \times 10^5\). AdamW achieves its lowest SFO at batch size 128 (SFO \(= 3.92 \times 10^6\)), with SFO increasing sharply at larger batch sizes. The resulting SFO advantage of Muon is \(\approx 4\times\), demonstrating that Muon’s efficiency gains transfer to a substantially different architecture and dataset.

15.1 Extended CBS Validation↩︎

15.1.0.1 CBS at different target precision levels: \(b^\star \propto 1/\varepsilon^2\).

To directly verify the \(\varepsilon\)-dependence of the CBS lower bound formula, we varied the target accuracy for ResNet-18/CIFAR-10 across four levels, \((75\%/80\%)\), \((80\%/85\%)\), \((85\%/90\%)\), and \((90\%/95\%)\) (test/train), and measured the SFO complexity at seven batch sizes (32–2048). As the target becomes more stringent (smaller \(\varepsilon\)), the best SFO increases sharply: from \(5.0 \times 10^4\) at the easiest target (75%/80%) to \(3.5 \times 10^5\) at the hardest target (90%/95%), consistent with the theoretical prediction that SFO scales with \(1/\varepsilon^2\). A log-log regression of best SFO against \(1/\varepsilon^2\) yields an exponent of \(1.02\) (\(R^2{=}0.984\)), close to the predicted linear relationship.

15.1.0.2 CBS lower bound formula decomposition.

Table 15 shows the gradient variance \(\sigma^2\) and its product with rank across widths, and Table 16 compares predicted vs.empirical CBS.

15.1.0.3 Interpretation of predicted/empirical ratios.

The predicted-to-empirical CBS ratios in Table 16 range from \(0.43\) to \(1.98\), indicating order-of-magnitude agreement rather than precise prediction. Because the theoretical CBS formula provides a lower bound (Proposition 6), ratios less than \(1\) (predicted \(<\) empirical) are theoretically expected: the bound may be tighter at some widths than others. Conversely, ratios greater than \(1\) indicate that the bound is loose for those configurations; that is, the empirical CBS is smaller than the lower bound predicts, which can occur because the bound involves worst-case constants. The practical utility of the formula is in predicting the scaling trend (how CBS varies with width, \(\beta\), and \(\sigma^2\)) rather than exact CBS values.

15.1.0.4 Choice of model scale.

Our Llama3.1 (320M) experiments use a 320M-parameter model, which is limited relative to frontier LLM scales. This choice is dictated by computational budget constraints: each batch-size sweep requires training multiple models to completion, and we sweep over eight batch sizes, four Muon variants, and multiple momentum values, totaling hundreds of training runs. Importantly, the theoretical predictions we validate are derived from properties of the optimizer and the loss landscape (smoothness, gradient variance) that are architecture-independent. The fact that these predictions hold consistently across both vision (ResNet-18/CIFAR-10) and language (Llama3.1 (320M)/C4) workloads provides evidence for their generality. Moreover, [13] and [18] have independently demonstrated that Muon’s efficiency gains persist at larger LLM scales (up to 3B parameters).

15.1.0.5 Ablation Study

Finally, we conducted an ablation study to examine the relationship between the weight decay parameter \(\beta\) and the learning rate across four batch sizes. Figure 15 illustrates how Muon’s weight decay and learning rate affect the loss for each batch size. For the training loss (top row), lower weight decay consistently corresponds to reduced loss. Similarly, the smallest learning rates within the explored range are preferred. For the test loss (bottom row), smaller learning rates yield better results. However, a clear inflection point emerges in weight decay at around \(10^{-1}\), or approximately \(10^{-2}\) for larger batch sizes, indicating that these weight decay settings minimize the test loss.

16 Empirical Behavior of \(r_1\) and the Convergence Floor (FashionMNIST)↩︎

The convergence bound in Theorems 1–2 contains two terms whose interpretation has been questioned: (a) the factor \(r_1 := \sup_t \operatorname{rank}(C_t - \nabla f(W_t))\), whose worst-case value is \(\min(m,n)\) and which could in principle grow with \(T\); and (b) the \(\sqrt{(1-\beta)\,r_1/b}\) term, which is independent of both \(\eta\) and \(T\) and therefore, taken in isolation, could be read as requiring a large batch size for convergence. This section provides controlled empirical evidence that addresses both points. The experiments use a 2-layer FashionMNIST MLP (fc1: \(m\times 784\) with \(m\in\{512, 2048\}\); fc2: \(10 \times m\)), Muon with Newton–Schulz orthogonalization (NS5), weight decay \(0.01\), and Nesterov momentum, trained with three seeds per configuration. To match the theoretical definition of \(r_1\) exactly, the gradient \(\nabla f(W_t)\) in Section 16.1 is computed as the full empirical-risk gradient over the entire 60,000-example FashionMNIST training set every 150 steps (not as a sub-sample probe). All code and configs are released under workspace/configs/rc2_e{1,2,3}*.json.

16.1 Saturation of \(r_1\) along the Training Trajectory↩︎

At every eval step we form \(E_t := C_t - \nabla f(W_t)\) for the fc1 parameter and compute its singular value decomposition. Importantly, \(\nabla f(W_t)\) here is the full empirical-risk gradient (averaged over all \(60{,}000\) training examples), not a mini-batch estimate, so that \(\operatorname{rank}(E_t)\) matches the theoretical \(r_1\) exactly. We report (i) the hard rank at tolerance \(10^{-3}\): \(\#\{i : \sigma_i(E_t) > 10^{-3}\,\sigma_{\max}(E_t)\}\), which is the object appearing in the theorem; and (ii) the stable rank \(r_{\mathrm{eff}}(E_t) := \|E_t\|_F^2 / \|E_t\|_{\mathrm{op}}^2\), a continuous proxy satisfying \(r_{\mathrm{eff}} \leq \operatorname{rank}\).

Two observations support the interpretation of \(r_1\) used in the analysis. First, the hard rank plateaus within the first \(\approx 500\) steps and remains flat thereafter (e.g., at \(m=2048, \beta=0.95\) it is \(186\) at step \(600\), \(198\) at step \(1200\), \(200\) at step \(1800\), and \(246\) at step \(3000\)), so \(r_1\) does not exhibit the unbounded growth with \(T\) that a worst-case reading of the bound might suggest. Second, the late-training ratio \(r_1/\min(m,n)\) is always \(\leq 0.42\) and frequently \(\leq 0.30\), which matches the empirical observation of [18] that rank \(r \approx n/4\) suffices for Muon in LLMs up to 3B parameters. The stable rank \(r_{\mathrm{eff}} \approx 1.2\) further indicates that the Frobenius energy of \(E_t\) concentrates on essentially a single direction, making the Frobenius-norm bound used in the proof considerably tighter than \(\sqrt{r_1}\,\|E_t\|_{\mathrm{op}}\) would suggest.

16.2 Batch-Size Dependence of the Gradient-Norm Plateau↩︎

Our bound predicts that the time-averaged gradient norm plateau has the structural form \((\mathrm{plateau})^2 = A + B/b\), where \(A = \mathcal{O}(\eta L \hat{r})\) collects the \(b\)-independent floor and \(B \propto (1-\beta)\, r_1\, \sigma^2\) captures the mini-batch variance. We evaluate this prediction by fixing \((\beta, \eta, m, \lambda) = (0.95, 0.05, 1024, 0.01)\) and sweeping \(b \in \{8, 16, 32, 64, 128, 256, 512\}\) with three seeds; the plateau is taken as the mean of \(\|\nabla f(W_t)\|_F\) over the last \(20\,\%\) of eval points.

The fitted coefficients are \(\sqrt{A} = 0.349\) and \(\sqrt{B} = 1.494\) with \(R^2 = 0.973\). Three interpretations follow. (i) The \(1/b\) variance component has effectively saturated by \(b \approx 64\); beyond that the plateau is controlled by the \(b\)-independent \(A\)-term. This directly refutes the reading that convergence requires a large batch size. (ii) In a log–log view, the naive expectation of a slope of \(-1/2\) is not observed (the empirical slope is \(-0.13\)) precisely because the plateau is a sum of an \(A\)-term and a \(B/b\)-term rather than a pure \(1/\sqrt{b}\). (iii) The magnitude of \(\sqrt{B}\) is consistent with a small effective \(r_1\): setting \(\sqrt{B}=\sqrt{(1-\beta)\,r_1\,\sigma^2}\) with \(\beta=0.95\) and the empirical \(\sigma^2\) recovers \(r_1 = \mathcal{O}(1)\)–\(\mathcal{O}(10)\), in line with the stable-rank measurements in Section 16.1.

16.3 Effect of a Diminishing Step-Size Schedule on the Floor↩︎

The remaining \(A\)-component is of order \(\eta L \hat{r}\) and, unlike the \(B/b\) term, is not independent of \(\eta\). The diminishing step-size regime of Theorems 1–2 (where \((1/T)\sum_{t}\eta_t \to 0\)) implies that replacing the constant schedule by \(\eta_t = \eta_0 / \sqrt{1+t}\) drives \(A \to 0\). We test this on FashionMNIST at the most floor-sensitive regime \(b=8\) (hidden \(m=1024\), \(\beta=0.95\), \(\lambda=0.01\), \(6000\) steps, three seeds).

The \(24\,\%\) plateau reduction at fixed batch size \(b=8\) is the direct empirical counterpart of the theoretical prediction that a decaying schedule eliminates the \(\mathcal{O}(\eta L \hat{r})\) floor. Together with the boundedness of \(r_1\) (Appendix 16.1) and the rapid saturation of the \(1/b\) component (Appendix 16.2), this shows that the lever for pushing the convergence floor down is the step-size schedule, not the batch size. This addresses the concern that the \(\sqrt{(1-\beta)r_1/b}\) term, which is by itself independent of \((\eta, T)\), forces the use of large batches: it does not, because the accompanying \(A\)-term is the binding constraint and is \(\eta\)-controllable.

17 Additional Results (Language Modeling)↩︎

We provide supplemental plots for the LLM experiments discussed in the main text. Figure 19 shows final loss and SFO complexity versus batch size. Muon is consistently better than AdamW and the gap widens at large batch sizes. Nesterov momentum and weight decay do not produce systematic gains in this setting.

To examine the role of momentum, we swept \(\beta\) under two configurations: with and without Nesterov (weight decay fixed at zero unless indicated). Figure 20 shows that the best trade-off is near \(\beta=0.95\). As \(\beta\) increases, the critical batch size decreases, but very small or very large values of \(\beta\) degrade both loss and SFO.

17.0.0.1 CBS robustness across stopping criteria.

To verify that the CBS conclusions on the LLM workload are not artifacts of the chosen loss target, we swept the stopping criterion across two metric types: gradient-norm thresholds (\(\|\nabla f\| \le \varepsilon\), \(\varepsilon \in [0.1, 0.7]\)) and training-loss targets (\(\mathcal{L} \le \ell\), \(\ell \in [3.3, 4.0]\)). Figure 21 shows the empirical CBS for Muon, Muon+Nesterov, and AdamW under each criterion. Although the exact CBS value shifts as the threshold tightens, the qualitative ranking (Muon attaining a CBS comparable to or larger than AdamW) is preserved across both metric types. This is consistent with the proxy-invariance observed on FashionMNIST (Table ¿tbl:tab:stopping95sensitivity?).

17.0.0.2 Gradient-norm CBS on C4.

To use a stopping criterion that directly corresponds to the quantity bounded by the theorems, we extracted the SFO complexity for Muon with Nesterov (\(\beta{=}0.95\)) at gradient-norm thresholds \(\|\nabla f\| \le \varepsilon\) across the full batch-size range (GBS \(= 16\) to \(8192\)). Table ¿tbl:tab:llm95cbs95target95sensitivity? reports the results. The CBS shifts from \(32\) at \(\varepsilon = 0.40\) to \(128\) at \(\varepsilon = 0.25\), consistent with the predicted \(b^\star \propto 1/\varepsilon^2\) scaling. At every threshold, the SFO curve exhibits a clear U-shape: for example, at \(\varepsilon = 0.20\), the SFO at GBS \(= 4096\) (\(1.33 \times 10^6\)) is \(10\times\) larger than at the CBS (\(1.30 \times 10^5\)).

17.0.0.3 Training Curves for Muon on Llama3.1 (320M)

To further illustrate Muon’s optimization dynamics on large-language-model workloads, we report full training curves for Llama3.1 (320M) trained on C4 at multiple batch sizes. Figure 22 plots the training loss against the number of steps for Muon and AdamW. Across all batch sizes, Muon exhibits faster loss reduction in early training and consistently reaches lower final loss within the same token budget. These results complement Figures 6 and 7, demonstrating that the advantages of Muon are visible not only in SFO complexity but also in the raw optimization trajectory.

References↩︎

---

1. https://huggingface.co/datasets/allenai/c4↩︎

2. For large-scale workloads, computing the full gradient is impractical; we therefore use the mini-batch gradient norm as a practical surrogate for the full gradient and report EMA-smoothed curves for stability. We emphasize that while the stochastic gradient is an unbiased estimator of the true gradient, its norm is not an unbiased estimator of the true gradient norm; accordingly, we treat this quantity as a calibrated surrogate rather than an exact proxy. In our controlled full-Muon experiments (Appendix 13), the Pearson correlation between the full-batch gradient norm and the mini-batch gradient norm exceeds \(0.99\) across all tested batch sizes and seeds (Table ¿tbl:tab:proxy95calibration?), confirming that the surrogate faithfully tracks the theoretical quantity. Theoretically, the CBS is the batch size that minimizes the computational cost required for the true gradient norm to reach a prescribed threshold. Empirically, however, the batch size that minimizes a given computational cost falls in the same regime whether measured via full-gradient norm, mini-batch gradient norm, or task-level metrics such as accuracy and loss (see Table ¿tbl:tab:stopping95sensitivity?). In a controlled toy setting (Appendix 13), we also compute the full-batch gradient norm to directly match the theoretical quantity. The EMA window (100 steps) is chosen to balance noise suppression and responsiveness; preliminary tests with windows of 50 and 200 steps yielded CBS estimates within the same power-of-two grid point.↩︎