Search papers, labs, and topics across Lattice.
This paper unifies popular divergences used in Shampoo-like optimizers under a Bregman divergence framework to analyze their impact on preconditioning. By empirically analyzing the spectral properties of gradient second moments and their interaction with finite-sample error, the authors demonstrate how specific divergences can better mitigate underestimation of the empirical second moment. These findings are validated through GPT-2 pretraining, offering principled guidance for future optimizer development.
Divergence choice in Shampoo-like optimizers fundamentally alters Kronecker approximation and finite-sample error compensation, explaining observed training behavior.
Understanding the principles behind Shampoo has recently guided the development of more effective neural network optimizers. These methods learn a preconditioner by optimizing the Frobenius or Kullback-Leibler (KL) divergence against the gradient second moment. In this work, we investigate how the choice of divergence shapes preconditioning, which remains unclear and blocks further improvements. To do so, we develop a unified Bregman divergence framework that connects all popular divergences, allowing us to study them jointly. Through empirical spectral analysis of gradient second moments, we examine how divergence choice shapes Kronecker approximation and interacts with finite-sample error in preconditioning. We find that some divergences can better compensate for finite-sample underestimation of the empirical second moment, helping explain the differing behavior of their corresponding Shampoo variants. We further validate this explanation through GPT-2 pretraining experiments. By connecting divergence choice to practical training behavior, we believe our framework provides principled guidance for understanding the foundations of, and further improving, Shampoo.
