Optimizer code to reproduce results on Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization (ICLR 2026).
We provide a prototype implementation of the KL-Shampoo/KL-SOAP.
We are developing a faster and more robust version of KL-Shampoo/KL-SOAP (at least 1.5-2 times faster than the current version). Stay tuned for new updates.