Two-Sided Low-Rank SOAP for Efficient LLM Training
Abstract
Matrix-aware adaptive optimizers such as Shampoo and SOAP exploit row and column structure in neural-network weights, but their optimizer states are expensive for large language models. Recent low-rank optimizers reduce memory by keeping adaptive statistics in a small subspace, yet most existing variants are one-sided: they compress only one matrix axis and therefore approximate only a restricted part of SOAP's two-sided geometry. We propose Adaptive Low-Dimensional Subspace SOAP Method (ALSO), a compact two-sided extension of Alice, an optimizer that estimates adaptive low-dimensional gradient subspaces and compensates discarded directions. For a matrix parameter, ALSO maintains low-dimensional row and column subspaces, applies Adam-like scaling in the resulting core coordinates, and compensates the three residual blocks outside the core. The method can be viewed as a two-sided low-rank approximation to SOAP that keeps full-parameter updates while reducing optimizer-state cost. We evaluate ALSO on C4 pre-training with LLaMA-style models.