Zeta：基于坐标自适应预条件的矩阵优化双白化方法

# Zeta：通过坐标自适应预条件实现矩阵优化的双重白化

来源：https://arxiv.org/html/2606.14187

陈凯文¹，张书海¹∗，陈秋武²∗，刘子墨²，李林晓²，孙颖²，李宇晨²，张一帆²，韩波³，谭明奎¹

¹华南理工大学，²AIGCode，³香港浸会大学

∗同等贡献。邮箱：[email protected]，[email protected]，[email protected]

†通讯作者。邮箱：[email protected]

###### 摘要

大规模神经网络训练日益依赖于能够利用权重参数结构（超越逐元素自适应）的矩阵感知优化器。然而，现有的矩阵感知方法（如 Muon）存在一个被低估的弱点：其核心操作 Newton–Schulz 迭代对输入条件化高度敏感，而原始动量矩阵却表现出严重的坐标维度尺度异质性。在本文中，我们首先通过卡方均匀性检验验证了这种尺度异质性，证明矩阵内部的尺度不均衡在 Transformer 各层中普遍存在，而坐标白化能够有效纠正这一问题。受此启发，我们提出 Zeta，一种双重白化优化器，它以严格有序的流水线方式依次应用坐标白化和谱白化。这种排序并非可调的选项，而是源于数学上的依赖关系：坐标白化建立了谱白化可靠运行所需的统计各向同性。我们进一步证明，与纯谱方法相比，这种双重流水线通过改善输入的条件数，严格降低了正交化误差。在实证方面，Zeta 在语言模型（0.6B 到 8B 参数）、混合专家架构和视觉任务中匹配甚至超越了强基线，表明在正交化之前解决尺度不均衡问题能够带来更快的收敛速度和更好的泛化性能。源代码请访问 https://gitcode.com/kevin259/MindSpeed。

## 1 引言

神经网络训练[1, 2]严重依赖于决定每个参数更新方式的优化算法。例如，自适应优化器[3, 4, 5]及其变体（如 AdamW[6]）按坐标操作，利用每个参数梯度的运行一阶和二阶矩来逐元素地重新缩放更新。这处理了网络不同部分[8]中可能出现的梯度幅值[7]的广泛范围。然而，通过将权重矩阵视为扁平化的向量，这些方法忽略了定义深度神经网络内部表示的空间相关性和谱属性——这些属性是有效学习的内在要素。

最近的几种优化器旨在捕捉标量方法所丢弃的矩阵结构。Muon[9]是一个典型例子。它通过谱白化对权重更新施加正交约束，该过程均衡了动量矩阵的奇异值，防止任何单一方向主导更新。这种设计在大规模训练中已被证明是有效的，最显著的例子是 Kimi-1.1T 模型[10]的训练。通过直接在矩阵上操作而非扁平化的条目，Muon 保留了标量方法无法捕捉的更新几何结构，从而带来更稳定的训练。

尽管有这些优势，Muon 仍存在一个被低估的弱点。其谱操作 Newton–Schulz 迭代[11,12,13]是到 Stiefel 流形[14,15]上的投影，其精度对输入条件化[16,17]很敏感。当奇异值分布良好时，该迭代效果良好。但深度网络训练中的原始动量矩阵通常被少数具有不成比例大值的坐标所主导。当这样的矩阵进入 Newton–Schulz 步骤时，大的条目会扭曲奇异谱，导致正交化效果下降[18]。更新方向不再真实反映梯度几何结构。简而言之，Muon 将一个强大的几何操作应用于一个坐标尺度未受控制的信号。这揭示了一个更深层次的错配，而非调参问题。谱白化需要一个尺度归一化的输入，但原始动量并不能提供这样的保证。缺少的一步是在进入 Newton–Schulz 迭代之前，先使坐标维度的尺度均匀化。

为了验证这种尺度异质性是否真实存在且可纠正，我们对 Transformer 中代表性权重矩阵的行能量应用了卡方均匀性检验[19]。原假设是矩阵的所有分块共享相同的尺度。如图 1 所示，未经任何预处理时的 p 值在各层中始终很小，证实了存在显著的矩阵内部尺度不均衡。经过坐标白化（每个条目由其运行二阶矩归一化）后，p 值明显趋近于 1，表明不均衡已被大幅消除。

参考图注

图 1：Muon 中矩阵内部尺度异质性的说明。我们对使用 Muon 训练的 Qwen3-0.6B 模型的动量矩阵应用了卡方均匀性检验，其中原假设是所有分块共享相同尺度，p 值越接近 1 表示均匀性越强。左图显示了网络浅层、中层和深层的代表性层次，而右图报告了所有层的平均值。白化前的 p 值始终很小，证实了原始动量矩阵存在显著的尺度不均衡。经过坐标白化后，p 值明显趋近于 1，表明这种不均衡可以被纠正。

受此发现启发，我们提出 Zeta，一种双重白化优化器。核心见解很直接：谱白化必须建立在统计稳定的基础之上。Zeta 将此实现为两个操作的顺序流水线。第一个，坐标白化，通过运行二阶矩归一化动量矩阵的每个条目，减少坐标间的尺度差异。第二个，谱白化，对现在条件良好的矩阵应用 Newton–Schulz 正交化，均衡其奇异值以产生平衡的更新方向。这两个操作是单一要求的两个互补部分：坐标白化建立了谱白化所依赖的统计各向同性，因此排序并非自由选择。只有在输入已经摆脱极端尺度扭曲之后，旋转才有意义。我们进一步通过理论分析支持了这种设计。具体来说，我们证明了在标准假设下，坐标白化起到近似对角预条件子的作用（定理 1），使第一阶段具有原则性的解释而非启发式步骤。然后我们证明，这种预条件改善了输入的条件数，从而相对于在原始动量上操作降低了 Newton–Schulz 迭代的正交化误差（定理 2）。最后，我们确认 RMS 缩放保留了更新的谱方向（定理 3），因此该流水线在保持更新几何结构不变的情况下产生了条件良好的更新。

