Reparametrizing Shampoo and SOAP

Shampoo-based methods, such as KL-Shampoo and SOAP, have demonstrated strong performance in neural network training, but their QR-based updates become costly for large preconditioning matrices. To address this, we propose a preconditioner reparametrization that accelerates full-basis QR computation and enables faster partial basis updates. By maintaining preconditioning factors in the current orthogonal basis, our approach exposes useful input structure to facilitate QR decomposition and allows dynamically selected basis vectors to be cheaply updated via subspace QR while preserving the full-basis orthogonality. To further accelerate QR computation, we provide a customized implementation supporting mixed-precision computation and batched execution. The resulting framework unifies full-basis and subspace updates and applies broadly to QR-based updates, including KL-Shampoo, SOAP, and KL-SOAP. Empirically, our approach reduces computational overhead and improves the training performance of SOAP-type methods in mixed precision, enabling KL-SOAP to match or exceed KL-Shampoo. Overall, our approach improves the time efficiency of Shampoo-based methods.