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! preview.redd.it/8pbtu9azgzrh1.png