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 (din) ℓ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) ℓ2 normalization are asymptotically equivalent in the case of the transformer. This substitution reduces the per-iteration computational complexity from O(mn⋅min(m,n)) to O(mn) for an m×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.
Figures & tables
Figure 1
Figure 5: Comparison among Transformer layerwise Hessian, Preconditioner for Muon , and Preconditioner for RMNP . The figure of Transformer layerwise Hessian is conceptual, the real case can be widely found in [ 52 , 51 , 8 ] . P=VtVtT , m and n are the number of rows and columns of the weight matrix, respectively. In Section 3.2 we further verified through experiments that the Muon preconditioner has such a certain diagonal dominance property.
Figure 6: Worst-case problems for an algorithm may not capture the essential properties of NN optimization problems.
Figure 7: Per-parameter diagonal dominance ratios ravg , rmin , rmax (rows) for three representative matrix parameters (columns) during GPT-2 Small (125M), GPT-2 Medium (355M) and GPT-2 Large (770M) pre-training. Transparent curves: raw values; solid curves: smoothed with window size 50. Red dashed line: y=1 threshold.
Figure 8: Global diagonal dominance ratios ravg , rmin , rmax (columns) averaged across all matrix parameters, comparing across model scales for two architectures: GPT-2 Small (125M), Medium (355M), and Large (770M) pre-trained on OpenWebText (top row), and LLaMA 60M, 130M, and 350M pre-trained on C4 (bottom row). The x-axis is rescaled to the relative training progress (%) so that all model scales within a row align on a shared horizontal range; the y-axis is in log scale. Transparent curves: raw values; solid curves: smoothed with window size 50. Red dashed line: y=1 threshold. For both architectures, the metrics quickly rise above 1 after warm-up and remain mostly above 1, and the magnitude of ravg , rmin , rmax tends to grow with model scale, confirming strong and progressively more pronounced diagonal dominance throughout training on both Transformer families.
Size
Time Cost (s)
Speedup (×)
Muon
RMNP
60M
1.480
0.115
12.9
125M
2.975
0.201
14.8
200M
4.140
0.260
15.9
355M
7.380
0.401
18.4
500M
15.720
0.462
34.0
Table 1: Efficiency comparison between Muon and RMNP ’s preconditioning cost on GPT-2 models. Time measured over 100 steps with batch size 16 on a single RTX Pro 6000 GPU.
Figure 9: Left: Results for GPT-2 on OpenWebText: Small (125M) trained with 5B tokens; Medium (355M) trained with 10B tokens, and Large (770M) trained with 20B tokens. Numeric values are reported in Table 16 . FineWeb-Edu-100B results are in Figure 15 and Table 17 . Right: Results for LLaMA: 60M trained with 1B tokens; 130M trained with 2B tokens, 350M trained with 6B tokens, and 1B trained with 9B tokens. Numeric values are reported in Table 18 .
Appendix figures & tables45 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Per-parameter diagonal dominance ratios ravg , rmin , rmax (rows) for three representative matrix parameters (columns) during GPT-2 Small (125M), GPT-2 Medium (355M), and GPT-2 Large (770M) pre-training. Transparent curves: raw values; solid curves: smoothed with window size 50. Red dashed line: y=1 threshold.
Figure 11: Per-parameter diagonal dominance ratios ravg , rmin , rmax (rows) for three representative matrix parameters (columns) during LLaMA 60M, LLaMA 130M, and LLaMA 350M pre-training. Transparent curves: raw values; solid curves: smoothed with window size 50. Red dashed line: y=1 threshold.
Figure 12: Cross-architecture, cross-scale comparison of the global diagonal dominance ratios ravg , rmin , rmax (columns). Top row: GPT-2 Small (125M), Medium (355M), and Large (770M) on OpenWebText. Bottom row: LLaMA 60M, 130M, and 350M on C4. The x-axis is rescaled to the relative training progress (%) so that all model scales within a row align on a shared horizontal range; the y-axis is in log scale. Transparent curves: raw values; solid curves: smoothed with window size 50. Red dashed line: y=1 threshold. The figure reproduces the main-body version (Figure 8 ) and makes explicit that for both architecture families, larger models exhibit progressively stronger diagonal dominance across all three statistics.
Figure 13: Per-parameter diagonal dominance ratios ravg , rmin , rmax (rows) for three representative matrix parameters (columns) on the largest scales evaluated in this paper: GPT-2 XLarge (1.5B) on FineWeb-Edu-100B (left block) and LLaMA 1B on C4 (right block). Transparent curves: raw values; solid curves: smoothed with window size 50. Red dashed line: y=1 threshold. All metrics remain comfortably above the threshold throughout training, confirming that the row-wise block-diagonal dominance of the Muon preconditioner persists at the largest scales we evaluate.
Size
Time Cost (s)
Memory (MB)
Muon
RMNP
Muon
RMNP
60M
1.480
0.115
7804
7804
125M
2.975
0.201
11797
11797
200M
4.140
0.260
15352
15352
355M
7.380
0.401
23225
23225
500M
15.720
0.462
30011
30011
Appendix
Table 2: Efficiency comparison between Muon (Newton-Schulz orthogonalization) and RMNP (row normalization) on GPT-2 models. Time and memory usage are measured over 100 steps with batch size 16 on one single RTX Pro 6000 GPU.
Model
Params
Layers
Heads
dmodel
GPT-2 60M
60M
6
10
640
GPT-2 Small
125M
12
12
768
GPT-2 200M
200M
16
14
896
GPT-2 Medium
355M
24
16
1024
GPT-2 500M
500M
28
18
1152
GPT-2 Large
770M
36
20
1280
Appendix
Table 3: GPT-2 Model Configurations for Preconditioning Time Cost
Model
Params
Layer
Heads
demb
Steps
Warm-up
Token Count
Batch Size
LR schedule
GPT-2 Small
125M
12
12
768
10K
1K
5B
480
Cosine
GPT-2 Medium
355M
24
16
1024
20K
2K
10B
480
Cosine
GPT-2 Large
770M
36
20
1280
40K
4K
20B
480
Cosine
Appendix
Table 4: GPT-2 Model Configurations and specified hyperparameters for OpenWebText experiments.
Model
Params
Layer
Heads
demb
Steps
Warm-up
Token Count
Batch Size
LR schedule
GPT-2 Small
125M
12
12
768
10K
1K
5B
480
Cosine
GPT-2 Medium
355M
24
16
1024
20K
2K
10B
480
Cosine
GPT-2 Large
770M
36
20
1280
40K
4K
20B
480
Cosine
GPT-2 XLarge
1.5B
48
25
1600
50K
5K
25B
480
Cosine
Appendix
Table 5: GPT-2 Model Configurations and specified hyperparameters for FineWeb-Edu-100B experiments.
Model
Optimizer
lrAdamW
lrMatrix
Weight Decay
β
Schedule
Small (125M)
AdamW
6×10−4
—
0.1
(0.9, 0.95)
Cosine
Muon
3×10−3
2×10−2
0.1
(0.9, 0.95)
Cosine
RMNP
3×10−3
2×10−2
0.1
(0.9, 0.95)
Cosine
Medium (355M)
AdamW
3×10−4
—
0.1
(0.9, 0.95)
Cosine
Muon
1.5×10−3
1×10−2
0.1
(0.9, 0.95)
Cosine
RMNP
1.5×10−3
1×10−2
0.1
(0.9, 0.95)
Cosine
Appendix
Table 6: Optimizer hyperparameters for GPT-2 pre-training on FineWeb-Edu-100B.
Params
Hidden
Intermediate
Heads
Blocks
Steps
Warm-up
Token Count
Batch Size
LR schedule
60M
512
1376
8
8
10K
1K
1B
512
Cosine
130M
768
2048
12
12
20K
2K
2B
512
Cosine
350M
1024
2736
16
24
60K
6K
6B
512
Cosine
1B
2048
5461
32
24
90K
9K
9B
512
Cosine
Appendix
Table 7: LLaMA Model Configurations and specified hyperparameters.
Matrix LR
0.01
0.015
0.02
0.025
Muon
23.62
26.74
22.86
22.87
Matrix LR
0.002
0.003
0.004
0.005
RMNP
23.58
22.95
22.82
26.42
Appendix
Table 8: Hyperparameter search on GPT-2 Small with AdamW learning rate fixed at 3×10−3 .
Matrix LR
0.005
0.01
0.02
0.03
Muon
18.33
18.26
17.38
17.44
Matrix LR
0.001
0.002
0.003
0.005
RMNP
18.58
17.31
17.42
17.88
Appendix
Table 9: Hyperparameter search on GPT-2 Medium with AdamW learning rate fixed at 1.5×10−3 .
Matrix LR
0.005
0.01
0.02
0.03
0.04
Muon
29.90
29.58
30.46
30.49
30.03
Matrix LR
0.001
0.004
0.005
0.01
0.02
RMNP
31.00
28.99
28.95
29.26
29.64
Matrix LR
0.005
0.01
0.02
0.03
0.04
Shampoo
31.04
30.69
31.07
29.74
30.61
Matrix LR
0.001
0.002
0.003
0.004
0.005
Appendix
Table 10: Hyperparameter search on LLaMA-60M (LM head and token-embedding parameters handled by AdamW ). Validation perplexity is reported.
Matrix LR
0.005
0.01
0.02
0.03
Muon
22.51
22.42
22.47
22.51
Matrix LR
0.01
0.02
0.03
0.04
RMNP
22.42
22.49
22.14
23.31
Matrix LR
0.005
0.01
0.03
0.04
Shampoo
23.22
22.70
22.69
23.49
Matrix LR
0.001
0.002
0.003
0.005
Appendix
Table 11: Hyperparameter search on LLaMA-130M (LM head and token-embedding parameters handled by AdamW ). Validation perplexity is reported.
Matrix LR
0.003
0.004
0.005
Muon
17.01
16.87
16.89
Matrix LR
0.003
0.004
0.005
RMNP
17.02
16.86
16.85
Appendix
Table 12: Hyperparameter search on LLaMA-350M (LM head and token-embedding parameters handled by AdamW ). Validation perplexity is reported.
Optimizer
LLaMA 60M
LLaMA 130M
GPT-2 Small (OWT)
AdamW
28.23
21.35
20.97
Muon
27.03
20.84
20.88
RMNP
26.44
20.53
20.41
Appendix
Table 13: Final validation PPL (↓) under an extended training budget ( 2× standard). Lower is better.
Matrix LR
0.005
0.01
0.02
0.03
0.04
Muon
30.41
29.49
29.63
29.38
30.57
Matrix LR
0.001
0.004
0.005
0.01
0.02
RMNP
34.92
29.56
29.28
29.03
31.45
Appendix
Table 14: Hyperparameter search on LLaMA-60M with LM head and embedding parameters optimized by the matrix optimizer. Validation perplexity is reported.
Matrix LR
0.005
0.01
0.02
0.03
Muon
22.89
22.55
22.80
22.87
Matrix LR
0.01
0.02
0.03
0.04
RMNP
22.16
22.11
22.06
23.62
Appendix
Table 15: Hyperparameter search on LLaMA-130M with LM head and embedding parameters optimized by the matrix optimizer. Validation perplexity is reported.
Figure 14: Final validation perplexity ( ↓ ) on OpenWebText for GPT-2 Small, Medium, and Large. Numeric values are reported in Table 16 .
Small (125M)
Medium (355M)
Large (770M)
AdamW
24.19
18.80
15.27
Muon
22.86
17.38
14.67
RMNP
22.82
17.31
14.43
Appendix
Table 16: Final validation perplexity ( ↓ ) on OpenWebText for GPT-2 models.
Figure 15: Final validation perplexity ( ↓ ) on FineWeb-Edu-100B for GPT-2 Small, Medium, Large, and XLarge. Numeric values are reported in Table 17 .
Small (125M)
Medium (355M)
Large (770M)
XLarge (1.5B)
AdamW
23.85
18.19
14.81
13.12
Muon
22.71
17.13
14.16
12.97
RMNP
22.60
17.07
13.75
12.58
Appendix
Table 17: Final validation perplexity ( ↓ ) on FineWeb-Edu-100B for GPT-2 models.
Figure 16: Final validation perplexity ( ↓ ) on C4 for LLaMA 60M, 130M, 350M, and 1B. Numeric values are reported in Table 18 .
60M
130M
350M
1B
AdamW
33.28
23.24
17.08
15.33
Muon
29.58
22.42
16.87
14.13
RMNP
28.95
22.14
16.85
13.75
Appendix
Table 18: Final validation perplexity ( ↓ ) on C4 for LLaMA models.
Figure 17: GPT-2 Small (125M) on OpenWebText. Training loss is smoothed with a 20-step rolling window. RMNP ends with the lowest training and validation loss.
Figure 18: GPT-2 Medium (355M) on OpenWebText. RMNP achieves the lowest validation loss among the three optimizers.
Figure 19: GPT-2 Large (770M) on OpenWebText. RMNP ’s lead over Muon grows with model scale.
Figure 20: GPT-2 Small (125M) on FineWeb-Edu-100B. RMNP attains the lowest training and validation loss.
Figure 21: GPT-2 Medium (355M) on FineWeb-Edu-100B. RMNP maintains a slight but consistent edge over Muon on validation loss while AdamW lags throughout training.
Figure 22: GPT-2 Large (770M) on FineWeb-Edu-100B. RMNP ’s lead over Muon grows with model scale, while AdamW converges to a noticeably higher validation loss.
Figure 23: GPT-2 XLarge (1.5B) on FineWeb-Edu-100B. RMNP continues to track Muon closely and surpasses it in late training, while delivering an order-of-magnitude reduction in preconditioning wall-clock cost (Appendix C ).
Figure 24: LLaMA-60M on C4. The available AdamW log extends beyond the canonical training horizon and has been clipped to match Muon and RMNP on a shared x-range. RMNP achieves the lowest validation loss; Muon is close behind, while AdamW converges to a clearly higher value.
Figure 25: LLaMA-130M on C4. RMNP outperforms both baselines in validation loss throughout training.
Figure 26: LLaMA-350M on C4. RMNP matches Muon on training loss and edges ahead on validation loss in late training.
Figure 27: LLaMA-1B on C4. The RMNP curve tracks Muon closely on both training and validation loss while delivering a substantially lower preconditioning cost.
Figure 28: Mamba on FineWeb-Edu. Validation perplexity is shown on a log scale. RMNP matches Muon throughout training and clearly outperforms AdamW , demonstrating that the row-wise normalized preconditioner generalizes beyond Transformer architectures to state-space models.
Figure 29: Diagonal dominance ratios for Mamba pre-training on FineWeb-Edu. Transparent curves: raw values; solid curves: smoothed with window size 50. Red dashed line: y=1 threshold. All metrics remain above the threshold throughout training, demonstrating that the row-wise block-diagonal dominance property holds for the Mamba state-space architecture both at the global aggregate level (panel (a)) and at the per-parameter level (panel (b)).
Matrix LR
6.67×10−4
0.008
0.009
Muon
36.55
32.95
33.02
Matrix LR
6.67×10−4
8×10−4
1×10−3
RMNP
32.56
32.32
32.33
Appendix
Table 19: Hyperparameter search on Mamba (FineWeb-Edu) with AdamW learning rate fixed at 1×10−4 . Validation perplexity is reported.
Figure 30: ResNet-18 on CIFAR-10 , comparing Muon and RMNP . The two matrix-aware optimizers track each other closely and converge to essentially identical final accuracy, demonstrating that RMNP extends to convolutional vision tasks without loss of optimization quality.
Figure 31: Diagonal dominance ratios for ResNet-18 training on CIFAR-10 . Transparent curves: raw values; solid curves: smoothed with window size 50. Red dashed line: y=1 threshold. All metrics remain above the threshold throughout training, demonstrating that the row-wise block-diagonal dominance property holds for the convolutional vision architecture both at the global aggregate level (panel (a)) and at the per-parameter level (panel (b)).
Matrix LR
0.01
0.04
0.05
Muon
94.57
94.65
94.39
Matrix LR
0.006
0.008
0.01
RMNP
94.33
93.93
94.31
Appendix
Table 20: Test accuracy (%) on CIFAR-10 for ResNet-18 : matrix LR search with AdamW learning rate fixed at 0.006 . Higher is better.
Figure 32: Gradient clip rate during GPT-2 pre-training on OpenWebText, one panel per model size. Transparent line: raw values; solid line: 50-step rolling mean. RMNP (red) is drawn last so it sits on top of AdamW (blue) and Muon (green).
Figure 33: Gradient clip rate during GPT-2 pre-training on OpenWebText, with x-axis rescaled to relative training progress (%). Each panel shows one optimizer; within a panel, lighter to darker shades encode Small / Medium / Large . The clip rate falls below 1.0 progressively later for larger models.
Figure 34: Gradient clip rate during GPT-2 pre-training on FineWeb-Edu-100B, one panel per model size. Transparent line: raw values; solid line: 50-step rolling mean. AdamW on the XLarge (1.5B) model has its gradients clipped at every step throughout the entire run; both Muon and RMNP progressively reduce the clip rate.
Figure 35: Gradient clip rate during GPT-2 pre-training on FineWeb-Edu-100B, with x-axis rescaled to relative training progress (%). Each panel shows one optimizer; lighter to darker shades within a panel encode Small / Medium / Large / XLarge . The size-dependent delay before the clip rate begins to drop is most pronounced under AdamW (left) and least pronounced under RMNP (right).
Orthogonality-based optimizers, such as Muon, have recently shown strong performance across large-scale training and community-driven efficiency challenges. However, these methods rely on a costly gradient orthogonalization step. Even efficient iterative approximations such as Newton-Schulz remain expensive, typically requiring dozens of matrix multiplications to converge. We introduce a pre-conditioning procedure that improves the initialization of the Newton--Schulz iterations while incurring negligible overhead. Furthermore, our pre-conditioning reduces the initial polar error and enables the removal of one Newton-Schulz iteration (out of the five iterations usually used in practice). The resulting implementation significantly reduces Muon's overhead. At the end-to-end training level, we observe consistent runtime improvements across speed-run and standard benchmarks, including ∼3% reductions in training time on multiple fast training benchmarks, while matching reference performance on both language and vision tasks. Crucially, these improvements require no hyperparameter tuning and can be adopted as a simple drop-in replacement. Beyond empirical gains, we provide theoretical insight into the geometry of the update and its potential robustness against feature collapse. Our code is publicly available on github, in optax and huggingface kernels.
Thibaut Boissin, Thomas Massena, Franck Mamalet +1
IRIT-MISFIT · Institut de Recherche Technologique Saint-Exupery, France · DTIPG - SNCF, IRIT-MISFIT +2
Matrix-based optimizers have demonstrated immense potential in training Large Language Models (LLMs), however, designing an ideal optimizer remains a formidable challenge. A superior optimizer must satisfy three core desiderata: efficiency, achieving Muon-like preconditioning to accelerate optimization; stability, strictly adhering to the scale-invariance inherent in neural networks; and speed, minimizing computational overhead. While existing methods address these aspects to varying degrees, they often fail to unify them, either incurring prohibitive computational costs like Muon, or allowing radial jitters that compromise stability like RMNP. To bridge this gap, we propose Nora, an optimizer that rigorously satisfies all three requirements. Nora achieves training stability by explicitly stabilizing weight norms and angular velocities through row-wise momentum projection onto the orthogonal complement of the weights. Simultaneously, by leveraging the block-diagonal dominance of the Transformer Hessian, Nora effectively approximates structured preconditioning while maintaining an optimal computational complexity of O(mn). Furthermore, we prove that Nora is a scalable optimizer and establish its corresponding scaling theorems. With a streamlined implementation requiring only two lines of code, our preliminary experiments validate Nora as an efficient and highly promising optimizer for large-scale training.
Jinghui Yuan, Jiaxuan Zou, Shuo Wang +2
1. School of Artificial Intelligence, Optics and Electronics (iOPEN), Northwestern Polytechnical University. · 2. School of Mathematics and Statistics, Xi’an Jiaotong University. · 4. Gaoling School of Artificial Intelligence, Renmin University of China. +1
Matrix optimizers have emerged as a promising direction, with Muon standing out as a prominent design. Revisiting Muon through its full-Gram representation, we observe that it jointly processes marginal-scale and interaction information. This opens an alternative way to organize geometric information hierarchically, motivating the Normalize-Then-Precondition framework. Specifically, it first uses diagonal-Gram information to construct a marginally normalized update, then applies spectral preconditioning to its directional interaction geometry. Building on this framework, we develop NormPre with NormPre-G and NormPre-L adopting global and localized spectral preconditioning, grounded in spectral-norm steepest descent and a regularized formulation followed by leading mode selection, respectively. To enable large-scale training, NormPre-G uses Newton-Schulz iterations and NormPre-L employs randomized sketching to approximate the leading interaction eigenspace. Theoretically, we establish O(T−1/2) convergence guarantees for simplified versions of NormPre. Across extensive pretraining experiments on GPT-2 Small, LLaMA and Qwen3, both variants consistently outperform AdamW, Muon and MANO under matched training budgets. Further efficiency and spectral analyses reveal the complementary strengths of two variants and characterize their performance-efficiency trade-off. We open-source our code through a GitHub repository at https://github.com/zx-gong/NormPre.
Zixuan Gong, Zeyu Gan, Jiaye Teng +1
Gaoling School of Artificial Intelligence Renmin University of China · School of Statistics and Management Shanghai University of Finance and Economics