Bregman 散度如何影响 Shampoo 预条件器

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