🤖 machine learning

Rethinking Bregman Divergences in Kronecker-Factored Optimizers

本文分析了不同的 Bregman 散度如何将 Kronecker 近似误差分布在协方差谱上，揭示了顶端特征空间是可靠的而尾部是噪声化的这一现象，从而启发了一种结合了基于特征值的预处理与自适应各向同性加速的新型子空间感知优化器。

原作者： Bing Liu, Wenjie Zhou, Chengcheng Zhao

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

原作者： Bing Liu, Wenjie Zhou, Chengcheng Zhao

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

核心理念：在迷雾缭绕的山峦中航行

想象你正在试图引导一个巨大的、沉重的机器人（AI 模型）走下山坡，去寻找最低的谷底（最佳性能）。这个机器人太庞大了，无法瞬间转向，因此它需要一张地图来告诉它哪个方向才是“向下”。

在 AI 世界中，这张地图被称为预处理器（Preconditioner）。它告诉机器人每个方向的地势有多陡峭，以便它能采取完美的步长。

问题在于？这座山太大了，要绘制一张包含每一块岩石和碎石的完美、详细的地图是不可能且太慢的。因此，工程师们使用了一种捷径：他们使用“克罗内克因子化（Kronecker-factored）”结构来近似这张地图。你可以把它想象成通过只观察两个独立的二维切片（比如从侧面看一个影子，再从正面看一个影子），并将它们结合起来，从而描述一个复杂的 3D 地形。

问题所在：“影子”并不完美

作者指出这种捷径存在一个根本性的缺陷：影子永远无法完美匹配真实物体。

因为真实的“大山”（数据）非常复杂，你永远无法仅通过两个二维切片就完美地重建它。总会存在“近似误差”——即地图中那些略有偏差的部分。

长期以来，研究人员一直认为这些快捷方法基本上是相同的，只是使用了略微不同的数学公式（称为 Bregman 散度）来衡量“影子”与原物的偏差。这篇论文提出了疑问：如果我们无法获得完美的地图，那么我们使用哪种公式来衡量误差是否真的重要？

发现：不同的公式，不同的盲区

作者发现，答案是肯定的，这非常重要。不同的公式对地图中“错误”的处理方式各不相同：

Frobenius（“大局观”爱好者）： 这个公式主要关注山的宏观特征（巨大的峰峦和深邃的谷底）。它忽略了细小的碎石。这就像一位制图师，只绘制主要的公路，而忽略了小巷。
LogDet（“显微镜”爱好者）： 这个公式痴迷于微小的细节。它对地图中微小的、带有噪声的部分极其敏感。它试图修复每一个微小的碎石，即使那个碎石可能只是一个并不代表真实地形形状的随机石块。
von Neumann（“中间派”）： 它介于两者之间，既关注宏观特征，又不像第一种方法那样激进。

核心洞察： 作者发现，“宏观特征”（数据的上谱/top spectrum）通常与山的真实形状（Hessian）高度契合。然而，“微小的细节”（下谱/bottom spectrum）往往只是噪声——即那些对机器人导航毫无帮助的随机静电信号。

如果你使用“显微镜”（LogDet）方法，你就会试图根据随机噪声来引导机器人，这会导致它左右摇晃且移动缓慢。如果你使用“大局观”方法，你则会专注于那些可靠的部分。

解决方案：面向子空间的优化器

作者没有尝试用一种公式去修复整张地图，而是提出了一种名为 BregTop 的新策略。他们将地图分为两个区域：

“高置信度”区域（上谱）：
- 定义： 山峦中清晰、宏大的特征。
- 策略： 在这里使用精确的、基于特征值的地图。这会根据真实的地面情况，准确地告诉机器人该如何转向以及移动多快。
- 类比： 这就像是在主要公路上使用具有高分辨率卫星图像的 GPS 导航。
“噪声”区域（下谱）：
- 定义： 那些主要是静态噪声的、微小且不可靠的细节。
- 策略： 不要再试图绘制这些细节了！相反，只需给机器人一个稳定的、均匀的推力。不要试图绕过每一颗微小的碎石，只需保持恒定的速度前进即可。
- 类比： 这就像是在一片大雾弥漫的田野中驾驶。与其试图躲避每一个看不见的石头，不如保持直线行驶，并以安全且恒定的速度前进。

实验结果

作者在语言模型训练任务中，将这种新的“分层策略”（BregTop）与旧的方法（如标准的 Shampoo 及其变体）进行了对比测试。

结果： 与其他方法相比，新方法能更快（用更少的步骤）达到目标性能。
原因： 通过信任可靠的“大局观”数据，并忽略掉那些带有噪声的“微小细节”，机器人不会因为左右摇晃而浪费能量。它在下山的过程中移动得更加高效。

总结

本文认为，当我们无法完美绘制复杂的 AI 数据地图时，我们不应该试图修复每一个微小的错误。相反，我们应该：

信任清晰、宏大的模式。
忽略那些嘈杂、不可靠的细节，通过简单且统一的方式对待它们。

这种“知晓信任什么，以及忽略什么”的方法，能够带来更快速、更高效的 AI 训练。

技术摘要：重新思考 Kronecker 分解优化器中的 Bregman 散度

