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

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