PaperlayerView exact revision on arXiv

Source-linked AI summary

RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization

Shenyang Deng, Zhuoli Ouyang, Tianyu Pang, Zihang Liu, Ruochen Jin, Shuhua Yu, Yaoqing Yang

arXiv:2603.20527v3cs.LG

TL;DR

Matrix-based preconditioning can capture curvature information, but balancing its effectiveness with computational cost remains a challenge. RMNP replaces Muon’s Newton-Schulz iteration with row-wise ℓ2 normalization; the paper reports faster preconditioning with competitive optimization performance and convergence guarantees matching Muon’s recent results.

  • Problem

    Preconditioned adaptive methods capture curvature information, but implementing effective matrix preconditioners efficiently remains challenging.

    Source: §1 · pp. 1–2

  • Method

    RMNP replaces Muon’s Newton-Schulz iteration with row-wise ℓ2 normalization motivated by Transformer curvature structure.

    Source: §3.1 · p. 6, §6 · p. 13

  • Results

    RMNP delivers 13–44× speedup on preconditioning over Muon across 60M to 1.5B parameter models while matching or outperforming Muon and AdamW in final performance.

    Source: §6 · p. 13

  • Takeaways & Limitations

    RMNP offers a more scalable preconditioning approach when Muon’s preconditioning process becomes a computational bottleneck in large-scale training.

    Source: §6 · p. 13

    The convergence analysis assumes a smooth loss, unbiased stochastic gradients, and bounded second moments.

    Source: §5 · p. 10

Abstract

from arXiv · show

Preconditioned adaptive methods have gained significant attention for training deep neural networks, as they capture rich curvature information of the loss landscape. The central challenge in this field lies in balancing preconditioning effectiveness with computational efficiency of implementing the preconditioner. Among recent advances, Muon stands out by using Newton-Schulz iteration to obtain preconditioned updates without explicitly constructing the preconditioning matrix. Despite its advantages, the efficiency of Muon still leaves room for further improvement. In this paper, we introduce RMNP (Row Momentum Normalized Preconditioning), an optimizer that replaces Newton-Schulz iteration with a simple row-wise ($d_{\text{in}}$) $\ell_2$ normalization operation, motivated by the empirically observed diagonal block structure of the Transformer layerwise Hessian. We empirically verified that orthogonalization and row-wise (on input dim) $\ell_2$ normalization are asymptotically equivalent in the case of the transformer. This substitution reduces the per-iteration computational complexity from ${O}(mn\cdot\min(m,n))$ to ${O}(mn)$ for an $m\times n$ weight matrix while maintaining comparable optimization performance. Theoretically, we establish convergence guarantees for RMNP in the non-convex setting that match recent results for Muon optimizers, achieving the minimax optimal complexity. Extensive experiments on large language model pretraining show that RMNP delivers competitive optimization performance compared with Muon while substantially reducing preconditioning wall-clock time. Our code is available at https://github.com/Dominator-Index/RMNP.

1 Introduction

The introduction frames matrix-based preconditioning as a way to capture parameter interactions, but its computational cost remains a barrier. RMNP uses row-wise normalization to reduce that cost while retaining matrix-level adaptivity, and the paper reports competitive performance and non-convex convergence guarantees.

1 Introduction

  • Diagonal preconditioners ignore parameter correlations, while full-matrix curvature methods can be prohibitively expensive for large models.Muon avoids explicitly constructing the preconditioner but its iterative matrix-polynomial evaluations cost O(mn·min(m,n)) for an m × n weight matrix.Source: §1 · pp. 1–2
  • RMNP replaces Muon’s Newton-Schulz iterations with row-wise ℓ2 normalization, reducing per-iteration complexity from O(mn·min(m,n)) to O(mn).The design is motivated by the observed block-diagonal dominance of Transformer curvature.Source: §1 · p. 3
  • RMNP matches or exceeds Muon’s final perplexity while reducing preconditioning wall-clock time by up to an order of magnitude.The introduction reports comparisons across model architectures and scales.Source: §1 · p. 3
  • The paper reports non-convex convergence guarantees on par with current Muon theory and minimax optimal complexity.The convergence analysis is presented for the non-convex smooth setting.Source: §1 · p. 3

