PC Layer: Polynomial Weight Preconditioning for Improving LLM Pre-Training Senmiao Wang∗ 1 , Tiantian Fang∗ 2 , Haoran Zhang∗ 1 , Yushun Zhang1,4 , Kunxiang Zhao1 , Alex Schwing3 , and Ruoyu Sun†1,4
arXiv:2606.06470v1 [cs.LG] 4 Jun 2026
1 The Chinese University of Hong Kong, Shenzhen, China 2 Google LLC, Mountain View, CA, United States
3 University of Illinois at Urbana-Champaign, Urbana, IL, United States
4 Shenzhen International Center for Industrial and Applied Mathematics, Shenzhen
Research Institute of Big Data {senmiaowang1, haoranzhang2, yushunzhang, kunxiangzhao}@link.cuhk.edu.cn {tf6, aschwing}@illinois.edu [email protected]
Abstract We propose a preconditioning (PC) layer — a weight parameterization via polynomial preconditioner that ensures stable weight conditioning throughout LLM training. The PC module reshapes the singularvalue spectrum of weight matrices via low-degree polynomial preconditioning. After training, the preconditioned weights can be merged back into the original architecture, incurring no inference overhead. We demonstrate the advantage of the proposed PC layer over standard transformers in Llama-1B pretraining, for both the AdamW and Muon optimizers. Theoretically, we justify this spectrum-control principle by proving that uniformly bounding each layer’s singular values ensures geometric convergence of gradient descent to global minima, for certain deep linear networks. Our code is available at https://github.com/Empath-aln/PC-layer.
1
Introduction
Training modern neural networks, especially at scale, relies on normalization techniques. These methods stabilize optimization by inserting layers that normalize intermediate quantities—such as features or parameters—within the network [Huang et al., 2023]. Representative examples include batch normalization [Ioffe and Szegedy, 2015], layer normalization [Ba et al., 2016], instance normalization [Ulyanov et al., 2016], group normalization [Wu and He, 2018], etc. In the era of large language models (LLMs), normalization methods remain critical. For instance, RMSNorm [Zhang and Sennrich, 2019], a variant of LayerNorm [Ba et al., 2016], has become a de facto standard in most prevalent LLM architectures, and Query–Key Normalization (QK-Norm) [Henry et al., 2020, Loshchilov et al., 2025] has been widely used to control attention logits and stabilize training. RMSNorm and QK-Norm are designed to control the layer outputs of neural networks. Another family of normalization methods acts directly in weight space: classical examples include weight normalization (WN) [Salimans and Kingma, 2016] and spectral normalization (SN) [Miyato et al., 2018]. While such weight control ∗ Equal contribution. † Corresponding author.
1
Figure 1: Illustration of the PC layer: a low-degree matrix polynomial reshapes the singular-value spectrum of a weight matrix, amplifying small singular values and saturating large ones, without explicit SVD. Unlike sign/polar-style maps (e.g., those used in Muon) that drive nearly all nonzero singular values toward one, PC performs soft spectrum shaping rather than near-orthogonalization.
methods have not become standard components of mainstream Transformer-based LLM pre-training, recent works have revived this line by explicitly normalizing or constraining weights, including row/column-wise normalizations on weights [Loshchilov et al., 2025, Franke et al., 2025, Fu et al., 2025], hyperball/hypersphere constraints on weights [Wen et al., 2025a, Ren et al., 2026], spectral constraints on weights [Newhouse et al., 2025, Xie et al., 2026a], etc. These developments suggest that weight geometry is becoming an increasingly relevant design axis for stable LLM training. This motivates two questions: (i) Is there a theoretical principle relating weight properties to the performance of the algorithm? (ii) If so, can such a principle inform the design of weight-based control techniques? Regarding the first question, we note that the spectrum of weight matrices can be understood from the perspective of signal propagation. For a linear layer, the largest and smallest singular values control the upper and lower bounds of the propagated signal magnitude, so their ratio characterizes the conditioning of signal propagation. A collapsed lower tail of singular values can therefore degrade signal propagation across depth. This suggests that controlling the weight spectrum can promote better signal propagation, which in turn makes the network easier to optimize. Moreover, we prove that under certain assumptions, for deep linear networks, if all weights of a multi-layer fully connected network have upper bounded condition number, then gradient descent converges to global minimizer at a rate dependent on the weight condition numbers. This gives a partial answer to the first question. Regarding the second question on the algorithm design, an immediate challenge is how to make the spectrum control method compatible with the original backpropagation method1 . A possible way to reshape the spectrum of a weight matrix W is to compute an explicit factorization (e.g., a singular value decomposition (SVD) W = U diag(σ1 , . . . , σr ) V ⊤ ) and then modify the singular values σ1 , . . . , σr . However, such matrix-factorization–style approaches may be rather costly. 1 Background: Existing normalization layers are embedded into the network so that auto-differentiation applies to normalization layers.
2
We propose to add polynomial preconditioning into the neural network. The key insight is that matrix polynomials provide a convenient mechanism for acting directly on the spectrum. Concretely, let W = U diag(σ1 , . . . , σr ) V ⊤ be the SVD of W, and let g be a certain feasible preconditioning polynomial. Then the transformed matrix g(W ) has singular values {| g(σ1 )|, . . . , | g(σr )|}, i.e., each singular value is reshaped by the same scalar function g. This yields an explicit and controllable modification of the weight spectrum without resorting to computationally heavy matrix decompositions. We choose the polynomial g to approximate a desired shaping function that amplifies small singular values relative to large ones, thereby narrowing the spread of the spectrum and promoting well-conditioned weights throughout training. Our main contributions are threefold: • Algorithm design. We propose a built-in-layer module, termed preconditioning (PC) layer, which encourages well-conditioned weight matrices via polynomial preconditioning. This method is lightweight during training and causes no additional inference cost after training. • Empirical validation in LLM pre-training. We evaluate PC layers on Llama-271M and Llama-1B models with both the AdamW [Kingma, 2015, Loshchilov and Hutter, 2019] and Muon [Jordan et al., 2024] optimizers. On Llama-1B, PC delivers consistent token-efficiency speedups under both optimizers—i.e., it reaches the same loss with fewer training tokens (2× with AdamW and 1.13× with Muon)—and improves zero-shot downstream accuracy in both cases. We further verify that PC improves weight-spectrum conditioning. • Theory for controlling weight spectrum. Our intuition is that a “good weight spectrum” helps training. To support this intuition, we prove that a bounded weight spectrum implies geometric convergence of the loss to global minimum on deep linear networks, providing theoretical backing for the proposed method. A preliminary version of this work, focusing on generative adversarial nets (GANs), is available online [Fang et al., 2021]. Compared to that version, the present version targets large language models and introduces several critical LLM-specific modifications to the method, including (i) magnitude adjustment operations (norm recovery and learnable scaling) and (ii) a specific choice of layers to apply preconditioning. These changes are essential for achieving strong performance in the LLM setting, which presents distinct optimization challenges not found in the earlier GAN-based experiments.
2
Motivation: From Signal Propagation to Weight-Spectrum Control
Signal propagation as a guiding lens. Signal propagation provides a classical lens for understanding trainability: it studies how activations and gradients are transformed as they pass through depth. Historically, this lens has been most influential in the design of weight initialization schemes. Xavier [Glorot and Bengio, 2010] and Kaiming [He et al., 2015] initializations aim to keep forward and backward signal magnitudes stable at the start of training. The line of work on orthogonal initialization further refines this principle by providing a more fine-grained analysis of how depth affects signal propagation [Saxe et al., 2014, Pennington et al., 2017, 2018, Xiao et al., 2018, Hu et al., 2020]. Recent LLM training methods have also begun to control weights or their updates beyond initialization, through normalized Transformers [Loshchilov et al., 2025], hyperspherical or hyperball constraints [Wen et al., 2025a, Ren et al., 2026], and spectral-sphere constraints [Xie et al., 2026a]. This viewpoint is also compatible with µP-style spectral analyses, which suggest that the spectral scale of weights and updates is central to stable feature learning under width scaling [Yang et al., 2021, 2023]. Together, these observations motivate treating the weight spectrum itself as a direct control knob for signal propagation. Weight spectrum as a propagation control knob.
Consider a linear transformation
y = Wx.
3
The singular values of W determine how much the layer can amplify or suppress different input directions: σmin (W )∥ x ∥2 ≤ ∥Wx ∥2 ≤ σmax (W )∥ x ∥2 . The same singular values also govern the adjoint transformation W ⊤ that appears in backpropagation. Thus, the spectrum of W gives a layer-local proxy for forward and backward signal propagation. In the ideal square case, if W is orthogonal, then ∥Wx ∥2 = ∥ x ∥2 for every input x: the layer is exactly norm-preserving. This connects weight-spectrum control to the classical dynamical-isometry and orthogonal-initialization literature, where near-isometric transformations are used to maintain well-conditioned propagation through depth [Saxe et al., 2014, Pennington et al., 2017, 2018, Xiao et al., 2018, Hu et al., 2020]. Why exact orthogonality may not be the right target. Although orthogonal weights provide a clean signal-propagation ideal, enforcing exact orthogonality throughout training is too restrictive. To see this, consider a depth-L linear network f θ ( x ) = WL WL−1 · · · W1 x, where each Wℓ ∈ Rd×d is orthogonal. Then the end-to-end map WL · · · W1 is also orthogonal. Even if we add a global scalar α and consider αWL · · · W1 , all singular values of the end-to-end map are still identical. Hence such a model cannot represent anisotropic maps that require direction-dependent amplification, such as x 7→ diag(1, 2) x. This example highlights a basic tension. Perfectly flat spectra are favorable for norm preservation, but they can remove useful spectral anisotropy and weaken representation power. Therefore, the goal should not be to collapse all singular values to one. Soft spectrum conditioning. The preceding discussion suggests the following design principle: the weight spectrum should be well-conditioned, but not perfectly flat. Equivalently, the condition number κ (W ) =
σmax (W ) σmin (W )
should be kept moderate, while still allowing nontrivial variation among singular values. This principle balances two desiderata. On the optimization side, reducing spectral spread mitigates excessive amplification or attenuation of signals. On the representation side, preserving controlled spectral variation allows the model to implement anisotropic transformations. Design implication. The discussion above leads to a training-time weight-spectrum control problem. We do not merely want an initialization whose spectrum is well-behaved at step zero. Instead, we want a mechanism that acts on the effective weight matrices used during training. Such a mechanism should reduce the relative spread between small and large singular values, while avoiding the over-restrictive regime where all singular values are forced to be identical. The next section develops a polynomial weight-preconditioning mechanism that satisfies these requirements.
3
Polynomial Weight Preconditioning
We now instantiate the training-time weight-spectrum control desideratum from Section 2. In analogy to normalization methods, we embed the control mechanism—preconditioning (PC) layer—directly into the network. Background on polynomial preconditioning. Polynomial preconditioning is a classical technique in numerical linear algebra, originally developed to accelerate iterative solvers for symmetric linear systems Qx = b [Johnson et al., 1983]. The idea is to replace Q by g( Q) := p( Q) Q, where p is a low-degree polynomial chosen so that the spectrum of g( Q) concentrates near 1, dramatically reducing its condition number and 4
speeding up iterative methods such as conjugate gradient. Johnson et al. [1983] formulate the design of p as a scalar approximation problem (minimax or least-squares) on the spectral interval, and observe that least-squares polynomials tend to yield better empirical convergence by improving the entire spectrum rather than just the extreme condition number. We will adapt this perspective to the deep-learning setting, where the matrices of interest are rectangular weight matrices and the goal is to control their singular-value spectrum during training. See Appendix B.1 for a self-contained review of polynomial preconditioning and Appendix B.2 for a discussion of the relation between preconditioning, the condition number, and the full spectrum. Outline. We first extend polynomial preconditioning to rectangular matrices in deep nets (§3.1), describe how to determine the preconditioning polynomials (§3.2), and finally propose the preconditioning (PC) layer method (§3.3).
3.1
Preconditioning Rectangular Matrices in Deep Nets
In deep networks, the objects we wish to control are typically rectangular weight matrices W ∈ Rn×m . Polynomial preconditioning via a Gram matrix. Polynomials are not directly defined on a rectangular W, so we work through a symmetric Gram matrix—either WW ⊤ ∈ Rn×n or W ⊤ W ∈ Rm×m —whose nonzero eigenvalues are the squared singular values of W. Concretely, we embed a preconditioning polynomial by left-multiplying W with a polynomial p in WW ⊤ : g(W ) ≜ p(WW ⊤ )W ∈ Rn×m , which preserves the shape of W and can be implemented efficiently using only repeated matrix multiplications. (Equivalently, one can use the right-preconditioned form g(W ) = W p(W ⊤ W ); in practice one can pick the smaller Gram matrix for efficiency.) The following proposition shows that controlling the conditioning of g(W ) reduces to shaping a one-dimensional map g(σ ) = p(σ2 )σ on the singular-value interval of W.2 See the proof in Appendix E. Proposition 1 (Singular-value mapping). Let W have singular values σ1 ≥ · · · ≥ σm ≥ 0 and define the scalar map g(σ ) = p(σ2 ) σ. Then the singular values of g(W ) = p(WW ⊤ )W are {| g(σi )|}im=1 . Proposition 1 shows that, despite W being rectangular, polynomial preconditioning still reduces spectral control to designing a scalar map g(σ ) = p(σ2 )σ acting on individual singular values. In particular, the spread (and hence the conditioning) of the singular values of g(W ) is determined entirely by how g deforms the singular values of W: by choosing g to expand small σi relative to large ones, we directly reduce the relative gap between σmin and σmax and thereby improve conditioning. This is exactly the mechanism we will exploit in Section 3.2 to determine the polynomial p and in Section 3.3 to build our preconditioning layer.
3.2
Finding Preconditioning Polynomials
Building on the singular-value mapping in Section 3.1, this subsection describes the criteria and procedure for choosing the preconditioning map g(σ ) = p(σ2 )σ. Optimization formulation. Suppose we are given a domain [γL , γU ], a target function f , and an integer k, we seek the best approximation to f (σ) on [γL , γU ] within the polynomial class Gk = { g(σ ) = p(σ2 )σ | deg( p) ≤ k }. 2We slightly abuse notation by using g to denote both the matrix map W 7 → p (WW ⊤ )W and the induced scalar mapping σ 7 → p ( σ2 ) σ.
5
A standard approach in polynomial preconditioning is to determine g by minimizing a continuous weighted least-squares objective [Johnson et al., 1983]. Concretely, writing g(σ ) = a0 + a1 σ2 + · · · + ak σ2k σ, a = ( a 0 , . . . , a k ) ∈ Rk +1 , the continuous weighted least-squares objective is min
a ∈Rk +1
Z γU
2
g(σ ) − f (σ ) w(σ ) dσ,
(1)
γL
where w(σ ) = σα is the weight function used by Johnson et al. [1983] and α is a numerical constant. In practice, we approximate the objective by finite-sample approximation n
2
∑ g(σi ) − f (σi ) w(σi ), a ∈Rk +1 min
(2)
i =1
where σ1 , . . . , σn ∈ [γ L , γU ] are the sample points. Note that g(σi ) is linear in the coefficient vector a: g(σi ) = ϕ(σi )⊤ a, where ϕ(σi ) = (σi , σi3 , . . . , σi2k+1 )⊤ . Therefore, (2) is a standard weighted linear leastsquares problem in a and can be solved efficiently. See Appendix F.1 for the formal polynomial fitting algorithm. Choosing the fitting interval [γL , γU ] via spectral normalization. The singular spectrum of W can vary substantially across layers and training steps. To make a single fitted polynomial applicable to matrices e = W/s(W ), where s(W ) is intended with different spectral scales, we design it for normalized weights W e into [0, 1]. In to approximate ∥W ∥2 . Ideally, taking s(W ) = ∥W ∥2 would map the singular values of W practice, during training, our method normalizes each selected weight matrix by an estimate s(W ) ≈ ∥W ∥2 , computed via streaming power iteration; this PC-layer implementation is described in Section 3.3, and the estimator is detailed in Appendix F.2. To make the polynomial fit robust to the small error of this estimate, we use the slightly enlarged fitting interval [γL , γU ] = [0, 1.1] instead of [0, 1]. Empirically, the streaming power-iteration estimator is accurate enough for this margin to cover the observed approximation error during training, keeping the polynomial evaluation within its design domain; see Appendix I for the estimator validation. A feasible desired mapping. Recall that driving all singular values to a common value would yield a perfectly conditioned matrix. However, any map of the form g(σ ) = p(σ2 )σ necessarily satisfies g(0) = 0, so it cannot send arbitrarily small singular values to a fixed positive constant. We therefore adopt a more realistic target: amplify small singular values (to improve conditioning) while saturating large ones near 1. This motivates the piecewise-linear target ( σ/b, σ < b, PLb (σ ) = 1, σ ≥ b, with cutoff b > 0. Intuitively, PLb enlarges the low end of the spectrum by a factor 1/b and caps the high end at 1, thereby reducing the spread of singular values. Choosing the cutoff: optimization vs. expressiveness. The cutoff b operationalizes the optimization– expressiveness trade-off introduced in Section 2 by controlling how aggressively singular values are driven toward 1. If b ≈ 1, then PLb (σ ) ≈ σ over most of [0, 1], so the spectrum is barely changed and conditioning is not improved. At the other extreme, if b is so small that b < σmin (W ), then PLb maps all singular values to 1; this is undesirable because it can substantially reduce the expressive capacity of a deep network. To balance these effects, we consider a set of targets with different strengths, e.g., b ∈ {0.8, 0.6, 0.4, 0.3} (see Figure 2(a)). Additionally, Appendix J provides a complementary over-flattening stress test, showing that near-perfect spectrum flattening can hurt performance and supporting the use of moderate rather than overly aggressive targets. 6
Fitted polynomials. Solving the weighted least-squares problem above yields a sequence of preconditioning polynomials with increasing strength. Concretely, we associate each pc_level with a cutoff b in the target PLb , using b ∈ {0.8, 0.6, 0.4, 0.3} for k ∈ {1, 2, 3, 4}, where smaller b corresponds to a more aggressive push of singular values toward 1. Using the polynomial-fitting algorithm detailed in Appendix F.1, we obtain one set of fitted polynomials gk (σ ) = pk (σ2 )σ (with overall degrees 3, 5, 7, 9) are: g1 (σ ) = 1.507σ − 0.507σ3 , g2 (σ ) = 2.083σ − 1.643σ3 + 0.560σ5 , g3 (σ ) = 2.909σ − 4.649σ3 + 4.023σ5 − 1.283σ7 , g4 (σ ) = 3.625σ − 9.261σ3 + 14.097σ5 − 10.351σ7 + 2.890σ9 . Figure 2(b) plots these four fitted mappings. Other valid polynomial sets can be obtained by re-solving (2). We refer to the index k as the PC level (pc_level). Concretely, pc_level = k means that pk has degree k in gk (σ ) = pk (σ2 )σ, so the induced scalar map gk has overall degree 2k + 1. A larger pc_level corresponds to a smaller cutoff b in the target PLb , and therefore a stronger spectrum-shaping operation: it more aggressively lifts small singular values and saturates large ones, while still avoiding exact spectrum flattening. Table 1 summarizes this correspondence. pc_level k
cutoff b
deg pk
deg gk
intuition
1 2 3 4
0.8 0.6 0.4 0.3
1 2 3 4
3 5 7 9
mild shaping moderate shaping strong shaping strongest in this work
1.0
1.0
0.8
0.8
0.6
0.6
g( )
f( )
Table 1: The PC level k as a spectrum-shaping strength knob: larger k uses a higher-degree polynomial fitted to a smaller cutoff b, giving a more aggressive map.
0.4 Piecewise linear, cutoff = 0.3 Piecewise linear, cutoff = 0.4 Piecewise linear, cutoff = 0.6 Piecewise linear, cutoff = 0.8
0.2 0.0 0.0
0.2
0.4
0.6
0.8
0.4 fitted polynomial on PL0.3 fitted polynomial on PL0.4 fitted polynomial on PL0.6 fitted polynomial on PL0.8
0.2 0.0
1.0
0.0
(a) Piecewise-linear target mappings. We show PLb (σ ) = min(σ/b, 1) for several cutoffs b ∈ {0.3, 0.4, 0.6, 0.8} in the fitting interval [0, 1.1]. Smaller b more aggressively enlarges small singular values and saturates large ones to 1.
0.2
0.4
0.6
0.8
1.0
(b) Fitted polynomial preconditioners. Weighted leastsquares solutions in Gk = { g(σ ) = p(σ2 )σ : deg p ≤ k } for k ∈ {1, 2, 3, 4} (overall degrees 3, 5, 7, 9), each fitted to the corresponding PLb target in (a) by Algorithm 2 in Appendix F.1. The fitted map g closely tracks the target while respecting the structural constraint g(0) = 0.
Figure 2: Finding preconditioning polynomials on a fixed spectral domain. (a) Target piecewise-linear mappings. (b) Low-degree polynomial preconditioners obtained by weighted least squares.
7
3.3
Methodology: Preconditioning (PC) Layer
In this subsection, we formally propose the Preconditioning (PC) layer, a built-in weight-space module that reshapes the singular-value spectrum of selected weight matrices. PC operates on a configurable set of blocks, acting inside the computation graph during training. Algorithm 1 summarizes the procedure. Algorithm 1 PC Layer for Transformers 1: Input: Transformer architecture A; subset of weight blocks PC_blocks ⊆ param_blocks; preconditioning
polynomial gk selected by pc_level= k; a per-block learnable scalar γ (initialized to 1) for each W ∈ PC_blocks. 2: for each weight matrix W in PC_blocks do 3: s(W ) ← estimation of ∥W ∥2 e = W/s(W ) 4: W # weight normalization e ) = p (W eW e ⊤ )W e 5: g (W # polynomial preconditioning on singular values e) 6: PC(W ) ← γ · [s(W )]stop-grad · g(W # norm recovery and learnable scaling 7: end for 8: Return: Transformer architecture APC = A[W → PC(W )]W ∈PC_blocks . In general, our weight preconditioning method (Algorithm 1) replaces each weight matrix W in selected blocks with its preconditioned form PC(W ). Importantly, this replacement is architectural rather than numerical — it modifies the computation graph itself, rather than simply altering the parameter values. In this work, we apply PC to the attention output projection matrices and the FFN projection matrices in Llama 2 architecture. We use WO to denote the attention output projection in each layer, and define ffn := {Wgate , Wup , Wdown }, where the three FFN matrices correspond to gate_proj, up_proj, and down_proj, respectively. Thus, unless otherwise specified, PC_blocks = {WO , ffn}. The following explains the key components of the weight preconditioning method: e = W/s(W ), where • Weight normalization (Line 4): We scale the original weight matrix by W s(W ) ≈ ∥W ∥2 is an estimate of the spectral norm. The goal is to confine the singular values of e to (approximately) [0, 1], the interval over which the polynomial g is designed to be effective in W preconditioning. We compute s(W ) by a streaming power-iteration procedure that avoids exact SVD; see Appendix F.2 for details. • Polynomial preconditioning (Line 5): As detailed in Section 3.1, the odd polynomial of degree 2k +1, e ) = p (W eW e ⊤ ) W, e maps each singular value σ of W e to g(σ ) = σ p(σ2 ), thereby reshaping the g (W spectrum for better conditioning. We implement this step using compute-efficient techniques; see Appendix F.3 for details. • Norm recovery and learnable scaling (Line 6): After the polynomial preconditioning, we rescale the preconditioned matrix by the spectral-norm estimate [s(W )]stop-grad (norm recovery)3 and by a learnable scalar γ (initialized to 1). Intuitively, multiplying s(W ) back restores a comparable overall magnitude after normalization and polynomial spectrum shaping. This rescaling is crucial for the loss (see Section 6.3). The learnable γ then allows the model to adaptively adjust this magnitude during training, leaving the final loss nearly unchanged but stabilizing signal-propagation metrics (see Section 6.4). Because both are scalar multipliers, they only rescale the singular values uniformly and do not alter their relative distribution. Hence the spectrum shape tailored by the PC polynomial g is preserved. 3 Here [·]
stop-grad indicates that the scalar recovery factor is detached from the computation graph.
8
Remark 1. Weight preconditioning is implemented as a reparameterization during training, thus it introduces no inference-time overhead. To be more specific, after training APC , we simply store PC(W ) as the new weight and use it as the parameter of the original architecture A during inference. Thus, there is no additional PC-polynomial computation at inference, and the FLOPs remain the same.
4
Experimental Results
In this section, we present an empirical evaluation of the proposed PC layer. We first describe the experimental settings (§4.1), then report pre-training results (§4.2), and finally evaluate the computation and memory cost (§4.3).
4.1
Experimental Settings
Training setup. We evaluate PC by pre-training Llama-2-style [Touvron et al., 2023] models at two scales, 271M and 1B parameters, on the FineWeb dataset [Penedo et al., 2024] using the TorchTitan codebase [Liang et al., 2025].4 Our main results are reported under two optimizers: AdamW [Kingma, 2015, Loshchilov and Hutter, 2019] and Muon [Jordan et al., 2024]. Detailed optimizer configurations are deferred to Appendix H. Guided by recent pre-training practice in data-rich regimes [Gadre et al., 2024, Sardana et al., 2024, Wen et al., 2025b], we use token-to-parameter ratios well above the Chinchilla compute-optimal guideline (≈20 tokens/parameter) [Hoffmann et al., 2022]: the 271M model is trained on 57B tokens (≈210 tokens/parameter, 10.5×Chinchilla), and the 1B model on 160B tokens (≈160 tokens/parameter, 8×Chinchilla). Both scales use sequence length 8192 and a fixed global batch of 2.62M tokens per optimization step. All runs are conducted on 8×NVIDIA H100 GPUs. Learning-rate schedule and grid search. We use a cosine learning-rate schedule with linear warmup: the LR is linearly warmed up over the first 1% of training steps to the peak value, and then decayed to 10% of the peak by the end of training following a cosine schedule. We tune the√peak LR only on the baseline configuration. √ For each √ model size and optimizer, we sweep the peak LR over a 2-spaced geometric grid (e.g., a/2, a/ 2, a, a 2, 2a). Among the runs that train stably (no spikes in the training loss curve), we pick the peak LR with the lowest final validation loss (see Table 2). The PC run then directly adopts this baseline-tuned peak LR (together with all other non-PC hyperparameters), without any further tuning. This procedure is conservative for PC, since the LR is tuned only on the baseline and is not re-optimized after adding PC. Model
Optimizer
Sweep range
Selected peak LR
Llama-271M
AdamW Muon
[3.0 × 10−4 , 3.394 × 10−3 ] [1.061 × 10−3 , 1.2 × 10−2 ]
8.486 × 10−4 4.243 × 10−3
Llama-1B
AdamW Muon
[5.303 × 10−5 , 4.243 × 10−4 ] [7.5 × 10−4 , 4.243 × 10−3 ]
3 × 10−4 1.5 × 10−3
Table 2: Peak learning-rate grid search on the baseline configuration.
PC configuration. Unless otherwise specified, PC is applied to the FFN projections and the attention output projection in each transformer layer, i.e., PC_blocks = {ffn, WO } with ffn = {Wgate , Wup , Wdown }. For spectral normalization, we estimate ∥W ∥2 using a 10-step streaming power iteration during each training step. All PC runs use norm recovery and a per-block learnable scalar γ, as described in Algorithm 1. The 4 https://github.com/pytorch/torchtitan. Our implementation is built on an earlier revision of TorchTitan.
9
only optimizer-dependent PC hyperparameter in our main experiments is pc_level, the degree k of pk in gk (σ ) = pk (σ2 )σ and hence a knob controlling the strength of spectrum shaping (Table 1); a larger pc_level applies a higher-degree, more aggressive PC map. We use pc_level = 4 with AdamW and pc_level = 2 with Muon, following the ablation in Section 6. Additional hyperparameters, model configurations, and implementation details are provided in Appendix H.
4.2
Main Results
We present pre-training results of PC under two optimizers, AdamW and Muon, on Llama-271M and Llama-1B models. For each optimizer, we compare a transformer equipped with PC layers against an optimizer-matched standard Llama-style transformer baseline without PC. PC improves optimization under AdamW. We first evaluate PC under AdamW on Llama-271M and Llama-1B. On Llama-271M (Figure 3(a)), under the same token budget, the Transformer with PC layer reaches a 0.055 lower final validation loss than the standard baseline. It attains the baseline loss with about 39% fewer training tokens, a 1.63× speedup. The advantage carries over to the 1B scale (Figure 3(b)). Here PC lowers the final validation loss by 0.070 and reaches the baseline loss with 50% fewer tokens, a 2× speedup. The optimization gain does not diminish at 1B scale and is even slightly larger.
2x faster
1.63x faster
-0.070
-0.055
(a) Llama-271M. Validation loss vs. training tokens.
(b) Llama-1B. Validation loss vs. training tokens.
Figure 3: PC performance under AdamW. Validation loss vs. training tokens on (a) Llama-271M and (b) Llama-1B. Under the same total token budget, adding the PC layer reduces the final validation loss by 0.055 on 271M (a 1.63× token-efficiency speedup) and by 0.070 on 1B (a 2× speedup). The gain does not diminish at the larger scale.
PC also helps under Muon. We further evaluate PC under Muon, a second widely used optimizer for LLM pre-training. Adding the PC layer again lowers the final validation loss on both scales (Figure 4(a), 4(b)). The reduction is 0.006 on Llama-271M, a 1.07× speedup, and 0.012 on Llama-1B, a 1.13× speedup. These suggest that the improvement brought by PC layer stays consistent across scales under Muon. Downstream task evaluation. We evaluate the pre-trained 1B-parameter model on a suite of standard zero-shot downstream tasks using the Language Model Evaluation Harness [Gao et al., 2024]. Specifically, we report results on LAMBADA [Paperno et al., 2016], HellaSwag [Zellers et al., 2019], WinoGrande [Sakaguchi et al., 2021], PIQA [Bisk et al., 2020], BoolQ [Clark et al., 2019], ARC-Easy/ARC-Challenge (ARC-E/ARC-C) [Clark et al., 2018], CommonsenseQA (CSQA) [Talmor et al., 2019], and OpenBookQA (OBQA) [Mihaylov et al., 2018]. Table 3 shows that adding the PC layer improves average downstream
10
1.07x faster 1.13x faster
-0.006
-0.012
(a) Llama-271M. Validation loss vs. training tokens.
(b) Llama-1B. Validation loss vs. training tokens.
Figure 4: PC performance under Muon. Validation loss vs. training tokens on (a) Llama-271M and (b) Llama-1B. Adding the PC layer consistently reduces the final validation loss across scales. The reduction is 0.006 on 271M (a 1.07× token-efficiency speedup) and 0.012 on 1B (a 1.13× speedup).
accuracy under both optimizers: by 0.0206 points under AdamW (0.4539 → 0.4745) and by 0.0125 points under Muon (0.4880 → 0.5005), winning on 8 out of 9 tasks in each case. Method
LMB.
Hella.
Wino.
PIQA
BoolQ
ARC-E
ARC-C
CSQA
OBQA
Avg.
Acc ↑
Accn ↑
Acc ↑
Acc ↑
Acc ↑
Accn ↑
Accn ↑
Acc ↑
Accn ↑
↑
AdamW baseline AdamW + PC
0.5026 0.5366
0.5137 0.5642
0.5438 0.5627
0.7122 0.7274
0.5645 0.5554
0.4769 0.5034
0.2645 0.2807
0.2088 0.2105
0.2980 0.3300
0.4539 0.4745
Muon baseline Muon + PC
0.5717 0.5812
0.6004 0.6071
0.5809 0.5967
0.7470 0.7486
0.5550 0.5847
0.5076 0.5370
0.2892 0.2986
0.1859 0.2048
0.3540 0.3460
0.4880 0.5005
(OpenAI)
Table 3: Zero-shot downstream evaluation on Llama-1B. Adding the PC layer improves average accuracy under both optimizers: from 0.4539 to 0.4745 under AdamW and from 0.4880 to 0.5005 under Muon, with consistent per-task improvements on 8 out of 9 tasks in each case. Best results within each optimizer block are in bold. To reduce length bias in candidate scoring, we report length-normalized accuracy (Accn , i.e. acc_norm) when available, namely for HellaSwag, ARC-E, ARC-C, and OBQA; the remaining tasks use plain accuracy (Acc). The average is computed over all nine reported task metrics.
4.3
Computational and Memory Cost Analysis
The polynomial preconditioner can be computed efficiently via Horner’s method (see Appendix F.3 for the implementation details). We provide a theoretical FLOPs analysis in Appendix G.1 based on matrixmultiplication counts. For Llama-1B training, the estimated relative FLOPs overhead is at most 0.39% under the AdamW default configuration (pc_level = 4) and 0.24% when PC is combined with Muon (pc_level = 2). In terms of memory, under the Llama-1B setting, the PC layer increases the per-GPU peak active memory by approximately 9.56% under AdamW and 8.73% under Muon (see Appendix G.2). As discussed in Remark 1, the learned preconditioned weights can be absorbed into the original architecture after training; hence inference incurs no additional computation or memory overhead.
11
5
PC Improves the Weight Spectrum
In this section, we quantify how PC improves the spectral conditioning of transformer weight matrices. We first propose a robust spectral metric, the modified condition number, and then extend it to a global model-level measure that we track during training. Modified condition number. First, we introduce a Modified Condition Number metric, denoted as κ̃ (W ), to robustly evaluate the conditioning of a weight matrix W. The traditional condition number σmax /σmin is numerically fragile, because the smallest singular value of a trained network is often extremely close to zero, which makes the ratio blow up. To mitigate this and capture the effective spectral spread, we define κ̃ (W ) as the ratio of the top singular value to the average of the bottom 10% singular values. Let n = min(m, d) be the number of singular values of W ∈ Rm×d , ordered such that σ1 ≥ σ2 ≥ · · · ≥ σn . The modified condition number is formally defined as: κ̃ (W ) =
σ1 σ̄bottom-10%
where σ1 denotes the largest singular value and σ̄bottom-10% denotes the average of the smallest ⌈0.1n⌉ singular values. Averaging over the bottom 10% rather than taking the single σmin smooths out the near-zero tail, giving a stable measure of how widely the spectrum is spread. Global modified condition number. To quantify the overall spectral health of the model, we further aggregate κ̃ (W ) computed on each critical weight block into a Global Modified Condition Number (GMCN, κ̃). We employ the geometric mean of the modified condition numbers across all critical weight matrices. The geometric mean is preferred over a simple product for numerical stability when aggregating across multiple blocks: ! 1/N
κ̃ =
∏ ∏ κ̃ (Wl,b′ )
l ∈ L b∈ Bl
where L represents the set of transformer layers, Bl = {WQ , WK , WV , WO , Wgate , Wup , Wdown } denotes the set of key weight blocks in layer l, and N = | L| · | Bl | is the total number of blocks.5 Crucially, the evaluated matrix ′ depends on the optimization strategy: for blocks targeted by our preconditioner (W , W Wl,b gate , Wup , Wdown ), O ′ = PC(W ), which we refer to throughout as the effective we compute κ̃ on the preconditioned matrix Wl,b l,b ′ =W . weight; for non-preconditioned blocks (WQ , WK , WV ), we use the original weights, Wl,b l,b Experimental analysis of κ̃. The same geometric mean can be restricted to any subset of blocks: we report it over the PC-targeted blocks (ffn and WO ), over the non-preconditioned attention blocks (WQ , WK , WV ), and over all blocks (the global GMCN). We monitor the modified condition number throughout training (Figure 5). Here the baseline curve uses the original weights from the baseline run, while the PC curve uses the effective weights PC(W ) from a separate PC run. On the blocks directly targeted by PC, namely ffn and WO , PC shows a short early transient and then steadily lowers κ̃. At the final checkpoint, this aggregate decreases from about 17.3 for the baseline to about 7.7 for PC. We also examine WQ , WK , and WV , which are not directly preconditioned. Even though these blocks are never preconditioned, in the PC run their aggregate κ̃ ends up clearly below the baseline near the end of training. This indicates that applying PC to ffn and WO not only avoids harming the untargeted weights but also improves their conditioning indirectly. For the global aggregate, the baseline GMCN grows to about 42.4, while PC stabilizes around 25.0, giving a roughly 41% reduction. This verifies that PC substantially improves the effective conditioning of the model weights during training. 5 In the Llama 2 architecture, W gate , Wup , and Wdown refer to the three linear projections in the feed-forward network (FFN): the gate (gate_proj), up (up_proj), and down (down_proj) matrices, respectively.
12
GMCN (PC blocks) vs Tokens
18
GMCN vs Tokens
AdamW baseline AdamW + PC
120
14
40 35 30
10 8
GMCN
100
12
GMCN
GMCN
GMCN (non-PC blocks) vs Tokens
140
16
80 60
6
AdamW baseline AdamW + PC
4 2
0
20
40
60
80
100
120
Total Tokens Trained (B)
140
20 15
40
160
25
AdamW baseline AdamW + PC
10 0
20
(a) FFN + WO
40
60
80
100
Total Tokens Trained (B)
120
140
5
160
0
20
40
(b) WQ , WK , WV
60
80
100
Total Tokens Trained (B)
120
140
160
(c) Global
Figure 5: Evolution of modified condition number under AdamW. We report κ̃ for the preconditioned blocks (ffn and WO ), the non-preconditioned attention-input blocks (WQ , WK , WV ), and their global aggregate. The baseline curves are computed on the original weights, while the PC curves are computed on the effective weights used during the PC trajectory. PC improves the conditioning of the blocks to which it is applied, while the non-preconditioned blocks also show a clear reduction in κ̃ rather than any deterioration. Together, these reductions give a lower model-level GMCN.
0.8
1.0
0.8
80
0.4
/ max
0.6
0.8
00.0
1.0
(e) Layer 10: WO
Frequency
150 100 50 0.2
0.4
0.6
/ max
0.8
(i) Layer 18: WO
1.0
0.6
0.8
1.0
AdamW baseline AdamW + PC
30 20 10
0.2
0.4
/ max
0.6
0.8
00.0
1.0
0.2
0.4
/ max
0.6
0.8
1.0
80
AdamW baseline AdamW + PC
60
/ max
0.6
0.8
1.0
(j) Layer 18: Wgate
00.0
/ max
0.6
0.8
1.0
AdamW baseline AdamW + PC
0.2
0.4
/ max
0.6
0.8
1.0
AdamW baseline AdamW + PC
50 40 30 20
20
0.4
0.4
(h) Layer 10: Wdown
40
0.2
60 50 40 30 20 10 00.0
(g) Layer 10: Wup
AdamW baseline AdamW + PC
0.2
(d) Layer 2: Wdown
40
Frequency
60 50 40 30 20 10 00.0
/ max
50
(f) Layer 10: Wgate
AdamW baseline AdamW + PC
0.4
(c) Layer 2: Wup
AdamW baseline AdamW + PC
60
0.2
Frequency
0.2
00.0
1.0
20
10
Frequency
Frequency 0.6
40
20
00.0
/ max
Frequency
Frequency
30
200
0.4
(b) Layer 2: Wgate
AdamW baseline AdamW + PC
40
0.2
AdamW baseline AdamW + PC
Frequency
0.6
/ max
60 50 40 30 20 10 00.0
Frequency
0.4
10
Frequency
0.2
50
250
20
20 00.0
AdamW baseline AdamW + PC
30
40
(a) Layer 2: WO
00.0
40
AdamW baseline AdamW + PC
60
Frequency
AdamW baseline AdamW + PC
Frequency
300 250 200 150 100 50 00.0
10 0.2
0.4
/ max
0.6
0.8
(k) Layer 18: Wup
1.0
00.0
0.2
0.4
/ max
0.6
0.8
1.0
(l) Layer 18: Wdown
Figure 6: Singular-value histograms at the final-step checkpoint (AdamW, Llama-1B). We visualize the singular-value spectra for representative layers (2, 10, and 18 out of 18) and PC_blocks (WO , Wgate , Wup , Wdown ). For the baseline, spectra are computed on the original weights W of the baseline-trained model; for PC, spectra are computed on the effective preconditioned matrices PC(W ) from the PC-trained model. Within each subplot, singular values are rescaled by the largest singular value of that matrix so that all spectra lie in [0, 1], which makes the spectral shape directly comparable across blocks and depths; the y-axis reports the bin frequency (raw counts per bin). Across depths and blocks, PC moves more singular-value mass away from the lower end of the normalized spectrum and into a more regular middle range.
Singular-value spectrum comparison. To further investigate how preconditioning reshapes the full singular-value spectrum across individual blocks and depths, we examine singular-value distributions at the final checkpoint. For the transformer baseline, we compute singular values from the original weights W; for the PC run, we compute singular values of the effective (preconditioned) weight matrices, i.e., PC(W ). To probe depth-wise behavior while avoiding boundary effects, we report three representative layers in our
13
18-layer Llama-1B model: layer 2 (shallow), layer 10 (middle), and layer 18 (deep). For each selected layer, we plot singular-value histograms (with bin frequencies on the y-axis) for all PC_blocks: WO , Wgate , Wup and Wdown , where each block’s singular values are first rescaled by their per-matrix maximum so that all spectra share the support [0, 1], making the spectral shapes directly comparable (see Figure 6). Across layers and blocks, the AdamW baseline places substantial histogram mass near the lower end of the normalized singular-value spectrum. PC moves more of this mass into the middle range. The spectra are therefore less concentrated at the lower end and have a narrower relative spread. The dominant visible effect is that PC lifts the lower part of the normalized spectrum and narrows the relative spread, consistent with the intended soft spectrum-shaping behavior of the PC polynomial and with the lower modified condition numbers reported above.
6
Ablation Studies
To better understand the contributions of individual components within the PC module, we conduct ablation studies. Unless otherwise specified, ablations are conducted on the Llama-271M model. The pc_level and PC-block ablations report results under both AdamW and Muon, while the norm-recovery and learnable-γ ablations focus on the AdamW setting. We examine four aspects in order: the polynomial degree (pc_level), the choice of PC blocks, the effect of norm recovery, and the use of the learnable scalar γ. Unless otherwise specified, each ablation varies only the factor under study and keeps all remaining settings at their default values described in Section 4.1. The default PC configuration is PC_blocks = {ffn, WO } and spectral-norm normalization with 10 power-iteration steps, with pc_level = 4 for AdamW and pc_level = 2 for Muon. The no-PC baseline has a final validation loss of 2.5944 under AdamW and 2.4831 under Muon, which serve as the references in this section.
6.1
PC Level Selection
Recall that pc_level is the degree k of pk , equivalently the strength knob of the PC map: larger values use a higher-degree, more aggressive spectral shaping (Table 1). We evaluate different PC levels by sweeping pc_level ∈ {1, 2, 3, 4} under the spectral-norm normalizer using 10 power-iteration steps. Figure 7(a) reports the final validation loss against the AdamW baseline (2.5944). pc_level = 1 is slightly worse than the baseline, and increasing the degree consistently improves validation performance, with pc_level = 4 achieving the lowest loss (2.5397). Within the tested range k ∈ {1, 2, 3, 4}, this monotone trend indicates that stronger polynomial preconditioning is more effective at reshaping the weight spectrum under AdamW. Under Muon (with the same spectral-norm normalizer), this monotone trend reverses: sweeping the same range, pc_level = 2 achieves the lowest final validation loss and higher degrees degrade performance (Figure 7(b)). Intuitively, Muon already exerts an implicit form of spectral control on the update matrices, so the residual spectral shaping required from PC is smaller and a lower-degree polynomial suffices; pushing pc_level higher risks over-constraining the weight spectrum. We therefore adopt pc_level = 2 as the default under Muon, and discuss the mechanism behind this difference in Appendix K.2.
6.2
Choice of PC Blocks
Figure 8 studies which transformer weight blocks should be equipped with PC. We compare configurations obtained by adding different attention-side projections to ffn, with candidates {(WQ , WK ), WO , WV }, where (WQ , WK ) denotes applying PC jointly to the query and key projections. We keep ffn enabled in all variants. Under AdamW, all tested PC variants improve over the no-PC baseline, and our default choice PC_blocks = {ffn, WO } is among the best-performing configurations, reducing the baseline loss by 5.47 × 10−2 . Although several configurations containing (WQ , WK ) achieve slightly lower validation loss, their advantage over our default is relatively small, within 6.6 × 10−3 in final validation loss, while requiring PC on two additional attention projections in each transformer block. Under Muon, the picture is more selective (Figure 8(b)): several configurations still improve and {ffn, WO } ties for the best-performing block choice, but adding 14
Transformer Transformer w/ PC Layer
2.6145
2.60 Baseline final val loss = 2.5944
2.58
2.5817
2.56 2.5542
2.54 2.52
2.5397
1
2
PC Level
3
Transformer Transformer w/ PC Layer
2.500
Final Validation Loss
Final Validation Loss
2.62
2.495 2.490 2.485
2.4853
2.4831
Baseline final val loss = 2.4831
2.480 2.4774
2.475 1
4
(a) AdamW. pc_level = 4 is best.
2.4975
2
PC Level
3
4
(b) Muon. Loss is non-monotone; pc_level = 2 is best.
Figure 7: pc_level ablation. Sweeping pc_level ∈ {1, 2, 3, 4} with {ffn, WO } fixed. (a) AdamW: larger degrees help. (b) Muon: non-monotone, pc_level = 2 optimal.
2.60
Transformer Transformer w/ PC Layer
2.58 2.56 2.54
2.495
Final Validation Loss
Final Validation Loss
certain attention-side projections can degrade performance relative to the Muon baseline. Since {ffn, WO } is among the best configurations under AdamW and tied for best under Muon, we adopt this simpler and more efficient configuration as the default across optimizers.
Baseline final val loss = 2.5944
2.5742 2.5623 Final choice 2.5468 2.5424 2.5397
2.5331 2.5353 2.5343
2.4924
2.490 2.485
Baseline final val loss = 2.4831
2.4840
2.4855
2.4841 2.4842
Final 2.480 2.4779 choice 2.4774 2.4774
2.52
2.475
FFN O V QK
FFN O V QK
PC Block Combinations
Transformer Transformer w/ PC Layer
PC Block Combinations
(a) AdamW.
(b) Muon.
Figure 8: PC block selection ablation. All variants include ffn; the hatched bar (“Final choice”) marks our default {ffn, WO }, which is among the best under both (a) AdamW and (b) Muon.
6.3
Effect of Weight Norm Recovery After Preconditioning
Recall from Algorithm 1 that PC first applies the polynomial to a spectrally normalized weight matrix. After this spectrum-shaping step, the PC module can further rescale the result in two ways: by recovering the spectral scale removed during normalization, and by applying a learnable scalar γ. We ablate these two components jointly, giving four variants with or without norm recovery and with or without γ. Table 4 reports the final-loss difference of these variants relative to the baseline, and Figure 9 shows the corresponding validation-loss curves. The dominant factor is norm recovery. Without norm recovery, both PC variants perform worse than the transformer baseline, regardless of whether γ is used. With norm recovery, the same polynomial spectrum shaping consistently improves the final validation loss. Thus, applying the polynomial to the normalized 15
weight alone is not sufficient; restoring the spectral scale after shaping is essential for turning PC into a consistent gain. The learnable scalar γ has a smaller effect on the final validation loss in this ablation, and its main role in stabilizing training dynamics is examined in the next subsection. 2.90
Transformer Baseline
poly(W/ W 2) poly(W/ W 2) W 2 poly(W/ W 2) W 2 poly(W/ W 2)
2.85
w/o norm recovery w/ norm recovery
w/o γ
w/ γ
+0.0447 ↑ −0.0480 ↓
+0.0430 ↑ −0.0547 ↓
Validation Loss
2.80
Setting
2.75 2.70 2.65 2.60 2.55 0
Table 4: Ablation on norm recovery and learnable γ. Loss difference denotes final validation loss minus the baseline final validation loss (2.5944). ↓ indicates an improvement over the baseline, and ↑ indicates worse performance than the baseline. The best result is bolded.
6.4
10
20
30
40
50
Total Tokens Trained (B)
60
Figure 9: Validation-loss curves for the norm-recovery ablation. The curves correspond to the four variants in Table 4. Norm recovery consistently moves PC from worsethan-baseline performance to clear improvement.
The Role of Learnable γ in the PC Module
We next isolate the role of the learnable scalar γ. Switching γ on or off changes the final validation loss only mildly compared with the change caused by norm recovery. Instead, its main effect appears in the training dynamics. Here we track the activation root mean square (RMS), computed on the outputs of the full Attention and FFN submodules, after the output projections and before residual addition. As shown by the RMS curves in Figure 10, γ stabilizes signal propagation. At the 1B scale, as shown in Figure 10, the run without γ exhibits sizable spikes in attention RMS and FFN RMS during training, whereas the γ-enabled run remains much smoother. We therefore adopt γ to stabilize training signal propagation.
10
1
10
0
10
1
Transformer Baseline
W poly( W ) W poly( W )
FFN Activation RMS
Attention Activation RMS
Transformer Baseline
0
20
40
60
80
100
Total Tokens Trained (B)
120
140
160
10
3
10
2
10
1
10
0
10
1
W poly( W ) W poly( W )
0
(a) Cross-layer average attention activation RMS.
20
40
60
80
100
Total Tokens Trained (B)
120
140
160
(b) Cross-layer average FFN activation RMS.
Figure 10: Learnable γ stabilizes activation RMS at the 1B scale. The run without γ shows many spikes in attention RMS and FFN RMS, while the run with γ remains stable.
16
7
Theory: Why Controlling the Spectrum?
We provide a theoretical justification for the spectrum-control principle underlying PC layers by analyzing deep linear networks. Specifically, if all weight matrices maintain uniformly bounded condition numbers during training, then the Jacobian is well-conditioned, leading to geometric convergence of gradient descent. Settings of theoretical analysis.
We consider an L-layer linear network that maps x ∈ Rdx to F (θ; x ) = WL WL−1 · · · W1 x ∈ Rdy ,
where θ = (W1 , . . . , WL ), and Wl ∈ Rdl ×dl −1 for l ∈ {1, . . . , L}, with d0 = d x and d L = dy . We assume a pyramidal architecture as follows. Assumption 1 (Pyramidal architecture). There exists some r ∈ {1, . . . , L}, such that d0 ≤ d1 ≤ · · · ≤ dr
and
d r ≥ d r +1 ≥ · · · ≥ d L .
Suppose that we have n training samples. We denote X = [ x1 , . . . , xn ] ∈ Rdx ×n as the input, F (θ; X ) = ( F (θ; x1 ); . . . ; F (θ; xn )) ∈ Rndy as the stacked model predictions, and y = (y1 ; . . . ; yn ) ∈ Rndy as the stacked labels. We further impose the following standard assumption on the input data matrix. Assumption 2 (Non-degenerate input data). The input matrix X ∈ Rdx ×n has full column rank, i.e., σmin ( X ) > 0. Assumption 2 implicitly places us in the over-parameterized regime, i.e., n ≤ d x . See Remark 2 in Appendix C for further discussion of the linear-network setting; see Remark 3 in Appendix C for the roles of Assumption 1 and 2 in our theoretical analysis. We consider the squared loss of the training data:
L(θ ) =
1 n 1 ∥ F (θ; xi ) − yi ∥2 = ∥ F (θ; X ) − y∥2 . ∑ 2 i =1 2
We analyze gradient descent applied to the loss L(θ ). Given an initial parameter θ (0), the iterates are updated as θ (t + 1) = θ (t) − η ∇L(θ (t)), where η > 0 denotes the learning rate. For given {(τl , µl )}lL=1 with τl ≥ 1 ≥ µl > 0, we define the region n o R := θ = (W1 , . . . , WL ) σmax (Wl ) ≤ τl , σmin (Wl ) ≥ µl , ∀l ∈ {1, 2, . . . , L} . Within R, each layer weight matrix is well-conditioned: its singular values are bounded away from both 0 and ∞. The next theorem shows that, as long as the training trajectory stays inside R, gradient descent achieves a geometric decrease. Theorem 1 (Geometric convergence within R). Suppose the iterates satisfy θ (t) ∈ R for all t ∈ {0, 1, . . . , T }. Define !2 q !2 L L β := ∏ τl 2L(θ (0)) + ∥ X ∥ F Lσmax ( X ), µ := ∏ µl σmin ( X )2 . l =1
l =1
Then for any learning rate η ∈ 0, 1/β , it holds that L(θ (t + 1)) ≤ 1 − ηµ L(θ (t)),
∀ t ∈ {0, 1, . . . , T }.
In particular, choosing η = 1/β yields the contraction factor 1 − µ/β: µ L(θ (t + 1)) ≤ 1 − L(θ (t)), ∀ t ∈ {0, 1, . . . , T }. β
17
The geometric decrease in Theorem 1 is closely related to global convergence. If the entire training trajectory satisfies θ (t) ∈ R for all t ≥ 0, then iterating the one-step contraction yields
L(θ (t)) ≤ (1 − ηµ)t L(θ (0)) −−→ 0. t→∞
Proof sketch. The argument follows a direct implication chain: (i) For θ ∈ R, the weight matrices stay well-conditioned (each layer satisfies σmin (Wl ) ≥ µl and σmax (Wl ) ≤ τl ), which implies that the Jacobian matrix G (θ ) is well-conditioned, with its spectrum bounded as µI ⪯ G (θ )⊤ G (θ ) ⪯ βI; (ii) The above spectral bounds imply a Polyak–Łojasiewicz (PL) inequality and (local) smoothness of the loss function; (iii) under the PL inequality and smoothness condition, a standard convergence proof for gradient descent yields the geometric decrease of the squared loss. See Appendix C for the full proof and a detailed discussion. Dependence on global conditioning. Theorem 1 reveals a direct relationship between the weight conditioning within R and the convergence rate 1 − µ/β. Let κ (Wl ) := σmax (Wl )/σmin (Wl ) denote the layer-wise condition number, and define the global condition-number bound over R by L τ 1/L . κR := ∏ l µ l =1 l 2L , where C depends only on the data ( X, y ) and the initialization θ (0), but not on the Then β/µ = C κR layer-wise condition-number bounds. Hence better weight conditioning (smaller κR ) directly translates to a faster geometric convergence rate.
Corollary 1 (Iteration complexity). Under the assumptions of Theorem 1, take η = 1/β. Then for all t ∈ {0, 1, . . . , T }, µ t L(θ (t)) ≤ 1 − L(θ (0)). β Consequently, reaching L(θ ( T )) ≤ ϵ is guaranteed after L(θ (0)) 2L T = O κR log ϵ iterations within the region R. Corollary 1 implies that smaller layer-wise upper bounds τℓ /µℓ lead to a smaller global conditioning bound κR , and hence a smaller upper bound on the iteration complexity T required to reach a given error level (see Appendix D for the proof). This result provides theoretical support for the design principle underlying PC layers: better weight conditioning during training leads to faster convergence. The preceding sections have instantiated this principle via polynomial weight preconditioning.
8
Related Work
Weight control during NN training. Several prior works improve optimization by controlling the weight parameterization or statistics during training. One line of work focuses on channel-level weight normalizations to improve signal propagation and training stability. For example, weight normalization (WN) [Salimans and Kingma, 2016] and related schemes [Karras et al., 2024, Loshchilov et al., 2025, Franke et al., 2025, Fu et al., 2025] reparameterize weights along rows, columns, or filters to control the norms. Relevant variants include but are not limited to centered WN [Huang et al., 2017] and weight standardization [Qiao et al., 2019, Kolesnikov et al., 2020], which further control mean and variance along weight channels. Along a conceptually related line, [Xie et al., 2026b] proposes Manifold-Constrained Hyper-Connections (mHC) , which projects the residual mixing matrices onto the Birkhoff polytope (i.e., the manifold of non-negative matrices with each row and each column summing to one) to stabilize signal propagation.
18
Another line of work controls matrix-level properties of weight matrices. For example, some methods project weights into or onto Frobenius-norm hyperballs/spheres [Wen et al., 2025a, Ren et al., 2026], while others constrain the spectral norm of weight matrices [Yoshida and Miyato, 2017, Miyato et al., 2018, Xie et al., 2026a, Newhouse et al., 2025, Jiang et al., 2026]. In contrast, our PC layer targets spectrum conditioning: it regulates the relative structure of the singular-value spectrum, rather than upper-bounding or shrinking the maximum singular value. Weight spectrum conditioning during NN training. There are many methods that target spectrum conditioning during training. Some of them are not “built-in layer” solutions: some require explicit or repeated SVD/eigen-decomposition computations [Jia et al., 2017, Huang et al., 2018, Jiang et al., 2019, Li et al., 2019], while others incorporate constraints via optimizer co-designs [Cisse et al., 2017, Bansal et al., 2018, Bernstein, 2025b]. In contrast, our PC layer is designed as a plug-and-play module that can be inserted into existing architectures without modifying the optimizer or the overall training pipeline. A few recent works have also explored built-in conditioning mechanisms that operate through weight matrices. Saratchandran et al. [2024] preconditions weights via diagonal row/column equilibration and Saratchandran and Lucey [2026b,a] modifies the query, key, and value projection weights with additive correction terms or conditioned initialization to improve attention-Jacobian conditioning. These methods provide complementary evidence for the usefulness of weight conditioning, with empirical studies largely in CNNs, vision transformers and BERT-style masked-language-modeling settings. In contrast, PC targets decoder-only LLM pre-training and performs (low-degree) polynomial weight preconditioning. This polynomial formulation provides a flexible mechanism for encouraging the desired reshaping of the effective weight spectrum during training. Some recent works maintain the singular-value spectrum of weight matrices during training via orthogonalequivalence transformations [Qiu et al., 2026a,b, Shi et al., 2026]; in contrast, our PC layer actively reshapes the spectrum through polynomial preconditioning. From a conceptual standpoint, our PC layer is related to the principle behind orthogonal weight initialization [Xiao et al., 2018, Hu et al., 2020]. Orthogonal initialization is motivated by the dynamical isometry theory [Saxe et al., 2014, Pennington et al., 2017, 2018], which suggests that (approximately) orthogonal transformations help maintain a well-conditioned input–output Jacobian, thereby improving training dynamics. Subsequent works have extended the orthogonality principle to the entire training process [Jia et al., 2017, Cisse et al., 2017, Bansal et al., 2018, Huang et al., 2018, Li et al., 2019]. In contrast, PC layer performs soft spectrum conditioning rather than collapsing all singular values to one, thereby preserving controlled spectral variation and model expressivity. Convergence of neural network optimization. The output Jacobian (i.e., the Jacobian of the mapping from network parameters to the output vector) serves as a foundational mathematical object in the theoretical analysis of neural network training [Jacot et al., 2018, Lee et al., 2019, Du et al., 2019a, Arora et al., 2019c], generalization [Jacot et al., 2018, Chen et al., 2020, Simon et al., 2021] and architecture-dependent inductive biases [Bietti and Mairal, 2019, Yang and Littwin, 2021]. Many global convergence analyzes [Jacot et al., 2018, Lee et al., 2019, Du et al., 2019a, Arora et al., 2019c] utilized the following simple fact: gradient descent converges to a global minimum if the Jacobian remains non-degenerate along the iterates. Nevertheless, to rigorously guarantee this condition is nontrivial. A common route to ensure this condition is to assume an ultra-wide (overparameterized) regime, in which parameters move only mildly from initialization and the Jacobian stays close to its initial value throughout training [Lee et al., 2019, Du et al., 2019a, Arora et al., 2019c, Allen-Zhu et al., 2019, Xiao et al., 2020]. In contrast, our analysis does not rely on an ultra-width assumption; instead, our convergence analysis is based on the weight spectrum remaining bounded during training. This permits nontrivial parameter movement, while still implies non-degenerate Jacobian. From a strict theoretical perspective, the spectral bound is a simplifying assumption. However, it serves as a well-defined and architecturally actionable anchor that allows us to move from abstract convergence analysis to concrete model design.
19
Matrix preconditioning and matrix polynomials. Preconditioning is a fundamental technique in numerical linear algebra for accelerating iterative methods [Chen, 2005]. In general, it improves convergence by reducing the effective condition number or clustering the spectrum of the underlying operator, and has been widely used for solving large-scale linear systems [Saad, 2003, Benzi, 2002]. Among the preconditioning strategies, polynomial preconditioning leverages matrix polynomials as an efficient yet effective mechanism for spectrum shaping. By applying a suitably chosen low-degree polynomial to a matrix, one can enhance eigenvalue clustering, thereby improving the convergence of iterative solvers such as the conjugate gradient method [Johnson et al., 1983]. Importantly, it can be implemented using only a few matrix multiplications, without explicit matrix inverses or decompositions. In deep learning, a notable recent success of matrix-polynomial preconditioning is the Muon optimizer (MomentUm Orthogonalized by Newton–Schulz) [Jordan et al., 2024]. Muon applies a few Newton–Schulz steps to the momentum matrix to approximate its polar (orthogonal) factor, i.e., pushing singular values toward 1. While both the Muon and our PC layer leverage matrix polynomials, they differ in three key ways: (i) Where the polynomial acts. Muon acts on momentum, whereas PC acts on weight matrices. Muon is an optimization algorithm, while PC layer is a new architecture component. (ii) Target spectral transform. Muon’s polynomial is designed to (approximately) realize a matrix sign / polar-decomposition map that drives singular values toward 1 [Amsel et al., 2025]. In contrast, our PC layer approximates a piecewise-linear spectral map that stays close to 1, while allowing a controlled deviation on small singular values to preserve model expressivity. See Section 3.2 for the details of PC polynomials. (iii) Theoretical analysis. Muon, as a general-purpose optimizer, is analyzed under broad non-convex settings [Bernstein and Newhouse, 2024, Bernstein, 2025a, Shen et al., 2025, Wang et al., 2025, Li and Hong, 2025]. Because the PC layer is explicitly designed for neural networks, our analysis leverages the specific structure of neural networks. This allows us to derive problem-specific guarantees, such as convergence to global minima under certain conditions. Notably, the PC layer can be orthogonally combined with Muon, since the two methods act at different parts of the training process: Muon manipulates the momentum matrices, whereas our PC layer preconditions the weight matrices themselves. Empirically, adding PC on top of Muon yields further improvement (§4.2), with additional spectral and ablation analysis in Appendix K.
9
Conclusion and Limitations
In this work, we introduce a Preconditioning (PC) layer, a weight-based built-in module that maintains healthy weight conditioning throughout LLM pre-training. We provide theoretical support in deep linear networks, showing that uniformly bounding each layer’s condition number yields geometric convergence of gradient descent, which motivates training-time control of the weight spectrum. Guided by this principle, we develop an efficient polynomial preconditioning algorithm that reshapes singular-value spectra using only lightweight matrix multiplications, avoids expensive decompositions, and can be merged into the original model after training with no inference overhead. We empirically validate that PC consistently improves optimization in Llama-1B pre-training. Our empirical evaluation is constrained by available computational resources, so we mainly focus on Llama pre-training at 271M and 1B parameters. We do not yet report results on a broader range of architectures (e.g., non-Llama transformer variants) or substantially larger model scales. While the observed gains are consistent across these two settings, more extensive studies are needed to characterize the generality of PC across different model families and parameter regimes. In addition, PC introduces design choices (e.g., pc_level) that may admit layer-specific optima in transformers. For instance, shallow and deep layers may benefit from different degrees of preconditioning. Systematically exploring such heterogeneous configurations and developing principled or adaptive rules for selecting PC settings are promising directions that we leave to future work. 20
References Z. Allen-Zhu, Y. Li, and Z. Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning, pages 242–252. PMLR, 2019. N. Amsel, D. Persson, C. Musco, and R. M. Gower. The polar express: Optimal matrix sign methods and their application to the Muon algorithm. arXiv preprint arXiv:2505.16932, 2025. S. Arora, N. Cohen, N. Golowich, and W. Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations, 2019a. S. Arora, N. Cohen, W. Hu, and Y. Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019b. S. Arora, S. S. Du, W. Hu, Z. Li, R. R. Salakhutdinov, and R. Wang. On exact computation with an infinitely wide neural net. Advances in Neural Information Processing Systems, 32, 2019c. J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. N. Bansal, X. Chen, and Z. Wang. Can we gain more from orthogonality regularizations in training deep networks? Advances in Neural Information Processing Systems, 31, 2018. P. Bartlett, D. Helmbold, and P. Long. Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks. In International Conference on Machine Learning, pages 521–530. PMLR, 2018. M. Benzi. Preconditioning techniques for large linear systems: a survey. Journal of Computational Physics, 182 (2):418–477, 2002. J. Bernstein. Deriving Muon, 2025a. URL https://jeremybernste.in/writing/deriving-muon. J. Bernstein. Modular manifolds. Thinking Machines Lab: Connectionism, 2025b. doi: 10.64434/tml.20250926. https://thinkingmachines.ai/blog/modular-manifolds/. J. Bernstein and L. Newhouse. Old optimizer, new norm: An anthology. In OPT 2024: Optimization for Machine Learning, 2024. A. Bietti and J. Mairal. On the inductive bias of neural tangent kernels. Advances in Neural Information Processing Systems, 32, 2019. Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. PIQA: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432–7439, 2020. T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877–1901, 2020. K. Chen. Matrix preconditioning techniques and applications. Number 19. Cambridge University Press, 2005. L. Chen, J. Li, and Q. Liu. Muon optimizes under spectral norm constraints. arXiv preprint arXiv:2506.15054, 2025. Z. Chen, Y. Cao, Q. Gu, and T. Zhang. A generalized neural tangent kernel analysis for two-layer neural networks. Advances in Neural Information Processing Systems, 33:13363–13373, 2020. M. Cisse, P. Bojanowski, E. Grave, Y. Dauphin, and N. Usunier. Parseval networks: Improving robustness to adversarial examples. In International Conference on Machine Learning, pages 854–863. PMLR, 2017.
21
C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 conference of the north American chapter of the association for computational linguistics: Human language technologies, volume 1 (long and short papers), pages 2924–2936, 2019. P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? Try ARC, the AI2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018. S. Du and W. Hu. Width provably matters in optimization for deep linear neural networks. In International Conference on Machine Learning, pages 1655–1664. PMLR, 2019. S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pages 1675–1685. PMLR, 2019a. S. S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations, 2019b. T. Fang, A. Schwing, and R. Sun. Precondition layer and its use for GANs, 2021. URL https://openreview. net/forum?id=1yXhko8GZEE. J. K. Franke, U. Spiegelhalter, M. Nezhurina, J. Jitsev, F. Hutter, and M. Hefenbrock. Learning in compact spaces with approximately normalized transformer. arXiv preprint arXiv:2505.22014, 2025. Y. Fu, X. Dong, S. Diao, H. Ye, W. Byeon, Y. Karnati, L. Liebenwein, H. Zhang, N. Binder, M. Khadkevich, et al. Nemotron-Flash: Towards latency-optimal hybrid small language models. Advances in Neural Information Processing Systems, 38, 2025. S. Y. Gadre, G. Smyrnis, V. Shankar, S. Gururangan, M. Wortsman, R. Shao, J. Mercat, A. Fang, J. Li, S. Keh, et al. Language models scale reliably with over-training and on downstream tasks. arXiv preprint arXiv:2403.08540, 2024. L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. Le Noac’h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou. The language model evaluation harness, 07 2024. URL https://zenodo.org/records/12608602. X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. In International Conference on Artificial Intelligence and Statistics, pages 249–256. JMLR Workshop and Conference Proceedings, 2010. K. He, X. Zhang, S. Ren, and J. Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE International Conference on Computer Vision, pages 1026–1034, 2015. A. Henry, P. R. Dachapally, S. S. Pawar, and Y. Chen. Query-key normalization for transformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4246–4253, 2020. J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022. W. G. Horner. XXI. A new method of solving numerical equations of all orders, by continuous approximation. Philosophical Transactions of the Royal Society of London, (109):308–335, 1819. W. Hu, L. Xiao, and J. Pennington. Provable benefit of orthogonal initialization in optimizing deep linear networks. In International Conference on Learning Representations, 2020.
22
L. Huang, X. Liu, Y. Liu, B. Lang, and D. Tao. Centered weight normalization in accelerating training of deep neural networks. In Proceedings of the IEEE International Conference on Computer Vision, pages 2803–2811, 2017. L. Huang, X. Liu, B. Lang, A. Yu, Y. Wang, and B. Li. Orthogonal weight normalization: Solution to optimization over multiple dependent stiefel manifolds in deep neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018. L. Huang, J. Qin, Y. Zhou, F. Zhu, L. Liu, and L. Shao. Normalization techniques in training DNNs: Methodology, analysis and application. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(8): 10173–10196, 2023. S. Ioffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pages 448–456. PMLR, 2015. A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems, 31, 2018. Z. Ji and M. Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019. K. Jia, D. Tao, S. Gao, and X. Xu. Improving training of deep neural networks via singular value bounding. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 4344–4352, 2017. H. Jiang, Z. Chen, M. Chen, F. Liu, D. Wang, and T. Zhao. On computation and generalization of generative adversarial networks under spectrum control. In International Conference on Learning Representations, 2019. X. Jiang, A. Semenov, and S. U. Stich. Enhancing llm training via spectral clipping. arXiv preprint arXiv:2603.14315, 2026. O. G. Johnson, C. A. Micchelli, and G. Paul. Polynomial preconditioners for conjugate gradient calculations. SIAM Journal on Numerical Analysis, 20(2):362–376, 1983. K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse, and J. Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://kellerjordan.github.io/posts/muon/. J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. T. Karras, M. Aittala, J. Lehtinen, J. Hellsten, T. Aila, and S. Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 24174–24184, 2024. K. Kawaguchi. Deep learning without poor local minima. Advances in Neural Information Processing Systems, 29, 2016. D. P. Kingma. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. A. Kolesnikov, L. Beyer, X. Zhai, J. Puigcerver, J. Yung, S. Gelly, and N. Houlsby. Big transfer (bit): General visual representation learning. In European conference on computer vision, pages 491–507. Springer, 2020. T. Laurent and J. Brecht. Deep linear networks with arbitrary loss: All local minima are global. In International Conference on Machine Learning, pages 2902–2907. PMLR, 2018. J. Lee, L. Xiao, S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in Neural Information Processing Systems, 32, 2019. 23
J. Li and M. Hong. A note on the convergence of Muon. arXiv preprint arXiv:2502.02900, 2025. S. Li, K. Jia, Y. Wen, T. Liu, and D. Tao. Orthogonal deep neural networks. IEEE transactions on pattern analysis and machine intelligence, 43(4):1352–1368, 2019. W. Liang, T. Liu, L. Wright, W. Constable, A. Gu, C.-C. Huang, I. Zhang, W. Feng, H. Huang, J. Wang, S. Purandare, G. Nadathur, and S. Idreos. TorchTitan: One-stop PyTorch native solution for production ready LLM pretraining. In International Conference on Learning Representations, 2025. D. Lin, R. Sun, and Z. Zhang. Faster directional convergence of linear neural networks under spherically symmetric data. Advances in Neural Information Processing Systems, 34, 2021. J. Liu, J. Su, X. Yao, Z. Jiang, G. Lai, Y. Du, Y. Qin, W. Xu, E. Lu, J. Yan, et al. Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982, 2025. I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. I. Loshchilov, C.-P. Hsieh, S. Sun, and B. Ginsburg. nGPT: Normalized transformer with representation learning on the hypersphere. In International Conference on Learning Representations, 2025. T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In Proceedings of the 2018 conference on empirical methods in natural language processing, pages 2381–2391, 2018. T. Miyato, T. Kataoka, M. Koyama, and Y. Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018. L. Newhouse, R. P. Hess, F. Cesista, A. Zahorodnii, J. Bernstein, and P. Isola. Training transformers with enforced Lipschitz constants. arXiv preprint arXiv:2507.13338, 2025. Q. Nguyen and M. Hein. The loss surface of deep and wide neural networks. In International Conference on Machine Learning, pages 2603–2612. PMLR, 2017. Q. N. Nguyen and M. Mondelli. Global convergence of deep networks with one wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33:11961–11972, 2020. D. Paperno, G. Kruszewski, A. Lazaridou, N.-Q. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th annual meeting of the association for computational linguistics (volume 1: Long papers), pages 1525–1534, 2016. G. Penedo, H. Kydlíček, A. Lozhkov, M. Mitchell, C. A. Raffel, L. Von Werra, T. Wolf, et al. The FineWeb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:30811–30849, 2024. J. Pennington, S. Schoenholz, and S. Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. Advances in Neural Information Processing Systems, 30, 2017. J. Pennington, S. Schoenholz, and S. Ganguli. The emergence of spectral universality in deep networks. In International Conference on Artificial Intelligence and Statistics, pages 1924–1932. PMLR, 2018. S. Qiao, H. Wang, C. Liu, W. Shen, and A. Yuille. Micro-batch training with batch-channel normalization and weight standardization. arXiv preprint arXiv:1903.10520, 2019. Z. Qiu, S. Buchholz, T. Xiao, M. Dax, B. Schölkopf, and W. Liu. Reparameterized LLM training via orthogonal equivalence transformation. Advances in Neural Information Processing Systems, 38:140775–140821, 2026a.
24
Z. Qiu, L. Liu, A. Weller, H. Shi, and W. Liu. POET-X: Memory-efficient LLM training by scaling orthogonal transformation. arXiv preprint arXiv:2603.05500, 2026b. L. Ren, Y. Liu, Y. Shen, and W. Chen. Rethinking language model scaling under transferable hypersphere optimization. arXiv preprint arXiv:2603.28743, 2026. Y. Saad. Iterative methods for sparse linear systems. SIAM, 2003. K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. WinoGrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106, 2021. T. Salimans and D. P. Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in Neural Information Processing Systems, 29, 2016. H. Saratchandran and S. Lucey. Conditioned initialization for attention. In The Fourteenth International Conference on Learning Representations, 2026a. URL https://openreview.net/forum?id=cKNOCYPo2W. H. Saratchandran and S. Lucey. Spectral conditioning of attention improves transformer performance. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2026b. URL https://openreview. net/forum?id=RBx1AfoL2J. H. Saratchandran, T. X. Wang, and S. Lucey. Weight conditioning for smooth optimization of neural networks. In European Conference on Computer Vision, pages 310–325. Springer, 2024. N. Sardana, J. Portes, S. Doubov, and J. Frankle. Beyond Chinchilla-optimal: Accounting for inference in language model scaling laws. In International Conference on Machine Learning, pages 43445–43460. PMLR, 2024. A. M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference on Learning Representations, 2014. W. Shen, R. Huang, M. Huang, C. Shen, and J. Zhang. On the convergence analysis of Muon. arXiv preprint arXiv:2505.23737, 2025. K. Shi, H. Li, Z. Qiu, Y. Wen, S. Buchholz, and W. Liu. Pion: A spectrum-preserving optimizer via orthogonal equivalence transformation. arXiv preprint arXiv:2605.12492, 2026. J. B. Simon, M. Dickens, and M. R. DeWeese. Neural tangent kernel eigenvalues accurately predict generalization, 2021. URL https://openreview.net/forum?id=lycl1GD7fVP. A. Talmor, J. Herzig, N. Lourie, and J. Berant. Commonsenseqa: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4149– 4158, 2019. H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. D. Ulyanov, A. Vedaldi, and V. Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv:1607.08022, 2016. S. Wang, F. Zhang, J. Li, C. Du, C. Du, T. Pang, Z. Yang, M. Hong, and V. Y. Tan. Muon outperforms Adam in tail-end associative memory learning. arXiv preprint arXiv:2509.26030, 2025. K. Wen, X. Dang, K. Lyu, T. Ma, and P. Liang. Fantastic pretraining optimizers and where to find them 2.1: Hyperball optimization, 12 2025a. URL https://tinyurl.com/muonh.
25
K. Wen, D. Hall, T. Ma, and P. Liang. Fantastic pretraining optimizers and where to find them. arXiv preprint arXiv:2509.02046, 2025b. Y. Wu and K. He. Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV), pages 3–19, 2018. L. Xiao, Y. Bahri, J. Sohl-Dickstein, S. Schoenholz, and J. Pennington. Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, pages 5393–5402. PMLR, 2018. L. Xiao, J. Pennington, and S. Schoenholz. Disentangling trainability and generalization in deep neural networks. In International Conference on Machine Learning, pages 10462–10472. PMLR, 2020. T. Xie, H. Luo, H. Tang, Y. Hu, J. K. Liu, Q. Ren, Y. Wang, W. X. Zhao, R. Yan, B. Su, et al. Controlled llm training on spectral sphere. arXiv preprint arXiv:2601.08393, 2026a. Z. Xie, Y. Wei, H. Cao, C. Zhao, C. Deng, J. Li, D. Dai, H. Gao, J. Chang, K. Yu, L. Zhao, S. Zhou, Z. Xu, Z. Zhang, W. Zeng, S. Hu, Y. Wang, J. Yuan, L. Wang, and W. Liang. mHC: Manifold-constrained hyper-connections, 2026b. URL https://arxiv.org/abs/2512.24880. G. Yang and E. Littwin. Tensor programs IIb: Architectural universality of neural tangent kernel training dynamics. In International Conference on Machine Learning, pages 11762–11772. PMLR, 2021. G. Yang, E. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, J. Pachocki, W. Chen, and J. Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems, 34:17084–17097, 2021. G. Yang, J. B. Simon, and J. Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310.17813, 2023. Y. Yoshida and T. Miyato. Spectral norm regularization for improving the generalizability of deep learning. arXiv preprint arXiv:1705.10941, 2017. R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th annual meeting of the association for computational linguistics, pages 4791–4800, 2019. B. Zhang and R. Sennrich. Root mean square layer normalization. Advances in neural information processing systems, 32, 2019. D. Zou, P. M. Long, and Q. Gu. On the global convergence of training deep linear ResNets. In International Conference on Learning Representations, 2020.
26
A
Notation Definition
Theory (deep linear network) L∈N number of layers number of training samples n∈N Wl ∈ Rdl ×dl −1 , l ∈ {1, . . . , L} weight matrix of the l th layer θ = (W1 , . . . , WL ) collection of all parameters x ∈ Rd x ×1 input of the neural network collection of inputs X ∈ Rd x × n ndy ×1 y∈R collection of labels gradient descent iteration t ∈ {0, 1, . . . , T } µl ∈ R lower bound of weight matrix Wl spectrum in region R τl ∈ R upper bound of weight matrix Wl spectrum in region R R region where every Wl is well-conditioned (σmin (Wl ) ≥ µl , σmax (Wl ) ≤ τl ) Polynomial weight preconditioning W ∈ Rn × m rectangular matrix e ∈ Rn × m W scaled W whose spectrum lies in [0, 1] σ1 ≥ · · · ≥ σm ≥ 0 singular values (ordered from largest to smallest) p, a = ( a0 , . . . , ak ) polynomial p of degree k and its coefficients g ( σ ) = p ( σ2 ) σ induced scalar preconditioning map PC level; gk (σ ) = pk (σ2 )σ has overall degree 2k + 1 pc_level = k PLb (σ ) = min(σ/b, 1) piecewise-linear target with cutoff b > 0 s (W ) ≈ ∥W ∥ 2 streaming power-iteration estimate of the spectral norm γ∈R per-block learnable scalar (initialized to 1) Spectrum metric κ̃ (W ) = σ1 /σ̄bottom-10% modified condition number Table 5: Main symbols used throughout the paper.
27
B
Background on Preconditioning
Preconditioning originated in numerical linear algebra as a technique for accelerating iterative solvers for large sparse linear systems, where the cost of direct methods (e.g., Gaussian elimination) is prohibitive. The central observation is that the convergence speed of iterative methods is governed by the conditioning of the system matrix, and that a problem-specific transformation can substantially improve this conditioning. The canonical setting in which this idea is developed is solving a symmetric linear system Qx = b via the conjugate gradient method; we briefly review this setting below before turning to polynomial preconditioners, which are the variant most relevant to our work. Consider a linear system of equations Qx = b, where Q ∈ Rn×n is real symmetric, and b ∈ Rn×1 . Conjugate gradientp(CG) is one of the most popular methods to solve the system of equations. It has iteration complexity O( κ ( Q) log 1/ϵ), where κ ( Q) is the condition number of Q. For ill-conditioned problems (i.e., large κ ( Q)), the convergence can be slow. Thus, in practice, preconditioned CG is commonly used instead of the original CG. Suppose there is a certain way to find a preconditioner M that reduces the condition number, i.e., κ ( MQ) < κ ( Q). Define Q̃ = MQ and b̃ = Mb, then we can solve the alternative problem Q̃x = b̃, for which CG (and other gradient methods) converge faster. One simple example is Jacobi preconditioning √ (closely related to whitening in machine learning) where M is a diagonal matrix with Mii = 1/ Qii .
B.1
Polynomial Preconditioner
We review the polynomial preconditioners proposed by Johnson et al. [1983]. However, we do not directly utilize the polynomials proposed by Johnson et al. [1983] since our setting differs. But we do borrow two lessons, which we will explain at the end of this subsection. Consider a linear system of equations Qx = b, where Q ∈ Rn×n is real symmetric, and b ∈ Rn×1 . To find a polynomial preconditioning g( Q) = p( Q) Q which is well-conditioned, we only need to find a polynomial g such that g(λ) = p(λ)λ maps [λ1 , λm ] to [1 − ϵ, 1]. This can be formulated as an approximation theory problem: find a polynomial g( x ) that approximates a function f where f (0) = 0, f ([λ1 , λm ]) = 1. Since the exact values of λ1 , λm vary for each Q, we loosen the range from [λ1 , λm ] to [γ L , γU ] such that [γ L , γU ] ⊇ [λ1 , λm ]. Define Pk to be the set of all polynomials with degree no more than k, i.e., k
Pk = { p(λ) | p(λ) = ∑ a j λ j , a j ∈ R for any j ∈ {0, 1, . . . , k }}. j =0
Johnson et al. [1983] consider two polynomial preconditioners: minimax and least-squares polynomials. Minimax polynomials are the solution to the following problem: min
max
p∈ Pk λ∈[γ L ,γU ]
|1 − λp(λ)|.
(3) (µ(λ))
T
Denote g(λ) = λp(λ), then there is a closed-form solution to the above problem: g∗ (λ) = 1 − Tk+1 (µ(0)) , k +1
where µ(λ) = γUγ+Uγ−Lγ−L2λ , and Tk is the (first-kind) Chebyshev polynomial satisfying Tk (cos(z)) = cos(kz). Johnson et al. [1983] also consider least-squares polynomials, which are the solutions to the problem min g
Z γU
|1 − g(λ)|w(λ)dλ.
γL
28
(4)
There is a closed-form solution to the above problem since it is a quadratic problem in the coefficients of g. For the Jacobian weight function w(λ) = (γU − λ)α (λ − γL ) β where α ≥ β ≥ −1/2, the optimal solution g∗ stays positive in [γ L , γU ]. As we will implement the preconditioning on rectangular weight matrices while Q is a square matrix, we cannot directly borrow the polynomials used by Johnson et al. [1983]. Nevertheless, we borrow two lessons for our design. First, the polynomial preconditioner can be designed by solving an optimization problem (either minimax or least squares). Second, they found that least-squares polynomials perform better than minimax polynomials for iterative algorithms. For this reason, we adopt the least-squares polynomials instead of the minimax polynomials.
B.2
Relation between Preconditioning and the Spectrum
For researchers less familiar with preconditioning, there may appear to be a discrepancy between theoretical results and practical implementation: while theory often relates the condition number to convergence speed, in practice, convergence frequently depends on the behavior of the entire spectrum. This distinction is wellrecognized within numerical linear algebra and optimization. We have incorporated several considerations to bridge this conceptual gap, as elaborated below. (1) Spectral metrics and practical impact. (a) We utilize a modified condition number, defined as the ratio of the largest singular value to the average of the bottom 10% of singular values (see Sec. 5), to assess the health of the singular value distribution. This metric accounts for the broader spectral distribution rather than being sensitive only to the extreme outliers. While one could further investigate the optimality of this metric following Chen [2005] (Sec. 1.5), we found our current approach sufficient for achieving robust empirical results. (b) Our preconditioner is designed to improve the entire spectrum, not merely the extreme condition number. For example, if an initial spectrum is [1, 0.1, 0.1, 10−8 ], our method aims to shift it toward a more favorable distribution, such as [1, 0.2, 0.2, 2 × 10−8 ], thereby enhancing overall convergence characteristics. (2) Historical context in numerical analysis. The distinction between the condition number and the full spectrum has long been documented. Johnson et al. [1983] observed that "the optimal polynomial preconditioner M (which achieves the best condition number in a certain set) may map small eigenvalues of A into large ones of M−1 A," which can degrade convergence. They argued that minimizing the condition number should not be the sole objective; instead, the goal should be the improvement of the whole spectrum. While the ideal metric remains a subject of discussion, we adopt their least-squares preconditioning approach as an effective alternative. Etymologically, "preconditioning" often evokes the "condition number," yet historically, the term encompasses the broader intent of spectral improvement. As noted in the preface of Chen [2005] (page xvi), the two most relevant terms in the field are "condition number and clustering." Section 1.5 of that text further emphasizes the importance of eigenvalues clustering. While a more descriptive name might be "spectral improver," we adhere to the standard terminology of "preconditioner." (3) Theoretical vs. empirical objectives. There remains an inherent disconnect in the optimization literature: researchers often prove results based on condition numbers (e.g., O(κ log 1/ϵ) iteration complexity) while designing algorithms that practically manipulate the full spectrum. This pattern is evident in foundational works; for instance, Johnson et al. [1983] aims to improve Conjugate Gradient bounds, and Nesterov’s acceleration is framed around improving the Gradient Descent bound. The community generally acknowledges that while the spectrum ultimately dictates performance, the condition number serves as a tractable, albeit weaker, proxy for theoretical analysis. Bridging this theoretical-practical divide remains an important area for future research. In summary, the standard paradigm, proving theorems on the condition number to justify methods that improve the spectrum, is a well-established practice in numerical analysis [Johnson et al., 1983]. We follow this convention and provide this discussion to clarify potential confusion regarding these related concepts.
29
C
Proof of Theorem 1
In this section, we will provide a detailed proof of Theorem 1. Note that, for the squared loss, the Gram matrix of the Jacobian coincides with the neural tangent kernel (NTK); therefore, conditioning of the Jacobian is equivalent to conditioning of the NTK. Throughout the proofs in this appendix, we will use the NTK-based formulation. We decompose the argument into three steps. • Step 1 (Appendix C.1) establishes weight conditioning → NTK conditioning: uniform lower and upper bounds on the singular values of all weight matrices imply, respectively, a lower and an upper bound on the eigenvalues of the empirical NTK along the optimization trajectory. • Step 2 (Appendix C.2) establishes NTK conditioning → optimization properties of the loss function: the NTK lower-eigenvalue bound implies a Polyak–Łojasiewicz (PL) inequality for the squared loss, and the NTK upper-eigenvalue bound implies (local) smoothness of the loss. • Step 3 (Appendix C.3) follows the standard global-convergence proof template for gradient descent: combining the descent lemma (from local smoothness) with the PL inequality yields a per-iteration contraction of the loss, and hence the usual geometric convergence conclusion. Throughout, we keep the constants explicit so that the final convergence rate can be read off directly from the weight spectral bounds. For the reader’s convenience, we restate the setting and Theorem 1 here before giving the proof. We consider an L-layer linear network that maps x ∈ Rdx to F (θ; x ) = WL WL−1 · · · W1 x ∈ Rdy , where θ = (W1 , . . . , WL ), and Wl ∈ Rdl ×dl −1 for l ∈ {1, . . . , L}, with d0 = d x and d L = dy . Remark 2 (Deep linear networks). The deep linear network is a widely used setting in theoretical analysis. While the input–output map is linear, the objective is generally nonconvex due to the matrix-factorization parameterization, making the setting non-trivial. Deep linear networks have been extensively studied for optimization landscape characterizations [Kawaguchi, 2016, Laurent and Brecht, 2018], convergence and training dynamics [Arora et al., 2019a, Du and Hu, 2019, Hu et al., 2020, Lin et al., 2021], implicit bias [Arora et al., 2019b, Ji and Telgarsky, 2019], etc. Even global convergence in this setting is non-trivial and has received extensive study [Arora et al., 2019a, Bartlett et al., 2018, Ji and Telgarsky, 2019, Zou et al., 2020]; the closely related theory of dynamical isometry was also pioneered in linear models [Saxe et al., 2014]. Our analysis is in the same spirit as a line of works that aim at narrowing the gap between theory and practice through linear networks [Pennington et al., 2017, 2018, Xiao et al., 2018, Lee et al., 2019, Xiao et al., 2020, Hu et al., 2020, Lin et al., 2021]: we use this setting to extract the actionable insight that a well-behaved weight spectrum facilitates training, which in turn motivates our algorithm design. We assume a pyramidal architecture as follows. Assumption 1 (Restated) (Pyramidal architecture). There exists some r ∈ {1, . . . , L}, such that d0 ≤ d1 ≤ · · · ≤ dr
and
d r ≥ d r +1 ≥ · · · ≥ d L .
Suppose that we have n training samples. We denote X = [ x1 , . . . , xn ] ∈ Rdx ×n as the input, F (θ; X ) = ( F (θ; x1 ); . . . ; F (θ; xn )) ∈ Rndy as the stacked model predictions, and y = (y1 ; . . . ; yn ) ∈ Rndy as the stacked labels. We further impose the following standard assumption on the input data matrix. Assumption 2 (Restated) (Non-degenerate input data). The input matrix X ∈ Rdx ×n has full column rank, i.e., σmin ( X ) > 0. Assumption 2 implicitly places us in the over-parameterized regime, which implies n ≤ d x . Remark 3 (Pyramidal architecture). The pyramidal architecture assumption (Assumption 1) is not particularly restrictive, and it has been adopted in prior theory as a convenient way to exclude rank-bottleneck pathologies while 30
still covering realistic width profiles. First, the widely used “equal-width” hidden-layer assumption is a special case of a pyramidal topology, and it is standard in over-parameterization/NTK-style (neural tangent kernel) analyses for fully-connected networks [Hu et al., 2020, Allen-Zhu et al., 2019, Du et al., 2019b]. Second, pyramidal architecture is a common structural assumption in prior optimization analyses to obtain benign loss landscape and global convergence guarantees [Nguyen and Hein, 2017, Nguyen and Mondelli, 2020, Laurent and Brecht, 2018]. In our theoretical analysis, together with the full-column-rank input assumption (Assumption 2), the pyramidal architecture provides a convenient sufficient condition for obtaining a non-degenerate Jacobian, which is exactly what enables the convergence results presented later. We consider the quadratic loss of the training data:
L(θ ) =
1 1 n ∥ F (θ; xi ) − yi ∥2 = ∥ F (θ; X ) − y∥2 . ∑ 2 i =1 2
We analyze gradient descent applied to the loss L(θ ). Given an initial parameter θ (0), the iterates are updated as θ (t + 1) = θ (t) − η ∇L(θ (t)), where η > 0 denotes the learning rate. For given {(τl , µl )}lL=1 with τl ≥ 1 ≥ µl > 0, we define the regions
Wlow (µ1 , . . . , µ L ) := {θ = (W1 , . . . , WL ) | σmin (Wl ) ≥ µl , ∀l ∈ {1, . . . , L}}, Wup (τ1 , . . . , τL ) := {θ = (W1 , . . . , WL ) | σmax (Wl ) ≤ τl , ∀l ∈ {1, . . . , L}}, R := Wup (τ1 , . . . , τL ) ∩ Wlow (µ1 , . . . , µ L ). Within R, each layer weight matrix is well-conditioned: its singular values are bounded away from both 0 and ∞. The next theorem shows that, as long as the training trajectory stays inside R, gradient descent achieves a geometric decrease. Theorem 1 (Restated) (Geometric convergence within R). Suppose the iterates satisfy θ (t) ∈ R for all t ∈ {0, 1, . . . , T }. Define !2 q !2 L L 2L(θ (0)) + ∥ X ∥ F Lσmax ( X ), β := ∏ τl µ := ∏ µl σmin ( X )2 . l =1
l =1
Then for any learning rate η ∈ 0, 1/β , it holds that L(θ (t + 1)) ≤ 1 − ηµ L(θ (t)),
∀ t ∈ {0, 1, . . . , T }.
In particular, choosing η = 1/β yields the contraction factor 1 − µ/β: µ L(θ (t + 1)) ≤ 1 − L(θ (t)), ∀ t ∈ {0, 1, . . . , T }. β
C.1
From Weight Matrices Conditioning to NTK Conditioning
In this subsection, we first review some basics of the neural tangent kernel (NTK), then derive the upper bound (Lemma 1) and the lower bound (Lemma 3) of NTK eigenvalues. The neural tangent kernel (NTK) was introduced by Jacot et al. [2018] (see also Du et al. [2019a]). For a ∂F (θ;x ) sample xi , the parameter Jacobian is ∂θ i ∈ RP×dy , where P is the number of parameters. Stacking the Jacobians for n samples gives the (global) Jacobian matrix ∂F (θ; X ) ∂F (θ; x1 ) ∂F (θ; xn ) G (θ ) := = ,..., ∈ RP×ndy . ∂θ ∂θ ∂θ The neural tangent kernel (NTK) is then defined as K (θ ) := G (θ )⊤ G (θ ) ∈ Rndy ×ndy . 31
For the L-layer linear network with weights {Wl }lL=1 and widths d0 , . . . , d L (where Wl ∈ Rdl ×dl −1 and d L = dy ), we define the shorthand Wi:j := Wi Wi−1 · · · Wj
( i ≥ j ).
Then the Jacobian G (θ ) admits the block form X ⊗ (WL:2 )⊤ G1 (θ ) .. .. . . ⊤ G (θ ) = Gl (θ ) = Wl −1,1 X ⊗ (WL:l +1 ) , . . . . . . GL (θ ) WL−1:1 X ⊗ Id L where ⊗ denotes the Kronecker product and Id L is the d L × d L identity matrix. Here G (θ ) is formed by stacking the layerwise Jacobian blocks Gl (θ ) ∈ Rdl −1 dl ×nd L for l = 1, . . . , L, so that in total G (θ ) has P = ∑lL=1 dl −1 dl rows and nd L columns. Moreover, the NTK has the form L
K (θ ) = G (θ )⊤ G (θ ) = ∑ Gl (θ )⊤ Gl (θ ). l =1
Lemma 1 (Upper bound of NTK eigenvalues). Consider a deep linear network of any shape. For any θ ∈ Wup (τ1 , . . . , τL ), it holds that λmax (K (θ )) ≤ Lσmax ( X )2 (τ1 . . . τL )2 , or equivalently,
∥ G (θ )∥2 ≤
√
Lσmax ( X )τ1 . . . τL .
Proof. By the submultiplicative property of matrix norms, we have
∥Wl −1:1 X ∥2 ≤ ∥Wl −1 ∥2 . . . ∥W1 ∥2 ∥ X ∥2 ≤ σmax ( X )τ1 . . . τl −1 , ∥WL:l +1 ∥2 ≤ ∥WL ∥2 . . . ∥Wl +1 ∥2 ≤ τL . . . τl +1 . Since τl ≥ 1 for any l ∈ {1, 2, . . . , L}, we have σmax ( Gl (θ )) = σmax (Wl −1:1 X ⊗ (WL:l +1 )⊤ ) = σmax (Wl −1:1 X ) · σmax ((WL:l +1 )⊤ )
≤ (σmax ( X )τ1 . . . τl −1 ) · (τl +1 . . . τL ) ≤ σmax ( X )τ1 . . . τL . Since this holds for every l, we have L
L
l =1
l =1
λmax (K (θ )) ≤ ∑ λmax ( Gl (θ )⊤ Gl (θ )) ≤ ∑ σmax ( Gl (θ ))2 ≤ Lσmax ( X )2 (τ1 . . . τL )2 . Thus,
∥ G (θ )∥2 =
q
λmax ( G (θ )⊤ G (θ )) =
q
λmax (K (θ )) ≤
√
Lσmax ( X )τ1 . . . τL .
Lemma 1 shows that an upper bound on the spectral norms of the weight matrices (i.e., θ ∈ Wup (τ1 , . . . , τL )) directly implies an upper bound on the largest NTK eigenvalue λmax (K (θ )). Notably, this argument does not depend on any pyramid-like constraint on the network widths and therefore applies to deep linear networks of arbitrary shapes. In contrast, to derive a meaningful lower bound on λmin (K (θ )), the pyramidalarchitecture assumption (Assumption 1) becomes essential. Lemma 2. Suppose Am , Am−1 , . . . , A1 are matrices of size nm × nm−1 , nm−1 × nm−2 , . . . , n1 × n0 , where nm ≥ nm−1 ≥ · · · ≥ n0 . Suppose σmin ( Ai ) ≥ µi , i = 1, . . . , m. Then the product M = Am Am−1 . . . A1 satisfies σmin ( M ) ≥ µ1 µ2 . . . µm .
32
Proof. We first prove the following result: for a k1 × k2 matrix A and a k2 × k3 matrix B, where k1 ≥ k2 ≥ k3 , σmin ( AB) ≥ σmin ( A)σmin ( B). This is proved by the following chain of inequalities: λmin ( B T A T AB) = min u T B T A T ABu ≥λmin ( A T A)∥ Bu∥2 = λmin ( A T A)u T B T Bu ∥u∥=1
≥λmin ( A T A)λmin ( B T B). Since matrices A, B, and AB are all tall matrices, we have q q q σmin ( AB) = λmin (( AB)⊤ AB) ≥ λmin ( A⊤ A) λmin ( B⊤ B) = σmin ( A)σmin ( B). Applying the result multiple times, we immediately obtain the desired result. Lemma 3 (Lower bound of NTK eigenvalues). Consider an L-layer pyramidal linear network satisfying Assumptions 1 and 2. For any θ ∈ Wlow (µ1 , . . . , µ L ), it holds that λmin (K (θ )) ≥ σmin ( X )2 (µ1 . . . µ L )2 . Proof. Recall the pyramidal-architecture assumption (Assumption 1), we have d0 ≤ d1 ≤ · · · ≤ dr
and
d r ≥ d r +1 ≥ · · · ≥ d L .
We focus on the r-th Jacobian block Gr (θ ) = (Wr−1:1 X ) ⊗ (WL:r+1 )⊤ , which is a (dl −1 dl × dy n) matrix. Since the number of training samples satisfies n ≤ d x = d0 (a consequence of Assumption 2), Wr−1:1 X = Wr Wr−1 . . . W1 X is a multiplication of many tall matrices. By Lemma 2, we have σmin (Wr−1:1 X ) ≥ σmin ( X )µ1 . . . µr−1 . ⊤ ⊤ ⊤ Similarly, (WL:r+1 ) = Wr+1 . . . WL is also a multiplication of many tall matrices, so σmin ((WL:r+1 )⊤ ) ≥ µr+1 . . . µ L . T and σmin (WL:r +1 ) ≥ µr +1 . . . µ L . Since that Gr ( θ ) is also a tall matrix, we have
σmin ( Gr (θ )) = σndy ( Gr (θ )) = σn (Wr−1:1 X ) · σdy ((WL:r+1 )⊤ ) = σmin (Wr−1:1 X ) · σmin ((WL:r+1 )⊤ )
= (σmin ( X )µ1 . . . µr−1 )(µr+1 . . . µ L ) ≥ σmin ( X )µ1 . . . µr−1 µr µr+1 . . . µ L , where the inequality follows from the assumption that µr ≤ 1. Therefore, L
K (θ ) = ∑ Gl⊤ (θ ) Gl (θ ) ⪰ Gr⊤ (θ ) Gr (θ ) ⪰ (σmin ( X )µ1 . . . µ L )2 In . l =1
C.2
From NTK Conditioning to PL Inequality and Smoothness of the Loss
In this subsection, we convert the NTK eigenvalue bounds from Appendix C.1 into two key analytic properties of the training loss function. More precisely, Lemma 4 uses the lower bound on the minimum NTK eigenvalue to derive a Polyak–Łojasiewicz (PL) inequality for the squared loss, while Lemma 7 uses the upper bound on the maximum NTK eigenvalue to establish the local smoothness of the loss function. We note that both results hold provided that the parameters remain within the well-conditioned region. Lemma 4 (PL inequality). Consider an L-layer pyramidal linear network satisfying Assumptions 1 and 2. For any θ ∈ Wlow (µ1 , . . . , µ L ), it holds that ∥∇L(θ; X )∥2 ≥ 2µL(θ; X ), 33
where µ := σmin ( X )2 (µ1 µ2 . . . µ L )2 . Proof. For the quadratic loss function, we have
∇L(θ; X ) = G (θ )e(θ ),
where e(θ ) = F (θ; X ) − y.
By Lemma 3, we have
∥∇L(θ; X )∥2 = e(θ )⊤ G (θ )⊤ G (θ )e(θ ) ≥ µ∥e(θ )∥2 = 2µL(θ; X ).
Remark 4 (Global minimum is zero). To interpret Lemma 4 as a standard PL inequality, We note that the model is realizable in our setting, i.e., the network can interpolate the data, and hence the global minimum of L(·; X ) is zero. Under the pyramidal-architecture assumption (Assumption 1), the widths satisfy the no-bottleneck condition min
l =1,...,L−1
dl ≥ min(d0 , d L ).
It is shown by Laurent and Brecht [2018] that, under this condition, the deep linear parametrization A = WL · · · W1 can represent any linear map A ∈ Rdy ×dx . Since the input matrix X has full column rank (Assumption 2), there exists a linear map A⋆ ∈ Rdy ×dx that interpolates the data, i.e., A⋆ X = Y. Combining these two facts, there exist weights (W1⋆ , . . . , WL⋆ ) such that WL⋆ · · · W1⋆ X = Y, and hence L(θ ⋆ ; X ) = 0. Therefore, the global minimum value of L(·; X ) is zero, and Lemma 4 indeed provides the usual PL condition ∥∇L(θ; X )∥2 ≥ 2µ (L(θ; X ) − L(θ ∗ ; X )) . Lemma 5 (Lipschitz continuity of F). Consider an L-layer linear network of any shape. For any θ, θ̂ ∈ Wup (τ1 , . . . , τL ), it holds that √ ∥ F (θ; X ) − F (θ̂; X )∥ ≤ L∥ X ∥F τL . . . τ1 · ∥θ − θ̂ ∥. Proof. For any θ, θ̂ ∈ Wup (τ1 , . . . , τL ), we have
∥ F (θ; X ) − F (θ̂; X )∥ = ∥WL · · · W1 X − ŴL · · · Ŵ1 X ∥F ≤ ∥WL · · · W1 − ŴL · · · Ŵ1 ∥2 ∥ X ∥F L
∑ ŴL · · · Ŵl+1 (Wl − Ŵl )Wl−1 · · · W1 2 ∥X ∥F .
=
l =1
By the assumption τl ≥ 1 and Cauchy’s inequality, we obtain L
RHS of (5) ≤ ∥ X ∥ F ∑ ∥ŴL ∥2 . . . ∥Ŵl +1 ∥2 ∥Wl − Ŵl ∥2 ∥Ŵl −1 ∥2 . . . ∥Ŵ1 ∥2 l =1 L
≤ ∥ X ∥ F ∑ τL . . . τl +1 ∥Wi − Ŵi ∥2 τl −1 · · · τ1 i =1
L
≤ τL · · · τ1 ∥ X ∥ F ∑ ∥Wi − Ŵi ∥2 i =1
≤
√
v u L u Lτ · · · τ ∥ X ∥ t ∥W − Ŵ ∥2 . L
F
1
∑
i
i 2
i =1
We can relax
v v u L u L u u t ∥W − Ŵ ∥2 ≤ t ∥W − Ŵ ∥2 = ∥θ − θ̂ ∥,
∑
i =1
i
i 2
∑
i =1
thus obtaining the desired result. 34
i
i F
(5)
Lemma 6 (Lipschitz continuity of G). Consider a deep linear network of any shape. Assume σmax ( X ) ≤ τ0 . For any θ, θ̂ ∈ Wup (τ1 , . . . , τL ), it holds that
∥ G (θ ) − G (θ̂ )∥2 ≤ Lσmax ( X )τL . . . τ1 ∥θ − θ̂ ∥. Proof. Recall that G (θ )⊤ = [ G1 (θ )⊤ ; G2 (θ )⊤ ; . . . ; GL (θ )⊤ ], So we have
Gl (θ ) = (Wl −1:1 X ) ⊗ (WL:l +1 )⊤ .
q
λmax ( G (θ ) − G (θ̂ ))⊤ ( G (θ ) − G (θ̂ )) v u L u = tλmax ∑ ( Gl (θ ) − Gl (θ̂ ))⊤ ( Gl (θ ) − Gl (θ̂ ))
∥ G (θ ) − G (θ̂ )∥2 =
≤
√
(6)
l =1
L max ∥ Gl (θ ) − Gl (θ̂ )∥2 . l
We denote Gl (θ̂ ) = (Ŵl −1:1 X ) ⊗ (ŴL:l +1
)⊤ . Then for any l ∈ {1, 2, . . . , L}, we have
∥ Gl (θ ) − Gl (θ̂ )∥2 = ∥(Wl −1:1 X ) ⊗ (WL:l +1 )⊤ − (Ŵl −1:1 X ) ⊗ (ŴL:l +1 )⊤ ∥2 = ∥(Wl −1:1 X − Ŵl −1:1 X ) ⊗ (WL:l +1 )⊤ + (Ŵl −1:1 X ) ⊗ (WL:l +1 − ŴL:l +1 )⊤ ∥2 ≤ ∥(Wl −1:1 X − Ŵl −1:1 X ) ⊗ (WL:l +1 )⊤ ∥2 + ∥(Ŵl −1:1 X ) ⊗ (WL:l +1 − ŴL:l +1 )⊤ ∥2
(7)
= ∥Wl −1:1 X − Ŵl −1:1 X ∥2 ∥WL:l +1 ∥2 + ∥Ŵl −1:1 X ∥2 ∥WL:l +1 − ŴL:l +1 ∥2 . Following the same proof strategy as in Lemma 5, we obtain l −1
∥Wl −1:1 X − Ŵl −1:1 X ∥2 ≤ τl −1 · · · τ1 σmax ( X ) ∑ ∥Wi − Ŵi ∥2 , i =1
L
∥WL:l +1 − ŴL:l +1 ∥2 ≤ τl +1 · · · τL
∑ ∥Wi − Ŵi ∥2 .
i = l +1
By the submultiplicative property of matrix norms, we have
∥Ŵl −1:1 X ∥2 ≤ τl −1 · · · τ1 σmax ( X ),
∥WL:l +1 ∥2 ≤ τl +1 · · · τL .
Since we assume that τl ≥ 1, we get RHS of (7) l −1
≤
τl −1 · · · τ1 σmax ( X ) ∑ ∥Wi − Ŵi ∥2
!
L
(τl +1 · · · τL ) + (τl −1 · · · τ1 σmax ( X )) τl +1 · · · τL
i =1
∑ ∥Wi − Ŵi ∥2
!
i = l +1
L
≤ σmax ( X )τ1 · · · τL ∑ ∥Wi − Ŵi ∥2 ≤
√
l =1
Lσmax ( X )τ1 · · · τL ∥θ − θ̂ ∥.
Here, the last inequality follows the proof strategy of Lemma 5. Combining it with (6) and (7), we get √ ∥ G (θ ) − G (θ̂ )∥2 ≤ L max∥ Gl (θ ) − Gl (θ̂ )∥2 ≤ Lσmax ( X )τ1 · · · τL ∥θ − θ̂ ∥. l
Lemma 7 (Local smoothness of loss function). Consider an L-layer linear network F (θ; x ) = WL . . . W1 x. Then
∥∇L(θ; X ) − ∇L(θ̂; X )∥ ≤ β(θ )∥θ − θ̂ ∥,
35
∀θ, θ̂ ∈ Wup (τ1 , . . . , τL ),
where β(θ ) = Lσmax ( X )
q
2L(θ ) + ∥ X ∥F (τ1 · · · τL )2 .
Proof. For the quadratic loss function, we have
∇L(θ; X ) = G (θ )e(θ ),
where e(θ ) = F (θ; X ) − y.
For any θ, θ̂, the norm of their gradient difference satisfies
∥∇L(θ; X ) − ∇L(θ̂; X )∥ = ∥ G (θ )e(θ ) − G (θ̂ )e(θ̂ )∥ ≤ ∥ G (θ ) − G (θ̂ )∥2 ∥e(θ̂ )∥ + ∥ G (θ )∥2 ∥e(θ ) − e(θ̂ )∥ = ∥ G (θ ) − G (θ̂ )∥2 ∥e(θ )∥ + ∥ G (θ )∥2 ∥ F (θ; X ) − F (θ̂; X )∥. Applying Lemma 1, Lemma 5, and Lemma 6, since τl ≥ 1 for any l ∈ {1, 2, . . . , L}, we get
∥∇L(θ; X ) − ∇L(θ̂; X )∥
√ √ ≤ Lσmax ( X )τ1 · · · τL ∥θ − θ̂ ∥ · ∥e(θ )∥ + ( Lσmax ( X )τ1 · · · τL ) · ( L∥ X ∥F τ1 · · · τL · ∥θ − θ̂ ∥) ≤ Lσmax ( X ) ∥e(θ )∥τ1 · · · τL + ∥ X ∥F (τ1 · · · τL )2 ∥θ − θ̂ ∥ q 2L(θ ) + ∥ X ∥F (τ1 · · · τL )2 ∥θ − θ̂ ∥. ≤ Lσmax ( X )
C.3
Geometric Convergence within the Well-Conditioned Region R
In this subsection, we complete the proof of Theorem 1 by applying a standard gradient-descent analysis under the PL inequality and smoothness established in Appendix C.2. We choose a stepsize satisfying the usual stability condition determined by the smoothness constant (i.e., η ≤ 1/β), and then combine (i) the descent lemma with (ii) the PL inequality to obtain a one-step contraction of the loss. Proof of Theorem 1. We prove this theorem by induction. Suppose that the result holds for 0, 1, . . . , k − 1. Then we have L(θ (k)) ≤ L(θ (0)). Recall the definition of β(θ ) and β, q β(θ ) = Lσmax ( X ) 2L(θ ) + ∥ X ∥F (τ1 · · · τL )2 , q 2L(θ (0)) + ∥ X ∥F (τ1 · · · τL )2 . β = Lσmax ( X ) This implies β(θ (k )) ≤ β(θ (0)) = β.
(8)
For the notational simplicity, let us denote g(θ ) = ∇L(θ; X ), and gk = ∇L(θ (k ); X ). Denote δk = −ηgk = θ (k + 1) − θ (k ). We follow the proof of the descent lemma. We obtain
L(θ (k + 1)) = L(θ (k)) + gk⊤ δk + 2
Z 1 0
≤ L(θ (k)) − η ∥ gk ∥ +
[ g(θ (k) + tδk ) − g(θ (k))]⊤ δk dt
Z 1 0
(9)
∥ g(θ (k) + tδk ) − g(θ (k))∥ · ∥δk ∥dt.
We denote (k)
(k)
θ (k) := (W1 , . . . , WL ),
( k +1)
θ (k + 1) := (W1
then θ (k) + tδk = t(θ (k) + δk ) + (1 − t)θ (k)
= tθ (k + 1) + (1 − t)θ (k) 36
( k +1)
, . . . , WL
),
(k)
(k)
(k)
(k)
( k +1)
(k)
= (tW1 + (1 − t)W1 , . . . , tWL + (1 − t)WL ). Since θ (k ), θ (k + 1) ∈ Wup (τ1 , . . . , τL ), we have ∥Wl ∥2 , ∥Wl function, for any l ∈ {1, 2, . . . , L}, we have ( k +1)
∥tWl
(k)
( k +1)
+ (1 − t)Wl ∥2 ≤ t∥Wl
∥2 ≤ τl . Since every norm is a convex
(k)
∥2 + (1 − t)∥Wl ∥2 ≤ tτl + (1 − t)τl = τl .
This implies θ (k ) + tδk ∈ Wup (τ1 , . . . , τL ). According to Lemma 7 and (8), we have
∥ g(θ (k) + tδk ) − g(θ (k))∥ ≤ β(θ (k))∥tδk ∥ ≤ βt∥δk ∥. Plugging this into (9), we get
L(θ (k + 1)) ≤ L(θ (k)) − η ∥ gk ∥2 + β∥δk ∥2 For any 0 ≤ η ≤ 1/β, it holds that η −
βη 2 tdt = L(θ (k )) − η − ∥ gk ∥2 . 2 0
Z 1
η βη 2 2 ≥ 2 , so
βη 2 η L(θ (k + 1)) ≤ L(θ (k)) − η − ∥ gk ∥2 ≤ L(θ (k)) − ∥ gk ∥2 . 2 2 This completes the proof of the descent lemma. According to the PL inequality (Lemma 4), we get
∥ gk ∥2 ≥ 2µL(θ (k)). Combining the two inequalities above, we obtain η L(θ (k + 1)) ≤ L(θ (k)) − ∥ gk ∥2 ≤ L(θ (k)) − ηµL(θ (k)) = (1 − ηµ)L(θ (k)), 2 for η ∈ (0, 1/β]. This completes the proof.
37
D
Proof of Corollary 1
Before proving Corollary 1, we expand the ratio β/µ to make the data-dependent constant C explicit. Recall from Theorem 1 that !2 q !2 L L σmin ( X )2 . β = ∏ τl 2L(θ (0)) + ∥ X ∥ F Lσmax ( X ), µ = ∏ µl l =1
l =1
Dividing β by µ and grouping data-/initialization-dependent factors separately from the layer-wise conditionnumber bounds, we obtain !2 p L 2L(θ (0)) + ∥ X ∥ F Lσmax ( X ) τl β 2L · ∏ = C κR . = µ µ σmin ( X )2 l l = 1 | {z } =: C
Here
p C :=
2L(θ (0)) + ∥ X ∥ F Lσmax ( X ) σmin ( X )2
depends only on the data ( X, y), the depth L, and the initialization θ (0), and is independent of the layer-wise condition-number bounds {τl , µl }lL=1 . Proof of Corollary 1. Take η = 1/β. By Theorem 1, for every t ∈ {0, 1, . . . , T − 1}, µ L(θ (t)). L(θ (t + 1)) ≤ 1 − β Iterating this inequality yields the geometric bound. Formally, by induction on t: for t = 0 the claim is trivial. If L(θ (t)) ≤ (1 − µ/β)t L(θ (0)), then µ µ t +1 L(θ (t + 1)) ≤ 1 − L(θ (t)) ≤ 1 − L(θ (0)). β β Hence, for all t ∈ {0, 1, . . . , T }, µ t L(θ (0)). L(θ (t)) ≤ 1 − β For the iteration complexity, use the standard inequality 1 − x ≤ e− x for x ∈ [0, 1]. Since µ/β ∈ (0, 1], we have µ µ T 1− ≤ exp − T . β β Therefore, µ L(θ ( T )) ≤ exp − T L(θ (0)). β If T≥
L(θ (0)) L(θ (0)) β 2L log = C κR log , µ ϵ ϵ
2L log(L( θ (0)) /ϵ )) then exp(− β T ) L(θ (0)) ≤ ϵ, which implies L(θ ( T )) ≤ ϵ. This yields the claimed T = O(κR bound. µ
38
E
Proof of Proposition 1
Proof. Let W = UΣV ⊤ be an SVD with Σ = diag(σ1 , . . . , σm ) and σi ≥ 0. Then WW ⊤ = UΣ2 U ⊤ , and hence 2 )). Therefore p(WW ⊤ ) = U p(Σ2 )U ⊤ with p(Σ2 ) = diag( p(σ12 ), . . . , p(σm m g(W ) = p(WW ⊤ )W = U p(Σ2 )Σ V ⊤ = U D V ⊤ , D ≜ diag p(σi2 )σi i=1 . Note that the diagonal entries of D may be negative, while singular values are by definition nonnegative. Let | D | ≜ diag(| p(σi2 )σi |)im=1 . Writing D = S | D | with S = diag(sign( p(σi2 )σi )), we have g (W ) = U S | D | V ⊤ . Since US is still orthogonal, this exhibits an SVD of g(W ) with singular values given by the diagonal entries of | D |, i.e., {| p(σi2 )σi |}im=1 = {| g(σi )|}im=1 .
39
F
Details of PC Layer Implementation
F.1
Polynomial Fitting Algorithm
The detailed procedure of identifying polynomial g is summarized in Algorithm 2. Additionally, note that this procedure is “offline”, i.e., it is not needed during deep net training. Algorithm 2 POLY-FITTING 1: Input: target map f : [0, 1] → R (default f (σ ) = PLb (σ )), fitting interval [γ L , γU ] (default [0, 1.1]). 2: Hyper-parameters: degree k and α appearing in w (σ ) = σ α . i.i.d. 3: Sampling: draw σ1 , . . . , σn−1 ∼ Unif[γ L , γU ] and set σn = 1. 4: Solve: compute a⋆ ∈ Rk+1 by minimizing the discrete weighted least-squares objective (2). k 5: Output: g (σ ) = σ ∑t=0 a⋆t σ2t .
In Algorithm 2, we include the deterministic sample σn = 1 to anchor the least-squares fit at the nominal unit scale after spectral normalization. The enlarged fitting interval [0, 1.1] is used only as a robustness margin for spectral-norm estimation error; σ = 1 remains the reference value for the intended normalized top singular scale. Since our target maps satisfy f (1) = 1, including this point encourages g(1) ≈ f (1) and prevents the finite-sample fit from unintentionally changing the top normalized scale.
F.2
Streaming Power Iteration for Spectral Normalization
Algorithm 1 normalizes each selected weight matrix by a scalar s(W ) that approximates the spectral norm ∥W ∥2 . This subsection describes how s(W ) is computed in practice. During training, each PC block maintains two auxiliary buffers u and v that estimate the top left and right singular vectors of the current weight matrix. At training step t, instead of running power iteration from scratch, we initialize the iteration from the buffers saved at the previous step, (ut−1 , vt−1 ). After a small number of power-iteration steps, the updated vectors are written back to the buffers and the Rayleigh quotient is used as the spectral-norm estimate. We refer to this procedure as streaming power iteration because the singular-vector estimates are carried over across consecutive training steps as the weights evolve. This design avoids exact SVD computation and exploits the fact that weights change smoothly across consecutive training steps, so that estimates from the previous step provide a high-quality initialization for the current step. Algorithm 3 Streaming power iteration for spectral-norm normalization 1: Input: weight matrix Wt ∈ Rn×m in PC_blocks at training step t; persistent buffers ut−1 ∈ Rn , vt−1 ∈ Rm
from the previous step; number of power-iteration steps q (default q = 10); stability constant ϵ > 0. 2: if ut−1 , vt−1 are uninitialized then 3: Initialize ut−1 and vt−1 as random unit vectors. 4: end if 5: u ← ut−1 , v ← v t −1 6: for i = 1, . . . , q do Wt⊤ u Wt v 7: v← . , u ← ∥W ∥Wt⊤ u∥2 t v ∥2 8: end for
# warm-start from the previous training step
9: Update buffers: ut ← u, vt ← v. 10: Estimate spectral-norm: st (Wt ) ← u⊤ Wt v + ϵ. 11: Return: st (Wt ), ut , vt .
In all experiments, we use q = 10 power-iteration steps per training step. The streaming power-iteration routine is used only during training to estimate s(Wt ). After training, we materialize the final effective weight 40
PC(W ) and store it as the corresponding weight in the original transformer architecture; the power-iteration buffers are then discarded and no spectral-norm estimation is needed during inference.
F.3
Tricks for Matrix Polynomial Calculation
We discuss a few computational tricks to implement a PC layer. These tricks can greatly reduce the computation cost. In the following, suppose we decide to use a polynomial g(σ ) = σ p(σ2 ) to implement the PC layer. We use
Trick 1: choose an efficient equivalent form. ⊤
g( A) = A p( A A)
if A is tall,
otherwise use
g( A) = p( AA⊤ ) A.
Note that g( A) = p( AA⊤ ) A = A p( A⊤ A). We call p( AA⊤ ) A the AA⊤ -form and A p( A⊤ A) the A⊤ A-form. Although the two forms give the same value, their implementation time can differ a lot. For A ∈ Rn×m with n ≥ m, forming A⊤ A costs O(m2 n), while AA⊤ costs O(n2 m); hence pick the A⊤ A-form when A is tall (n ≥ m), and the AA⊤ -form when A is wide (n < m). Trick 2: cache and reuse B.
Store
( B=
A⊤ A,
if A is tall,
AA⊤ ,
if A is wide,
and reuse it. We will compute A p( B) (tall case) or p( B) A (wide case). Trick 3: evaluate p via Horner’s method [Horner, 1819].
Write a degree-k polynomial as
k p ( σ ) = ∑ a i σ i = a 0 + σ a 1 + σ a 2 + · · · + σ ( a k −1 + a k σ ) . i =0
A naive implementation k
p ( B ) = ∑ ai Bi i =0
k ( k −1) requires ∑ik=1 (i − 1) = matrix multiplications and k additions. Using Horner’s method, we only 2
need k matrix multiplications. For example, for a degree-4 polynomial p (corresponding to a degree-9 g(σ ) = σp(σ2 )), Horner’s method reduces the work from 6 matrix multiplications to 4. For a degree-15 g (i.e., k = 7), it reduces 21 multiplications to 7. Summary of the three tricks. • Use the A⊤ A-form A p( A⊤ A) if A is tall, and the AA⊤ -form p( AA⊤ ) A if A is wide. • Store B = AA⊤ or B = A⊤ A, and compute p( B) A or A p( B), respectively. • Use Horner’s method to evaluate p( B).
Example. Suppose we choose g( x ) = a0 x + a1 x3 + a2 x5 + a3 x7 and want to apply it to a wide matrix A. There are two steps: 1. Compute B = AA⊤ . 2. Compute g ( A ) = a0 + B a1 + B ( a2 + a3 B )
41
A.
If we apply it to a tall matrix A, the two steps are: 1. Compute B = A⊤ A. 2. Compute g ( A ) = A a0 + B a1 + B ( a2 + a3 B ) .
42
G
Computational and Memory Cost Analysis
G.1
Computational Cost Analysis
We estimate the additional FLOPs introduced by the PC layer. Following standard practice in numerical computing, we count one matrix multiplication of an ( a × b) matrix with a (b × c) matrix as ≈ 2abc FLOPs. We only account for matmul (matrix multiplication) FLOPs and omit lower-order elementwise/scalar operations. e ∈ Rn×m . Without loss of Per-matrix overhead (forward). Consider a normalized weight matrix W generality we assume the wide-form e ) = p (W eW e ⊤ ) W, e g (W e p (W e ⊤W e ) is fully where p is a degree-k polynomial evaluated by Horner’s method. Note that the tall-form W analogous; one always chooses the smaller Gram matrix to reduce cost. eW e ⊤ ∈ Rn×n be the chosen Gram matrix in this form. In practice, we always choose the smaller Let B = W Gram; hence the Gram dimension is s = min(n, m) and the other dimension is ℓ = max(n, m). The dominant additional matrix multiplications in the forward pass are: (i) One Gram construction (≈ 2ℓs2 FLOPs); (ii) Horner calculation of a degree-k polynomial p( B), which requires (k − 1) multiplications of s × s matrices (≈ 2(k − 1)s3 FLOPs); e (≈ 2ℓs2 FLOPs). (iii) One final multiplication to apply p( B) to W Therefore, ignoring lower-order terms, the forward overhead is approximately 4ℓ s2 + 2( k − 1) s3 . Spectral-norm estimation. To approximate the weight spectral norm, we use streaming power iteration with persistent buffers. Each power-iteration step uses two matrix-vector multiplications, W ⊤ u and Wv, and therefore costs approximately 4ℓs FLOPs. Including the final Rayleigh quotient, the spectral-norm estimation cost is approximately (4q + 2)ℓs, where q is the number of power iterations in each training step. Combining the polynomial-preconditioning cost and the spectral-norm estimation cost, the total additional forward FLOPs for one PC-applied matrix in a training step are approximately ∆FLOPsPC,fwd ≈ 4ℓs2 + 2(k − 1)s3 + (4q + 2)ℓs. Since s3 ≤ ℓs2 , this admits the upper bound ∆FLOPsPC,fwd ≤ 2(k + 1)ℓs2 + (4q + 2)ℓs. For the matrix sizes used in our experiments, s is large, so the streaming power-iteration term is lower-order compared with the matrix-polynomial term. Training-step overhead. A common rule of thumb is that, in typical deep networks, the backward pass costs about twice the forward pass in FLOPs [Brown et al., 2020], so a full training cost (forward + backward) is about 3× the forward cost. Thus, for the blocks where the PC layer is applied, the FLOPs overhead in one training step satisfies ∆FLOPsPC,train ≈ 3 ∆FLOPsPC,fwd ≤ 6(k + 1) ℓs2 + 6(2q + 1) ℓs.
43
For a standard block without PC, the baseline training per-step FLOPs (one forward and backward pass) is approximately 6nmB, where B is batch size (tokens) [Kaplan et al., 2020]. Using nm = ℓs, the relative FLOP overhead of the PC layer is bounded by ∆FLOPsPC,train 6(k + 1)ℓs2 + 6(2q + 1) ℓs (k + 1)s + 2q + 1 ≤ = . (10) 6nmB 6ℓ s B B Note that this bound is block-wise; the network-level overhead is even smaller since PC is implemented only for a subset of blocks. This overhead is small whenever the effective tokens-per-step B is large compared to the model width s and the PC polynomial degree k. Overhead ≜
Instantiating with our Llama-1B setting. For transformer with PC layer under AdamW, we use pc_level= 4, i.e., a degree-k polynomial with k = 4, and q = 10 streaming power-iteration steps for spectral-norm estimation (see Appendix F.2). Our training processes B = 2.62 × 106 tokens per step. For any PC-applied weight W ∈ Rn×m , let s = min(n, m) be the Gram dimension. From Inequality (10), the relative FLOPs overhead is bounded by 5s + 21 (k + 1)s + 2q + 1 = . B B In our 1B model, PC is applied to WO ∈ Rd×d and the FFN matrices Wgate ∈ Rmffn ×d , Wup ∈ Rmffn ×d , Wdown ∈ Rd×mffn , where d = 2048 and mffn = 5632 (see Appendix H for model configurations). All these matrices have the same smaller dimension s = d = 2048, hence 5 · 2048 + 21 Overhead ≤ ≈ 3.92 × 10−3 ≈ 0.39%. 2.62 × 106 Therefore, the training-FLOP overhead brought by PC layer is acceptable in our setting. Overhead ≤
Remark 5. When PC is combined with Muon, we empirically find that a lower polynomial degree pc_level= 2, i.e., k = 2, is optimal (see §4.2 for the results and Appendix K.2 for an intuitive explanation). Keeping q = 10, our overhead bound becomes 3s + 21 (k + 1)s + 2q + 1 = . Overhead(W ) ≤ B B Instantiating with s = d = 2048 and B = 2.62 × 106 tokens/step, this yields 3 · 2048 + 21 ≈ 2.35 × 10−3 ≈ 0.24%, 2.62 × 106 leading to even lower computational cost overhead. Overhead(W ) ≤
G.2
Memory Cost Analysis
To evaluate the memory efficiency of our proposed approach, we measured the peak memory footprint during training. All Llama-1B runs use 8×NVIDIA H100 GPUs; the reported footprint is the per-GPU peak active memory on a single H100, taken as the maximum across the 8 ranks (in practice the largest value is on rank 0). Under the Llama-1B AdamW configuration, the transformer baseline incurs a per-GPU peak footprint of 65.90 GiB, whereas the PC layer consumes 72.20 GiB, an increase of approximately 9.56%. Under Muon, the corresponding footprints are 64.08 GiB for the baseline and 69.67 GiB with the PC layer, an increase of approximately 8.73%. The extra peak memory observed with the PC layer likely comes from additional differentiable intermediate tensors that PyTorch autograd caches during the PC forward computation for use in the backward pass, such as the Gram matrix, the streaming power-iteration buffers, and the intermediate matrices produced when evaluating the polynomial p(·) (see Algorithm 1). Our current implementation does not explicitly control which of these intermediates are retained, so the observed overhead may overestimate the minimum amount of state actually required for backpropagation through the PC layer. In principle, this footprint could be reduced by (i) activation recomputation (checkpointing) at the PC-block boundary, which trades 44
additional FLOPs for lower memory, or (ii) a custom fused PC kernel with hand-managed backward, which re-materializes selected intermediates instead of relying on autograd to cache them. A more aggressive kernel-level memory optimization is left as future work.
G.3
Summary
In summary, Table 6 provides an overview of the computational and memory overheads introduced by the PC layer. Type
Optimizer
Overhead
Note
FLOPs
AdamW Muon
≤ 0.39% ≤ 0.24%
k = 4, q = 10 k = 2, q = 10
Peak memory (per GPU)
AdamW Muon
≈ 9.56% ≈ 8.73%
65.90 GiB → 72.20 GiB 64.08 GiB → 69.67 GiB
Table 6: Overhead summary of PC layer in Llama-1B experiments
45
H
Additional Experimental Details
Model configurations. The model configurations of the Llama 2 architecture used in our experiments are summarized in Table 7. Model Llama-271M Llama-1B
d_model
n_layers
n_heads
hidden_dim
1024 2048
16 18
16 16
2816 5632
Table 7: Model configurations used in our experiments.
Batchsize realization. Recall from Section 4.1 that we use sequence length 8192 and a fixed global batch of 2.62M tokens per optimization step, which corresponds to 320 sequences (8192 × 320 ≈ 2.62M). On 8 GPUs, this global batch is realized by setting the per-GPU micro-batch size and gradient accumulation to (8, 5) for the 271M model and (4, 10) for the 1B model, so that the total number of sequences per step satisfies micro × accum × 8 = 320. Optimizers. We train all models with either AdamW [Kingma, 2015, Loshchilov and Hutter, 2019] or the Muon optimizer [Jordan et al., 2024]. For AdamW, we set ( β 1 , β 2 ) = (0.9, 0.95) and weight decay to 0.1. For Muon, we follow the standard practice of applying Muon only to matrix-shaped (2D) weight tensors, while using AdamW for lower-dimensional parameters; we also adopt the root mean square (RMS) matching trick [Liu et al., 2025]. Muon hyperparameters are momentum= 0.95 and ns_steps= 5. Additionally, to ensure training stability, both optimizers apply gradient-norm clipping at 1.0. The learning-rate schedule and the peak-LR grid search are described in Section 4.1.
46
I
Accuracy of the Streaming Power-Iteration Spectral-Norm Estimator
We empirically validate the streaming power-iteration estimator s(W ) used by the PC layer to approximate ∥W ∥2 (Algorithm 3). On the LLaMA2-271M AdamW run we compare s(W ) against the ground-truth ∥W ∥2 from a full SVD of the same weight, and report the relative error relerr(W ) =
| s (W ) − ∥ W ∥ 2 | . ∥W ∥ 2
For each of the four preconditioned matrix types (WO , Wgate , Wup , Wdown ) we aggregate relerr across the 16 transformer blocks at every step, plotting per-step median and per-step maximum over the full 22,000-step trajectory (Figure 11). median (16 layers) max (16 layers)
Relative err of powiter approx
Relative err of powiter approx
10 1 10 2 10 3 10 4 0
5000
10000
step
15000
median (16 layers) max (16 layers)
10 2 10 3 10 4 10 5
20000
0
(a) WO (attention output projection)
10 3 10 4
0
5000
10000
step
15000
step
10 1
median (16 layers) max (16 layers)
10 2
10 5
10000
15000
20000
(b) Wgate (FFN gate projection)
Relative err of powiter approx
Relative err of powiter approx
10 1
5000
20000
(c) Wup (FFN up projection)
median (16 layers) max (16 layers)
10 2 10 3 10 4 10 5
0
5000
10000
step
15000
20000
(d) Wdown (FFN down projection)
Figure 11: Relative error of the streaming power-iteration estimator s(W ) on LLaMA2-271M (AdamW, 22,000 steps). Per-step median and per-step maximum of relerr(W ) aggregated over the 16 transformer blocks.
Initial transient and steady state. During the first few hundred steps the weights drift rapidly while the warm-started (u, v) buffers have not yet aligned with the dominant singular subspace, so the estimator exhibits a short transient, during which the per-step maximum across all four matrix types stays below 0.08. After roughly 3,000 steps the estimator enters a tight steady state: the per-step maximum stays below 4 × 10−3 and the per-step median drops below 1.5 × 10−4 for the remainder of training, i.e., s(W ) tracks ∥W ∥2 to three to four significant digits across all 16 layers.
47
Behavior across matrix types. The output projection WO is the slowest to converge: its per-step maximum keeps producing intermittent spikes up to ∼ 10−2 until step ≈ 3,000, whereas the corresponding maxima for Wgate , Wup , Wdown settle to their steady-state floor by step ≈ 1,000. This is consistent with WO having a tighter relative gap between its two leading singular values, which slows the geometric rate of power iteration. Even so, WO enters the same steady-state band as the other three matrices, confirming that ten power-iteration sweeps with warm-starting suffice even for the hardest case. Implication for the [0, 1.1] fitting interval. Both regimes above are comfortably absorbed by the 10% safety margin built into the fitting interval [γ L , γU ] = [0, 1.1] (Section 3.2): in the worst case observed, the normalized spectrum W/s(W ) can overshoot 1 by at most ∼ 0.08 < 0.1, so the polynomial g is still evaluated inside its design domain throughout training.
48
J
Overly Aggressive Spectrum Flattening
Our PC polynomial is designed to improve conditioning without collapsing the entire singular-value spectrum. To test the limiting case, we replace the default PC polynomial with an over-flattening polynomial whose scalar map nearly sends every nonzero normalized singular value to one, as shown in Figure 12. Concretely, we use the Polar Express composite of degree-5 Newton–Schulz polynomials [Amsel et al., 2025], which approximates the matrix-sign map. On the normalized domain σ ∈ [0, 1], the over-flattening map is the composition of T = 8 odd degree-5 polynomials, p(σ ) = p8 ◦ p7 ◦ · · · ◦ p1 (σ ), where p1 (σ ) = 7.2086 σ − 15.5131 σ3 + 9.0178 σ5 ,
p2 (σ ) = 3.9623 σ − 2.5813 σ3 + 0.4542 σ5 ,
p3 (σ ) = 3.9466 σ − 2.5765 σ3 + 0.4544 σ5 ,
p4 (σ ) = 3.8991 σ − 2.5671 σ3 + 0.4566 σ5 ,
p5 (σ ) = 3.7186 σ − 2.5308 σ3 + 0.4653 σ5 ,
p6 (σ ) = 3.1390 σ − 2.3073 σ3 + 0.4733 σ5 ,
p7 (σ ) = 2.1715 σ − 1.5246 σ3 + 0.3885 σ5 ,
p8 (σ ) = 1.8648 σ − 1.2224 σ3 + 0.3577 σ5 .
All other PC settings are kept identical to the per-optimizer default (PC_blocks = {ffn, WO }, with pc_level = 4 under AdamW and pc_level = 2 under Muon).
2 X W S X W
3 R O D U ( [ S U H V V L W H U V 7 D U J H W 1
, Q L W L D O
Figure 12: Over-flattening scalar map. The Polar Express composite of degree-5 Newton–Schulz polynomials [Amsel et al., 2025] nearly maps every nonzero normalized singular value to one.
Figure 13 shows that this aggressive flattening is harmful under both optimizers: in each case the overflattened polynomial gives higher validation loss than the transformer baseline, whereas the softer default PC polynomial remains better than the baseline. This suggests that PC should not aim to perfectly orthogonalize the selected weight matrices at every training step. Preserving some spectral anisotropy appears important for expressiveness and optimization. We therefore use moderate piecewise-linear targets, which improve conditioning while retaining spectral flexibility.
49
2.90 2.85
2.65
Validation Loss
2.80
Validation Loss
2.70
Transformer Baseline Transformer w/ PC Layer (level 4) Transformer w/ PC Layer (overflatten poly.)
2.75 2.70 2.65 2.60
2.60 2.55 2.50
Transformer Baseline Transformer w/ PC Layer (level 2) Transformer w/ PC Layer (overflatten poly.)
2.55 0
10
20
30
40
Total Tokens Trained (B)
50
2.45
60
(a) AdamW.
0
10
20
30
40
Total Tokens Trained (B)
50
60
(b) Muon.
Figure 13: Overly aggressive spectrum flattening hurts under both optimizers. Validation loss on Llama-271M under (a) AdamW and (b) Muon. In both cases the over-flattened polynomial performs worse than the transformer baseline, while the default PC layer (pc_level = 4 for AdamW, pc_level = 2 for Muon) stays below the baseline, supporting soft spectrum shaping over near-perfect flattening.
50
K
Additional Results for PC Layer with Muon
Notably, since the PC layer and the Muon optimizer operate on different aspects of training (update matrices vs. weight matrices), they are largely orthogonal and thus compatible. The main-text validation-loss curves and downstream comparison (Figure 4, Table 3) already establish that PC improves training under Muon, and the pc_level/PC_blocks ablations under Muon are reported alongside the AdamW ablations in Section 6. This section provides the remaining supporting details: we first examine the effects on spectral conditioning (§K.1), and then offer additional discussions on the relationship between PC layer and Muon, including the computational template shared by both methods and their respective impact on weight-spectrum control (§K.2).
K.1
Spectral Conditioning
We then assess PC’s effect on weight spectral conditioning under Muon by tracking the global modified condition number over training and visualizing final-checkpoint singular-value spectra across representative depths and blocks. Global modified condition number (GMCN). Figure 14 shows the evolution of the Global Modified Condition Number (GMCN) κ̃ throughout training under Muon. Enabling PC shifts the curve downward on the blocks to which PC is applied (ffn and WO ), while the non-preconditioned WQ , WK , and WV blocks do not deteriorate and are slightly improved. The global aggregate therefore remains lower under PC, indicating improved spectral conditioning of weights. GMCN (non-PC blocks) vs Tokens
GMCN (PC blocks) vs Tokens
11 10
80
22.5
7
GMCN
8
60
20.0 17.5
6
50
15.0
5
40
12.5
Muon baseline Muon + PC
4 0
20
40
60
80
100
Total Tokens Trained (B)
120
140
160
Muon baseline Muon + PC
25.0
70
GMCN
GMCN
9
GMCN vs Tokens
27.5
Muon baseline Muon + PC
90
10.0
30 0
(a) FFN + WO
20
40
60
80
100
Total Tokens Trained (B)
(b) WQ , WK , WV
120
140
160
0
20
40
60
80
100
Total Tokens Trained (B)
120
140
160
(c) Global
Figure 14: Block-wise and aggregate modified condition numbers under Muon. We track κ̃ separately for the PC-targeted blocks (ffn and WO ), the attention-input blocks left outside PC (WQ , WK , WV ), and the resulting global aggregate. The targeted blocks show the clearest conditioning gain, and the untargeted QKV blocks also exhibit a noticeable improvement relative to the baseline, further contributing to the lower aggregate curve.
Singular-value spectrum comparison. To complement the scalar GMCN curves, we visualize the finalcheckpoint singular-value distributions following the identical protocol as in Section 5 (same three representative depths, baseline spectra on W versus PC spectra on PC(W ), with per-matrix max-normalization to [0, 1]); see Figure 15. Figure 15 corroborates the GMCN trend. Across depths and blocks, PC pulls the lower part of the spectrum upward toward σmax , and this is most visible in the ffn blocks. As a result, the relative spectral spread narrows and the conditioning improves.
K.2
Supplementary Discussions
In this section, we provide further discussions on the similarities between Muon and PC in terms of their computational templates and their approaches to spectral conditioning.
51
/ max
0.6
0.8
00.0
1.0
(a) Layer 2: WO
0.8
40
Muon baseline Muon + PC
30
/ max
0.6
0.8
00.0
1.0
(e) Layer 10: WO
0.2
0.4
/ max
0.6
30 25 20 15 10 5 00.0
0.8
1.0
(f) Layer 10: Wgate 50
Muon baseline Muon + PC
Frequency
Frequency
0.4
/ max
0.6
0.8
(i) Layer 18: WO
1.0
00.0
0.2
0.4
/ max
0.6
0.8
1.0
(j) Layer 18: Wgate
00.0
Frequency
0.2
0.4
/ max
0.6
0.8
1.0
(d) Layer 2: Wdown Muon baseline Muon + PC
50 40 30 20 10
0.2
0.4
/ max
0.6
0.8
00.0
1.0
0.2
0.4
/ max
0.6
0.8
1.0
(h) Layer 10: Wdown
Muon baseline Muon + PC
50 40
Muon baseline Muon + PC
30 20
20
10 0.2
00.0
1.0
40
20
10
0.8
Muon baseline Muon + PC
60
30
20
/ max
0.6
(g) Layer 10: Wup
Muon baseline Muon + PC
40
30
0.4
Frequency
0.4
10 0.2
(c) Layer 2: Wup
10
0.2
00.0
1.0
20
20
00.0
/ max
0.6
Frequency
Frequency
40
40
0.4
(b) Layer 2: Wgate
Muon baseline Muon + PC
60
0.2
Frequency
0.4
Frequency
0.2
20
10
10
Muon baseline Muon + PC
40 30
20
20
10
Muon baseline Muon + PC
30
30
20
00.0
Muon baseline Muon + PC
40
Frequency
Frequency
30
00.0
50
Frequency
Muon baseline Muon + PC
40
Frequency
50
10 0.2
0.4
/ max
0.6
0.8
(k) Layer 18: Wup
1.0
00.0
0.2
0.4
/ max
0.6
0.8
1.0
(l) Layer 18: Wdown
Figure 15: Singular-value histograms at the final checkpoint (Muon, Llama-1B). Setup as in Figure 6 (three depths, all PC_blocks, per-matrix max-normalization to [0, 1]). Across depths and blocks, PC shifts the bulk of the spectrum upward toward σmax , narrowing the relative spectral spread.
A shared computational template. Despite these conceptual differences, both methods share a computational template that makes them GPU-friendly. Concretely, both can be expressed using matrix-polynomial transforms of the form T ( X ) = q( XX ⊤ ) X, i.e., a specific polynomial in XX ⊤ applied to X. This structure reduces to a small number of matrix multiplications and avoids explicit spectral decompositions, while still inducing a controllable map on singular values. Notably, Muon employs a 25th-degree polynomial (comprising 5 iterations of degree-5 polynomials) to precondition the matrix, while PC layer uses a 9th-degree polynomial (pc_level = 4) under AdamW and a 5th-degree polynomial (pc_level = 2) when combined with Muon. A detailed FLOPs analysis of the PC layer overhead is provided in Appendix G.1. Implicit vs. explicit control of the weight spectrum. Although Muon acts on updates rather than weights, it can still be viewed as exerting an implicit form of weight-spectrum control through the optimizer dynamics. Recent work suggests that Muon with decoupled weight decay corresponds to an implicit spectral constraint (notably affecting the top singular value) [Chen et al., 2025], and empirical studies report that Muon-trained weights tend to be more isotropic, with higher effective rank / spectral entropy [Wang et al., 2025, Liu et al., 2025]. In contrast, PC provides explicit weight-spectrum control by directly shaping the singular values of selected weight blocks toward a healthy range throughout training. Since Muon already exerts implicit spectral control on the updates, the additional spectrum shaping required from PC under Muon might be smaller. As a result, a lower pc_level (i.e., a lower polynomial degree) already suffices to shape the spectrum without overly interfering with the model’s expressiveness under Muon. Higher pc_level values, on the other hand, may risk compromising the model’s ability to represent complex patterns by over-constraining the weight spectrum. This explains why the optimal pc_level for Muon is smaller than that for AdamW.
52