MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training

Summary (Overview)

• MONA is a novel optimizer that integrates curvature-aware acceleration into the Muon optimizer's matrix orthogonalization framework, enabling escape from sharp minima while preserving Muon's geometric benefits.
• The method augments gradients with an exponential moving average of gradient differences before orthogonalization, allowing the momentum buffer to capture curvature-aware directions.
• Theoretical analysis proves convergence at $O(K^{-1/2})$ rate under standard assumptions, matching SGD and AdamW, while formalizing how acceleration selectively resists sharp minima.
• Empirically, MONA outperforms both Muon and AdamW across three MoE model scales (1B to 68B parameters) on code, mathematical reasoning, and general text benchmarks.
• MONA-Lite, a memory-efficient variant using BF16 quantization and streaming gradient computation, reduces extra memory overhead by ~75% without sacrificing training quality.

Introduction and Theoretical Foundation

Problem Setting

The paper considers the stochastic non-convex optimization problem:

$$\min_{\theta \in \mathcal{D}} f(\theta), \quad f(\theta) = \mathbb{E}_{\zeta}[\ell(\theta, \zeta)]$$

where $\ell(\theta, \zeta)$ is the loss on sample $\zeta$ from distribution $\mathcal{P}$. Parameters partition into matrix-valued parameters $\{W^{(i)} \in \mathbb{R}^{m_i \times n_i}\}$ and vector-valued parameters.

Muon Optimizer

Muon updates matrix parameters via:

$$M_k = \mu M_{k-1} + G_k$$ $$O_k = \text{Newton-Schulz}(M_k)$$ $$W_{k+1} = W_k - \eta(\gamma O_k + \lambda W_k)$$

The Newton-Schulz iteration approximates polar decomposition with coefficients $a = 3.4445, b = -4.7750, c = 2.0315$, ensuring convergence for singular values in $[0, 1]$. The scaling factor is:

$$\gamma = 0.2 \cdot \sqrt{\max(m, n)}$$

Acceleration Mechanism

The acceleration term computes:

$$d_k = g_k - g_{k-1}$$ $$a_k = \beta_a a_{k-1} + (1 - \beta_a) d_k$$ $$\tilde{g}_k = g_k + \alpha a_k$$

The theoretical motivation derives from:

$$-\nabla\|\nabla f(\theta_k)\|^2 = -2H_k\bar{g}_k \approx \bar{g}_k - \bar{g}_{k-1}$$

where $H_k$ is the Hessian. This direction points away from sharp minima, and gradient differences serve as a computationally efficient proxy.

Methodology

MONA Algorithm

The key innovation is applying acceleration before momentum accumulation:

Algorithm 1: MONA Optimizer

1. Initialize $M_0 \leftarrow 0$, $A_0 \leftarrow 0$, $G_0 \leftarrow 0$
2. For $k = 1, 2, 3, \ldots$:
   • $G_k \leftarrow \nabla_W \ell(W_k, \zeta_k)$
   • $D_k \leftarrow G_k - G_{k-1}$
   • $A_k \leftarrow \beta_a A_{k-1} + (1 - \beta_a) D_k$
   • $\tilde{G}_k \leftarrow G_k + \alpha A_k$
   • $M_k \leftarrow \mu M_{k-1} + \tilde{G}_k$
   • $O_k \leftarrow \text{Newton-Schulz}(M_k, T)$
   • $W_{k+1} \leftarrow W_k - \eta(\gamma O_k + \lambda W_k)$
   • $G_{k-1} \leftarrow G_k$

For vector-valued parameters, MONA falls back to AdamW, following Muon's convention.

Geometric Intuition

MONA's effectiveness stems from two complementary mechanisms:

1. Spectral normalization (Muon): Newton-Schulz ensures $O_k$ has singular values near 1, preventing over-commitment to large-gradient directions.
2. Curvature-aware acceleration: The term $A_k$ encodes gradient changes—near sharp minima, $\|D_k\|$ is large, pushing toward flatter regions; in flat regions, $A_k$ is small, allowing stable convergence.

Theoretical Analysis

Key Assumptions

