Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage

Abstract: Shampoo-based methods, such as KL-Shampoo and SOAP, have demonstrated strong performance in training neural networks and rely on QR decomposition. Because existing QR implementations require single-precision (FP32) arithmetic and remain computationally expensive, these methods become time- and memory-intensive when their preconditioning matrices are large. Moreover, using BFloat16 (BFP16) storage to reduce memory usage can degrade the performance of Shampoo-based methods. We propose a reparametrization of the preconditioner that supports BFP16 storage and forms a complete basis by combining updated basis vectors with unchanged ones. By updating only part of the basis through QR decomposition in a subspace, our approach reduces computational overhead while mitigating the performance degradation caused by BFP16 storage. Our approach applies broadly to Shampoo-based methods that employ QR decomposition, including KL-Shampoo, SOAP, and KL-SOAP. In particular, it improves the performance of SOAP and KL-SOAP under BFP16 storage, enabling KL-SOAP to match or exceed KL-Shampoo. Overall, our approach makes Shampoo-based methods more memory- and time-efficient.

• The paper reparametrizes Shampoo optimizers to store products of basis vectors, reducing the computational overhead of QR and eigendecomposition.
• It demonstrates that subspace QR updates with BFloat16 storage preserve performance, achieving near-identical test loss while reducing runtime.
• The method supports flexible full and partial basis updates via greedy subspace selection, ensuring efficient, stable, and faster convergence.

Reparametrization of Shampoo-Based Preconditioners for Subspace Updates and BFloat16 Storage

Motivation and Scope

Shampoo-based optimizers—including Shampoo, SOAP, KL-Shampoo, and KL-SOAP—employ sophisticated matrix preconditioning techniques founded on Kronecker factorization and matrix decompositions to maximize neural network training efficiency. However, two principal bottlenecks persist: the computational burden from QR and eigendecomposition on large matrices, and performance degradation when memory-saving BFloat16 storage is used for optimizer state. Current decomposition algorithms are highly optimized for single-precision, resulting in increased time and memory footprints. Attempts to trade precision for efficiency result in instability and reduced convergence.

This paper proposes a novel reparametrization of Shampoo preconditioners, moving away from storing direct orthogonal bases to maintaining products of old and new basis vectors. This approach enables efficient subspace QR updates, supports half-precision storage, and unifies full-basis and partial-basis updates without compromising numerical stability or optimizer performance. The methodology generalizes across Shampoo-based variants utilizing QR decomposition, broadening its applicability.

Technical Contribution

The central technical advance is the reparametrization of preconditioning factors. Instead of persisting explicit eigenbases, the method stores matrices of the form $P_i := Q_i^\top M_i Q_i$, where $M_i$ is the matrix being preconditioned, and $Q_i$ is the approximated orthogonal basis. QR decomposition is applied directly in selected subspaces of $Q_i$, drastically reducing cost, while the full eigenbasis can be dynamically reassembled. This allows updates to only subsets of basis vectors, mitigating computational load.

With this architecture, decomposition intervals can be shortened (i.e., QR is performed more frequently), narrowing performance gaps typically induced by stale eigenbases in classical implementations while amortizing decomposition overhead.

Figure 1: Top left: KL-Shampoo and Muon performance with BFP16 storage. Top right: QR runtime cost for a $6000 \times 6000$ matrix and proportional subspaces. Bottom: Test loss and parameterization gaps, showing improved performance with reparametrization under BFP16.

Algorithmic Details and Implementation

The method is formalized as follows:

• The optimizer state is stored as $\{Q_i, \Lambda_i, P_i\}$, where $P_i = Q_i^\top M_i Q_i$.
• At each iteration, gradients and covariance factors are updated as in classical KL-Shampoo.
• Preconditioning is performed using updated bases and eigenvalues, unchanged from previous approaches.
• Every $T$ steps, QR decomposition is applied only to a selected block/subspace of $Q_i$. The subspace can be identified through random selection or greedy heuristics based on off-diagonal Frobenius norm.
• The updated subspace is merged with unchanged basis vectors to form a complete orthogonal basis.
• Eigenvalue tracking relies on exponential moving averages, compatible with BFP16 storage.
• The method seamlessly supports switching between full-basis and block-wise updates, and is compatible with both QR and eigendecomposition.

Subspace updates involve only local QR/Eig computation and block-diagonal augmentation, a process mathematically described in the paper as direct sums (see Eq. (lowrank_Q), (lowrank_P)). The update remains orthogonal due to the structure of the operation.

(Figure 2)

Figure 2: Side-by-side comparison of KL-Shampoo original and reparameterized variants, showing differences in optimizer state and update sequence.

Subspace selection can be random or greedy. The latter involves identifying pairs or blocks with maximum off-diagonal couplings and expanding the block, implemented efficiently on GPUs with vectorized operations.

(Figure 3)

Figure 3: KL-Shampoo subspace QR update step using the reparametrized optimizer state.

Empirical Evaluation

Extensive experiments on nanoGPT and several Llama3 variants (119M, 313M, 598M) confirm the practical value. Five core findings are highlighted:

• Preservation of Performance under BFP16: The reparametrized approach achieves nearly identical performance between FP32 and BFP16 storage, while classical parameterizations suffer degradation. This results in state-of-the-art test loss for KL-Shampoo and KL-SOAP (see Table shown in content; KL-Shampoo achieves 3.344 loss for Llama3 119M, delta 0.001 when switching to BFP16).
• Effectiveness of Subspace Update Selection: Greedy subspace selection consistently outperforms random selection, particularly with aggressive learning rates (nanoGPT, Llama3).
• Equivalence of QR and Eigendecomposition in Subspaces: Both decomposition methods converge to similar test losses, with QR providing substantial runtime advantages.
• Block-Jacobi Inner Loops: Multi-step subspace updates incrementally improve convergence, with trade-offs between runtime and accuracy transparent due to color-coded benchmarking (see blocking tables).
• Runtime Reduction: Subspace QR updates reduce wall-clock time, enabling frequent basis updates with minimal performance penalty. Optimal configuration achieves notable runtime savings on large Llama3 models with near-best accuracy.

(Figure 4)

Figure 4: Performance and runtime benchmarks for multi-subspace updates—runtime savings with moderate accuracy trade-off.

Theoretical and Practical Implications

The new parametrization addresses two major theoretical challenges: extending decomposition frequency without increasing computational cost, and supporting half-precision storage without introducing roundoff instability. This unlocks the potential for Shampoo-type optimizers in resource-constrained environments (memory, GPU compute).

Practically, optimizer pipelines can now:

• Perform block QR more frequently and adaptively, accelerating convergence.
• Employ memory-efficient BFP16 storage, crucial for scaling to larger models.
• Use subspace selection heuristics to balance computational cost and precision.
• Integrate parallel block updates (Jacobi-inspired) efficiently on accelerator hardware.
• Reduce decomposition overhead while maintaining up-to-date preconditioner states for fast convergence.

Future directions include hybrid full/subspace update schedules, adaptive decomposition interval selection, and parallelized subspace QR/Eig updates to further reduce runtime and memory usage.

Conclusion

The proposed reparametrization for Shampoo-based optimizers enables both subspace QR decomposition and reliable BFloat16 storage, addressing the dual bottlenecks of time and memory efficiency. Experimental and theoretical analyses confirm the method’s robustness and applicability. This enhances the scalability of Shampoo-style preconditioners and broadens their deployment in modern neural network optimization. The approach aligns with recent block-Jacobi paradigms and presents avenues for further acceleration through adaptive, parallelized updates.