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

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

cs.LGMar 20, 2026

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

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

Organizations: Dartmouth College · International Computer Science Institute · University of California, Berkeley · Meta

Abstract

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.

Figures & tables

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=V_{t}V_{t}^{T}$ , $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 $r_{\text{avg}}$ , $r_{\min}$ , $r_{\max}$ (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 $\overline{r}_{\text{avg}}$ , $\overline{r}_{\min}$ , $\overline{r}_{\max}$ (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 $\overline{r}_{\text{avg}}$ , $\overline{r}_{\min}$ , $\overline{r}_{\max}$ 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 $(\times)$
Size	Muon	RMNP	Speedup $(\times)$
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 $r_{\text{avg}}$ , $r_{\min}$ , $r_{\max}$ (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 $r_{\text{avg}}$ , $r_{\min}$ , $r_{\max}$ (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 $\overline{r}_{\text{avg}}$ , $\overline{r}_{\min}$ , $\overline{r}_{\max}$ (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 $r_{\text{avg}}$ , $r_{\min}$ , $r_{\max}$ (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)
Size	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	$d_{\text{model}}$
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	$d_{\text{emb}}$	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	$d_{\text{emb}}$	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	$\text{lr}_{\text{AdamW}}$	$\text{lr}_{\text{Matrix}}$	Weight Decay	$\beta$	Schedule
Small (125M)	AdamW	$6\times 10^{-4}$	—	0.1	(0.9, 0.95)	Cosine
	Muon	$3\times 10^{-3}$	$2\times 10^{-2}$	0.1	(0.9, 0.95)	Cosine
	RMNP	$3\times 10^{-3}$	$2\times 10^{-2}$	0.1	(0.9, 0.95)	Cosine
Medium (355M)	AdamW	$3\times 10^{-4}$	—	0.1	(0.9, 0.95)	Cosine
	Muon	$1.5\times 10^{-3}$	$1\times 10^{-2}$	0.1	(0.9, 0.95)	Cosine
	RMNP	$1.5\times 10^{-3}$	$1\times 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\times 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\times 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\times

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 (

\downarrow

) 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 (

\downarrow

) on OpenWebText for GPT-2 models.

Figure 15: Final validation perplexity (

\downarrow

) 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 (

\downarrow

) on FineWeb-Edu-100B for GPT-2 models.

Figure 16: Final validation perplexity (

\downarrow

) 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 (

\downarrow

) on C4 for LLaMA models.

Panel of Figure 17 — 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.

Panel of Figure 18 — Figure 18: GPT-2 Medium (355M) on OpenWebText. RMNP achieves the lowest validation loss among the three optimizers.

Panel of Figure 19 — Figure 19: GPT-2 Large (770M) on OpenWebText. RMNP ’s lead over Muon grows with model scale.

Panel of Figure 20 — Figure 20: GPT-2 Small (125M) on FineWeb-Edu-100B. RMNP attains the lowest training and validation loss.

Panel of Figure 21 — 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.

Panel of Figure 22 — 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.

Panel of Figure 23 — 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 ).

Panel of Figure 24 — 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.

Panel of Figure 25 — Figure 25: LLaMA-130M on C4. RMNP outperforms both baselines in validation loss throughout training.

Panel of Figure 26 — Figure 26: LLaMA-350M on C4. RMNP matches Muon on training loss and edges ahead on validation loss in late training.

Panel of Figure 27 — 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.

Panel of Figure 28 — 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.

Panel of Figure 29 — 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\times 10^{-4}$	0.008	0.009
Muon	36.55	32.95	33.02
Matrix LR	$6.67\times 10^{-4}$	$8\times 10^{-4}$	$1\times 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\times 10^{-4}

. Validation perplexity is reported.

Panel of Figure 30 — 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.

Panel of Figure 31 — 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).

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 331 papers

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates
Dec 4, 2025Thibaut Boissin, Thomas Massena, Franck Mamalet +1Newton-Schulz Iteration Orthogonal Optimization
IRIT-MISFIT · Institut de Recherche Technologique Saint-Exupery, France · DTIPG - SNCF, IRIT-MISFIT +2
Nora: Normalized Orthogonal Row Alignment for Scalable Matrix Optimizer
May 5, 2026Jinghui Yuan, Jiaxuan Zou, Shuo Wang +2Neural Network Optimization Matrix Optimization
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
Normalize-Then-Precondition: A Hierarchical Approach to Marginal Scale and Interaction Geometry for LLM Training
Sep 29, 2026Zixuan Gong, Zeyu Gan, Jiaye Teng +1Language Model Pretraining Matrix Optimization
Gaoling School of Artificial Intelligence Renmin University of China · School of Statistics and Management Shanghai University of Finance and Economics

Dec 4, 2025J/K move · Enter open · S save

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates

Thibaut Boissin, Thomas Massena, Franck Mamalet, Mathieu Serrurier

IRIT-MISFIT · Institut de Recherche Technologique Saint-Exupery, France · DTIPG - SNCF, IRIT-MISFIT+2

Newton-Schulz Iteration Orthogonal Optimization Neural Network Optimization Muon Optimizer Second-Order Optimization

Open paper AskPDF arXiv

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

\sim

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.

Abstract

Neural Network Optimization

Explore similar work

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates

Nora: Normalized Orthogonal Row Alignment for Scalable Matrix Optimizer

Normalize-Then-Precondition: A Hierarchical Approach to Marginal Scale and Interaction Geometry for LLM Training