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

Introduction

• Muon’s full-Gram transformation jointly uses marginal row scales and pairwise interactions, while diagonal-Gram normalization removes scale weighting to expose directional interactions.The introduction motivates this separation with an example where a dominant row scale can dominate the leading Gram eigenvector despite unchanged directional interactions.Source: §1 · p. 2

The Normalize-Then-Precondition Framework

• Normalize-Then-Precondition first constructs a marginally normalized update, then refines it using interaction information through global or localized spectral preconditioning.The framework formalizes global preconditioning through spectral-norm steepest descent and localized preconditioning through a regularized formulation with leading-mode selection.Source: §1 · pp. 2–4

The NormPre Optimizer

• NormPre-G uses Newton-Schulz iterations for global preconditioning, while NormPre-L uses randomized sketching for localized eigenspace extraction; simplified variants have O(T^-1/2) convergence guarantees.The framework is presented as addressing the unresolved choice between full-spectrum transformation and targeted spectral modification.Source: §1 · pp. 2, 4

Optimization Performance and Trade-offs

• Under matched training budgets, NormPre-G and NormPre-L consistently outperform AdamW, Muon, and MANO across the evaluated architectures and model scales.The variants offer different performance-cost balances: NormPre-L provides strong optimization gains at lower cost, while NormPre-G further reduces validation loss.Source: §1 · pp. 3–4

2 Geometric Foundations of Normalize-Then-Precondition

• The framework first establishes a full-Gram view, derives diagonal-Gram normalization, and then formulates spectral preconditioning of the normalized update’s interactions.The section presents the row-side analysis, with the column-side counterpart obtained by applying the same construction to the transpose.Source: §2 · p. 4

2.1 Orthogonalized Updates and Full-Gram Transformation

• Muon’s full row Gram matrix encodes row scales on its diagonal and pairwise interactions off-diagonal, so its transformation jointly processes both.The section uses this full-Gram view as the geometric starting point for the hierarchical construction.Source: §2.1 · p. 5

2.2 Diagonal-Gram Normalization

• Diagonal-Gram normalization uses only row scales to form a marginally normalized update, while its normalized interaction matrix retains cosine similarities between rows.The resulting interaction matrix has unit diagonal; its off-diagonal entries encode the rows’ directional correlations.Source: §2.2 · pp. 5–6

2.3 Spectral Interaction Preconditioning

• Global preconditioning orthogonalizes the normalized update across its spectrum; localized preconditioning retains it as a reference and modifies selected leading active modes.A spectral budget r limits the number of transformed modes, while the localized update remains dense.Source: §2.3 · pp. 6–7

3 NormPre

• NormPre comprises NormPre-G and NormPre-L, which use global and localized spectral preconditioning, respectively.The optimizer alternates row and column orientation, uses relaxed tangent momentum processing, and rescales updates to a target RMS of 0.2.Source: §3 · p. 7, §3.1 · pp. 7–8

3.2 Implementation of Spectral Preconditioning

• NormPre-G approximates full-spectrum preconditioning with Newton-Schulz iterations, while NormPre-L can use exact eigenspace extraction or randomized sketching of leading interaction modes.The sketch-based implementation restricts extraction to a subspace of dimension ℓ=min{m, r+o}, reducing the stated cost to O((p+1)mnℓ+(m+n)ℓ^2+ℓ^3).Source: §3.2 · p. 8
• The convergence guarantee is for simplified deterministic NormPre variants and relies on stated smoothness, lower-boundedness, and geometric conditions.The analysis excludes momentum, RMS scaling, and weight decay; its sufficient condition is not necessary for convergence.Source: §3.2 · p. 9

4.1 Scaling Experiments

• Both NormPre variants outperform AdamW, Muon, and MANO across GPT-2 Small, LLaMA, and Qwen3; NormPre-G achieves the lowest loss.On GPT-2 Small, they reduce validation loss versus Muon by 0.0397 and 0.0238, respectively; at LLaMA-1.3B, the reductions are 0.0392 and 0.0301.Source: §4.1 · pp. 9–10, Table 1 · p. 10

4.2 Training Efficiency and Spectral Dynamics

• NormPre-L reduces optimizer latency through sketching, while end-to-end step-time changes are smaller because forward and backward computation dominate.Relative to Muon, NormPre-G increases step time by about 0.44% and 0.08%, while NormPre-L reduces it by 1.16% and 3.13% on the two reported models.Source: §4.2 · pp. 10–11, Table 2 · p. 10
• After marginal normalization, the update remains anisotropic; the top 32 modes account for approximately 62%–67% of global spectral transformation energy across checkpoints.Localized preconditioning focuses on these modes and captures a substantial portion of the global transformation.Source: §4.2 · p. 11, Figure 4 · p. 11

4.3 Comparisons of NormPre-G and NormPre-L

• NormPre-G achieves the lowest validation loss by transforming the full spectrum, while NormPre-L trades some performance for lower optimizer latency and step time at r = 32.On GPT-2 Small, larger r lowers validation loss at higher optimizer cost; r = 32 is the default balance.Source: §4.3 · pp. 11–12

