Bregman散度如何塑造Shampoo优化器

中文摘要 AI 辅助

理解Shampoo背后的原理近期指导了更有效的神经网络优化器的发展。这些方法通过优化Frobenius或Kullback-Leibler（KL）散度相对于梯度二阶矩来学习预条件矩阵。在这项工作中，我们研究了散度的选择如何影响预条件化，这一问题目前尚不清楚，并阻碍了进一步的改进。为此，我们开发了一个统一的Bregman散度框架，连接了所有流行的散度，使我们能够共同研究它们。通过对梯度二阶矩的经验谱分析，我们考察了散度选择如何塑造Kronecker近似，并如何与预条件化中的有限样本误差相互作用。我们发现，某些散度能更好地补偿经验二阶矩的有限样本低估，这有助于解释其对应Shampoo变体的不同行为。我们通过GPT-2预训练实验进一步验证了这一解释。通过将散度选择与实际训练行为联系起来，我们相信我们的框架为理解Shampoo的基础并进一步改进提供了原则性指导。

英文摘要

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.

发表机构

• Zhejiang University（浙江大学）
• University of the Chinese Academy of Sciences（中国科学院大学）
• Peking University（北京大学）
• Concordia University（康考迪亚大学）
• Mila（米拉研究所）
• University of Central Florida（中佛罗里达大学）

机构由 AI 辅助整理，请以论文原文为准。