2 Related Work

Related work spans structured matrix preconditioners, Muon’s iterative orthogonalization, and newer row- or column-normalized optimizers. The paper motivates RMNP from Transformer curvature structure and argues that worst-case steepest-descent analyses do not explain why a particular structure suits neural-network optimization.

Discussion with Recent Row Normalization Optimizers

  • Row-normalized methods include SRON, SWAN, MNGD, Mano, and MOGA, while SCALE uses column-wise normalization.Zhang et al. [21] introduced row-wise normalization for optimizer design, and Pethick et al. [22] proposed the LMO framework.Source: §2 · p. 4

Why steepest-descent analyses cannot explain NN-specific benefits

  • Worst-case steepest-descent guarantees provide a floor but cannot explain why a particular norm suits neural-network optimization.The paper instead examines neural-network curvature and reports asymptotic equivalence between full orthogonalization and row ℓ2 normalization for Transformers.Source: §2 · p. 4

Preconditioned Optimization Algorithms

  • Structured matrix methods approximate curvature with lower-dimensional factors, whereas Muon uses Newton-Schulz iterations for matrix-level adaptation.Muon’s iterative evaluations have complexity O(mn·min(m,n)) for an m × n weight matrix.Source: §2 · p. 4

Hessian Properties of Neural Networks

  • Transformer layerwise Hessians exhibit row-wise block-diagonal dominance, motivating row-level preconditioning that focuses on within-row interactions.The paper states that full orthogonalization and row ℓ2 normalization are asymptotically equivalent for Transformers.Source: §2 · pp. 4–5

3 Method

RMNP approximates Transformer curvature by retaining row-wise diagonal blocks, motivating a preconditioner implemented as row-wise ℓ2 normalization. The authors assess Muon’s Gram-matrix diagonal dominance and report strong dominance across representative GPT-2 matrices and training scales.

3.1 RMNP Preconditioner

  • RMNP retains row-wise diagonal blocks of the preconditioner, based on the hypothesis that these capture dominant curvature while cross-row interactions contribute negligibly.The resulting structure applies scaling independently to each row and reduces to row-wise ℓ2 normalization.Source: §3.1 · p. 6

3.2 Analysis of Muon Preconditioner

  • The analysis measures diagonal dominance in V_tV_t^T using per-row ratios and their average, minimum, and maximum across rows.A ratio above 1 means a diagonal entry exceeds the average off-diagonal magnitude in its row.Source: §3.2 · p. 6
  • For GPT-2 Small, representative matrices had rmin above 2, ravg above 5, and rmax near 25; global values were about 1.6, 4.9, and 60, respectively.The ratios exceeded 1 throughout training in the three representative matrices, and dominance grew with model size.Source: Figure 4 · p. 7, §3.2 · p. 7

4 Main Experimental Results

Across the reported GPT-2 and LLaMA pretraining experiments, RMNP matches or improves on Muon’s perplexity while reducing preconditioning time. The experiments compare these optimizers across model scales and datasets, with RMNP’s preconditioning speed advantage reported across GPT-2 sizes.

4.1 Experimental Setup

  • The experiments cover GPT-2 scales from 125M to 770M on OpenWebText and LLaMA scales from 60M to 1B on C4.GPT-2 is also evaluated on FineWeb-Edu-100B across four scales, including 1.5B.Source: §4 · p. 7, §4.1 · p. 8
  • Across GPT-2 and LLaMA scales, global diagonal-dominance ratios rise above 1 after warm-up and tend to grow with model size.Figure 5 reports GPT-2 models pretrained on OpenWebText and LLaMA models pretrained on C4.Source: Figure 5 · p. 8

4.2 Preconditioning Time Cost

  • RMNP row normalization is approximately 13–44× faster than Muon’s Newton–Schulz orthogonalization across GPT-2 model sizes.For GPT-2 1.5B, RMNP achieves a 42.9× speedup over Muon over 100 steps.Source: Table 2 · p. 9, §4.2 · p. 9

