Large-scale neural network training increasingly relies on matrix-aware optimizers that exploit the structure of weight parameters beyond element-wise adaptation. However, existing matrix-aware methods such as Muon have an underappreciated vulnerability: their core operation, Newton-Schulz iteration, depends critically on input conditioning, yet the raw momentum matrices exhibit severe coordinate-wise scale heterogeneity. In this paper, we first verify this scale heterogeneity through a chi-square uniformity test, showing that intra-matrix scale imbalance is prevalent across Transformer layers and that coordinate whitening effectively corrects it. Motivated by this finding, we propose Zeta, a dual whitening optimizer that applies coordinate whitening and spectral whitening in a strictly ordered pipeline. The ordering is not a tunable choice but follows from a mathematical dependency: coordinate whitening establishes the statistical isotropy that spectral whitening requires to function reliably. We further prove that this dual pipeline strictly reduces orthogonalization error relative to pure spectral methods by improving the condition number of the input. Empirically, Zeta matches or surpasses strong baselines across language modeling (0.6B to 8B parameters), mixture-of-experts architectures, and vision tasks, demonstrating that resolving scale imbalance before orthogonalization leads to faster convergence and better generalization. Code is available at https://gitcode.com/kevin259/MindSpeed.
翻译:大规模神经网络训练日益依赖于能够利用权重参数结构(超越逐元素自适应)的矩阵感知优化器。然而,现有矩阵感知方法(如Muon)存在一个未被充分认识的风险:其核心操作——Newton-Schulz迭代——关键依赖于输入条件化,而原始动量矩阵呈现出严重的坐标方向尺度异质性。本文首先通过卡方均匀性检验验证了这种尺度异质性,表明Transformer层中存在普遍的矩阵内尺度失衡现象,且坐标白化能够有效纠正该问题。基于这一发现,我们提出Zeta——一种严格按序执行坐标白化与谱白化的双白化优化器。该顺序并非可调参数,而是源于数学依赖性:坐标白化建立了谱白化可靠运行所需的统计各向同性。我们进一步证明,通过改善输入条件数,该双白化管道相较于纯谱方法可严格降低正交化误差。实验表明,在语言建模(0.6B至8B参数)、混合专家架构及视觉任务中,Zeta的性能可匹配甚至超越强基线方法,证明在正交化前解决尺度失衡问题能带来更快的收敛速度与更优的泛化性能。代码开源地址:https://gitcode.com/kevin259/MindSpeed。