Airborne Muon: Spectral Matrix Orthogonal Optimizer for LLMs

Muon is a production-grade, drop-in replacement for torch.optim.AdamW designed specifically for modern Large Language Models and Transformer architectures.

By replacing Adam's coordinate-wise normalization with Spectral Matrix Sign Orthogonalization via Quintic Newton-Schulz iterations, Muon enables models to converge in half the steps of AdamW while maintaining 100% compute efficiency (zero extra backward passes).

Benchmark Highlights: The 6-Way LLM Optimizer Tournament

Tested head-to-head on autoregressive Transformer pre-training under identical seeds, architecture (MiniGPT), and cosine decay schedules:

Key Takeaway: Airborne Muon reached 8.42 PPL at Step 250, defeating AdamW's final 500-step score in half the optimization steps. On full 124M GPT-2 on an NVIDIA A100-80GB, Airborne Muon matched AdamW's final loss in $\le 180$ steps (64% step reduction).

Why Muon Outperforms AdamW on LLMs

1. Matrix-Aware Geometry: Adam normalizes each parameter scalar independently: $u_i = m_i / \sqrt{v_i}$. For 2D linear weight matrices ($W_q, W_k, W_v, W_{\text{out}}$), this is blind to coordinate rotations and causes severe zigzagging across anisotropic valley walls.
2. Spectral Orthogonalization (Newton-Schulz-5): Muon computes the approximate matrix sign of the momentum tensor: $$O = M (M^T M)^{-1/2}$$ This normalizes all singular values to $1.0$, ensuring that all principal feature directions update at a uniform, optimal velocity.
3. Automatic Hybrid Routing:
   • 2D Weight Matrices $\to$ Spectral Orthogonal Momentum ($O$)
   • 1D Vectors (LayerNorm, RMSNorm, Biases) & Token Embeddings $\to$ Decoupled AdamW
4. Zero Compute Overhead: Unlike second-order curvature methods (AdaHessian, Sophia) that require extra backward passes or double-backward autograd graphs, Muon runs in zero extra forward/backward passes.

Quickstart

Installation

pip install airborne-muon
# Or from source:
pip install git+https://github.com/AirBorneAI/airborne-muon.git

Usage (Drop-in Replacement for AdamW)

import torch
from muon import Muon

# Initialize your Transformer / LLM
model = MyTransformer().to('cuda')

# Initialize Muon with hybrid routing
optimizer = Muon(
    model.parameters(),
    lr=1e-3,           # Learning rate for 1D vectors / embeddings (AdamW)
    matrix_lr=2e-2,    # Spectral learning rate for 2D weight matrices (Muon)
    momentum=0.95,     # Momentum coefficient
    weight_decay=0.01  # Decoupled weight decay
)

# Standard training loop (no closures, no extra passes!)
for x, y in dataloader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast('cuda', dtype=torch.bfloat16):
        loss = criterion(model(x), y)
    loss.backward()
    optimizer.step()

Empirical Research Artifacts in this Repository

• muon.py: Production-grade Muon optimizer class with integrated AdamW routing.
• cs_adam.py: Curvature-Steered Adam implementing directional Gram-Schmidt projection.
• sophia_fd.py: Sophia-FD optimizer with finite-difference Hutchinson probing.
• benchmark_grand_showdown.py: Deterministic pretraining benchmark script on TinyShakespeare.
• models/architectures.py: Causal Transformer (MiniGPT) and ConvNet architectures.
• test_curvature_unit.py: Mathematical proof of the $1/\eta$ minibatch secant collapse.

Citation & Acknowledgments

Built by the research team at AirBorne AI. Inspired by spectral matrix momentum discoveries by Keller Jordan and the Sophia optimization framework by Liu et al. (Stanford).

@software{airborne_muon_2026,
  author = {Singh, Suryaansh Prithvijit and AirBorne AI},
  title = {Airborne Muon: Production Spectral Matrix Orthogonal Optimizer for LLMs},
  year = {2026},
  url = {https://github.com/AirBorneAI/airborne-muon}
}