4.3 Pretraining Performance

  • GPT-2 evaluation perplexity improves over Muon by 0.04, 0.07, and 0.24 for Small, Medium, and Large, respectively.The corresponding improvements over AdamW are 1.37, 1.49, and 0.84.Source: §4.3 · p. 9
  • LLaMA perplexity decreases over Muon by 0.63 on 60M, 0.28 on 130M, and 0.02 on 350M.The reported decreases over AdamW are 4.33 on LLaMA-60M and 1.10 on LLaMA-130M.Source: §4.3 · p. 10
  • RMNP outperforms Muon and AdamW on GPT-2 evaluation perplexity and achieves comparable perplexity with a slight edge over Muon on LLaMA.Figure 6 covers GPT-2 and LLaMA pretraining results across model scales.Source: §4.3 · pp. 9–10, Figure 6 · p. 10

5 Non-Convex Convergence

Under non-convex smoothness assumptions, RMNP has convergence guarantees whose complexity depends on the matrix geometry: O(m²) under Frobenius smoothness and O(m) under ∥·∥∞,2 smoothness. The latter matches Muon’s favorable dimension dependence under a different geometry.

5.2 Assumptions

  • The analysis assumes a smooth loss, unbiased stochastic gradients, bounded second moments, and an objective bounded below.The paper studies convergence under non-convex smooth settings.Source: §5 · p. 10, §5.2 · p. 11

5.1 Notation

  • The notation defines matrix rows and Frobenius, ∥·∥1,2, and ∥·∥∞,2 norms, including the duality between the latter two.Source: §5.1 · p. 11

5.3 Main Results

  • O(m²LFσ²∆ϵ^-4) complexity yields ϵ-stationarity in both the Frobenius norm and the ∥·∥1,2 norm under Frobenius smoothness.The two convergence measures have the same complexity.Source: §5.3 · pp. 11–12
  • O(mL∞,2σ²∆ϵ^-4) complexity yields an ϵ-stationary point in ∥·∥1,2 under ∥·∥∞,2 smoothness.This improves dimension dependence from O(m²) to O(m).Source: §5.3 · p. 12, §5.4 · p. 13

5.4 Comparison with Related Work

  • Under Frobenius smoothness, RMNP’s O(m²) complexity matches Muon, while its ∥·∥∞,2-smoothness result improves dimension dependence to O(m), as does Muon under nuclear-norm smoothness.The favorable geometries differ: RMNP uses ∥·∥1,2 geometry, while Muon uses nuclear-norm geometry.Source: §5.4 · p. 13

6 Conclusion

RMNP replaces Muon’s Newton–Schulz iteration with row-wise ℓ2 normalization, reducing per-iteration complexity while retaining matrix-level preconditioning. The conclusion also reports experimental, theoretical, and practical contributions.

6 Conclusion

  • RMNP replaces Muon’s Newton–Schulz iteration with row-wise ℓ2 normalization, reducing per-iteration complexity from O(mn · min(m, n)) to O(mn).The design is motivated by diagonal block dominance observed in Transformer Hessians.Source: §6 · p. 13
  • RMNP provides a more scalable option when Muon’s preconditioning process becomes a computational bottleneck in large-scale training.The conclusion also reports 13–44× preconditioning speedups over Muon across 60M–1.5B parameter models, comparable memory usage, and final performance matching or exceeding Muon and AdamW in the reported experiments.Source: §6 · p. 13

A.1 Notation; A.2 Assumptions

  • The notation defines W_i,: as row i of W and gives the Frobenius and mixed norms, while the assumptions include smoothness, unbiased gradient estimates, bounded variance, and a lower-bounded objective.The initial optimality gap is defined as Δ = f(W0) − f*.Source: §A.1 · p. 17, §A.2 · pp. 17–18

A.3 Proof of Lemmas

The lemma proofs use row normalization properties and control the discrepancy between momentum estimates and true gradients. They derive supporting bounds under the stated smoothness conditions.

