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。

0
下载
关闭预览

相关内容

【论文笔记】自注意力图池化
专知
82+阅读 · 2019年11月18日
从动力学角度看优化算法:GAN的第三个阶段
PaperWeekly
11+阅读 · 2019年5月13日
最新|深度离散哈希算法,可用于图像检索!
全球人工智能
14+阅读 · 2017年12月15日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
9+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2014年12月31日
国家自然科学基金
8+阅读 · 2014年12月31日
VIP会员
最新内容
《边缘计算关键技术分析及美军作战实践应用》
专知会员服务
0+阅读 · 今天14:08
边缘计算的军事应用
专知会员服务
1+阅读 · 今天13:50
一种考虑资源机动性的武器目标分配混合算法
专知会员服务
4+阅读 · 8月8日
《多域冲突比较支持模型》60页
专知会员服务
10+阅读 · 8月7日
相关VIP内容
相关基金
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
9+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2014年12月31日
国家自然科学基金
8+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员