arXiv:2605.03769v1 [cs.LG] 5 May 2026
Nora: Normalized Orthogonal Row Alignment for Scalable Matrix Optimizer Jinghui Yuan1 ∗, Jiaxuan Zou2,4 ∗, Shuo Wang3 ∗, Yong Liu4 †, Feiping Nie1 1. School of Artificial Intelligence, Optics and Electronics (iOPEN), Northwestern Polytechnical University. 2. School of Mathematics and Statistics, Xi’an Jiaotong University. 3. Institute for Interdisciplinary Information Sciences, Tsinghua University. 4. Gaoling School of Artificial Intelligence, Renmin University of China. [email protected], [email protected], [email protected], [email protected], [email protected]
Abstract 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 row orthogonal complement of the weights. Simultaneously, by leveraging the row 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.
1
Introduction
The training of Large Language Models (LLMs) relies heavily on adaptive optimizers such as Adam [1, 2]. However, such methods treat parameters as flat vectors, thereby overlooking the rich structural interactions inherent in the loss landscape [3]. Recently, matrix-based optimizers like Muon have emerged to bridge this gap, achieving state-of-the-art data efficiency through orthogonalized updates [4]. Despite its remarkable performance, Muon relies on the Newton-Schulz iteration [5], which introduces significant computational overhead of O(m2 n), where the weight matrix w ∈ Rm×n . Without loss of generality, we assume m ≤ n. Subsequent research has sought to approximate this preprocessor. Dong et al.[6] identified the row block diagonal dominance of the Transformer Hessian, a phenomenon also noted by [7], who proposed a simplified preprocessor approximating Muon via diagonal matrices. This led to the development of RMNP (Row-wise Momentum Normalization). Although RMNP is computationally ∗ †
These authors contributed equally. Corresponding Author.
Preprint.
superior, it overlooks a important property of neural networks: weight space symmetry [8]. Due to the pervasive use of BatchNorm [9], RMSNorm [10] and LayerNorm [11], network representations exhibit scale invariance [12]. Radial updates—those aligning with the weight vector—do not alter the functional output [13]. However, radial momentum noise can interfere with the preprocessor’s output and the direction of weight learning. Simultaneously, it silently perturbs the weight norm, leading to chaotic oscillations in the effective learning rate and ultimately undermining network training [14]. Beyond these, other classical or state-of-the-art optimizers typically suffer from one or more deficiencies: they either neglect the use of preconditioners [15, 16, 17], fail to account for the scale-invariance of neural networks [18, 19], or incur excessive computational complexity [20, 21]. In fact, many existing methods simultaneously lack several of these essential properties. Our motivation stems from a commitment to the three core principles of optimizer design: efficiency, stability, and speed. To this end, we introduce Nora. The design of Nora relies on a simple row-wise orthogonality property: row-wise normalization preserves row-wise perpendicularity to the weight. This property enables us to unify preconditioning estimation, scale-invariance, and rapid computation. Specifically, we first project the momentum vt onto its component that is row-wise perpendicular to wt to obtain vtr⊥ , which ensures training stability through orthogonality. Subsequently, we apply a diagonal preconditioning estimation to vtr⊥ to derive the update dt . This step efficiently approximates the Muon-like preconditioning for enhanced efficiency, yet computationally simplifies to a mere row-wise normalization, thereby simultaneously satisfying both the speed and stability requirements. Furthermore, leveraging the Maximal Update Parametrization (µP) framework [22], we demonstrate that Nora is a scalable optimizer and derive its optimal learning rate scaling laws. We also provide a rigorous convergence analysis, establishing theoretical guarantees for Nora in non-convex optimization landscapes. Remarkably, the core logic of Nora can be implemented in just two lines of matrix-based code, ensuring its seamless compatibility as a plug-and-play module for any existing optimizer. We evaluate Nora by training LLaMA models [23] of various scales across a range of learning rates. Our experimental results demonstrate that Nora achieves superior performance in both convergence speed and wall-clock execution time. Our contributions are three-fold: • We propose Nora, a novel optimizer that simultaneously achieves efficiency in preconditioning estimation, stability by respecting the scale invariance of neural networks, and high computational speed during training. • We evaluated Nora by training LLaMA models across various learning rates. Extensive experiments demonstrate that our algorithm achieves better results in terms of both training efficiency and computational wall-clock time. • We rigorously prove through theorems that Nora is a scalable optimizer and provide scaling criteria. We also provide rigorous convergence guarantees for Nora in non-convex optimization, demonstrating that our approach is characterized by theoretical completeness.
2
Notations
Throughout this paper, w ∈ Rm×n denotes the learnable matrix parameters, and f (w) denotes the loss function. Let wt , gt , and vt denote the parameters, gradients, and momentum at iteration t, respectively. The operator diag(·) keeps the diagonal entries of a square matrix and sets all offdiagonal entries to zero. We denote the (i, j)-th entry of the momentum matrix by vt,ij and its i-th row by vt,i: . We decompose momentum into v ∥ and v ⊥ , the components parallel and orthogonal to w. Similarly, v r∥ and v r⊥ denote row-wise parallel and orthogonal components, while v c∥ and v c⊥ denote the corresponding column-wise components. The symbol ⊗ denotes the Kronecker product. The notation (·)T denotes matrix transpose, and H −1 [·] denotes the action of a preconditioning operator on a vector or matrix.
3
Related Work
3.1
Matrix-Based Optimizers and Muon
Adaptive optimizers such as AdamW [24] treat parameters as vectors and scale gradients componentwise, thereby overlooking the intricate parameter interactions within the loss landscape f (w). Matrix2
based optimizers instead treat the weight w ∈ Rm×n , typically with m ≤ n, as a structured matrix [25]. Let gt = ∇f (wt ), and let vt = βvt−1 + (1 − β)gt denote the exponential moving average of T − 12 the gradient. Muon orthogonalizes this momentum and uses the update direction dM vt , t = (vt vt ) thereby applying a global structural preconditioner. To avoid the exact inverse square root, Muon uses the Newton–Schulz iteration Xk+1 = 12 Xk (3I − XkT Xk ) from a scaled initialization X0 ∝ vt . This update improves the conditioning of the search direction, but the dense matrix multiplications in the Newton–Schulz loop cost O(m2 n) [26], limiting scalability in large LLM layers. 3.2
Row-Momentum Normalized Preconditioning
Recent studies show that the layer-wise Hessian matrices of Transformers exhibit strong row-wise block diagonal dominance [27]. Muon can be written as applying the full-spectral preconditioner 1 HM = (vt vtT ) 2 ⊗ In , where vt is the momentum matrix. Row-Momentum Normalized Preconditioning (RMNP) uses the row block dominant structural prior by retaining only the diagonal blocks of 1 the Gram matrix: HR = (diag(vt vtT )) 2 ⊗ In . This approximation reduces matrix orthogonalization vt,i: −1 to row-wise ℓ2 normalization of the momentum: HR [vt ]i: = ∥vt,i: ∥2 . By eliminating the Newton– Schulz loop, RMNP reduces the per-iteration complexity from O(m2 n) to O(mn). However, this algebraic simplification neglects the scale symmetry of neural networks and allows unconstrained radial updates that inject noise into the training dynamics. 3.3
Scale-invariance in Neural Networks
The widespread use of BatchNorm, RMSNorm, and LayerNorm makes scale invariance a central property of modern neural networks [28]. For scale-invariant parameters w, the loss satisfies f (w) = f (λw) for any λ > 0. This property implies that the gradient is orthogonal to the weights: ⟨∇f (w), w⟩ = 0. Thus, the magnitude of w is decoupled from the loss value, and learning changes the angular orientation of the parameters [29, 30]. In this setting, the radial momentum component v ∥ interferes with angular optimization and perturbs the weight norm. This perturbation destabilizes the effective learning rate, defined as the relative angular step size, and hinders stable training [31]. Our goal is therefore to update along directions dt satisfying ⟨dt , wt ⟩ = 0, so that the update remains orthogonal to wt .
4
Nora Optimizer
4.1
Design Concept
Modern neural networks contain many scale-invariant parameters: changing ∥wt ∥F does not change the loss. Therefore, effective motion is angular motion rather than radial motion, a point sometimes overlooked by designers. Based on the principle of stability, the update direction should strictly reside within the tangent space of wt . This approach effectively filters out meaningless radial noise while preventing oscillations in the weight norm ∥wt ∥. Furthermore, tangential updates allow the learning rate schedule to maintain exclusive control over the angular velocity (i.e., the effective learning rate), since the parameter norm grows at a negligible rate and the preconditioning typically yields an update magnitude determined by the matrix dimensions Rm×n . Existing preconditioners do not satisfy this requirement. Even when applied to tangential momentum ⊥ ⊥ T − 12 ⊥ vt⊥ , the Muon operator produces an update dM vt that deviates from the tangent t = (vt (vt ) ) direction, so ⟨dM t , wt ⟩ ̸= 0. RMNP has the same issue: its preconditioner can destroy the orthogonality of the update to wt . SSO [32] preserves this geometric property, but it relies on a computationally intensive bisection method and is too slow in practice. Inspired by prior work on the row block diagonal dominance of the Transformer Hessian, we observe that preconditioning can be made compatible with tangential updates by considering vtr⊥ . Here, r⊥ vtr⊥ is obtained by orthogonalizing vt against each row of wt , ensuring that ⟨vt,i: , wt,i: ⟩ = 0 for every row i. Based on the principle of efficiency, we seek a preconditioner for vtr⊥ . Given the row block diagonal dominance of the Transformer Hessian, it follows that vtr⊥ (vtr⊥ )T also exhibits rowdiagonal dominance. We therefore adopt the diagonal elements of vtr⊥ (vtr⊥ )T as the preconditioner, 3
1
specifically HN = (diag(vtr⊥ (vtr⊥ )T )) 2 ⊗ In . Consequently, the preconditioned update simplifies 1 −1 r⊥ to HN [vt ] = (diag(vtr⊥ (vtr⊥ )T ))− 2 vtr⊥ . Through algebraic simplification, we arrive at: r⊥ vt,i:
1
−1 r⊥ HN [vt ]i: = ((diag(vtr⊥ (vtr⊥ )T ))− 2 vtr⊥ )i: =
r⊥ ∥ ∥vt,i: 2
.
(1)
In other words, applying the preconditioner HN [·] to vtr⊥ is mathematically equivalent to performing a simple row-wise normalization. Furthermore, the following Equation (2) leads to our long-sought objective: by preconditioning vtr⊥ (which is already row-orthogonal to wt ), we simultaneously achieve optimal efficiency while preserving the orthogonality between vtr⊥ and wt . This ensures that the update consistently adheres to the principle of stability. m X −1 r⊥ −1 r⊥ ⟨HN [vt ], wt ⟩ = ⟨HN [vt ]i: , wt,i: ⟩ i=1
=
m X i=1
⟨
1 r⊥ ∥ ∥vt,i: 2
r⊥ vt,i: , wt,i: ⟩ =
(2)
m X
1
i=1
r⊥ ∥ ∥vt,i: 2
−1 r⊥ ⊥ r⊥ ⟨vt,i: , wt,i: ⟩ = ⟨(HN [vt ]) , wt ⟩ = 0.
From the perspective of Speed, Nora maintains high efficiency and stability while introducing minimal computational overhead. The additional computation consists solely of row-wise projections of vt onto the row tangent space of wt , which are element-wise operations that consume virtually no additional time compared to RMNP. Furthermore, by exploiting the row block diagonal dominance of the Transformer Hessian, Nora replaces the computationally intensive Newton-Schulz (NS) iterations used in Muon with simple row-wise normalization. This effectively reduces the overall time complexity and ensures optimal training throughput. 1
Furthermore, updating along the direction of (diag(vtr⊥ (vtr⊥ )T ))− 2 vtr⊥ offers additional advantages. First, the discrete tangential updates ensure that the norm of each row increases only at a second-order rate. This growth is both monotonic—guaranteeing steady progression without erratic fluctuations—and gentle, remaining strictly controllable throughout training. Moreover, the row-wise orthogonality of dt relative to wt ensures a uniform growth rate across all rows of wt . By upholding the geometric consistency of scale-invariant parameters, this mechanism effectively mitigates internal covariate shift and maintains the representation balance of the neural network. It is also noteworthy that, due to the row-block diagonal dominance of the Transformer Hessian, Nora consistently performs row-wise normalization on the momentum vtr⊥ regardless of whether m ≤ n or m ≥ n. This stands in stark contrast to Muon, which restricts its Newton-Schulz iterations to the square matrix of the smaller dimension. The comprehensive procedure of the Nora algorithm is detailed in Algorithm 1. Furthermore, Table 1 presents a comparative summary of how various optimizers adhere to the three core design principles. Algorithm 1: The Nora Optimizer Require: Layer Weight wt ∈ Rm×n , momentum vt ∈ Rm×n , learning rate ηt at step t, momentum coefficient β, and weight decay coefficient λ (Default λ = 0). Initialize v0 ← 0 ∈ Rm×n , t ← 0. for each step do gt ← ∇f (wt ) vt ← βvt−1 + (1 − β)gt ⟨vt,i: ,wt,i: ⟩ r⊥ vt,i: ← vt,i: − ∥w {For all Row} 2 wt,i: t,i: ∥ 2
1
dt ← (diag(vtr⊥ (vtr⊥ )T ))− 2 vtr⊥ wt+1 ← wt − ηt (dt + λwt ) end for 4.2
{Row Normalization}
Theoretical Analysis
We now analyze Nora as a scalable optimizer and establish non-convex convergence guarantees. For a matrix x ∈ Rm×n , define: m X ∥x∥1,2 := ∥xi: ∥2 , ∥x∥∞,2 := max ∥xi: ∥2 . 1≤i≤m
i=1
4
Table 1: Comparison of Optimizers Based on the Three Design Principles Optimizer
Nora
Muon [4]
RMNP [7]
SSO [32]
SGD [33]
Adam [1]
AdamP [16]
Efficiency
✓
✓
✓
✓
✗
✗
✗
Stability
✓
✗
✗
✓
✗
✗
✓
Speed
✓
✗
✓
✗
✓
✓
✓
For any matrix w with nonzero rows, a matrix z is row-wise perpendicular to w if ⟨zi: , wi: ⟩ = 0 for every row i. The row-wise perpendicular projection is: [Pwr⊥ (x)]i: := xi: −
⟨xi: , wi: ⟩ wi: . ∥wi: ∥22
(3)
We use row-wise normalization with the convention 0/0 = 0: xi: , xi: ̸= 0, [RN(x)]i: := ∥xi: ∥2 0, x = 0.
(4)
i:
Throughout this subsection, the scalar m denotes the number of rows of w ∈ Rm×n , and vt denotes the momentum sequence. For t = 0, . . . , T − 1, we analyze the core Nora update without decoupled weight decay: gt = ∇f (wt ; ξt ), r⊥ vt+1 = Pwr⊥ (vt+1 ), t
4.2.1
vt+1 = βvt + (1 − β)gt , r⊥ dt = RN(vt+1 ),
v0 = 0, wt+1 = wt − ηdt .
(5)
Scaling for Nora
In this section, we address the most fundamental engineering question in LLM training: How should the learning rate η of Nora scale with the model width? According to the core principles of Maximal Update Parametrization (µP), an optimal optimizer must ensure that the change in hidden layer activations, ∆h, remains on the scale of Θ(1) as the network width increases (n → ∞). This ensures the network neither fails to learn features due to vanishing updates nor suffers from numerical explosion. We derive the rigorous learning rate formula required for Nora to satisfy this limit through the following theorem: Theorem 4.1 (Nora Scaling under the Scaling Hypothesis). Consider a neural network layer defined by h = wx, where the weights w ∈ Rm×n . Assume the input activation x ∈ Rn satisfies the scaling √ hypothesis under standard deep learning initialization: ∥x∥2 ≤ γ n, where γ = Θ(1) is a constant. √ Suppose the parameters are updated using Nora: wt+1 = wt − ηt dt . Then |∆hi | ≤ ηt γ n. To achieve the stable feature learning limit required by µP theory—specifically, to ensure the activation update magnitude √ satisfies |∆hi | = Θ(1)—the learning rate ηt for Nora must follow the width-scaling rule, ηt ∝ 1/ n. √ The theorem provides an upper bound on the forward update, |∆hi | ≤ ηt γ n. This bound ensures training stability but raises a more precise question: do Nora’s projection and normalization operations substantially weaken the effective gradient components in practice, causing ∆hi to fall significantly √ below n or even vanish? To show that Nora achieves Maximal Update Parametrization (µP), we must evaluate the true order of ∆hi in the high-dimensional limit n → ∞. To this end, we use the following theorem. Theorem 4.2 (Asymptotic Convergence). Consider a hidden layer h = wx, where w ∈ Rm×n . Under the standard infinite-width random initialization hypothesis, assume the row vectors of w 2 follow wi: ∼ N (0, σw In /n), and the components of the input activation x ∈ Rn are independent with zero mean and variance σx2 = Θ(1). Let the error signal backpropagated to this layer be δi = Θ(1), and let the vanilla gradient without momentum be gi: = δi x⊤ . If Nora generates the update direction di: = RN(Pwr⊥ (gi: )), then as the network width n → ∞, the inner product between √ p the update direction and the input converges in probability to ⟨di: , x⟩ − → sgn(δi√ )σx n. Therefore, to obtain non-trivial feature learning with ∆hi = Θ(1), Nora must use η = η0 / n. 5
4.2.2
Convergence Analysis
We next prove non-convex convergence guarantees for Nora. Define the natural filtration as: Ft := σ(w0 , ξ0 , . . . , ξt−1 ). It is the sigma-algebra generated before sampling ξt . Equivalently, wt and vt are Ft -measurable, and the conditional expectations below are taken only over the current stochastic gradient. We measure stationarity by the row-wise projected gradient: Gt := Pwr⊥ ∇f (wt ) . (6) t This is the relevant first-order signal for Nora because each row of dt is perpendicular to the corresponding row of wt . Hence its descent inner product depends only on the row-wise perpendicular component of the true gradient, not on the radial component removed by the projection. Under row-wise scale invariance, Gt coincides with ∇f (wt ), so the projected stationarity measure reduces to the standard first-order one. Compared with RMNP, Nora inserts one row-wise perpendicular projection before row-wise normalization. This additional step preserves the RMNP proof structure; the main new ingredient is the non-expansiveness of Pwr⊥ in the norms used below. t Assumption 4.3 (Smoothness). The objective f : Rm×n → R satisfies one of the following conditions: (a) (Frobenius smoothness) There exists LF > 0 such that, for all w, w′ , ∥∇f (w) − ∇f (w′ )∥F ≤ LF ∥w − w′ ∥F .
(7)
(b) (Matched (∞, 2)-smoothness) There exists L∞,2 > 0 such that, for all w, w′ , ∥∇f (w) − ∇f (w′ )∥1,2 ≤ L∞,2 ∥w − w′ ∥∞,2 . (8) Theorem 4.4 (Nora under matched (∞, 2)-smoothness). Suppose Assumptions 4.3(b) hold. If Nora uses a constant step size ηt = η, then: s " # √ T −1 1 X ∆ 1 L∞,2 ηβ mσ 1 − β L∞,2 η E ∥Gt ∥1,2 ≤ +2 1− + √ + . (9) T t=0 Tη T 1−β 1 + β 2 B Here m is the row dimension of w. With the choice: s (p ) L∞,2 ∆ (1 − β)∆ , 1 − β = min (10) η= √ √ ,1 , L∞,2 T 2 mσ T Nora reaches an ϵ-stationary point in the projected ∥ · ∥1,2 sense after: T = O mL∞,2 σ 2 ∆ ϵ−4 (11) iterations. Proposition 4.5 (Frobenius-smooth counterparts). Under the same conditions of Theorem 4.4, if Nora uses a constant step size ηt = η, then: s " # √ T −1 √ ∆ 1 LF η mβ 1 X σ 1−β LF ηm E ∥Gt ∥F ≤ + ( m + 1) 1 − +√ + . (12) T t=0 Tη T 1−β 2 B 1+β The corresponding projected ∥ · ∥1,2 bound is: s " # √ T −1 1 X ∆ 1 LF ηmβ mσ 1 − β LF ηm E ∥Gt ∥1,2 ≤ +2 1− + √ + . T t=0 Tη T 1−β 1+β 2 B
(13)
Consequently, Nora reaches an ϵ-stationary point in either projected measure after: T = O m2 LF σ 2 ∆ ϵ−4 (14) iterations. Corollary 4.6 (Standard first-order stationarity under row-wise scale invariance). Assume, in addition, that f is row-wise scale invariant: f (Dw) = f (w), ∀w ∈ Rm×n , ∀D ≻ 0 diagonal. (15) Then ∇f (wt ) is row-wise perpendicular to wt for all t, i.e., Pwr⊥ ∇f (wt ) = ∇f (wt ). (16) t Therefore, Theorem 4.4 and Proposition 4.5 hold with Gt replaced by ∇f (wt ). 6
Proof sketch. The proof follows the RMNP descent argument and adds one projection lemma: for fixed wt , the row-wise perpendicular projector Pwr⊥ is non-expansive in ∥ · ∥F , ∥ · ∥1,2 , and ∥ · ∥∞,2 . t Because each row of dt is perpendicular to the corresponding row of wt , ⟨∇f (wt ), dt ⟩ = Pwr⊥ (∇f (wt )), dt = ⟨Gt , dt ⟩. t
(17)
Define the projected momentum tracking error as: r⊥ et := vt+1 − Gt .
The row-normalization identities give: ⟨Gt , dt ⟩ ≥ ∥Gt ∥1,2 − 2∥et ∥1,2 ,
√ ⟨Gt , dt ⟩ ≥ ∥Gt ∥F − ( m + 1)∥et ∥F .
(18)
Combining these inequalities with smoothness-based descent and the standard momentum recursion yields the stated bounds. Full proofs and standard assumptions are deferred to Appendix A. 4.3
Compared with Mano
In this section, we delineate the distinctions between Mano [34] and Nora. Mano alternates between row-wise and column-wise perpendicular projections, followed by normalization. Specifically, Mano accumulates momentum vt in the standard manner. During odd iterations, it removes the row-wise radial component to obtain vtr⊥ and applies row normalization: dt = RN(vtr⊥ ). During even iterations, it removes the column-wise radial component to obtain vtc⊥ and applies column normalization: dt = CN(vtc⊥ ). The design philosophy of Nora diverges significantly from that of Mano. Nora originates from the properties expected by the optimizer itself, while Mano comes from Riemannian optimization [35, 36, 37]. In practical implementation, Nora is more streamlined as it eliminates the need to track iteration parity (odd vs. even steps). More importantly, Nora is explicitly designed to leverage the row block diagonal dominance characteristic of the Transformer Hessian. By removing the heuristic column-wise orthogonal projection and normalization used in Mano, Nora demonstrates superior properties in ablation studies. These empirical results suggest that the row block diagonal dominance of the Transformer Hessian is a critical structural prior that must be prioritized.
5
Experiments
We evaluate Nora on autoregressive language modeling with LLaMA-style Transformer models at two scales: 60M and 135M parameters. We compare Nora with three matrix-based optimizers, Muon, Mano, and RMNP, under the same data pipeline, tokenizer, context length, global batch size, learning-rate schedule, precision, evaluation cadence, and checkpointing cadence. Detailed model and training configurations are deferred to Appendix B. The compared optimizers differ only in the optimizer rule, the matrix learning-rate sweep grid, and weight decay. For all non-matrix parameter groups, we use the same auxiliary Adam setting within each model size. For matrix-shaped parameters, we tune the matrix learning rate over optimizerspecific grids and select the best run by validation loss. Muon uses a slightly different sweep range from the other matrix optimizers because it is more sensitive to large matrix learning rates in our preliminary runs. The full sweep grids are listed in Appendix B. Importantly, Nora uses weight decay 0 in all runs, whereas Muon, Mano, and RMNP use weight decay 0.1. We treat this as part of Nora’s optimizer configuration and report it explicitly in all result tables. 5.1
Best Results from Different Optimizers
Table 2 reports the main language-modeling results. For each optimizer and model size, we report the matrix learning rate selected from the sweep, the validation loss at the selected setting, and validation perplexity. This format separates peak validation performance from the final checkpoint quality. Figure 1 shows the 135M training dynamics from 10k to 20k steps. RMNP has the strongest earlystage loss and perplexity, while Muon improves more slowly and plateaus at a higher value. Mano and Nora start from higher loss and perplexity, but both continue to improve after 15k steps. Nora shows the clearest late-stage improvement: it overtakes the other methods near the end of training 7
Table 2: Main language-modeling results. For each optimizer and model size, the matrix learning rate is selected by validation loss over the sweep grid in Appendix B. Lower validation loss and perplexity are better. Best results are shown in bold. Optimizer Muon Mano RMNP Nora
60M
Weight decay
135M
Best matrix LR
Val. loss
Val. ppl.
Best matrix LR
Val. loss
Val. ppl.
0.01 0.004 0.004 0.004
3.44 3.39 3.41 3.37
31.09 29.55 30.12 28.94
0.01 0.003 0.01 0.003
3.142 3.097 3.112 3.079
23.17 22.13 22.46 21.74
0.1 0.1 0.1 0.0
Table 3: Comparison of Mano and Nora (Model: 135M, Weight decay = 0) 0.003 0.005 0.01 0.02 Optimizer ppl loss ppl loss ppl loss ppl loss Mano 22.13 3.097 22.14 3.098 23.34 3.150 25.07 3.222 Nora 21.74 3.079 21.86 3.085 22.43 3.111 23.39 3.152
and reaches the lowest final loss and perplexity. This indicates that Nora uses the same training budget more efficiently, achieving better optimization quality without requiring additional training steps. This trend is consistent with Table 2, where Nora achieves the best 135M validation loss and perplexity among the compared optimizers. 5.2
Runtime Comparison
We also compare the cost of row normalization with Newton–Schulz orthogonalization on representative matrix shapes from LLaMA-style models. The benchmark uses CUDA with bfloat16 precision, 20 warmup iterations, and 200 measured iterations. Newton–Schulz uses five iterations. As shown in Table 4, row normalization is consistently cheaper than Newton–Schulz. The gap is already about one order of magnitude for small and medium matrices, and becomes much larger for the 1B-scale MLP matrices, where Newton–Schulz is more than 70 times slower. This highlights the substantial speed advantage of row normalization in practical training settings. Results support the practical motivation for replacing matrix-level orthogonalization with a row-wise normalization step. Table 4: Runtime comparison between row normalization and five-step Newton–Schulz orthogonalization on representative LLaMA-style matrix shapes. Times are mean kernel runtimes in milliseconds under CUDA and bfloat16 precision. Model scale
Matrix shape
Representative layer
Row normalization (ms)
NS(5) (ms)
NS / row-norm
60M 60M 60M 135M 135M 135M 350M 350M 350M 1B 1B 1B
512 × 512 1376 × 512 512 × 1376 768 × 768 2048 × 768 768 × 2048 1024 × 1024 2816 × 1024 1024 × 2816 2048 × 2048 5461 × 2048 2048 × 5461
attention: hidden × hidden MLP: intermediate × hidden MLP: hidden × intermediate attention: hidden × hidden MLP: intermediate × hidden MLP: hidden × intermediate attention: hidden × hidden MLP: intermediate × hidden MLP: hidden × intermediate attention: hidden × hidden MLP: intermediate × hidden MLP: hidden × intermediate
0.0689 0.0682 0.0688 0.0686 0.0692 0.0691 0.0674 0.0687 0.0675 0.0684 0.0985 0.1084
0.6554 0.6853 0.6623 0.6520 0.6871 0.6534 0.6397 0.8625 0.6941 2.0552 6.9985 7.9678
9.52× 10.05× 9.63× 9.50× 9.93× 9.46× 9.49× 12.56× 10.28× 30.06× 71.02× 73.51×
5.3
Ablation Experiment between Nora and Mano
To verify that row-normalization indeed captures the row-diagonal dominance inherent in the Transformer Hessian, we conducted an ablation study. Given that Nora and Mano share similar algorithmic frameworks, it is crucial to distinguish their structural advantages from hyperparameter effects. While the previous experiments utilized Mano’s recommended weight decay = 0.1 against Nora’s default weight decay = 0, we eliminated this discrepancy in the ablation study by setting Mano’s 8
3.4
Mano RMNP Muon Nora
27 26
Perplexity
3.3
Loss
28
Mano RMNP Muon Nora
3.2
3.1
25 24 23 22
3.0 10k
12k
15k
17k
10k
20k
12k
15k
17k
20k
Training Steps
Training Steps
Figure 1: Training dynamics on the 135M model. Left: loss over training steps. Right: perplexity over training steps. Nora continues to improve late in training and finishes with the lowest loss and perplexity. weight decay to 0. As shown in Table 3, the results demonstrate that Nora’s superiority stems from its intrinsic algorithmic architecture rather than an advantage gained from weight decay settings. Mano wd=0
Loss, lr=0.003 3.4
Nora
Loss, lr=0.01
Loss, lr=0.02
3.4
3.3
Loss
Mano wd=0.1
Loss, lr=0.005 3.4
3.3
3.2
3.4
3.2 3.2
3.1
3.1
3.0
3.0 10k
12k
17k
20k
10k
12k
Perplexity, lr=0.003
28
Perplexity
15k
3.2
15k
17k
20k
10k
12k
Perplexity, lr=0.005
15k
17k
20k
Perplexity, lr=0.01
28
30.0
26
27.5
24
25.0
10k
12k
15k
17k
20k
17k
20k
Perplexity, lr=0.02
35
26 30
24 25 22
22 10k
12k
15k
17k
Training Steps
20k
22.5 10k
12k
15k
17k
20k
Training Steps
10k
12k
15k
Training Steps
17k
20k
10k
12k
15k
Training Steps
Figure 2: Training dynamics on the 135M model. This figure illustrates the perplexity and loss decay curves for Nora (default weight decay = 0) in comparison with Mano (under weight decay = 0 and weight decay = 0.1). Furthermore, Figure 2 illustrates the perplexity and loss decay curves for both Mano and Nora at the 135M model scale, under configurations of weight decay = 0 and weight decay = 0.1. It is evident that Nora consistently outperforms Mano across all settings. For a comprehensive overview, all remaining experimental results and corresponding convergence curves are provided in Appendix B, where we focus on the sensitivity of hyper-parameters and stability of our proposed optimizer.
6
Conclusion
We introduce Nora, a normalized orthogonal row-alignment optimizer for scalable LLM training. Nora is motivated by the geometry of scale-invariant neural networks, where effective learning should mainly occur along angular rather than radial directions. By projecting momentum onto the row-wise orthogonal complement of the weights and applying row-wise normalization, Nora preserves stable scale-invariant dynamics while retaining efficient Muon-like preconditioning under the row-diagonal structure of Transformer Hessians. We further establish Nora’s scalability through µP-based width-scaling analysis and provide non-convex convergence guarantees. Experiments on LLaMA-style models show that Nora achieves the best validation loss and perplexity among the compared matrix-based optimizers, while its row-wise normalization brings a clear speed advantage over Newton–Schulz orthogonalization. Overall, Nora demonstrates that efficiency, stability, and speed can be unified through a simple row-wise geometric principle, offering a practical and principled direction for scalable optimizer design.
9
References [1] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. [2] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024. [3] Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for llm training. arXiv preprint arXiv:2502.16982, 2025. [4] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. [5] Günther Schulz. Iterative berechung der reziproken matrix. ZAMM-Journal of Applied Mathematics and Mechanics/Zeitschrift für Angewandte Mathematik und Mechanik, 13(1):57–59, 1933. [6] Zhaorui Dong, Yushun Zhang, Jianfeng Yao, and Ruoyu Sun. Towards quantifying the hessian structure of neural networks. In OPT 2025: Optimization for Machine Learning. [7] Shenyang Deng, Zhuoli Ouyang, Tianyu Pang, Zihang Liu, Ruochen Jin, Shuhua Yu, and Yaoqing Yang. Rmnp: Row-momentum normalized preconditioning for scalable matrix-based optimization. arXiv preprint arXiv:2603.20527, 2026. [8] Johanni Brea, Berfin Simsek, Bernd Illing, and Wulfram Gerstner. Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape. arXiv preprint arXiv:1907.02911, 2019. [9] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448–456. pmlr, 2015. [10] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in neural information processing systems, 32, 2019. [11] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. [12] Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics: Learning dynamics of normalized neural network using sgd and weight decay. Advances in Neural Information Processing Systems, 34:6380–6391, 2021. [13] Jh Yuan and Feiping Nie. Spherical cautious optimizers. In Workshop on Scientific Methods for Understanding Deep Learning, 2026. [14] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016. [15] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International conference on machine learning, pages 4596–4604. PMLR, 2018. [16] Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung-Woo Ha. Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights. arXiv preprint arXiv:2006.08217, 2020. [17] Hao Chen, Jh Yuan, and Hanmin Zhang. Decoupled orthogonal dynamics: Regularization for deep network optimizers. In Workshop on Scientific Methods for Understanding Deep Learning, 2026. 10
[18] Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(12):9508–9520, 2024. [19] Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 33:18795–18806, 2020. [20] Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Twelfth International Conference on Learning Representations, 2024. [21] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842–1850. PMLR, 2018. [22] Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems, 34:17084–17097, 2021. [23] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023. [24] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. [25] Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. arXiv preprint arXiv:2509.02046, 2025. [26] Nicholas J Higham. Functions of matrices: theory and computation. SIAM, 2008. [27] Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why transformers need adam: A hessian perspective. Advances in neural information processing systems, 37:131786–131823, 2024. [28] Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro. Path-sgd: Path-normalized optimization in deep neural networks. Advances in neural information processing systems, 28, 2015. [29] Kaiyue Wen, Xingyu Dang, Kaifeng Lyu, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them 2.1: Hyperball optimization, 12 2025. [30] Liliang Ren, Yang Liu, Yelong Shen, and Weizhu Chen. Rethinking language model scaling under transferable hypersphere optimization. arXiv preprint arXiv:2603.28743, 2026. [31] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020. [32] Tian Xie, Haoming Luo, Haoyu Tang, Yiwen Hu, Jason Klein Liu, Qingnan Ren, Yang Wang, Wayne Xin Zhao, Rui Yan, Bing Su, et al. Controlled llm training on spectral sphere. arXiv preprint arXiv:2601.08393, 2026. [33] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. nature, 323(6088):533–536, 1986. [34] Yufei Gu and Zeke Xie. Mano: Restriking manifold optimization for llm training, 2026. [35] Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023. 11
[36] Jh Yuan, Fangyuan Xie, Feiping Nie, and Xuelong Li. Riemannian optimization on relaxed indicator matrix manifold. In The Fourteenth International Conference on Learning Representations, 2026. [37] Jh Yuan, Zhuo Liu, and Feiping Nie. Riemannian fuzzy k-means on product manifolds. In Non-Euclidean Foundation Models: Advancing AI Beyond Euclidean Frameworks, 2025. [38] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1–67, 2020.
12
Contents 1
Introduction
1
2
Notations
2
3
Related Work
2
3.1
Matrix-Based Optimizers and Muon . . . . . . . . . . . . . . . . . . . . . . . . .
2
3.2
Row-Momentum Normalized Preconditioning . . . . . . . . . . . . . . . . . . . .
3
3.3
Scale-invariance in Neural Networks . . . . . . . . . . . . . . . . . . . . . . . . .
3
4
Nora Optimizer
3
4.1
Design Concept . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
3
4.2
Theoretical Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
4
4.2.1
Scaling for Nora . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
5
4.2.2
Convergence Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . .
6
Compared with Mano . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
7
4.3 5
6
Experiments
7
5.1
Best Results from Different Optimizers . . . . . . . . . . . . . . . . . . . . . . .
7
5.2
Runtime Comparison . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
8
5.3
Ablation Experiment between Nora and Mano . . . . . . . . . . . . . . . . . . . .
8
Conclusion
9
A Proof of Theorem
14
A.1 Proof of Theorem 4.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
14
A.2 Proof of Theorem 4.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
14
A.3 Preliminaries for Theorem 4.7–Corollary 4.9 . . . . . . . . . . . . . . . . . . . . .
15
A.3.1 Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
15
A.3.2 Auxiliary Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
16
A.4 Proof of Theorem 4.7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
20
A.5 Proof of Proposition 4.8 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
21
A.6 Proof of Corollary 4.9 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
23
B Additional Experimental Details
24
B.1 Model and Training Configurations . . . . . . . . . . . . . . . . . . . . . . . . . .
24
B.2 Matrix Learning-rate Sweep Grids . . . . . . . . . . . . . . . . . . . . . . . . . .
24
B.3 Other Experimental Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
24
C Running and Reference Code
26
C.1 Run Quickly . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
26
C.2 Reference Code . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
26
13
A
Proof of Theorem
A.1
Proof of Theorem 4.1
Theorem 4.1 (Nora Scaling under the Scaling Hypothesis) Consider a neural network layer defined by h = wx, where the weights w ∈ Rm×n . Assume the input activation x ∈ Rn satisfies the scaling √ hypothesis under standard deep learning initialization: ∥x∥2 ≤ γ n, where γ =√Θ(1) is a constant. Suppose Nora updates the parameters by wt+1 = wt − ηt dt . Then |∆hi | ≤ ηt γ n. To achieve the stable feature learning limit required by µP theory—specifically, to ensure the activation update magnitude satisfies |∆hi | = Θ(1), the learning rate ηt for Nora must follow the width-scaling rule, ηt ∝ √1n . Proof. For any output neuron i, the change in the hidden activation ∆hi during the forward pass is: ∆hi = (wt+1,i: − wt,i: )x = −ηt ⟨dt,i: , x⟩
(19)
By taking the absolute value and applying the Cauchy-Schwarz inequality, we obtain: |∆hi | = ηt |⟨(dt )i: , x⟩| ≤ ηt ∥(dt )i: ∥2 ∥x∥2
(20)
Nora’s row-wise normalization strictly ensures ∥(dt )i: ∥2 ≤ 1. Combined with the input assumption √ ∥x∥2 ≤ γ n, substituting these into the inequality yields a rigorous upper bound for the forward update: √ √ (21) |∆hi | ≤ ηt · 1 · γ n = ηt γ n To ensure that the feature update magnitude |∆hi | neither diverges nor vanishes as n → √ ∞ (i.e., satisfies the Maximal Update condition |∆hi | = Θ(1)), we require the upper bound ηt γ n to also be of order Θ(1). Given that γ is a constant, this necessarily implies ηt = √cn , where c is a constant independent of n. This completes the proof. A.2
Proof of Theorem 4.2
Theorem 4.2 (symptotic Convergence)) Consider a hidden layer h = wx, where w ∈ Rm×n . Under the standard infinite-width random initialization hypothesis, assume the row vectors of w σ2 follow wi: ∼ N (0, nw In ), and the components of the input activation x ∈ Rn are independent with zero mean and variance σx2 = Θ(1). Let the error signal backpropagated to this layer be δi = Θ(1), and the vanilla gradient (without momentum) be gi: = δi x⊤ . If Nora is applied to generate the update direction di: = RN(Pwr⊥ (gi: )), then as the network width n → ∞, the inner product between √ p the update direction and the input converges in probability to ⟨di: , x⟩ − → sgn(δi )σx n. To achieve non-trivial feature learning such that ∆hi = Θ(1) as n → ∞, the Nora learning rate must scale as η = √η0n . Proof. For output neuron i, Nora first removes the row-wise radial component of the gradient. Writing row vectors as columns inside inner products, this gives: ui: = δi x −
δi ⟨x, wi: ⟩ wi: . ∥wi: ∥22
(22)
2 We now compute the high-dimensional orders. Since xj has variance σx2 and wi,j has variance σw /n, the weak law of large numbers gives: √ √ p ∥x∥22 − → nσx2 =⇒ ∥x∥2 = σx n + op ( n). (23)
Similarly: p
2 ∥wi: ∥22 − → σw = Θ(1).
(24)
Because x and w are independent and centered, ⟨x, wi: ⟩ is a sum of n independent products with 2 total variance σx2 σw . The central limit theorem gives: d
2 ⟨x, wi: ⟩ − → N (0, σx2 σw ) = Op (1).
14
(25)
Let ri: denote the radial component removed by the projection. Its Euclidean norm satisfies: δi ⟨x, wi: ⟩ |δi ||Op (1)| ∥wi: ∥2 = σw = Op (1). (26) 2 ∥wi: ∥22 σw √ The primary gradient term δi x has norm Θ( n), while the removed radial component is Op (1). Hence the projected vector satisfies: √ ∥ui: ∥2 = ∥δi x − ri: ∥2 = |δi |∥x∥2 + Op (1) = |δi |σx n(1 + op (1)). (27) ∥ri: ∥2 =
Row-wise normalization gives: ui: δ x − ri: √i = . (28) ∥ui: ∥2 |δi |σx n(1 + op (1)) The forward update depends on the inner product: δi ⟨x, x⟩ − ⟨ri: , x⟩ √ ⟨di: , x⟩ = . (29) |δi |σx n(1 + op (1)) √ Since ⟨x, x⟩ = nσx2 + op (n) and |⟨ri: , x⟩| ≤ ∥ri: ∥2 ∥x∥2 = Op ( n), √ √ √ δi nσx2 − Op ( n) √ ⟨di: , x⟩ = = sgn(δi )σx n + op ( n). (30) |δi |σx n(1 + op (1)) √ √ p Thus ⟨di: , x⟩ − → sgn(δi )σx n holds. The actual forward update is ∆hi = −η⟨di: , x⟩ = ∓ησx n. To achieve ∆hi = Θ(1) for non-trivial feature learning, it is necessary that η = Θ(n−1/2 ). di: =
A.3 A.3.1
Preliminaries for Theorem 4.7–Corollary 4.9 Setup
For the convergence analysis, let m denote the number of rows in w ∈ Rm×n . We reserve vt for the momentum sequence. For any matrix w with nonzero rows, define the row-wise perpendicular projection: ⟨xi: , wi: ⟩ [Pwr⊥ (x)]i: := xi: − wi: . (31) ∥wi: ∥22 This projection removes the row-wise radial component: ⟨[Pwr⊥ (x)]i: , wi: ⟩ = 0 for every row i. Equivalently, a matrix z is row-wise perpendicular to w when ⟨zi: , wi: ⟩ = 0 for all i. For t = 0, . . . , T − 1, Nora is indexed as: gt = ∇f (wt ; ξt ),
vt+1 = βvt + (1 − β)gt , v0 = 0, (32) r⊥ dt = RN vt+1 , wt+1 = wt − ηdt . (33) Thus dt is the row-wise normalization of the component of the momentum that is row-wise perpendicular to wt . We use the convention 0/0 = 0 in RN, applied row-wise: xi: , xi: ̸= 0, [RN(x)]i: = ∥xi: ∥2 (34) 0, xi: = 0. r⊥ vt+1 := Pwr⊥ (vt+1 ), t
For a matrix x ∈ Rm×n , define: ∥x∥1,2 :=
m X
∥xi: ∥2 ,
∥x∥∞,2 := max ∥xi: ∥2 . 1≤i≤m
i=1
(35)
We also define the natural filtration: Ft := σ(w0 , ξ0 , . . . , ξt−1 ),
(36)
so that wt is Ft -measurable and gt is sampled conditionally on Ft . Finally, define the projected gradient and momentum-tracking errors: r⊥ Gt := Pwr⊥ ∇f (wt ) , et := vt+1 − ∇f (wt ), ēt := vt+1 − Gt = Pwr⊥ (et ). (37) t t The projected gradient Gt is the relevant stationarity measure because Nora uses directions that are row-wise perpendicular to wt . Therefore radial components of ∇f (wt ) do not contribute to the descent inner product with dt . 15
A.3.2
Auxiliary Lemmas
Lemma A.1 (Row-wise projection is non-expansive). For any w ∈ Rm×n with nonzero rows and any x ∈ Rm×n , ∥Pwr⊥ (x)∥F ≤ ∥x∥F ,
∥Pwr⊥ (x)∥1,2 ≤ ∥x∥1,2 ,
∥Pwr⊥ (x)∥∞,2 ≤ ∥x∥∞,2 .
(38)
Proof. For each row i, let Pi := In −
wi:⊤ wi: . ∥wi: ∥22
(39)
Then [Pwr⊥ (x)]i: = xi: Pi . Since Pi is an orthogonal projector on Rn , ∥xi: Pi ∥2 ≤ ∥xi: ∥2 .
(40)
Summing over rows gives the Frobenius and ∥ · ∥1,2 bounds, while taking the maximum over rows gives the ∥ · ∥∞,2 bound. Lemma A.2 (Basic geometry of Nora). Let z = Pwr⊥ (x) and d = RN(z). Then: 1. ⟨di: , wi: ⟩ = 0 for every row i; √ 2. ∥d∥F ≤ m; 3. ∥d∥∞,2 ≤ 1; 4. ⟨z, d⟩ = ∥z∥1,2 ; 5. ⟨z, d⟩ ≥ ∥z∥F . Proof. Each row zi: is orthogonal to wi: . Every nonzero row of d is a scalar multiple of zi: , so it is also orthogonal to wi: . Thus d is row-wise perpendicular to w. For every row, ∥di: ∥2 ≤ 1. Thus: ∥d∥2F =
m X
∥di: ∥22 ≤ m,
∥d∥∞,2 = max ∥di: ∥2 ≤ 1. i
i=1
Moreover, ⟨z, d⟩ =
m X i=1
zi: zi: , ∥zi: ∥2
∥zi: ∥2 ≥
X
=
m X
∥zi: ∥2 = ∥z∥1,2 .
(41)
(42)
i=1
Finally, !1/2 ∥z∥1,2 =
X i
∥zi: ∥22
= ∥z∥F ,
(43)
i
which gives ⟨z, d⟩ ≥ ∥z∥F . Lemma A.3 (Descent under Frobenius smoothness). Under Assumption 4.3(a), for every t, f (wt ) − f (wt+1 ) ≥ η⟨∇f (wt ), dt ⟩ −
LF η 2 m . 2
Proof. Let ∆t := wt+1 − wt . By the fundamental theorem of calculus, Z 1 f (wt + ∆t ) − f (wt ) = ⟨∇f (wt + s∆t ), ∆t ⟩ ds.
(44)
(45)
0
Subtracting ⟨∇f (wt ), ∆t ⟩ and applying Cauchy–Schwarz gives: Z 1 f (wt + ∆t ) − f (wt ) − ⟨∇f (wt ), ∆t ⟩ = ⟨∇f (wt + s∆t ) − ∇f (wt ), ∆t ⟩ ds
(46)
0
Z 1 ≤
∥∇f (wt + s∆t ) − ∇f (wt )∥F ∥∆t ∥F ds. 0
16
(47)
Assumption 4.3(a) implies: ∥∇f (wt + s∆t ) − ∇f (wt )∥F ≤ LF s∥∆t ∥F .
(48)
Hence the standard quadratic upper bound follows from the stated smoothness assumption: LF ∥∆t ∥2F . 2
(49)
LF η 2 ∥dt ∥2F . 2
(50)
f (wt + ∆t ) ≤ f (wt ) + ⟨∇f (wt ), ∆t ⟩ + Substituting ∆t = −ηdt gives: f (wt+1 ) ≤ f (wt ) − η⟨∇f (wt ), dt ⟩ + Using Lemma A.2(ii), ∥dt ∥2F ≤ m, which yields the claim.
Lemma A.4 (Descent under matched (∞, 2)-smoothness). Under Assumption 4.3(b), for every t, f (wt ) − f (wt+1 ) ≥ η⟨∇f (wt ), dt ⟩ −
L∞,2 η 2 . 2
Proof. Let ∆t := wt+1 − wt = −ηdt . By the fundamental theorem of calculus, Z 1 f (wt + ∆t ) − f (wt ) = ⟨∇f (wt ), ∆t ⟩ + ⟨∇f (wt + s∆t ) − ∇f (wt ), ∆t ⟩ ds.
(51)
(52)
0
For any matrices A and B, the (1, 2)–(∞, 2) duality bound gives: |⟨A, B⟩| ≤ ∥A∥1,2 ∥B∥∞,2 .
(53)
|⟨∇f (wt + s∆t ) − ∇f (wt ), ∆t ⟩| ≤ ∥∇f (wt + s∆t ) − ∇f (wt )∥1,2 ∥∆t ∥∞,2
(54)
Therefore, by Assumption 4.3(b), ≤ L∞,2 s∥∆t ∥2∞,2 . Lemma A.2(iii) gives ∥∆t ∥∞,2 = η∥dt ∥∞,2 ≤ η. Hence: Z 1 L∞,2 η 2 . f (wt+1 ) − f (wt ) ≤ −η⟨∇f (wt ), dt ⟩ + L∞,2 sη 2 ds = −η⟨∇f (wt ), dt ⟩ + 2 0
(55)
(56)
Rearranging proves the claim. Lemma A.5 (Projected inner-product lower bound in Frobenius norm). For every t, √ ⟨∇f (wt ), dt ⟩ ≥ ∥Gt ∥F − ( m + 1)∥et ∥F .
(57)
Proof. Since dt is row-wise perpendicular to wt by Lemma A.2(i), and ∇f (wt ) − Gt is row-wise parallel to wt , we have ⟨∇f (wt ), dt ⟩ = ⟨Gt , dt ⟩. (58) r⊥ Since Gt = vt+1 − ēt , r⊥ ⟨Gt , dt ⟩ = ⟨vt+1 , dt ⟩ − ⟨ēt , dt ⟩.
(59)
By Lemma A.2(v), r⊥ r⊥ ⟨vt+1 , dt ⟩ ≥ ∥vt+1 ∥F . By Cauchy–Schwarz and Lemma A.2(ii), √ |⟨ēt , dt ⟩| ≤ ∥ēt ∥F ∥dt ∥F ≤ m ∥ēt ∥F .
Also, Combining the above,
(60) (61)
r⊥ ∥vt+1 ∥F = ∥Gt + ēt ∥F ≥ ∥Gt ∥F − ∥ēt ∥F .
(62)
√ ⟨∇f (wt ), dt ⟩ ≥ ∥Gt ∥F − ( m + 1)∥ēt ∥F .
(63)
(et )∥F ≤ ∥et ∥F . ∥ēt ∥F = ∥Pwr⊥ t
(64)
Finally, by Lemma A.1, This proves the claim. 17
Lemma A.6 (Projected inner-product lower bound in (1, 2)-norm). For every t, ⟨∇f (wt ), dt ⟩ ≥ ∥Gt ∥1,2 − 2∥et ∥1,2 .
(65)
r⊥ ⟨∇f (wt ), dt ⟩ = ⟨Gt , dt ⟩ = ⟨vt+1 , dt ⟩ − ⟨ēt , dt ⟩.
(66)
r⊥ r⊥ ⟨vt+1 , dt ⟩ = ∥vt+1 ∥1,2 .
(67)
|⟨ēt , dt ⟩| ≤ ∥ēt ∥1,2 ∥dt ∥∞,2 ≤ ∥ēt ∥1,2 .
(68)
r⊥ ∥vt+1 ∥1,2 = ∥Gt + ēt ∥1,2 ≥ ∥Gt ∥1,2 − ∥ēt ∥1,2 .
(69)
⟨∇f (wt ), dt ⟩ ≥ ∥Gt ∥1,2 − 2∥ēt ∥1,2 .
(70)
∥ēt ∥1,2 ≤ ∥et ∥1,2 ,
(71)
Proof. As in Lemma A.5,
By Lemma A.2(iv), By duality and Lemma A.2(iii),
Also, Hence, Using Lemma A.1, and the proof is complete. Assumption A.7 (Unbiased stochastic gradients). For all t, E[gt | Ft ] = ∇f (wt ).
(72)
Assumption A.8 (Bounded variance). There exists σ > 0 such that, for all t, σ2 , E ∥gt − ∇f (wt )∥2F | Ft ≤ B
(73)
where B is the batch size. Assumption A.9 (Lower bounded objective). The objective is bounded below by f ⋆ . We write ∆ := f (w0 ) − f ⋆ . Lemma A.10 (Momentum tracking under Frobenius smoothness). Under Assumptions 4.3(a), A.7, and A.8, s √ T −1 X 1−β LF η mβ σ E ∥et ∥F ≤ (T − 1) +T√ . (74) 1 − β 1+β B t=0 Proof. Let ζt := gt − ∇f (wt ). Then E[ζt | Ft ] = 0
(75)
σ2 E ∥ζt ∥2F | Ft ≤ B
(76)
et = vt+1 − ∇f (wt ) = βvt + (1 − β)gt − ∇f (wt ) = β vt − ∇f (wt−1 ) + β ∇f (wt−1 ) − ∇f (wt ) + (1 − β)ζt = βet−1 + β ∇f (wt−1 ) − ∇f (wt ) + (1 − β)ζt ,
(77)
by Assumption A.7, and
by Assumption A.8. For t ≥ 1,
while e0 = (1 − β)ζ0 . Unrolling the recursion gives et =
t X j=0
β t−j (1 − β)ζj +
t X
β t−j+1 ∇f (wj−1 ) − ∇f (wj ) .
j=1
18
(78)
Therefore, ∥et ∥F ≤
t X
β
t−j
(1 − β)ζj
j=0
+
t X
β t−j+1 ∥∇f (wj−1 ) − ∇f (wj )∥F .
(79)
j=1
F
Using Assumption 4.3(a), wj − wj−1 = −ηdj−1 , and Lemma A.2(ii), t X
β t−j+1 ∥∇f (wj−1 ) − ∇f (wj )∥F ≤
j=1
t X
β t−j+1 LF ∥wj−1 − wj ∥F
j=1
=
t X
β t−j+1 LF η∥dj−1 ∥F (80)
j=1
√
≤ LF η m
t X
βk
k=1
√
≤ LF η m
β . 1−β
For the noise term, Jensen’s inequality implies
E
t X
v u 2 u X t u β t−j (1 − β)ζj . ≤ tE
β t−j (1 − β)ζj
j=0
j=0
F
(81)
F
Since {ζt } is a martingale difference sequence, the cross terms vanish, and thus
E
t X
2
β t−j (1 − β)ζj
j=0
=
t X
β 2(t−j) (1 − β)2 E∥ζj ∥2F
j=0
F
t
≤
X σ2 (1 − β)2 β 2k B
(82)
k=0
σ2 1 − β ≤ . B 1+β Hence, E
t X
σ ≤√ B
β t−j (1 − β)ζj
j=0
s
1−β . 1+β
(83)
F
Combining the previous bounds and summing over t = 0, . . . , T − 1 yields s √ T −1 X LF η mβ σ 1−β E ∥et ∥F ≤ (T − 1) +T√ . 1−β B 1+β t=0
(84)
Lemma A.11 (Momentum tracking under matched (∞, 2)-smoothness). Under Assumptions 4.3(b), A.7, and A.8, s √ T −1 X mσ 1 − β L∞,2 ηβ E ∥et ∥1,2 ≤ (T − 1) +T √ . (85) 1−β 1+β B t=0 Proof. The recursion for et is the same as in Lemma A.10. Hence, et =
t X j=0
β t−j (1 − β)ζj +
t X
β t−j+1 ∇f (wj−1 ) − ∇f (wj ) .
j=1
19
(86)
By the triangle inequality, ∥et ∥1,2 ≤
t X
β
t−j
(1 − β)ζj
t X
+
j=0
β t−j+1 ∥∇f (wj−1 ) − ∇f (wj )∥1,2 .
(87)
j=1
1,2
Using Assumption 4.3(b), wj − wj−1 = −ηdj−1 , and Lemma A.2(iii), t X
β t−j+1 ∥∇f (wj−1 ) − ∇f (wj )∥1,2 ≤
j=1
t X
β t−j+1 L∞,2 ∥wj−1 − wj ∥∞,2
j=1
=
t X
β t−j+1 L∞,2 η∥dj−1 ∥∞,2 (88)
j=1 t X
≤ L∞,2 η
βk
k=1
≤ L∞,2 η For the noise term, we use ∥A∥1,2 ≤ E
t X
√
β . 1−β
m∥A∥F , hence
β t−j (1 − β)ζj
√
≤
j=0
mE
t X
β t−j (1 − β)ζj
j=0
1,2
.
(89)
F
Applying the Frobenius-noise bound from the proof of Lemma A.10, s √ t X mσ 1−β E β t−j (1 − β)ζj ≤ √ . 1+β B j=0
(90)
1,2
Summing over t = 0, . . . , T − 1 proves the lemma. A.4
Proof of Theorem 4.7
Theorem 4.7. Suppose Assumptions 4.3(b), A.7, A.8, and A.9 hold, and Nora uses a constant step size ηt = η. Then: s " # √ T −1 1 X ∆ 1 L∞,2 ηβ mσ 1 − β L∞,2 η E ∥Gt ∥1,2 ≤ +2 1− + √ + . (91) T t=0 Tη T 1−β 1+β 2 B With the parameter choice: s η=
(p
(1 − β)∆ , L∞,2 T
1 − β = min
) L∞,2 ∆ √ √ ,1 , 2 mσ T
Nora reaches an ϵ-stationary point in the projected ∥ · ∥1,2 sense with complexity: T = O mL∞,2 σ 2 ∆ ϵ−4 .
(92)
(93)
Proof. Summing Lemma A.4 over t = 0, . . . , T − 1, we obtain f (w0 ) − f (wT ) ≥ η
T −1 X
⟨∇f (wt ), dt ⟩ −
t=0
T L∞,2 η 2 . 2
(94)
By Assumption A.9, ∆≥η
T −1 X
⟨∇f (wt ), dt ⟩ −
t=0
20
T L∞,2 η 2 . 2
(95)
Applying Lemma A.6, ∆≥η
T −1 X
∥Gt ∥1,2 − 2η
t=0
T −1 X t=0
∥et ∥1,2 −
T L∞,2 η 2 . 2
(96)
Taking expectations and using Lemma A.11, η
T −1 X
T −1 X T L∞,2 η 2 E ∥Gt ∥1,2 ≤ ∆ + 2η E ∥et ∥1,2 + 2 t=0 t=0 s # " √ T L∞,2 η 2 L∞,2 ηβ mσ 1 − β + ≤ ∆ + 2η (T − 1) +T √ . 1−β 1+β 2 B
(97)
Dividing both sides by T η yields T −1 1 X ∆ E ∥Gt ∥1,2 ≤ +2 T t=0 Tη
"
1 1− T
s # √ L∞,2 ηβ L∞,2 η mσ 1 − β + √ + . 1−β 1+β 2 B
(98)
This proves the explicit bound. To derive the rate, set q := 1 − β. For B = 1, the preceding bound implies, up to universal constants, RT ≲
∆ L∞,2 η √ √ + + mσ q + L∞,2 η. Tη q
(99)
Since q ≤ 1, the last term is dominated by L∞,2 η/q. For a fixed q, we balance the descent term and the momentum-tracking term: s ∆ L∞,2 η q∆ ≍ , which gives η= . (100) Tη q L∞,2 T With this choice, the bound reduces to s RT ≲
L∞,2 ∆ √ √ + mσ q. qT
(101)
The first term whereas the stochastic-noise term decreases. Balancing them √ p increases√as q decreases, gives q ≍ L∞,2 ∆/( mσ T ). Thus we take s ) (p L∞,2 ∆ (1 − β)∆ η= . (102) q = 1 − β = min √ √ ,1 , L∞,2 T 2 mσ T This is a constructive parameter choice used to obtain the complexity bound; it is not meant to be a unique tuning rule for practical training. A.5
Proof of Proposition 4.8
Proposition 4.8. Suppose Assumptions 4.3(a), A.7, A.8, and A.9 hold, and Nora uses a constant step size ηt = η. Then: s " # √ T −1 √ LF ηm 1 X ∆ 1 LF η mβ σ 1−β E ∥Gt ∥F ≤ + ( m + 1) 1 − +√ + , (103) T t=0 Tη T 1−β 2 B 1+β and similarly, T −1 1 X ∆ E ∥Gt ∥1,2 ≤ +2 T t=0 Tη
"
1 1− T
s # √ LF ηmβ LF ηm mσ 1 − β + √ + . 1−β 1 + β 2 B
(104)
Consequently, Nora reaches an ϵ-stationary point in either projected measure with complexity: T = O m2 LF σ 2 ∆ ϵ−4 . (105) 21
Proof. We first prove the Frobenius-norm bound. Summing Lemma A.3 over t = 0, . . . , T − 1 gives f (w0 ) − f (wT ) ≥ η
T −1 X
⟨∇f (wt ), dt ⟩ −
t=0
T LF η 2 m . 2
(106)
By Assumption A.9, f (wT ) ≥ f ⋆ , hence ∆≥η
T −1 X
⟨∇f (wt ), dt ⟩ −
t=0
T LF η 2 m . 2
(107)
Applying Lemma A.5, ∆≥η
T −1 X
√
∥Gt ∥F − η( m + 1)
t=0
T −1 X
∥et ∥F −
t=0
T LF η 2 m . 2
(108)
Taking expectation and using Lemma A.10, T −1 X T LF η 2 m √ E ∥et ∥F + E ∥Gt ∥F ≤ ∆ + η( m + 1) 2 t=0 t=0 s " # √ √ LF η mβ σ 1−β T LF η 2 m ≤ ∆ + η( m + 1) (T − 1) +T√ + . 1−β 2 B 1+β (109) Dividing both sides by T η yields s # " √ T −1 √ 1 X ∆ σ LF ηm 1 LF η mβ 1−β E ∥Gt ∥F ≤ + ( m + 1) 1 − +√ + . (110) T t=0 Tη T 1−β 2 B 1+β
η
T −1 X
We next prove the ∥ · ∥1,2 bound under the same Frobenius smoothness assumption. Starting again from Lemma A.3, T −1 X T LF η 2 m ∆≥η ⟨∇f (wt ), dt ⟩ − . (111) 2 t=0 Using Lemma A.6, ∆≥η
T −1 X
∥Gt ∥1,2 − 2η
t=0
Since ∥A∥1,2 ≤
√
T −1 X t=0
∥et ∥1,2 −
T LF η 2 m . 2
(112)
m∥A∥F , Lemma A.10 implies
s √ −1 √ TX LF ηmβ mσ 1 − β E ∥et ∥1,2 ≤ m E ∥et ∥F ≤ (T − 1) +T √ . 1−β 1+β B t=0 t=0
T −1 X
Substituting the above bound and dividing by T η gives s " # √ T −1 1 X ∆ 1 LF ηmβ mσ 1 − β LF ηm E ∥Gt ∥1,2 ≤ +2 1− + √ + . T t=0 Tη T 1−β 1+β 2 B It remains to justify the stated complexity. Take B = 1 and choose: s √ (1 − β)∆ LF ∆ √ ,1 . η= , 1 − β = min LF mT 2σ T Substituting this choice into either of the two preceding bounds gives ! r r T −1 2 2 1 X mLF ∆ 4 m LF σ ∆ E ∥Gt ∥F = O + , T t=0 T T 22
(113)
(114)
(115)
(116)
and similarly, T −1 1 X E ∥Gt ∥1,2 = O T t=0
r 4
m2 LF σ 2 ∆ + T
r
mLF ∆ T
! .
(117)
Thus Nora reaches an ϵ-stationary point in either projected measure with iteration complexity: T = O m2 LF σ 2 ∆ ϵ−4 . (118)
A.6
Proof of Corollary 4.9
Corollary 4.9. Assume, in addition, that f is row-wise scale invariant, namely: f (Dw) = f (w),
∀w ∈ Rm×n , ∀D ≻ 0 diagonal.
Then ∇f (wt ) is row-wise perpendicular to wt for all t, i.e., Pwr⊥ ∇f (wt ) = ∇f (wt ). t
(119)
(120)
Therefore, Theorem 4.4 and Proposition 4.5 hold with Gt replaced by ∇f (wt ). Proof. Fix any row index i, and define the diagonal matrix Di (c) := I + (c − 1)ei e⊤ i ,
c > 0.
(121)
By row-wise scale invariance, f (Di (c)w) = f (w),
∀c > 0.
(122)
Differentiating both sides with respect to c at c = 1, we obtain 0=
d = ⟨∇f (w)i: , wi: ⟩ . f (Di (c)w) dc c=1
(123)
Since i is arbitrary, every row of ∇f (w) is orthogonal to the corresponding row of w. Hence Pwr⊥ ∇f (w) = ∇f (w). (124) Applying this identity to every iterate wt shows that Gt = ∇f (wt ),
∀t.
(125)
Therefore, the statements of Theorem 4.4 and Proposition 4.5 hold with Gt replaced by ∇f (wt ).
23
Table 5: Detailed model and training configurations. Configuration
60M
135M
Layers Hidden size Attention heads MLP size Maximum context length Training steps GPUs Per-GPU batch size Gradient accumulation Global batch size Nominal token budget Base learning rate Auxiliary Adam learning rate Warmup steps Learning-rate schedule Gradient clipping Precision Evaluation frequency Checkpoint frequency
8 512 8 1376 256 10,000 2 64 4 512 1.31B 1 × 10−3 5 × 10−3 1,000 cosine 1.0 bf16 every 1,000 steps every 5,000 steps
12 768 12 2048 256 20,000 4 64 2 512 2.62B 1 × 10−3 3 × 10−3 2,000 cosine 1.0 bf16 every 1,000 steps every 5,000 steps
Table 6: Matrix learning-rate sweep grids. Model size
Optimizer
Weight decay
Matrix learning-rate grid
60M 60M 60M 60M
Muon Mano RMNP Nora
0.1 0.1 0.1 0.0
{0.005, 0.01, 0.02, 0.03, 0.04} {0.001, 0.004, 0.005, 0.01, 0.02} {0.001, 0.004, 0.005, 0.01, 0.02} {0.001, 0.004, 0.005, 0.01, 0.02}
135M 135M 135M 135M
Muon Mano RMNP Nora
0.1 0.1 0.1 0.0
{0.005, 0.01, 0.02, 0.03} {0.003, 0.005, 0.01, 0.02} {0.003, 0.005, 0.01, 0.02} {0.003, 0.005, 0.01, 0.02}
B
Additional Experimental Details
B.1
Model and Training Configurations
Table 5 reports the detailed model and training configurations used in the 60M and 135M experiments. Both settings use context length 256, global batch size 512, cosine learning-rate decay, bf16 precision, gradient clipping at 1.0, evaluation every 1,000 steps, and checkpointing every 5,000 steps. B.2
Matrix Learning-rate Sweep Grids
Table 6 lists the matrix learning-rate sweep grids for all optimizers. Nora uses the same matrix learning-rate grid as Mano and RMNP, but differs in weight decay: Nora uses weight decay 0, while Muon, Mano, and RMNP use weight decay 0.1. B.3
Other Experimental Results
In Tables 7, 8, 9, and 10, we present the perplexity and loss results across 60M and 135M model scales, evaluated under various algorithms and learning rate configurations. As illustrated, Nora demonstrates a significant advantage over all compared baselines across these diverse settings.
24
Table 7: Perplexity results of all baselines under different learning rates. (60M) Perplexity 0.001 0.004 0.01 0.02 Mano
34.09
29.55
31.53
34.39
RMNP
35.42
30.12
30.30
93.26
Muon
41.65
31.44
31.09
31.99
Nora
31.24
28.94
29.89
31.17
Table 8: Loss results of all baselines under different learning rates. (60M) Loss 0.001 0.004 0.01 0.02 Mano
3.529
3.386
3.451
3.538
RMNP
3.567
3.405
3.411
4.535
Muon
3.729
3.448
3.437
3.466
Nora
3.442
3.365
3.398
3.440
Table 9: Perplexity results of all baselines under different learning rates. (135M) Matrix LR 0.003 0.005 0.01 0.02 Mano Nora
22.13 21.74
22.14 21.86
23.34 22.43
25.07 23.39
Matrix LR
0.005
0.01
0.02
0.03
RMNP Muon
22.62 23.40
22.46 23.17
22.64 24.06
22.67 23.23
Table 10: Loss results of all baselines under different learning rates. (135M) Matrix LR 0.003 0.005 0.01 0.02 Mano Nora
3.097 3.079
3.098 3.085
3.150 3.111
3.222 3.152
Matrix LR
0.005
0.01
0.02
0.03
RMNP Muon
3.119 3.153
3.112 3.143
3.120 3.180
3.121 3.146
25
C
Running and Reference Code
C.1
Run Quickly
We provide a dedicated repository containing the Nora3 source code, along with detailed rationales and instructions for replacing Adam with Nora. Additionally, a separate reproduction repository4 is made available to facilitate the rapid replication of all experimental results presented in this paper. Please set up the C4 [38] dataset yourself and download the reproduction repository, then run: SCRIPT_DIR = " $ ( cd " $ ( dirname " $ { BASH_SOURCE [0]} " ) " && pwd ) " exec " $SCRIPT_DIR / train_universal . sh " \ -- model_size 135 m \ -- optimizer nora \ -- num_gpus 4 \ -- lr_matrix 0.005 \ -- lr_adam 0.02 \ -- num_steps 20000 \ -- batch_size 64 \ -- total_batch_size 512 \ -- warmup_steps 2000 \ -- weight_decay 0.0 \ -- save_every 10000 \ -- eval_every 1000 \ " $@ "
It should be noted that the results for Muon, RMNP, and Mano can be reproduced using the exact same experimental setup as Nora. C.2
Reference Code
import math import torch import torch . nn . functional as F
L O W _PRECISION_DTYPES = ( torch . float16 , torch . bfloat16 )
class Nora ( torch . optim . Optimizer ) : " " " Normalized Orthogonal Row Alignment optimizer for scalable matrix training . " " " def __init__ ( self , param_groups , lr_nora =0.005 , lr_adam =0.001 , momentum =0.95 , beta =0.95 , weight_decay =0.0 , betas =(0.9 , 0.95) , eps =1 e -10 , ): defaults = dict ( lr_nora = lr_nora , lr_adam = lr_adam , momentum = momentum , beta = beta , 3 4
https://github.com/Yuan-Jinghui/Nora https://github.com/JiaxuanZou0714/Lrp
26
weight_decay = weight_decay , betas = betas , eps = eps , ) super () . __init__ ( param_groups , defaults ) def step ( self , closure = None ) : loss = None if closure is not None : loss = closure () for group in self . param_groups : lr = group [ " lr " ] momentum = group . get ( " momentum " , 0.95) beta = group . get ( " beta " , 0.95) weight_decay = group . get ( " weight_decay " , 0.0) betas = group . get ( " betas " , (0.9 , 0.95) ) eps = group . get ( " eps " , 1e -10) is_nora = group . get ( " is_nora " , True ) for p in group [ " params " ]: if p . grad is None : continue grad = p . grad . data param_state = self . state . setdefault (p , {}) use_master_param = p . data . dtype in LOW_PRECISION_DTYPES if use_master_param : if " fp32_param " not in param_state : param_state [ " fp32_param " ] = p . data . detach () . float () . clone () elif param_state [ " fp32_param " ]. dtype != torch . float32 : param_state [ " fp32_param " ] = param_state [ " fp32_param " ]. float () param_data = param_state [ " fp32_param " ] grad_data = grad . float () else : param_data = p . data grad_data = grad if is_nora and grad . dim () >= 2: if " momentum_buffer " not in param_state : buf = torch . zeros_like ( grad_data ) else : buf = param_state [ " momentum_buffer " ] if use_master_param and buf . dtype != torch . float32 : buf = buf . float () buf . lerp_ ( grad_data , 1 - beta ) m_t = grad_data . lerp ( buf , momentum ) theta_hat = F . normalize ( param_data , p =2 , dim = -1 , eps = eps ) dot_product = torch . sum ( m_t * theta_hat , dim = -1 , keepdim = True ) v = m_t - dot_product * theta_hat v_hat = F . normalize (v , p =2 , dim = -1 , eps = eps )
27
scale = max (1 , math . sqrt ( grad_data . size ( -2) / grad_data . size ( -1) ) ) update_direction = v_hat * scale if weight_decay > 0: param_data . mul_ (1 - lr * weight_decay ) param_data . add_ ( update_direction , alpha = - lr ) if use_master_param : p . data . copy_ ( param_data . to ( dtype = p . data . dtype ) ) param_state [ " momentum_buffer " ] = buf else : if " exp_avg " not in param_state : param_state [ " exp_avg " ] = torch . zeros_like ( grad_data ) param_state [ " exp_avg_sq " ] = torch . zeros_like ( grad_data ) param_state [ " step " ] = 0 elif use_master_param and param_state [ " exp_avg " ]. dtype != torch . float32 : param_state [ " exp_avg " ] = param_state [ " exp_avg " ]. float () param_state [ " exp_avg_sq " ] = param_state [ " exp_avg_sq " ]. float () exp_avg , exp_avg_sq = param_state [ " exp_avg " ] , param_state [ " exp_avg_sq " ] param_state [ " step " ] += 1 exp_avg . mul_ ( betas [0]) . add_ ( grad_data , alpha =1 betas [0]) exp_avg_sq . mul_ ( betas [1]) . addcmul_ ( grad_data , grad_data , value =1 - betas [1]) bias_correction1 = 1 - betas [0] ** param_state [ " step " ] bias_correction2 = 1 - betas [1] ** param_state [ " step " ] step_size = lr * math . sqrt ( bias_correction2 ) / bias_correction1 denom = exp_avg_sq . sqrt () . add_ ( eps ) adam_update = exp_avg / denom if weight_decay > 0: param_data . mul_ (1 - step_size * weight_decay ) param_data . add_ ( adam_update , alpha = - step_size ) if use_master_param : p . data . copy_ ( param_data . to ( dtype = p . data . dtype ) ) return loss
def get_nora_optimizer ( model , lr_nora =0.005 , lr_adam =0.001 , weight_decay =0.1 ,
28
momentum =0.95 , beta =0.95 , ): nora_params = [] adam_params = [] for name , param in model . named_parameters () : if param . requires_grad : if param . ndim >= 2 and " embed " not in name and " lm_head " not in name : nora_params . append ( param ) else : adam_params . append ( param ) param_groups = [ dict ( params = nora_params , lr = lr_nora , lr_nora = lr_nora , lr_adam = lr_adam , weight_decay = weight_decay , momentum = momentum , beta = beta , is_nora = True , ), dict ( params = adam_params , lr = lr_adam , lr_nora = lr_nora , lr_adam = lr_adam , weight_decay = weight_decay , momentum = momentum , beta = beta , is_nora = False , ), ] optimizer = Nora ( param_groups ) return optimizer
29