Newton-Muon vs Muon vs AdamW

Reproduction of Newton-Muon Optimizer (Du & Su, 2026) benchmarked on Tiny Shakespeare.

Results

Key Finding

Train loss makes Muon look dominant (0.73 vs 1.09). Val loss tells the opposite story — Muon memorises, Newton-Muon generalises. Muon peaks at step 1000 then its val loss climbs for 3000 more steps. Newton-Muon peaks later and overfits far less.

Generated Text — ROMEO:

AdamW: "Far a white tidings of laments — lappiness" — grammar breaks down fast

Muon: "Dignitions messerity" — memorised rhythm, invented words

Newton-Muon: "Titus, Aufidius, Bolingbroke" — real Shakespeare characters, coherent scene

Setup

• Model: TinyGPT (~3M params, 4 layers, 4 heads, 256 dim)
• Dataset: Tiny Shakespeare (~1MB, character-level)
• Tokenizer: Character-level (65 tokens)
• Steps: 4000 each, same random seed
• Hardware: Kaggle free T4 GPU

Run It Yourself

Full runnable notebook on Kaggle

Reference

Du & Su (2026) — arXiv:2604.01472