At the end of 2024, independent researcher Keller Jordan proposed a new neural network optimizer—Muon (MomentUm Orthogonalized by Newton‑Schulz). This optimizer is specifically designed for 2D weight matrices in neural networks (the weight matrices of hidden layers). By combining Nesterov momentum with Newton‑Schulz orthogonalization iterations, Muon precisely optimizes the geometric structure of weight matrices. Between 2025 and 2026, Muon quickly gained popularity in the AI community, especially in Andrej Karpathy's minimal projects such as nanoGPT, nanochat and autoresearch, becoming a core component of current speedrun training setups. PyTorch 2.11 has integrated it into torch.optim.Muon, marking its entry into mainstream frameworks.
Muon emerged as a response to structural shortcomings of traditional optimizers like AdamW when training large Transformer models. Optimizers such as AdamW treat weight matrices as collections of independent scalars and perform element‑wise updates, ignoring the matrix as a geometric object and causing gradient interference across directions. Muon treats the weight matrix as an integrated object: after momentum update it immediately orthogonalizes the update, preserving better geometric properties of the update direction and thus improving training stability and convergence speed.
Muon PyTorch implementation and core mechanism
The official Muon implementation is concise and efficient. Keller Jordan's GitHub repository (https://github.com/KellerJordan/Muon) provides the full code, and PyTorch has also integrated it. The core idea is "momentum update first, then orthogonalize."
Below is a simplified core implementation (based on the official code, PyTorch style):
Python
import torch
from torch.optim import Optimizer
def newtonschulz5(G, steps=5, eps=1e-7):
"""Newton‑Schulz iteration implementing matrix orthogonalization"""
assert G.ndim == 2
a, b, c = (3.4445, -4.7750, 2.0315) # Keller Jordan's preferred polynomial coefficients
X = G.bfloat16()
X = X / (X.norm() + eps)
if G.size(0) > G.size(1):
X = X.T
for _ in range(steps):
A = X @ X.T
B = b * A + c * A @ A
X = a * X + B @ X
if G.size(0) > G.size(1):
X = X.T
return X
class Muon(Optimizer):
def __init__(self, params, lr=0.1, momentum=0.95, nesterov=True, steps=5):
defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, steps=steps)
super().__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
g = p.grad
if g.ndim != 2: # Apply Muon only to 2D weight matrices
continue
state = self.state[p]
if len(state) == 0:
state['momentum'] = torch.zeros_like(g)
# Nesterov momentum update
state['momentum'].mul_(group['momentum']).add_(g, alpha=1 - group['momentum'])
g = state['momentum'] if group['nesterov'] else g
# Newton‑Schulz orthogonalization
g_ortho = newtonschulz5(g, steps=group['steps'])
# Apply update
p.data.add_(g_ortho, alpha=-group['lr'])
return loss
Key insights：
- Newton‑Schulz iteration: Uses five iterations to quickly approximate an orthogonal matrix (operator norm near 1), avoiding gradient explosion or vanishing.
- Hybrid usage: In practice, a Muon + AdamW hybrid strategy is common—use Muon for 2D weight matrices (attention and FFN layers) and AdamW for 1D parameters such as biases and embeddings.
- Hyperparameters: Typical learning rate around 0.1 (much higher than AdamW's 0.001), momentum 0.95, and five Newton‑Schulz steps are usually sufficient.
Detailed comparison between Muon and AdamW
AdamW is the de facto standard optimizer for LLM training today, relying on element‑wise adaptive learning rates (based on first and second moment estimates of gradients). Muon is a matrix‑aware optimizer. They contrast across several dimensions:
Figure 1. Optimizer comparison by sample efficiency. [reproducible logs]
1. Optimization target：
- AdamW: Treats weights as independent scalars and updates element‑wise.
- Muon: Treats 2D weight matrices as geometric wholes—momentum update followed by orthogonalization.
2. Convergence speed：
- Experiments show Muon can speed up training by about 1.35x–2x on tasks like nanoGPT and CIFAR‑10 speedruns.
- In LLM training, Muon + AdamW hybrids often reach the same loss in fewer steps, saving up to ~50% compute in some cases.
3. Memory usage：
- Muon maintains only a single momentum state, reducing memory overhead by about 33% compared with AdamW.
- Hybrid schemes further lower peak memory while preserving performance.
4. Stability and generalization：
- Muon reduces spectral bias via orthogonalization, making updates more uniform and performing better on heavy‑tailed distributions.
- AdamW can exhibit gradient spikes during long runs, while Muon's updates are smoother.
5. Scope of application：
- AdamW: Highly general, suitable for all parameters.
- Muon: Designed for 2D weights; best used in combination with AdamW (use AdamW for embeddings/biases, Muon for hidden weights).
Usage in specific projects
1. Karpathy's autoresearch project：
- In train.py, Muon is offered as one of the default optimizers.
- Agents can modify the optimizer configuration (Muon or AdamW) and run experiments within a fixed 5‑minute budget, evaluating improvements via the val_bpb metric.
- Typical code snippet:
Python
from muon import Muon optimizer = Muon(model.parameters(), lr=0.1, momentum=0.95)
2. nanoGPT / nanochat：
- In Karpathy's minimal GPT training projects, Muon + AdamW hybrids have helped break multiple speed records.
- Typically, Muon is applied to QKV and FFN weight matrices, while embedding and bias parameters use AdamW.
3. Large‑scale LLM training：
- Teams such as Moonshot and Kimi have validated Muon's scalability on hundred‑billion‑parameter models.
- The hybrid approach has become a best practice: Muon handles geometric optimization of hidden layers while AdamW handles adaptive updates for the remaining parameters.
Practical effects and community feedback
In nanoGPT speedruns, switching to Muon increased training speed by about 35%. Muon also shows significant advantages on benchmarks like CIFAR‑10. The community generally regards Muon as one of the most important optimizer innovations since Adam, particularly well suited to Transformer‑based LLM training.
Summary and outlook
With its matrix‑aware design and efficient orthogonalization mechanism, Muon offers a new option for LLM training. It does not fully replace AdamW but complements it efficiently. Through deployments in Karpathy's minimal projects, Muon has demonstrated practical value. With PyTorch integration and more large‑scale validations, Muon could become a standard component of next‑generation AI training stacks.
From element‑wise updates to matrix‑level geometric optimization, Muon represents a shift from "general‑purpose" toward "structure‑aware" optimizers. In an era of high AI training costs, this innovation not only accelerates research cycles but also helps pave the way for more efficient and greener AI infrastructure.
▶️ Enjoyed this article? Creating it takes real effort!
Subscribe to our YouTube channel:
AI Rosetta
Your likes, subscriptions and shares are our greatest support! 🙏
For more AI news, visit:
https://www.airosetta.com