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. Stay tuned for new updates.