Technical Summary: MONA – Muon Optimizer with Nesterov Acceleration

1. Problem Statement

Training large language models (LLMs) increasingly demands optimizers with superior sample efficiency. While AdamW has been the standard for over a decade, recent alternatives like Muon have emerged. Muon improves upon first-order methods by treating weight matrices as single geometric units and applying matrix orthogonalization (via Newton-Schulz iteration) to the momentum buffer. This approach aligns with steepest descent under the spectral norm, offering approximately $2\times$ computational efficiency over AdamW in large-scale Mixture-of-Experts (MoE) training.

However, Muon shares a fundamental limitation with all first-order gradient methods: it lacks explicit mechanisms to explore the loss landscape beyond local gradient information. Consequently, Muon can become trapped in sharp local minima, particularly in large-batch training scenarios. Sharp minima are associated with poor generalization, whereas flatter minima tend to yield better downstream performance. While the ALTO optimizer introduced an acceleration term based on the exponential moving average (EMA) of gradient differences to escape sharp minima, its design as a general adaptor for Adam and LAMB presents integration challenges with Muon's highly nonlinear orthogonalization process.

2. Methodology: MONA

The authors propose MONA (Muon Optimizer with Nesterov Acceleration), which seamlessly integrates a curvature-aware acceleration term directly into Muon's orthogonalization framework.

Core Algorithm

MONA modifies the standard Muon pipeline by introducing an acceleration term before momentum accumulation and orthogonalization. The update process for a matrix parameter $W$ at iteration $k$ proceeds as follows:

1. Gradient Difference Calculation: Compute the difference between the current stochastic gradient $G_k$ and the previous gradient $G_{k-1}$:
   $$D_k = G_k - G_{k-1}$$
2. Acceleration Buffer Update: Maintain an exponential moving average of these differences:
   $$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$$
   Here, $\beta_a$ controls the memory of the acceleration.
3. Curvature-Aware Gradient: Construct an augmented gradient $\tilde{G}_k$ by adding the scaled acceleration term to the raw gradient:
   $$\tilde{G}_k = G_k + \alpha A_k$$
   The coefficient $\alpha$ scales the acceleration. Theoretical motivation suggests that $G_k - G_{k-1}$ approximates $H_k(\theta_k - \theta_{k-1})$, where $H_k$ is the Hessian. Thus, the term points away from sharp minima (where gradient changes are large) and remains small in flat regions.
4. Muon Orthogonalization: The augmented gradient is fed into the standard Muon momentum and orthogonalization steps:
   $$M_k = \mu M_{k-1} + \tilde{G}_k$$
   $$O_k = \text{Newton-Schulz}(M_k)$$
   $$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$$

Design Rationale

Applying the acceleration term before orthogonalization is critical. If applied after, it would disrupt the orthogonal structure that Muon relies on for spectral-norm regularization. By transforming the raw gradient into a curvature-aware direction first, the momentum buffer captures these enriched directions, which are then processed through Muon's standard pipeline.

For vector-valued parameters (e.g., embeddings, biases), MONA defaults to AdamW, adhering to Muon's existing conventions.

3. Theoretical Analysis

The paper provides a convergence analysis for MONA under standard assumptions (L-smoothness, unbiased gradients with bounded variance, and expected directional alignment).

• Convergence Rate: The authors prove that MONA achieves a convergence rate of $O(K^{-1/2})$ in non-convex settings, matching the asymptotic rates of SGD and AdamW. The acceleration term affects the constants in the bound but not the asymptotic rate.
• Sharp Minimum Escape: The analysis formalizes how the acceleration term selectively resists sharp minima. Near a stationary point, the expected acceleration $E[A_k]$ is proportional to the Hessian-weighted sum of past directions. For directions with large eigenvalues (sharp minima), the acceleration magnitude is large, pushing the optimizer away. For small eigenvalues (flat directions), the acceleration is small, allowing stable convergence. This mechanism preserves Muon's spectral-norm regularization while adding curvature awareness.

4. Empirical Results

MONA was evaluated across three scales of MoE pretraining models (1B, 6B, and 68B parameters) and a supervised fine-tuning (SFT) stage.

Pretraining Performance

• Models: MOE-1B-A0d2B (400B tokens), MOE-6B-A0d5B (1.2T tokens), and MOE-68B-A3B (700B tokens).
• Metrics: Validation loss was monitored across code, mathematical reasoning, general English, and Chinese academic text domains.
• Findings: MONA consistently achieved lower validation loss than both Muon and AdamW across all three model scales and all domains. The performance gap was particularly evident in the larger 68B model, where MONA demonstrated more stable convergence.

Downstream Benchmarks (MOE-68B-A3B)

The 68B model trained with MONA was evaluated on a comprehensive suite of benchmarks:

• General Capability: MONA achieved the highest average score (0.4557) on benchmarks including MMLU, MMLU-Pro, CMMLU, CEVAL, and BBH, outperforming Muon (0.4478) and AdamW (0.4382).
• Mathematical Reasoning & Code Generation: MONA showed superior performance on GSM8K, GPQA, and various code generation benchmarks (HumanEval+, MBPP+, BigCodeBench, LiveCodeBench). Notably, MONA delivered significant gains on tasks requiring transferable code representations.
• Supervised Fine-Tuning (SFT): When the MONA-pretrained checkpoint underwent SFT on code data, it maintained a lower training loss throughout the process compared to the Muon-pretrained baseline. Post-SFT evaluations on the BigCode suite confirmed that MONA-initialized models achieved higher scores on 6 out of 7 tasks compared to Muon-initialized models.

Efficiency and Deployment

To address the memory overhead of storing the previous gradient ($G_{k-1}$) and the acceleration buffer ($A_k$), the authors introduced MONA-Lite:

• Quantization: Storing auxiliary states in bfloat16 (BF16) instead of float32.
• Streaming Computation: Computing gradient differences in-place to eliminate the need for a dedicated $G_{k-1}$ buffer.
• Result: These strategies reduce the extra memory overhead by approximately 75% without sacrificing training quality. Computational overhead analysis showed MONA is only ~1% slower at the optimizer step level, with no practical difference in end-to-end iteration time compared to Muon.

5. Significance and Claims

The paper claims that MONA successfully bridges the gap between Muon's geometric efficiency and the need for curvature-aware acceleration. Its primary significance lies in:

1. Algorithmic Integration: Demonstrating that acceleration terms can be integrated into matrix orthogonalization frameworks without disrupting their geometric properties.
2. Scalability: Validating the method across a wide range of model sizes (1B to 68B parameters) and token counts (up to 1 trillion tokens), showing consistent improvements over both AdamW and Muon.
3. Generalization: Providing empirical evidence that escaping sharp minima during pretraining leads to better generalization in downstream tasks, particularly in complex reasoning and code generation.
4. Practicality: Offering a memory-efficient variant (MONA-Lite) that makes the optimizer viable for resource-constrained environments.

The authors conclude that MONA represents a state-of-the-art optimizer for large-scale LLM training, achieving superior convergence and downstream performance while maintaining the computational efficiency of the Muon framework.