1. L-smoothness: $\|\nabla_W \ell(W, \zeta) - \nabla_W \ell(W', \zeta)\|_F \leq L\|W - W'\|_F$
2. Unbiased gradient with bounded variance: $\mathbb{E}[G_k | W_k] = \nabla f(W_k)$ and $\mathbb{E}[\|G_k - \nabla f(W_k)\|_F^2] \leq \sigma^2$
3. Bounded gradient: $\|G_k\|_F \leq G$ a.s.
4. Expected directional alignment: $\mathbb{E}[\langle \nabla f(W_k), O_k \rangle | W_k] \geq \rho\|\nabla f(W_k)\|_F^2$

Key Results

Lemma 1 (Boundedness of acceleration): $\|A_k\|_F \leq 2G$ and $\|\tilde{G}_k\|_F \leq G(1 + 2|\alpha|)$

Lemma 2 (Momentum bound): $\|M_k\|_F \leq \frac{G(1 + 2|\alpha|)}{1 - \mu}$

Theorem 1 (Non-convex convergence): With $\eta \leq \min\{1/L, C_3/C_2\}$, after K iterations:

$$\frac{1}{K}\sum_{k=0}^{K-1}\mathbb{E}\left[\|\nabla f(W_k)\|_F^2\right] \leq \frac{f(W_0) - f^*}{\eta C_3 K} + \frac{\eta L C_4}{C_3}$$

With $\eta = O(K^{-1/2})$, this gives $O(K^{-1/2})$ convergence, matching SGD and AdamW.

Proposition 1 (Sharp minimum escape): Near a sharp minimum with $\lambda_{max}(H^*) \gg \lambda_{min}(H^*) > 0$:

$$\mathbb{E}[A_k] \approx -\eta\gamma H^* \sum_{j=0}^{k}(1 - \beta_a)\beta_a^{k-j}O_{j-1}$$

For large eigenvalues (sharp directions), acceleration is large, promoting escape; for small eigenvalues (flat directions), it's small, permitting convergence.

Empirical Validation

Pretraining Results

The paper evaluates three MoE models based on the LongCat/ScMoE architecture with MLA:

MOE-1B-A0d2B: 10 layers, 768 hidden dims, 128 experts, top-8 routing, ~400B tokens MOE-6B-A0d5B: 10 layers, 1536 hidden dims, 128 experts, top-6 routing, ~1.2T tokens MOE-68B-A3B: 14 layers, 3072 hidden dims, 256 experts, top-12 routing, ~700B tokens

MONA consistently achieves lower validation loss than Muon across all scales and domains (code, general English, math, Chinese academic text).

Benchmark Results for MOE-68B-A3B

Table 1: General Capability (mean ± std)

Table 2: Code & Math (mean ± std)

Supervised Fine-Tuning Results

After code-specific SFT on the MOE-68B-A3B model, MONA-pretrained checkpoints achieve higher scores on 6 out of 7 BigCode tasks compared to Muon-pretrained checkpoints, with notable gains on MBPP+ (0.7857 vs 0.7778), HumanEval+ (0.8841 vs 0.8720), and DS1000 (0.3703 vs 0.3395).

MONA-Lite Efficiency

MONA-Lite combines BF16 quantization with streaming gradient computation, reducing extra memory overhead by ~75% (from 2 full-precision buffers to ~0.5 gradient's worth). Validation loss curves show MONA-Lite closely tracks FP32 MONA while maintaining clear advantage over Muon, with negligible training speed overhead (~1% in optimizer step, disappearing at iteration level).

Theoretical and Practical Implications

Key Findings

1. Convergence guarantee: MONA maintains $O(K^{-1/2})$ convergence rate matching SGD and AdamW, with acceleration affecting only constants, not asymptotic behavior.
2. Sharp minimum escape: The acceleration term provides Hessian-dependent scaling—sharp directions amplify the acceleration while flat directions suppress it, enabling selective escape from sharp minima.
3. Practical scalability: MONA demonstrates consistent improvements across three orders of magnitude in model scale (1B to 68B parameters), suggesting robust generalization.
4. Transfer learning benefit: MONA-pretrained models adapt better to downstream tasks, achieving lower SFT loss and higher benchmark scores after fine-tuning.

Practical Considerations

• Hyperparameter tuning: MONA introduces only two new hyperparameters $(\beta_a, \alpha)$ with a consistent relationship $\alpha = -1/(2(1-\beta_a))$. Recommended values: $\beta_a = 0.99$ (1B), 0.98 (6B), 0.975 (68B).
• Memory overhead: Even with MONA-Lite optimizations, the acceleration buffer $A_k$ requires ~half a gradient's worth of extra memory. Users should balance this against training gains.
• Deployment flexibility: MONA-Lite makes the approach practical for memory-constrained environments without sacrificing training quality.

Conclusion

MONA successfully bridges Muon's orthogonalization framework with curvature-aware acceleration, providing:

• Theoretical guarantees: Convergence under standard assumptions with sharp minimum escape properties
• Empirical superiority: Consistent improvements over Muon and AdamW across model scales (1B-68B) and task domains
• Practical efficiency: MONA-Lite variant reduces memory overhead by ~75% while maintaining performance

Future work could explore adaptive hyperparameter selection for $\beta_a$ and $\alpha$, application to other architectures beyond MoE language models, and further memory optimizations for the acceleration buffer.