🤖 machine learning

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

This paper introduces Zeta, a dual whitening optimizer that addresses the critical vulnerability of existing matrix-aware methods by applying coordinate whitening before spectral whitening to correct scale heterogeneity, thereby improving convergence and generalization across large-scale language and vision models.

Original authors: Kaiwen Chen, Shuhai Zhang, Qiuwu Chen, Zimo Liu, Linxiao Li, Ying Sun, Yuchen Li, Yifan Zhang, Bo Han, Mingkui Tan

Published 2026-06-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kaiwen Chen, Shuhai Zhang, Qiuwu Chen, Zimo Liu, Linxiao Li, Ying Sun, Yuchen Li, Yifan Zhang, Bo Han, Mingkui Tan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a massive, complex robot (a neural network) to learn a new skill, like writing poetry or recognizing cats. To do this, you need a "coach" (an optimizer) that tells the robot how to adjust its brain cells (parameters) after every mistake.

For a long time, the best coaches treated every single brain cell as an individual. They looked at each cell, saw how much it messed up, and gave it a tiny nudge. This worked okay, but it ignored the fact that brain cells work in teams.

Recently, a new type of coach called Muon arrived. Instead of looking at cells one by one, Muon looked at whole teams (matrices) at once. It tried to make sure the team moved in a perfectly balanced, coordinated way, like a synchronized swimming team. This worked great for huge robots, but the paper authors found a hidden flaw in Muon's logic.

The Problem: The "Loudmouth" Teammates

The authors discovered that in Muon's synchronized swimming team, some teammates were shouting so loudly that they drowned out everyone else.

In technical terms, the "momentum" (the memory of past mistakes) for some parts of the matrix was huge, while others were tiny. When Muon tried to perform its fancy "synchronized dance" (called Newton–Schulz iteration), the loudmouths threw off the rhythm. The team couldn't coordinate properly because the input was unbalanced. It's like trying to conduct an orchestra where the trumpets are blasting at full volume while the violins are whispering; the music sounds terrible, no matter how good the conductor is.

The authors proved this by running a "uniformity test" (a statistical check). They found that before any fixing, the "volume levels" across the matrix were all over the place.

The Solution: Zeta (The Two-Step Coach)

The authors propose a new coach called Zeta. Zeta realizes that you can't fix the orchestra's rhythm until you first get the volume levels right. So, Zeta uses a strict two-step process:

Step 1: The Volume Knob (Coordinate Whitening)
Before the team tries to dance, Zeta walks around and turns down the volume on the shouting trumpets and turns up the volume on the whispering violins. It normalizes the "loudness" of every single entry in the matrix. Now, everyone is on an equal footing. The team is no longer dominated by a few loud voices.
Step 2: The Synchronized Dance (Spectral Whitening)
Now that the volume is balanced, Zeta lets the team perform its fancy synchronized dance. Because the input is now calm and balanced, the dance works perfectly. The team moves in a beautiful, coordinated geometric pattern that Muon tried to do but failed at because the volume was unbalanced.

Why the order matters: The paper emphasizes that you cannot swap these steps. If you try to dance first (Step 2) and then fix the volume (Step 1), the dance will still be ruined by the initial shouting. The volume fixing must happen first to create the stable foundation the dance needs.

The Results: Faster and Smarter

The authors tested Zeta on various "robots" (AI models) ranging from small ones (0.6 billion parameters) to massive ones (8 billion parameters), and even on models that use a "mixture of experts" (where different parts of the brain handle different tasks).

Faster Learning: Zeta learned faster than the old coaches (AdamW and Muon). It reached the same level of skill in fewer training steps.
Better Generalization: The models trained with Zeta were better at tasks they hadn't seen before, like answering tricky questions or recognizing images.
Efficiency: Zeta didn't slow things down. It learned faster without taking much extra time per step.

The Bottom Line

Think of Zeta as a coach who realized that before you can teach a team to move in perfect harmony, you first have to make sure everyone is speaking at the same volume. By fixing the "volume imbalance" first, the team can finally perform the "harmony dance" correctly, leading to a smarter, faster-learning AI.

The paper claims this is a fundamental fix to a structural problem in how these advanced optimizers work, proving that balancing the scales before trying to optimize the geometry is the key to success.

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

Problem Statement

Large-scale neural network training has increasingly shifted toward matrix-aware optimizers that exploit the structural properties of weight parameters, moving beyond element-wise adaptation. A leading example is Muon, which enforces orthogonal constraints on weight updates via spectral whitening (specifically, Newton–Schulz iteration) to equalize singular values and prevent any single direction from dominating the update.

However, this paper identifies a critical, underappreciated vulnerability in Muon and similar spectral methods: the Newton–Schulz iteration is highly sensitive to input conditioning. While the iteration functions well when singular values are well-distributed, raw momentum matrices in deep network training often exhibit severe coordinate-wise scale heterogeneity. In these matrices, a few coordinates possess disproportionately large magnitudes, skewing the singular spectrum. When such unconditioned matrices enter the Newton–Schulz step, the orthogonalization process degrades, leading to distorted update directions that no longer faithfully reflect the gradient geometry. The authors argue that spectral whitening implicitly requires a scale-normalized input, a condition that raw momentum systematically violates.