A.3 Proof of Lemmas

  • Row normalization gives each direction row unit ℓ2 norm, hence Frobenius norm √m and mixed norm ∥D∥∞,2 = 1.The proof states D_i,: = V_i,:/∥V_i,:∥2.Source: §A.3 · p. 18
  • The lemma arguments combine descent inequalities with lower bounds on the gradient-direction inner product that account for momentum estimation error.For the mixed-norm case, the inner-product bound subtracts twice the error’s 1,2 norm.Source: §A.3 · pp. 19, 24, 25
  • The proofs decompose momentum error into gradient changes and stochastic noise, then bound its accumulated contribution using smoothness and noise conditions.The error recursion follows from the momentum update; the lemmas apply both Frobenius and mixed-norm smoothness cases.Source: §A.3 · pp. 20, 24, 28

A.4 Proof of Theorem

The theorem proofs analyze constant-step-size updates with momentum under the stated assumptions. They sum descent inequalities across iterations and use prior lemmas to control accumulated error.

A.4 Proof of Theorem

  • Theorems 5.5 and 5.7 use Assumptions 5.1(a), 5.2, 5.3, and 5.4; Theorem 5.9 uses 5.1(b), 5.2, 5.3, and 5.4; all specify Algorithm 2 with constant step size ηt = η and momentum parameter β ∈ [0, 1).Source: §A.4 · pp. 30, 33, 34
  • The proofs sum per-iteration descent inequalities, use telescoping objective differences, and apply lemmas to bound accumulated error.The proof text notes that the initial update from W0 to W1 must be accounted for in the telescoping sum.Source: §A.4 · pp. 31–33
  • The argument isolates the sum of gradient norms and then takes expectations and averages across T iterations.The supplied theorem passages do not include the resulting bound itself.Source: §A.4 · p. 32

B Analysis of Muon Preconditioner

The appendix describes how it measures diagonal dominance in Muon’s momentum matrices and visualizes the resulting ratios across models and parameters.

B Analysis of Muon Preconditioner

  • The analysis computes Gram matrices and compares each diagonal entry with the average magnitude of off-diagonal entries to measure row-wise dominance.It aggregates these ratios into average, minimum, and maximum statistics per parameter and across parameters.Source: B Analysis of Muon Preconditioner · pp. 36–37
  • The analysis uses GPT-2 and LLaMA models, with model scales, training schedules, and Muon hyperparameters specified for the experiments.The listed setup uses Newton–Schulz iteration steps of 5, momentum 0.95, and weight decay 0.1.Source: B Analysis of Muon Preconditioner · p. 37
  • The plots show raw and smoothed ratios, with a dashed y = 1 threshold indicating when diagonal entries exceed average off-diagonal magnitude.The reported views include representative parameters, cross-scale and cross-architecture global ratios, and per-parameter ratios for the largest models.Source: B Analysis of Muon Preconditioner · p. 37

C Preconditioning Process Wall-Clock Time

RMNP substantially reduces preconditioner computation time versus Muon while maintaining identical memory usage. The appendix also reports the configurations, hyperparameter protocols, and model settings used in the timing and pretraining experiments.

C.1 Model Configuration for Preconditioning Time Cost

  • Timing measurements use the model configurations specified for preconditioning-cost evaluation.The section identifies Table 4 as the source of those configurations.Source: §C.1 · p. 38, Table 4 · p. 39

D Hyperparameter Search for Pretraining Performance

  • The hyperparameter search varies matrix learning rates across GPT-2 Small/Medium and LLaMA 60M/130M/350M; LLaMA-1B instead uses a fixed configuration.RMNP decouples AdamW and matrix learning rates, except LLaMA runs use a shared learning rate; the LLaMA-1B setting was fixed because of compute constraints.Source: D Hyperparameter Search for… · pp. 38, 40

C Preconditioning Process Wall-Clock Time

D.1 Hyperparameter Settings

  • Muon and RMNP share RMS learning-rate scaling and AdamW hyperparameters; GPT-2 matrix optimization covers all matrix parameters, while LLaMA leaves the LM head and embeddings to AdamW.The LLaMA parameter-group choice is examined in an appendix ablation.Source: §D.1 · pp. 40–41

