Pro-KLShampoo：通过正交化恢复白化的投影 KL-Shampoo

利用梯度矩阵结构的优化器是现代 LLM 预训练的核心，目前有两个不同的前沿方向：显式 Kronecker 分解预处理——最新的是 KL-Shampoo，它通过 KL 散度最小化来估计预条件器——以及梯度动量的正交化，以 Muon 为代表，并被分析为谱范数下的最速下降。这两条路线通常是孤立发展的。我们对 KL-Shampoo 的 Kronecker 预条件器提出了一个结构性观察：它们的特征值谱呈现出“尖峰-平坦”形状——少数主导特征值后接近似均匀的尾部——这在各层和训练阶段均成立，并在秩为 $\rho$ 的信号加噪声梯度模型下精确成立。我们利用这种结构，将 KL-Shampoo 的一个 Kronecker 因子限制为与“尖峰-平坦”形状对齐的参数化族：在被跟踪的 $r$ 维子空间上保留完整谱结构，在剩余的 $n-r$ 个方向上使用单一共享特征值。在这些方向上，我们应用正交化。一个恒等式表明，这种正交化恢复了完整 KL-Shampoo 预条件器的代数形式。在四种预训练规模（GPT-2 124M / 350M，LLaMA 134M / 450M）上，Pro-KLShampoo 在我们测试的每一个子空间秩上，在验证损失、每 GPU 峰值内存以及达到每个损失水平的实际时间方面，始终优于 KL-Shampoo。