• Matrix optimizers range from coordinate-wise and block-wise methods to matrix-level preconditioners, while this work focuses on update geometry through orthogonalization or row/column normalization.Matrix-level methods such as K-FAC, Shampoo, and SOAP build preconditioners from curvature or accumulated statistics.Source: §5 · p. 12
• Muon applies a matrix-sign transformation, while MANO and related methods use row/column normalization or adaptive scaling; prior normalization approaches do not establish a systematic hierarchy.The cited approaches include normalization after orthogonalization and normalization before orthogonalization for numerical stability.Source: §5 · p. 12
• Normalize-Then-Precondition formalizes a hierarchy of marginal normalization followed by spectral preconditioning, with global and localized realizations grounded in spectral steepest descent and its regularized form.Localized preconditioning uses leading-mode selection.Source: §5 · p. 12

6 Conclusion

• The paper formalizes Normalize-Then-Precondition to hierarchically organize marginal scales and interactions in matrix optimizers.NormPre outperforms established baselines in LLM pretraining, with a performance-efficiency trade-off between variants.Source: §6 · p. 13
• Future work includes hardware-aware spectral transformations, empirical validation beyond 1.7B parameters, and adaptive spectral schemes to narrow the localized-global performance gap while preserving efficiency.Source: §6 · p. 13

AI use statement

• The authors report using OpenAI Codex and Anthropic Claude Code for implementation, evaluation, translation, proof refinement, literature search, and manuscript work.They state that the research questions, conceptual framework, proof strategies, and experimental design were developed without AI assistance, and that authors reviewed the AI-assisted work.Source: AI use statement · p. 13

Ethics statement

• The work uses established language-modeling datasets and introduces or releases no new training data; it does not involve human-subject experiments or personal-information collection.Source: Ethics statement · p. 13

Reproducibility statement

• The paper points to appendices for implementation details, complexity analyses, theoretical assumptions and proofs, and additional experimental materials; its abstract links the code repository.Source: Reproducibility statement · p. 13

• The evaluated settings are GPT-2 Small on OpenWebText, LLaMA models from 130M to 1.3B on C4, and Qwen3 models from 0.6B to 1.7B on Pile.Source: §A.1 · p. 18
• C4 and Pile use fixed datasets of 5,120,000 training sequences and 2,048 validation sequences, each of length 1024; data order is fixed across optimizer runs.Documents are processed independently, with long documents truncated and shorter ones padded without cross-document packing.Source: §A.1 · p. 18
• NormPre-G uses five Newton–Schulz iterations, while NormPre-L uses sketch rank r = 32, oversampling 8, and one power iteration; both target update RMS 0.2.Muon, MANO, and both NormPre variants operate on hidden two-dimensional attention and MLP matrices, with other listed parameters optimized by AdamW.Source: §A.1 · p. 18, Table 4 · p. 19
• All models train for 10,000 steps with sequence length 1024 and effective batch size 512, using matched initialization, data order, budget, learning-rate schedule, and validation protocol.GPT-2 Small uses 200 sampled validation mini-batches; LLaMA and Qwen3 use the complete fixed validation split of 2,048 sequences.Source: §A.1 · pp. 18–19, Table 3 · p. 18

Training Efficiency Configurations

• Efficiency profiling uses 100 optimizer steps after 20 warmup steps, with optimizer latency isolating optimizer.step() and end-to-end time covering the full training loop.The reported configurations include sequence length 1024, effective batch size 512, and gradient clipping.Source: §A.2 · p. 19

Training and Validation Loss for Qwen3 Models

• Figure 5 presents training and validation loss curves for Qwen3-0.6B and Qwen3-1.7B on Pile across five optimizers.Dark training curves are 50-step moving averages; light curves show raw trajectories.Source: §A.2 · p. 19, Figure 5 · p. 20

Training Efficiency on Qwen3-1.7B

• NormPre-L reduces optimizer latency by 57.16% and step time by 4.21%, with a 4.41% throughput increase versus Muon on Qwen3-1.7B.NormPre-G increases optimizer latency by 6.77% and end-to-end step time by 0.53%; all evaluated optimizers use 53.17 GiB peak memory per GPU.Source: §A.2 · p. 20

Eigenspectra Before and After Normalization

• Marginal normalization removes row-norm weighting, increasing smaller-norm rows’ relative influence and potentially raising eigenvalues across multiple ranks.The spectra are compared after trace matching to isolate spectral shape from overall update magnitude.Source: §A.3 · pp. 21–22

Global and Localized Spectral Preconditioning

• The top 32 modes account for approximately 55%–60% of total spectral transformation energy across checkpoints in the column-active results.Localized preconditioning transforms selected modes, while global preconditioning also contracts remaining modes above one and amplifies positive modes below one.Source: Figure 6 · p. 21, §A.3 · p. 22

Training Dynamics

• Both NormPre variants show transient early gradient-norm peaks and lower descent alignment than Muon; later, their norms stabilize while directional differences persist.Descent alignment averages per-matrix cosine similarities between the update direction and the current negative gradient.Source: §A.3 · pp. 22–23, Figure 7 · p. 23

Ablations on Marginally Normalized Update