Methodology: Zeta Optimizer

To address this structural gap, the authors propose Zeta, a dual whitening optimizer that organizes coordinate whitening and spectral whitening into a strictly ordered, non-tunable pipeline. The core insight is that spectral whitening must be built upon statistical stabilization; coordinate whitening establishes the statistical isotropy required for spectral whitening to function reliably.

The Zeta pipeline operates as follows:

Coordinate Whitening (Preconditioning):
Before entering the spectral stage, Zeta applies a coordinate-wise normalization to the momentum matrix. It maintains exponential moving averages of the first moment ( $M_t$ ) and the second moment ( $V_t$ ), similar to Adam. The momentum matrix is then normalized element-wise by the running second moment:
$\tilde{G}_t = \frac{M_t}{\sqrt{V_t} + \epsilon}$
This step homogenizes the scale disparity across coordinates, transforming the matrix into a statistically isotropic state. Theoretically, this acts as an approximate diagonal preconditioner, addressing the scale imbalance that would otherwise skew the singular spectrum.
Spectral Whitening (Orthogonalization):
The homogenized matrix $\tilde{G}_t$ is then fed into the Newton–Schulz iteration. Because the input is now free of extreme entry-wise scale disparities, its singular spectrum is better conditioned. The iteration converges to a faithful orthogonal approximation, equalizing the singular values to produce a balanced update direction.
Update Construction:
The resulting orthogonal matrix $U_t$ is scaled using an RMS-normalized factor to ensure the update magnitude is compatible with standard optimizer configurations (e.g., AdamW branches) without altering the spectral direction.
Dual-Path Execution:
For matrix-valued parameters (e.g., attention and feed-forward projections), the full dual pipeline is applied. For non-matrix parameters (e.g., biases, LayerNorm scales), the spectral step is skipped, and the optimizer falls back to an AdamW-style update.

Theoretical Contributions

The paper provides three key theoretical guarantees supporting the design:

Approximate Whitening: The coordinate normalization step is proven to be an efficient diagonal approximation to the ideal whitening transform (Theorem 1), moving beyond heuristic scaling to a principled preconditioning step.
Reduced Orthogonalization Error: The authors prove that coordinate whitening strictly improves the condition number of the input to the Newton–Schulz iteration. Consequently, the orthogonalization error of the spectral stage is strictly reduced compared to operating on raw momentum (Theorem 2).
Scale Preservation: The final RMS scaling preserves the spectral direction of the update while synchronizing its global magnitude with the rest of the network (Theorem 3).

Empirical Results

Zeta was evaluated across language modeling (0.6B to 8B parameters), mixture-of-experts (MoE) architectures, and vision tasks, consistently matching or surpassing strong baselines (AdamW, Muon, and AdaMuon).

Convergence Speed: In pretraining tasks, Zeta demonstrated faster convergence than Muon and AdamW. For instance, on a 1.7B parameter model, Zeta achieved a 1.67× speedup (relative to AdamW) to reach the same loss level, outperforming Muon (1.52×) and AdaMuon (1.61×). On an 8B model, Zeta achieved a 1.25× speedup versus Muon's 1.19×.
Efficiency: The per-iteration time cost of Zeta remained comparable to AdamW (e.g., 25.47s vs. 24.46s for the 8B model on 16 GPUs), indicating that the dual whitening process does not incur significant runtime penalties.
MoE Performance: In MoE architectures, where sparse expert activation creates substantial gradient heterogeneity, Zeta maintained a clear advantage, converging to lower loss with less fluctuation than competitors.
Downstream Tasks: On 12 language understanding benchmarks (e.g., MMLU, GSM8k), Zeta-trained models achieved the highest average accuracy across both 1.7B and 8B scales. Similarly, on CIFAR-100 vision tasks, Zeta achieved the highest Top-1 validation accuracy for both ViT-Tiny and ViT-Base.
Hyperparameter Robustness: A grid search on momentum coefficients ( $\beta_1, \beta_2$ ) showed that Zeta is robust to hyperparameter choices, with final training loss varying by less than 1.2% across the tested range.

Significance and Claims

The paper claims that Zeta resolves a fundamental mismatch in existing matrix-aware optimizers: the application of powerful geometric operations (spectral whitening) to signals with uncontrolled coordinate-wise scales. By verifying the prevalence of intra-matrix scale imbalance through chi-square uniformity tests and demonstrating that coordinate whitening effectively corrects it, the authors establish that resolving scale imbalance before orthogonalization is a prerequisite for stable spectral methods.

The significance of Zeta lies in its unification of two previously competing optimization paradigms (coordinate-wise adaptive scaling and spectral orthogonalization) into a single, logically necessary pipeline. The ordering of operations is not a tunable hyperparameter but a mathematical dependency: coordinate whitening must precede spectral whitening to ensure the input is well-conditioned. This approach leads to faster convergence and better generalization without sacrificing the geometric benefits of matrix-level optimization.