问题陈述
Shampoo 类优化器使用 Kronecker 分解结构 $L \otimes R$ 来近似梯度协方差矩阵 $C = E[\text{vec}(G)\text{vec}(G)^\top]$ ，以降低计算复杂度。近期的研究（Lin et al., 2026）通过将各种 Shampoo 变体框架化为 $C$ 与 $L \otimes R$ 之间的 Bregman 矩阵散度最小化问题，实现了对它们的统一。然而，一个关键的空白仍然存在：对于通用的稠密协方差矩阵，精确的 Kronecker 分解是不可能的，这意味着近似误差是不可避免的。目前尚不清楚不同的 Bregman 散度如何将这种误差分布到协方差的光谱中，以及由此产生的预处理器是否对光谱的所有部分都同样信任。此外，由于梯度协方差通常与真实的 Hessian 矩阵显著不同，并非所有的谱分量都适合进行预处理。核心问题在于：当精确的 Kronecker 匹配无法实现时，优化器应该近似并信任协方差光谱中的哪些部分？

方法论与分析
作者通过两个主要的分析视角来研究这个问题：Bregman 近似的谱几何学以及协方差特征空间的经验可靠性。

Bregman 近似的谱几何学：
本文表征了 Frobenius、von Neumann 和 LogDet 散度如何以不同的方式惩罚近似误差。通过将散度分解为边缘谱项和交叉谱耦合项 $g_F(\lambda_i, \omega_j)$ ，作者表明：
- Frobenius 散度： 耦合项 $-\lambda_i \omega_j$ 施加了一个双线性量级惩罚，强烈强调顶端（大特征值）谱。
- von Neumann 散度： 项 $-\lambda_i \log \omega_j$ 保留了对顶端谱的偏好，但对其进行了对数级的平滑处理。
- LogDet 散度： 项 $\lambda_i / \omega_j$ 产生了相对比例失真，使得近似对小特征值（尾部谱）高度敏感。
定理 2 建立在这些最小化问题的驻点并非由原始残差 $(S^* - C)$ 决定，而是由散度加权残差 $\nabla^2 F(S^*)[S^* - C]$ 决定。因此，生成的 Kronecker 因子与所选散度的特定谱偏好相一致（例如，LogDet 因子对小特征值更敏感，而 Frobenius 因子则侧重于主子空间）。
协方差预处理的谱可靠性：
作者通过在 SST-2 数据集上训练的 Transformer 模型上，对梯度协方差与 Hessian 特征空间之间的对齐情况进行了经验性研究。结果（表 3）表明，顶端协方差特征空间与 Hessian 的对齐程度明显优于低特征值分量。观察发现，尾部谱在预处理方面更加嘈杂且不够可靠，这与先前的研究（Kunstner et al., 2019; Thomas et al., 2020）一致。
提出的解决方案：子空间感知 Kronecker 优化器 (BregTop)：
受散度诱导的谱偏好与实际谱可靠性之间不匹配的启发，作者提出了 BregTop。该优化器结合了两种观察结果的优势：
- 它使用对顶端敏感的散度（如 von Neumann 或 Frobenius）来选择 Kronecker 因子。
- 它仅在顶端子空间（联合 Kronecker 光谱中最大的 $\lceil \rho mn \rceil$ 个条目）内应用基于特征值的预处理。
- 它将互补的底端子空间视为自适应各向同性加速常数（ $\chi_{\rho, q}$ ）的处理对象，而不是试图用可能具有误导性的曲率信息去预处理嘈杂的尾部。
从理论上讲，这将被预处理器分解为一个曲率缩放的顶端组件和一个用于底端空间的均匀加速组件（定理 4）。

核心贡献

谱几何表征： 本文提供了关于 Frobenius、von Neumann 和 LogDet 散度如何诱导不同的依赖于特征值的惩罚的正式表征，表明驻 Kronecker 因子是由散度加权残差而非原始误差所控制。
经验谱可靠性： 研究证明，虽然全谱协方差在预处理方面并不可靠，但顶端特征空间与 Hessian 保持着强对齐，从而证明了子空间感知方法的合理性。
子空间感知优化器设计： 作者引入了 BregTop，这是一种能够根据特征值在可靠的顶端子空间中进行缩放，并在嘈杂的底端子空间中进行各向同性加速的自适应切换算法。

实验结果
作者在训练全量 SST-2 数据集的 4 层 Transformer 上，将 BregTop 与标准 Shampoo、KL-Shampoo、VN-Shampoo 和 F-Shampoo 进行了对比评估。使用的指标是 EMA 平滑后的训练损失降至 0.05 所需的步数。

BregTop-VN-v2 在 1003 步内达到了目标损失，优于标准 Shampoo（1119 步）和基准 VN-Shampoo（1127 步）。
BregTop 变体一致地改进了其对应的散度基准（例如，BregTop-VN 优于 VN-Shampoo；BregTop-F 优于 F-Shampoo），验证了隔离顶端子空间进行预处理的益处。

意义与主张
本文主张，结构化优化器的设计应当考虑 Kronecker 分解中不可避免的近似误差以及协方差谱的变化可靠性。通过认识到不同的 Bregman 散度本质上会优先考虑不同的谱区域，并且只有顶端谱与 Hessian 具有可靠的对齐，所提出的子空间感知方法提供了一种更鲁棒的预处理策略。作者指出，该设计与 Song et al. (2025) 中观察到的训练动态是一致的，即学习发生在底端子空间；因此，在那里应用各向同性加速可以在保留必要的梯度流的同时，通过顶端子空间预处理来抑制高曲率方向的嘈杂更新。结果表明，这种混合方法在不需要完整二阶信息的情况下提高了步骤效率。