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

Zixuan Gong; Zeyu Gan; Jiaye Teng; Yong Liu

cs.LGSep 29, 2026

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

Authors: Zixuan Gong, Zeyu Gan, Jiaye Teng, Yong Liu

Organizations: Gaoling School of Artificial Intelligence Renmin University of China · School of Statistics and Management Shanghai University of Finance and Economics

Abstract

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 $\mathcal{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.

Figures & tables

Panel of Figure 1 — Figure 1: GPT-2 Small on OpenWebText and LLaMA Models on C4. Training and validation loss for GPT-2 Small and LLaMA-{130M, 350M, 1.3B} with AdamW, Muon, MANO, NormPre-G and NormPre-L. Dark and light curves denote 50-step moving averages and raw training trajectories.

Figure 2: Geometric Workflow of the NormPre Optimizer (Algorithm 1 ).

Figure 3: Spectral Interaction Preconditioning.

Setting		Validation Loss $\downarrow$
Model	Dataset	AdamW	Muon	MANO	NormPre-G	NormPre-L
GPT-2 Small	OpenWebText	3.1444 $\pm$ 0.0054	3.1064 $\pm$ 0.0054	3.1156 $\pm$ 0.0057	3.0667 $\bm{\pm}$ 0.0049	3.0826 $\pm$ 0.0034
LLaMA-130M	C4	3.1363	3.1019	3.1102	3.0737	3.0930
LLaMA-350M	C4	3.0378	2.9999	2.9946	2.9678	2.9758
LLaMA-1.3B	C4	2.9385	2.9037	2.8963	2.8571	2.8662
Qwen3-0.6B	Pile	2.8956	2.8335	2.8382	2.7967	2.8066

Table 1: Validation loss across model scales, architectures and pretraining datasets. GPT-2 Small results show mean

\pm

standard deviation over three matched seeds, and larger models report single-run results under matched training budgets. Shaded columns indicate our variants. The best baseline and our outperforming variants are shown in underline and bold , respectively.

Model	Dataset	Optimizer	Optimizer Latency $\downarrow$	E2E Step Time $\downarrow$	Throughput $\uparrow$	Peak Memory $\downarrow$
			(ms/step)	(ms/step)	(k tokens/s)	(GiB)
GPT-2 Small	OpenWebText	AdamW	2.7	3364.8	155.82	14.93
		Muon	112.8	3479.8	150.67	14.61
		MANO	8.8	3365.6	155.78	14.61
		NormPre-G	119.3 $\uparrow$ 5.76%	3495.0 $\uparrow$ 0.44%	150.01 $\downarrow$ 0.44%	14.62
		NormPre-L	69.8 $\downarrow$ 38.12%	3439.4 $\downarrow$ 1.16%	152.44 $\uparrow$ 1.17%	14.61

Table 2: Training Efficiency Comparisons. Measurements use the same training configurations as the main experiments and are averaged over 100 optimizer steps after 20 warmup steps.

Panel of Figure 4 — Figure 4: (a) Eigenspectra before and after marginal normalization; (b) Spectral transformation percentages under global and localized preconditioning; (c) Performance-efficiency trade-off.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Model Size	Dataset	Layers	Hidden	FFN	Heads	Seq. Len.	Eff. Batch	Base Peak LR	Warmup	Steps
GPT-2 Small	124M	OpenWebText	12	768	3,072	12	1,024	512	$6\times 10^{-4}$	2,000	10,000
LLaMA-130M	130M	C4	12	768	2,048	12	1,024	512	$6\times 10^{-4}$	1,000	10,000
LLaMA-350M	350M	C4	24	1,024	2,736	16	1,024	512	$3\times 10^{-4}$	1,000	10,000
LLaMA-1.3B	1.3B	C4	24	2,048	5,461	32	1,024	512	$3\times 10^{-4}$	1,000	10,000
Qwen3-0.6B	0.6B	Pile	28	1,024	3,072	16	1,024	512	$3\times 10^{-4}$	1,000	10,000
Qwen3-1.7B	1.7B	Pile	28	2,048	6,144	16	1,024	512	$3\times 10^{-4}$	1,000	10,000

Appendix

Table 3: Model and Training Configurations.

Hyperparameter	AdamW	Muon	MANO	NormPre-G	NormPre-L
$\beta_{1}$	0.9	–	–	–	–
$\beta_{2}$	0.95	–	–	–	–
Momentum $\mu$	–	0.95	0.95	0.95	0.95
Newton-Schulz steps	–	5	–	5	–
Weight decay	0.1	0.1	0.1	0.1	0.1
Interaction rank $r$	–	–	–	–	32

Appendix

Table 4: Default Optimizer Configurations. A dash indicates that the corresponding hyperparameter is not applicable.

Panel of Figure 5 — Figure 5: Qwen3 Models on Pile. Training and validation loss for Qwen3-{0.6B, 1.7B} with AdamW, Muon, MANO, NormPre-G and NormPre-L. For training loss, dark curves denote 50-step moving averages and light curves the raw trajectories.

Model	Dataset	Optimizer	Optimizer Latency $\downarrow$	E2E Step Time $\downarrow$	Throughput $\uparrow$	Peak Memory $\downarrow$
			(ms/step)	(ms/step)	(k tokens/s)	(GiB)
Qwen3-1.7B	Pile	AdamW	18.3	12102.7	43.32	53.17
		Muon	940.0	13050.4	40.17	53.17
		MANO	118.7	12217.4	42.91	53.17
		NormPre-G	1003.6 $\uparrow$ 6.77%	13119.7 $\uparrow$ 0.53%	39.96 $\downarrow$ 0.52%	53.17
		NormPre-L	402.7 $\downarrow$ 57.16%	12501.4 $\downarrow$ 4.21%	41.94 $\uparrow$ 4.41%	53.17

Appendix

Table 5: Training Efficiency on Qwen3-1.7B. Measurements use the same training configuration as the main Qwen3-1.7B experiment and are averaged over 100 optimizer steps after 20 warmup steps. We report global throughput across 4 GPUs and per-GPU peak allocated memory.

Panel of Figure 6 — Figure 6: Spectral Dynamics. (a) Eigenspectra before and after marginal normalization in the column-active orientation. (b) Spectral transformation percentages across modes under global and localized preconditioning (column-active orientation).

Panel of Figure 7 — Figure 7: Training Dynamics. (a) Pre-clipping global gradient norm over all trainable parameters. (b) Descent alignment between the gradient and optimizer update direction over matrix-optimized parameters. Dark curves show 50-step moving averages and light curves show the corresponding raw trajectories.

Variant	Validation Loss
Variant	NormPre-G	NormPre-L
No Tangent	3.0945	3.1139
Strict Tangent	3.1869	3.1196
Relaxed Tangent	3.0711	3.0864

Appendix

Table 6: Ablations on Marginally Normalized Update. Results are reported on GPT-2 Small with OpenWebText using seed

1337

under the same

10

k-step training budget.

Interaction Source	NormPre-G		NormPre-L
Interaction Source	Training Loss	Validation Loss	Training Loss	Validation Loss
Update before normalization ( $X$ )	3.0844	3.0900	3.1375	3.1305
Update after normalization ( $\Psi$ )	3.0655	3.0711	3.0810	3.0864

Appendix

Table 7: Ablations on Interaction Geometry Source. Results on GPT-2 Small with OpenWebText using seed 1337 and the same 10k-step training budget.

Rank $r$	Training Loss	Validation Loss
8	3.0911	3.0962
16	3.0857	3.0908
32 (default)	3.0810	3.0864
64	3.0768	3.0827

Appendix

Table 8: Spectral Budget in NormPre-L. Results are reported on GPT-2 Small with OpenWebText using seed

1337

under the same

10

k-step training budget.

Method	Computational Complexity
Muon	$\mathcal{O}(mn+qmns)$
NormPre-G	$\mathcal{O}(mn+qmns)$
NormPre-L (Exact)	$\mathcal{O}(m^{2}n+m^{3})$
NormPre-L (Sketch)	$\mathcal{O}\!\left((p+1)mn\ell+(m+n)\ell^{2}+\ell^{3}\right)$

Appendix

Table 9: Computational Complexity of Muon and NormPre Variants.

Symbol	Dimension	Description
$\Psi$	$m\times n$	Marginally normalized update
$\Omega$	$n\times\ell$	Gaussian random matrix
$Y$	$m\times\ell$	Gaussian sketch
$Q$	$m\times\ell$	Candidate subspace
$\Gamma_{Q}$	$\ell\times\ell$	Rayleigh–Ritz matrix
$Z$	$\ell\times r$	Top- $r$ eigenvectors of $\Gamma_{Q}$

Appendix

Table 10: Dimensions of the Main Quantities in Algorithm 2

Explore similar work

MeqMuon: Matrix-Equilibrating Muon for LLM Pretraining
Sep 28, 2026Chang-Wei Shi, Xu Wang, Wu-Jun LiLarge Language Model Pretraining Muon
National Key Laboratory for Novel Software Technology, School of Computer Science, Nanjing University, P. R. China
MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning
Aug 5, 2026Tongle Wu, Huanyu Dong, Ying Sun +1Spectral Preconditioning Muon
School of Electrical Engineering and Computer Science, The Pennsylvania State University, University Park, PA, USA · Department of Computer Science, City University of Hong Kong, Hong Kong SAR, China
Muown: Row-Norm Control for Muon Optimization
May 11, 2026Kai Lion, Florian Hübler, Bingcong Li +2Muon Optimizer Design
Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, School of Computation, Information and Technology; Technical University of Munich, Germany · ELLIS Institute Tübingen, MPI-IS, Tübingen AI Center, Germany

Sep 28, 2026, cs.LGJ/K move · Enter open · S save

MeqMuon: Matrix-Equilibrating Muon for LLM Pretraining

Chang-Wei Shi, Xu Wang, Wu-Jun Li

National Key Laboratory for Novel Software Technology, School of Computer Science, Nanjing University, P. R. China

Large Language Model Pretraining Muon Normalization Convergence

Open paper AskPDF arXiv

The success of large language models (LLMs) has been accompanied by continued growth in model size and pretraining costs. Muon offers high accuracy and training efficiency in LLM pretraining. Recent work introduces row-wise normalization into Muon to balance update magnitudes and improve pretraining performance. However, row-wise normalization alone cannot accommodate different imbalance patterns in update matrices. In this paper, we propose an improved Muon optimizer, called \underline{m}atrix-\underline{eq}uilibrating Muon~(MeqMuon), for LLM pretraining. MeqMuon balances both row and column magnitudes through normalization that can be automatically tailored to different imbalance patterns without manual intervention. Moreover, MeqMuon eliminates the need to store AdamW's second-moment estimates, reducing optimizer-state memory usage. Empirical results demonstrate that MeqMuon achieves better convergence performance than AdamW, Muon, and other baselines in LLM pretraining.

Abstract

Explore similar work

MeqMuon: Matrix-Equilibrating Muon for LLM Pretraining

MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

Muown: Row-Norm Control for Muon Optimization