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

> MONA adds curvature-aware acceleration to Muon's orthogonalization, achieving consistent gains over Muon and AdamW across 1B-68B parameter MoE models with proven convergence guarantees.

- **Source:** [arXiv](https://arxiv.org/abs/2605.26842)
- **Published:** 2026-08-29
- **Permalink:** https://picx.dev/p/Ybe9ds
- **Whiteboard:** https://picx.dev/p/Ybe9ds/image

## Summary

# 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)

| Benchmark | AdamW | Muon | MONA |
|-----------|-------|------|------|
| MMLU-FewShot | 0.6218 ± 0.0098 | 0.6281 ± 0.0098 | **0.6373 ± 0.0097** |
| MMLU-Pro-FewShot | 0.3150 ± 0.0169 | 0.3250 ± 0.0170 | **0.3375 ± 0.0172** |
| CMMLU-FewShot | 0.6464 ± 0.0090 | 0.6723 ± 0.0088 | **0.6756 ± 0.0088** |
| CEVAL-FewShot | 0.6274 ± 0.0263 | 0.6586 ± 0.0255 | **0.6717 ± 0.0251** |
| BBH-FewShot | 0.3966 ± 0.0125 | 0.3962 ± 0.0126 | **0.4049 ± 0.0126** |
| **Average** | 0.4382 | 0.4478 | **0.4557** |

**Table 2: Code & Math** (mean ± std)

| Benchmark | AdamW | Muon | MONA |
|-----------|-------|------|------|
| Multiple | 0.2519 ± 0.0326 | 0.2239 ± 0.0313 | **0.2848 ± 0.0343** |
| BigCodeBench | 0.3123 ± 0.0269 | 0.3061 ± 0.0268 | **0.3377 ± 0.0275** |
| LiveCodeBench | 0.0358 ± 0.0219 | 0.0430 ± 0.0238 | **0.0502 ± 0.0257** |
| Bigmath Mathematics | 0.4627 ± 0.0173 | 0.4493 ± 0.0182 | **0.4733 ± 0.0173** |
| **Average** | 0.3353 | 0.3397 | **0.3495** |

### 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.

---

_Markdown view of https://picx.dev/p/Ybe9ds, served by PicX — AI-generated visual whiteboard summaries of research papers._