在实证方面，我们在从 0.6B 到 8B 参数的语言模型预训练、混合专家架构以及视觉任务上评估了 Zeta。在这些设置中，Zeta 在收敛速度和下游性能上始终匹配或超越 AdamW[6]、Muon[9] 和 AdaMuon[20]。这些结果表明，在正交化之前解决尺度不均衡问题能够带来更快的收敛和更好的泛化，同时不损失使得谱白化有价值的矩阵级几何结构。我们的贡献总结如下：

- • 我们识别并通过实验验证了像 Muon 这样的谱白化方法隐式要求尺度归一化的输入，而原始动量矩阵系统地违反了这一条件。通过卡方均匀性检验，我们证明了跨 Transformer 层存在显著的矩阵内部尺度异质性，并且坐标白化能够有效纠正它，揭示了此前被忽视的稳定谱正交化的先决条件。
- • 我们提出了 Zeta，一种双重白化优化器，它将坐标白化和谱白化组织成一个严格有序的流水线。这种排序并非可调的选项，而是源于数学上的依赖关系：坐标白化建立了谱白化可靠运行所需的统计各向同性。这种设计将那些被视为竞争优化范式的思想统一为单一要求的互补阶段。
- • 我们证明，通过展示坐标白化阶段改善了 Newton–Schulz 迭代输入的条件数，双重白化流水线相对于纯谱方法严格降低了正交化误差。我们在语言建模（0.6B 到 8B 参数）、混合专家架构和视觉任务上评估了 Zeta，显示其在收敛速度和下游性能上相对于 AdamW、Muon 和 AdaMuon 均有持续提升。

## 2 相关工作

深度学习中的优化算法主要遵循两条轨迹：坐标维度的自适应缩放和利用参数矩阵内部结构相关性的方法。

一阶自适应优化器。一阶方法因其效率在大规模训练中仍处于核心地位。除了经典的随机梯度下降（SGD[21]），自适应方法如 Adam[22] 和 AdamW[6] 是 Transformer 的标准选择，它们使用二阶矩估计来独立缩放坐标。Lion[23] 提供了基于动量符号的内存高效替代方案，而 Sophia[24] 结合了对角 Hessian 估计以捕捉损失曲率。尽管这些方法在坐标维度方差上表现有效，但它们将参数视为扁平化向量，因此忽略了权重矩阵的空间和谱结构。

矩阵结构和二阶优化器。为了捕捉参数间的相关性，另一条研究路线聚焦于矩阵结构化优化。Shampoo[25] 及其变体如 SOAP[26] 使用 Kronecker 因子化预条件子来近似完整二阶矩矩阵的逆，从而有效地对梯度进行跨行和跨列白化。然而，这些方法通常涉及昂贵的矩阵平方根或求逆操作。Muon[9] 通过使用 Newton–Schulz 迭代进行谱白化提供了一种更高效的替代方案，它迫使更新方向为正交矩阵。此后，像 AdaMuon[20] 这样的混合方法尝试将自适应缩放与谱方法结合起来。相比之下，我们的 Zeta 将自己定位在这两条轨迹的交汇处。通过将坐标白化与谱正交化耦合在统一的分层流水线中，Zeta 解决了纯谱方法仍然面临的梯度尺度不均衡问题，同时保留了基于矩阵优化的结构优势。

## 3 提出方法

Muon 优化器[9]通过谱白化对权重更新施加正交约束，这是一种在大规模训练中已被证明非常有效的矩阵感知设计[10]。尽管取得如此成功，Muon 仍存在一个关键的弱点。它所依赖的 Newton–Schulz 迭代[11,12,13]是到 Stiefel 流形[14,15]上的投影，其精度对输入条件化很敏感[16,17]。当动量矩阵被少数具有大值的坐标所主导时，其奇异谱变得歪斜，正交化效果下降[18]，产生扭曲的更新。本质上，Muon 将一个强大的几何操作应用于一个条件不充分的信号，而没有一种机制来预先稳定输入。这不是实现缺陷，而是一种结构性差距：谱白化需要一个条件良好的输入，但原始动量并不能提供这样的保证。

为了解决上述问题，我们提出 Zeta，一种双重白化优化器。我们的核心见解是：谱白化应该建立在统计稳定的基础之上：动量矩阵的坐标维度尺度必须在进入 Newton–Schulz 迭代之前均匀化。这一原则导致了一种设计，其中两个互补的操作顺序应用：坐标白化，它归一化每个条目以减少跨坐标的尺度差异；谱白化，它应用 Newton–Schulz 正交化来均衡更新的奇异值。这两个操作不是竞争性替代方案，而是顺序的必需品。一个在坐标基上操作以建立统计各向同性，另一个在谱基上操作以强制执行几何结构。关键的是，Zeta 不是现有技术的启发式组合，而是基于双重白化原则。我们强调三个关键方面。1) 这两个操作是互补的。坐标白化减少了条目的尺度差异（定理 1）；谱白化纠正了奇异值之间的方向性不均衡。两者单独都无法覆盖全部，合起来则形成了对更新的完整归一化。2) 顺序排序是逻辑上的必然，而非可调选项。谱白化执行旋转，其保真度取决于输入是否尺度归一化。颠倒顺序会将未归一化的矩阵送入 Newton–Schulz 步骤，回到 Zeta 所要消除的不稳定性。这种排序在理论上是有必要的：定理 2 表明坐标白化改进了