• Relaxed tangent projection achieves the lowest validation loss for both variants, while strict projection particularly degrades NormPre-G.The authors suggest fully removing the radial component can discard useful information before spectral refinement; alternating row/column normalization is the default for axis symmetry.Source: §A.4 · p. 24

Ablations on Interaction Geometry Source

• Extracting interaction geometry from Ψ rather than X yields lower training and validation loss for both variants.The authors suggest marginal scales in X can alter spectral ordering, making localized refinement more sensitive to the geometry source.Source: §A.4 · pp. 24–25

Spectral Budget in NormPre-L

• For NormPre-L, Table 8 reports training and validation losses across spectral budgets; both decrease as r increases over the evaluated range, and r = 32 is the default balancing gains against optimizer cost.Source: §A.4 · p. 25

Global Spectral Preconditioning

• NormPre-G applies Newton–Schulz iterations to the normalized update Ψ to approximate the global transformation msign(Ψ).The implementation follows the standard Newton–Schulz approach used by Muon.Source: B Detailed Implementations on Spectral… · p. 26

Exact Eigenspace Extraction

• Exact localized extraction forms Γ = ΨΨ⊤ and selects up to r largest eigenvalues above one for the preconditioner.The selected eigenpairs define the modes used by localized preconditioning.Source: B Detailed Implementations on Spectral… · p. 26

Sketch-Based Eigenspace Extraction

• The randomized sketch probes Ψ with a Gaussian matrix, uses power iterations and Rayleigh–Ritz extraction, then retains estimated top-r modes above one.This approximates the leading interaction eigenspace without explicitly forming and decomposing the full interaction matrix.Source: B Detailed Implementations on Spectral… · pp. 26–27

C.1 Shared Components in NormPre-G and NormPre-L

• Shared momentum, relaxed tangent projection, diagonal-Gram normalization, RMS scaling, and parameter updates each contribute O(mn) computation to both variants.These operations are included in the total costs reported for the variants.Source: §C.1 · p. 28

C.3 NormPre-L

• For fixed q Newton-Schulz iterations, Muon and NormPre-G have the same leading asymptotic complexity O(mn + qmns), while NormPre-L (Sketch) costs O((p + 1)mnℓ + (m + n)ℓ² + ℓ³).Here s = min{m, n}, ℓ = min{m, r + o}, and p is the number of power iterations.Source: C Detailed Computational Complexity… · p. 28, §C.2 · p. 29, §C.3 · p. 30

C.2 NormPre-G and Muon

• NormPre-G applies q Newton-Schulz iterations to the normalized update, giving total complexity O(mn + qmns), the same leading asymptotic order as Muon.The extra tangent projection and diagonal-Gram normalization do not change the asymptotic order.Source: §C.2 · pp. 28–29

C.3 NormPre-L

• Exact NormPre-L forms Γ = ΨΨ⊤ and computes its dense eigendecomposition, giving complexity O(m^2n + m^3).Its localized spectral application costs O(mnr); the simplification uses r ≤ m.Source: §C.3 · p. 29
• Sketch-based NormPre-L uses a randomized sketch and Rayleigh–Ritz extraction to estimate leading interaction eigenpairs without forming and decomposing the full m × m interaction matrix.The complexity uses ℓ = min{m, r + o} and p power iterations.Source: C Detailed Computational Complexity… · p. 28, §C.3 · pp. 29–30

D.1 Optimization Formulation of Localized Spectral Refinement

• Regularized spectral steepest descent clips singular values above one, while a rank budget retains at most r leading super-unit modes for localized refinement.The rank-constrained selection keeps the active modes with the largest eigenvalues; unselected modes remain unchanged.Source: §D.1 · pp. 31–32

D.2 Convergence Analysis of NormPre

• The proof chain analyzes a fixed orientation, forms a relaxed tangent gradient and its row-normalized update, then applies spectral preconditioning to obtain the refined update.The row normalization is assumed to be well defined.Source: §D.2 · pp. 33–34

D.2.1 Lemmas

• The lemmas show that row normalization retains a first-order gradient signal and exact spectral preconditioning retains at least a λmin(Pt) fraction of that signal.These bounds support the descent analysis alongside a bound on the refined-update norm.Source: §D.2.1 · pp. 34–36

D.2.2 Convergence of NormPre-G and NormPre-L (Exact)

• In deterministic settings without momentum, RMS scaling, or weight decay, NormPre-G with matrix-sign transformation and NormPre-L with exact implementation each achieve O(T^-1/2) convergence under their corresponding conditions.A sufficient condition limits worst-case radial cancellation to at most a (1 − ϵ) fraction of the positive first-order term, leaving at least an ϵ fraction in the full gradient-update inner product.Source: §D.2.2 · pp. 38–40

D.2.3 Convergence of NormPre-L (Sketch)

• A sketch-based NormPre-L preserves the O(T^-1/2) rate when its preconditioner error satisfies δ max_t λt,1 < εγ with a positive margin independent of T.The sketch and exact updates are constructed from the same normalized update and differ in interaction eigenspace extraction.Source: §D.2.3 · pp. 40–42