解决神经网络优化中权重矩阵尺度不均问题,提升训练速度与泛化性能。
Zeta: Dual Whitening for Matrix Optimization via Coordinate-Adaptive Preconditioning
- 通过坐标白化与谱白化两级处理,先统一各维度尺度再进行正交化。
- 在0.6B到8B参数模型上,收敛速度更快,精度更高。
- 适合大规模语言模型和视觉任务的高效训练,代码已开源。
大规模神经网络训练依赖于能捕捉权重结构的矩阵感知优化器,但现有方法如Muon存在隐患:其核心的牛顿-舒尔茨迭代对输入条件敏感,而原始动量矩阵存在严重的坐标尺度异质性。本文通过卡方均匀性检验验证了这一现象——Transformer层内尺度失衡普遍存在,坐标白化可有效纠正。受此启发,我们提出Zeta,一种严格按序执行坐标白化与谱白化的双重白化优化器。该顺序由数学依赖决定:坐标白化建立统计各向同性,使谱白化可靠运行。理论证明,该双阶段流程相比纯谱方法显著降低正交化误差,因改善了输入条件数。实验表明,Zeta在0.6B至8B参数的语言建模、混合专家架构及视觉任务中表现匹敌或超越强基线,证实先消除尺度不均衡再正交化可加速收敛并提升泛化能力。代码见https://github.com/AIGCodeOS/aigcode_zeta_optimizer。
原文摘要 · Abstract (English)
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://github.com/AIGCodeOS/aigcode_zeta_optimizer.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。