From Geometry to Generalization: Why Row Normalization Can Beat Adam and Muon

Detailed Research Summary: Geometry to Generalization: Why Row Normalization Can Beat Adam and Muon

This research paper presents a rigorous theoretical analysis demonstrating that row-wise normalization methods achieve strictly higher population accuracy in high-dimensional multiclass classification compared to several established optimizers, including full-batch Adam, exact-SVD Muon, and various spectral/coordinate-wise alternatives. The core mechanism hinges on how these different optimization geometries preserve or distort the underlying population decision boundaries.

Core Thesis and Primary Findings

The central argument is that row normalization's classwise Euclidean geometry asymptotically preserves the population decision-boundary directions, whereas the coordinate-wise geometry of Adam and the spectral geometry of Muon introduce systematic, nonvanishing distortions. This geometric preservation leads to superior generalization performance, quantified by higher population accuracy.

Key Comparative Results (Deterministic Setting):

The analysis establishes a strict population-accuracy ordering in the deterministic setting:

RowGD and RowGDM > SpecGD, Muon, SignGD, Signum, and Adam

The mechanism for this superiority is that row-normalized methods effectively preserve each class-mean direction, while the other methods introduce directional distortions.

Robustness and Extension Beyond Isotropy (Anisotropy)

The advantage of row normalization is not confined to isotropic data distributions. The authors extend the findings to more complex scenarios:

1. Deterministic Training on Class Means: The advantage persists even when training deterministically on class means, provided the class-mean and test-noise covariances are independent and possess random eigenbases, or when dealing with power-law spectra where the class-mean exponent is below one.

2. Anisotropic Test Noise: The advantage remains robust under heavily anisotropic test noise conditions (Theorem 4.2).

3. Coordinate Reversal and Restoration: A specific reversal of the ordering between RowGDM and Adam can occur when the class-mean and test-noise covariances are diagonal and sufficiently close, aligning with Adam's normalization coordinates. Crucially, this coordinate-specific reversal is restored by jointly applying a common Haar rotation to both the class means and the test-noise covariance.

Stochastic Training Robustness

The analysis also investigates robustness under stochastic training scenarios involving random reshuffling (Proposition 3.4). For sufficiently large momentum, methods like RR-RowSGDM, RR-Muon, and RR-Signum retain enough of their deterministic geometric separation to maintain the strict population-accuracy ordering. Furthermore, the analysis shows that RowGDM achieves strictly higher population accuracy than AdamProxy under certain conditions (Proposition 3.6).

Mathematical Formalism and Convergence Analysis

The paper utilizes sophisticated implicit-bias results derived from full-batch training trajectories to establish long-term convergence properties:

• Direction Set Convergence: The deterministic full-batch implicit-bias results show direction set convergence to UX row for RowGD/RowGDM, UX sp for SpecGD/Muon, and UX for SignGD/Signum/Adam.

• Error Exponent Gaps: The analysis introduces the Row–Spectral exponent-gap constant (alpha), which quantifies the separation between the error exponents of RowGDM and Muon. This gap is shown to be positive, leading to a quantifiable advantage:

d to infinity r squared t e d E pi,d(W Adam t) - (pi - 2/4 pi) = o(P(1))

(alpha) - o(P(1)) at most d to infinity r squared t e d E pi,d(W Muon t)

This demonstrates that RowGDM maintains a strictly smaller error exponent than Muon.

• Muon vs. Adam Comparison (Corollary C.28): For 0 < alpha < alpha (where alpha about 1.225), the long-term deterministic trajectories show that exact-SVD Muon has asymptotically smaller population error than full-batch Adam:

t to infinity r squared t e d E pi,d(W Muon t) at most-cMA(alpha) + o(P(1))

where cMA(alpha) > 0.

Practical Implications and Caveats

The paper provides a crucial practical caveat in Remark C.29: the theoretical comparison is idealized. The results are based on deterministic full-batch training, exact-SVD Muon, and Adam without numerical stability constants. In real-world last-layer training practice, Adam is often preferred over Muon. The mismatch observed in the theory suggests that the current analysis compares only asymptotic convergence directions and does not fully capture finite-time trajectories or convergence speed.

Conclusion

The paper provides a strong theoretical foundation for row-normalized optimizers (RowGD, RowGDM) as superior alternatives to spectral (SpecGD), coordinate-wise (Adam), and exact-SVD Muon in high-dimensional classification tasks. This superiority stems from the geometric property of row normalization which better preserves the essential class-mean directions of the data manifold, leading to demonstrably lower population error exponents under a wide range of data distribution assumptions, including anisotropy.

1. Always use row-wise normalization as a base optimizer for large language model training, as it can achieve strictly higher population accuracy than full-batch Adam in high-dimensional multiclass classification by asymptotically preserv[ing] the population decision-boundary directions.

2. Implement row-normalized gradient descent (RowGD) or row-normalized gradient descent with momentum (RowGDM) for deterministic training, as Theorem 3.1 establishes the deterministic optimizer population-accuracy ordering over competitors like Adam and Muon in Gaussian data models.

3. For stochastic training scenarios, use RR-RowSGDM when momentum is sufficiently large to ensure the strict population-accuracy ordering to persist, as shown in Proposition 3.4 regarding stochastic optimizer accuracy ordering.

4. When training on anisotropic or complex data (where class means and test-noise covariances are not isotropic), apply the row-normalized methods, as Theorem 4.2 shows that the advantage of row-normalized methods persists for arbitrary test-noise anisotropy.

5. Design a system that dynamically adjust its optimization strategy based on the covariance structure; specifically, if the high-variance class-mean and test-noise directions align with each other and with Adam’s normalization coordinates, switch to RowGDM or RowGD to maintain accuracy, as described in Proposition 4.5 concerning Coordinate dependence of the RowGDM–Adam ordering.

6. For models where the number of classes is large relative to the input dimension (i.e., when k/d → α ∈ (0, ∞)), use row-normalized methods to ensure generalization by leveraging the result that they maintain a separation between margin geometries even when class frequencies follow Zipf’s law.

7. In high-dimensional settings where the test noise is highly anisotropic, use RowGDM over Adam to achieve exponentially smaller population error as quantified in Corollary 3.3, demonstrating superior generalization performance on unseen data.

Different optimizers can fit the same training data while selecting classifiers with substantially different geometries, but whether this difference provably affects population performance remains unclear. We show that row-wise normalization can achieve strictly higher population accuracy than full-batch Adam, a proxy for random-reshuffling Adam, and exact-SVD Muon in high-dimensional multiclass classification. Under an isotropic Gaussian-cloud data model, this advantage arises because row normalization's class-wise Euclidean geometry asymptotically preserves the population decision-boundary directions, whereas Adam's coordinate-wise geometry and Muon's spectral geometry introduce nonvanishing distortions. Beyond isotropy, the advantage persists for full-batch training on class means with independently oriented class-mean and test-noise covariances. It holds for power-law spectra with class-mean exponent below one, even under heavily anisotropic test noise. When both covariances are diagonal and sufficiently close, the advantage over Adam can reverse, while applying the same random rotation to both restores it by changing only their alignment with Adam's coordinate axes. Synthetic and last-layer language-model experiments support the predicted advantage.