📌 Executive Summary

Modern large-scale models are bottlenecked by standard first-order adaptive gradient descent ($g / \sqrt{v}$). OmniMuon provides a unified toolkit implementing both Spectral Matrix Momentum and Finite-Difference Second-Order Curvature Probing:

• Muon: Specialized for Autoregressive LLMs & Transformers. Replaces coordinate-wise scaling with Quintic Newton-Schulz matrix sign orthogonalization on 2D weight matrices, paired with decoupled AdamW on 1D vectors and embedding tables.
• OmniMuon: The Universal Multi-Modal formulation. Adds Canonical Tensor Matricization (unfolding 3D/4D/5D convolution tensors) and Dynamic RMS Energy Matching for Vision Transformers (ViT) and ConvNets.
• SophiaFD: Second-Order Stochastic optimization utilizing 2-pass Hutchinson diagonal Hessian probing with scale-aware numerical stability.
• Polaris: First-principles Riemannian Soft-Polar geodesic optimizer with parameter-free self-consistent regularization $\tau = \text{RMS}(M)$, running on a single unified learning rate.

🚀 Installation

Install the production package directly from PyPI:

pip install omnimuon

Or install the bleeding-edge source from GitHub:

pip install git+https://github.com/AirBorneAI/airborne-muon.git

System Requirements

• Python $\ge 3.9$
• PyTorch $\ge 2.0.0$
• Recommended: CUDA $\ge 11.8$ with native torch.bfloat16 or torch.float16 support.

⚡ Quickstart

1. Training Large Language Models with Muon (Drop-in for AdamW)

import torch
import torch.nn as nn
from omnimuon import Muon

model = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(d_model=768, nhead=12, batch_first=True),
    num_layers=12
).cuda()

# Zero boilerplate! Defaults automatically use matrix_lr=0.02, lr=1e-3, momentum=0.95, weight_decay=0.01:
optimizer = Muon(model.parameters())

# Or customize any parameter whenever needed:
# optimizer = Muon(model.parameters(), lr=1e-3, matrix_lr=0.02, momentum=0.95, weight_decay=0.01)

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

2. Multi-Modal Vision & ConvNet Training with OmniMuon

from omnimuon import OmniMuon

# Works seamlessly across Conv2D, Vision Transformers, and Multi-Modal Models
optimizer = OmniMuon(
    model.parameters(),
    lr=3e-4,              # Vector learning rate
    matrix_lr=0.02,       # Spectral tensor learning rate
    rms_scaling=True,     # Scales updates by parameter RMS (prevents patch-embed overshoot)
    spectral_blend=0.8,   # 80% Orthogonal Matrix Sign + 20% AdamW residual
    weight_decay=0.01
)

3. Second-Order Curvature Probing with SophiaFD

from omnimuon import SophiaFD

optimizer = SophiaFD(
    model.parameters(),
    lr=3e-3,
    betas=(0.96, 0.99),
    gamma=5.0,            # Curvature clipping threshold
    rho=1.0,              # Maximum parameter step bound
    delta=1e-3,           # Finite-difference perturbation magnitude
    k=10                  # Re-compute diagonal Hessian every 10 steps
)

for step, (inputs, targets) in enumerate(dataloader):
    optimizer.zero_grad(set_to_none=True)
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    
    # Sophia-FD updates diagonal Hessian via 2-pass Hutchinson probing every k steps
    if step % optimizer.param_groups[0]['k'] == 0:
        base_grads = [p.grad.clone() if p.grad is not None else None for p in model.parameters()]
        probes = optimizer.sample_probe_vectors()
        delta = optimizer.param_groups[0]['delta']

        with torch.no_grad():
            for p, u in zip(model.parameters(), probes):
                if u is not None:
                    p.add_(u, alpha=delta)

        model.zero_grad(set_to_none=True)
        criterion(model(inputs), targets).backward()
        pert_grads = [p.grad.clone() if p.grad is not None else None for p in model.parameters()]

        with torch.no_grad():
            for p, u in zip(model.parameters(), probes):
                if u is not None:
                    p.sub_(u, alpha=delta)

        for p, bg in zip(model.parameters(), base_grads):
            if bg is not None:
                p.grad = bg

        optimizer.update_hessian(pert_grads, probes)
        
    optimizer.step()

🏛️ Optimizer Architecture Guide

📊 Empirical Benchmarks

1. The 6-Way LLM Pre-Training Tournament (MiniGPT on TinyShakespeare)

All optimizers evaluated under identical deterministic parameter initializations, token streams, and cosine decay learning rate schedules:

2. Full-Scale 124M GPT-2 on NVIDIA A100-80GB (FineWeb-Edu)

Tested under native torch.bfloat16 FlashAttention on real streaming educational web text:

🛠️ Production Runbooks & Best Practices

Recommended Learning Rates & Schedules

1. Cosine Decay with Warmup: Always use a short linear warmup ($2-5%$ of total training steps) followed by half-period cosine decay down to $10%$ of peak learning rate.

   Model Family	Muon (matrix_lr)	Muon (lr - AdamW)	Weight Decay
   Small LLMs (< 500M params)	0.02	1e-3	0.01
   Medium LLMs (1B - 7B params)	0.015	6e-4	0.01
   Large LLMs (7B - 70B params)	0.01	3e-4	0.05
   Vision Transformers (ViT)	0.01 (OmniMuon)	3e-4	0.05

2. Distributed Training (DDP / FSDP / DeepSpeed): Muon and OmniMuon operate entirely locally on the parameter gradients computed during backpropagation. They require no global cross-GPU communication during the Newton-Schulz polynomial step. You can use standard DistributedDataParallel or PyTorch FSDP without modifying your distributed training harness.

3. Mixed Precision: Newton-Schulz iterations automatically run in native torch.bfloat16 on Ampere, Hopper, and Blackwell architectures for maximum throughput.

🔬 Scientific Foundations & Prior Art

• Matrix Orthogonalization: Utilizes the degree-5 quintic Newton-Schulz iteration with optimal convergence coefficients $(a=3.4445, b=-4.7750, c=2.0315)$ discovered by Keller Jordan (2024).
• Curvature Optimization: Builds upon finite-difference Hutchinson stochastic Hessian probing formulated by Liu et al. (Stanford, 2023).
• The $1/\eta$ Minibatch Secant Collapse: Proved by AirBorne AI Research, demonstrating why consecutive-batch secant estimations collapse into learning rate artifacts.

📜 Citation

If you use OmniMuon in your academic research or production infrastructure, please cite:

@software{omnimuon2026,
  author = {Singh, Suryaansh Prithvijit and AirBorne AI Research},
  title = {OmniMuon: Universal Spectral Matrix and Second-Order Optimizer Suite for Deep Learning},
  year = {2026},
  publisher = {PyPI and GitHub},
  url = {https://github.com/AirBorneAI/airborne-muon}
}