Zeta: Dual Whitening for Matrix Optimization via Coordinate-Adaptive Preconditioning

2026-06-12

AI Summary

This work addresses the instability of existing matrix-sensing optimizers—such as Muon—whose momentum matrices exhibit pronounced coordinate-wise scale heterogeneity, rendering Newton–Schulz iterations highly sensitive to input conditioning. To resolve this, we propose Zeta, an optimizer that introduces a novel dual whitening mechanism with strict sequential dependency: it first applies coordinate whitening to rectify internal scale imbalances and then performs spectral whitening to enforce statistical isotropy, thereby yielding well-conditioned inputs for subsequent orthogonalization. This approach substantially reduces orthogonalization error and consistently outperforms strong baselines across diverse settings—including language models ranging from 0.6B to 8B parameters, mixture-of-experts architectures, and vision tasks—achieving faster convergence and improved generalization.

Abstract

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.

Problem

• matrix-aware optimization
• scale heterogeneity
• coordinate whitening
• conditioning
• momentum matrices

Innovation

• dual whitening
• coordinate-adaptive preconditioning
• matrix optimization
• scale heterogeneity
• spectral whitening