Tauon: A new optimizer outperforming Muon on GPT-Mini (lower loss, ~8.5% faster step time) [P]

Reddit r/MachineLearning Tools
optimizer machine-learning benchmark pytorch open-source

Summary

Tauon is a new optimizer using polynomial and orthogonalization techniques that outperforms Muon and AdamW in initial benchmarks on a small GPT-Mini model, showing lower loss and faster step time.

Hey r/MachineLearning! I’ve been working on a new optimizer called Tauon (turns out there is already "teon" but well if you have better idea, - i will gladly accept it! Anyway the core idea of optimizer is about polynomials and orthogonalization just like muon, the whole difference is that i managed to lower total number of steps (first through spectral filtering down to 3 steps then through coeff scheduling down to 2) + reduced matrix size (through dct-2). And I wanted to share some initial benchmark results... Benchmark Setup: Trained a GPT-Mini (d_model=512, 6 Layers) on TinyShakespeare against Muon and AdamW. Tauon: LR = 0.02 Muon: LR = 0.02 AdamW: LR = 0.0006 Results: Validation Loss: Tauon converged to a lower final loss (~1.6) compared to Muon (~1.65) and AdamW (~1.8). Stability: AdamW started overfitting/diverging around step 1200, whereas Tauon maintained stable progress throughout the 3000 steps. Compute Cost: On identical hardware, Tauon ran at 391.5 ms/step vs Muon’s 427.7 ms/step (~8.5% faster) and close to AdamW's baseline of 382.9 ms/step. And yeah i know that its hilariously tiny benchmark but well i have only 2 hours left on my kaggle free T4 so i really couldnt more + i hope someone would be able test it on a bigger setup! Links & Code: 📂 GitHub: erj2231/ai-projects/tree/main/tauon 📦 PyPI: pip install tauon-optimizer Would love to get your thoughts on the optimizer! If you have any ideas, suggestions - please tell me. Cheers, everyone! https://preview.redd.it/8pbtu9azgzrh1.png?width=5370&format=png&auto=webp&s=b2c694fd05a27919a41103e7f7d960eef41b2407
Original Article

Similar Articles

How Much Orthogonalization Does Muon Need?

arXiv cs.LG

This paper studies how much orthogonalization the Muon optimizer requires, proposing a five-step cubic Newton-Schulz schedule that reduces computational cost while achieving training quality similar to more expensive methods across GPT-2 Small and hybrid MoE/Mamba models.

Can Muon Fine-tune Adam-Pretrained Models?

Hugging Face Daily Papers

Research paper investigating performance degradation when using the Muon optimizer instead of Adam for fine-tuning pretrained models, demonstrating that parameter-efficient methods like LoRA effectively mitigate this optimizer mismatch across language and vision tasks.

Reassessing Muon for Matrix Factorization

arXiv cs.LG

This paper evaluates the Muon optimizer on low-rank matrix factorization, finding it does not consistently outperform AdamW, challenging earlier claims about its advantages in large-scale deep learning.

Dion3: Full-Stack Orthogonal Updates

arXiv cs.LG

Dion3 is a revised Muon optimizer that reduces computational and communication overhead via Gram Newton-Schulz, symmetric GEMM kernels, and megabatching, achieving up to 6x faster optimizer steps while matching or improving loss.