D.2 Model Configurations

  • GPT-2 runs use sequence length 1024 and batch size 480, whereas LLaMA runs use sequence length 256 and batch size 512.GPT-2 XLarge (1.5B) is trained only on FineWeb-Edu-100B.Source: §D.2 · pp. 41–42

D.3 Extended Training Budget

The authors double the standard training budget for three model-dataset combinations to test whether RMNP’s advantage persists at longer horizons. RMNP achieves the lowest final validation perplexity in every tested case.

D.3 Extended Training Budget

  • RMNP achieves the lowest final validation perplexity in every tested case at 2× the standard training budget.The extended runs are GPT-2 Small on OpenWebText (20K steps), LLaMA-60M on C4 (20K steps), and LLaMA-130M on C4 (40K steps).Source: §D.3 · p. 44, Table 14 · p. 45

D.4 LM Head and Embedding Ablation

Across the three main pretraining settings, RMNP attains the lowest final perplexity in every reported cell. The appendix also reports an LM-head and embedding ablation and full training curves across model and dataset combinations.

D.4 LM Head and Embedding Ablation

  • Including LM-head and token-embedding parameters in the matrix-optimizer group changes final perplexity by at most 0.13 PPL, with no consistent trend across scales or optimizers.The ablation covers LLaMA-60M and LLaMA-130M.Source: Table 15 · p. 45, Table 16 · p. 45, §D.4 · p. 45

E Full Training Curves

  • The appendix presents training-loss, validation-loss, and gradient clip-rate curves for every evaluated model-dataset combination using canonical hyperparameters.RMNP is drawn on top in the curves so it is not occluded.Source: E Full Training Curves · p. 45

E.1 Final Validation Perplexity Summary

E.2 GPT-2 on OpenWebText

E.3 GPT-2 on FineWeb-Edu-100B

  • Across four GPT-2 scales on FineWeb-Edu-100B, RMNP matches or surpasses Muon and clearly outperforms AdamW.This extends the reported trend to a more competitive corpus and larger token budget.Source: §E.3 · p. 48, Figure 16 · p. 49, Figure 17 · p. 49

E.4 LLaMA on C4

  • Across four LLaMA scales on C4, RMNP gives a slight, stable improvement over Muon, and the gap between matrix-aware optimizers and AdamW widens with scale.The curves cover training and validation loss.Source: Figure 21 · p. 51, Figure 22 · p. 51, §E.4 · p. 51

E.5 Mamba on FineWeb-Edu

  • On Mamba trained on FineWeb-Edu, RMNP tracks Muon while both outperform AdamW; all three diagonal-dominance ratios rise above 1 shortly after warm-up and remain there.The reported learning-rate sweep fixes AdamW at 1 × 10^-4 and varies the matrix learning rate.Source: §E.5 · pp. 51, 53, Figure 25 · p. 52, Figure 26 · p. 53

E.6 ResNet-18 on CIFAR-10

On ResNet-18/CIFAR-10, RMNP tracks Muon and reaches essentially identical final accuracy, with diagonal dominance also holding in the convolutional setting. The gradient clip-rate analysis finds that larger GPT-2 models remain clipped for longer, while RMNP releases the threshold earliest.

E.6 ResNet-18 on CIFAR-10

  • RMNP closely tracks Muon on ResNet-18/CIFAR-10 and reaches essentially identical final accuracy; row-wise block-diagonal dominance also holds for this convolutional architecture.The comparison omits AdamW and includes training/test loss and accuracy.Source: §E.6 · pp. 53, 55, Figure 27 · p. 54, Figure 28 · p. 55

E.7 Gradient Clip-Rate Trajectories

  • Across both datasets, larger GPT-2 models stay clipped longer; AdamW clips every step on XLarge, while RMNP releases the threshold earliest, indicating lower gradient-norm volatility than Muon and AdamW.The plots show clip rate by model size and relative training progress.Source: §E.7 · p. 56, Figure 30 · p. 56, Figure 31 · p. 57

Related papers

Loading 2603.20527v3…