PC Layer: Polynomial Weight Preconditioning for Improving LLM Pre-Training

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 singular-value 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 pre-training, 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.

• The paper introduces a deterministic PC layer that explicitly controls the singular value spectrum of weight matrices to enhance LLM training.
• It employs low-degree polynomial maps to reshape the weight distribution, achieving up to 39–50% token efficiency gains and improved zero-shot accuracy.
• Empirical results demonstrate a significant reduction in the global modified condition number with negligible computational overhead during training.

Polynomial Weight Preconditioning for LLM Training

Introduction and Motivation

The paper "PC Layer: Polynomial Weight Preconditioning for Improving LLM Pre-Training" (2606.06470) introduces a deterministic, built-in parameterization module—termed the "PC layer"—for actively controlling the singular value spectrum of weight matrices during LLM pre-training. Traditional normalization techniques (e.g., RMSNorm, QK-Norm, spectral normalization) mitigate issues with training instability, but they typically operate on either activations or coarse statistics (mean, variance, or maximum singular value) and rarely provide explicit control over the full spectrum of weight matrices. Recent developments highlight the intrinsic link between the weight spectrum and trainability—informing both forward and backward signal propagation and, ultimately, optimization convergence. The PC layer targets explicit, architecture-level spectrum shaping, aiming to ameliorate signal degradation and poor conditioning without incurring inference overhead.

The method builds on two pillars: (1) a theoretical justification that global optimization convergence in deep linear networks is governed by the uniform conditioning of layer weights (i.e., bounded singular value ratios) and (2) a polynomial preconditioning routine that acts on weight matrices in a manner conducive to gradient-based optimization and hardware efficiency.

Method: Polynomial Preconditioning Layer

Spectrum Shaping via Matrix Polynomials

Spectral control in deep nets is achieved by parameterizing certain weight matrices $W$ as $g(W)$, where $g$ is a low-degree polynomial map defined so that, if $\sigma$ is a singular value of $W$, then $g(\sigma) = p(\sigma^2)\sigma$ modulates that singular value. The relevant computational pathway is

$g(W) = p(WW^\top)W$

for left multiplication (or the analogous right multiplication on $W^\top W$), ensuring that spectrum shaping acts on the effective singular value interval. This approach contrasts with recently proposed almost-orthogonal transformations (as in Muon) which aggressively collapse the spectrum towards isotropy—a practice that, as empirically demonstrated, restricts model expressiveness and degrades performance when over-applied.

Figure 1: A low-degree polynomial reshapes the singular-value spectrum, amplifying small singular values and saturating large ones, achieving "soft" spectrum shaping rather than strict orthogonalization.

Determining Spectrum-Conditioning Polynomials

The target spectral manipulation is informed by a piecewise-linear function ("PL" mapping) that aggressively amplifies small singular values when below a threshold $b$, saturating large ones near unity:

$$\mathrm{PL}_b(\sigma) = \begin{cases} \sigma / b, & \sigma < b \ 1, & \sigma \geq b \end{cases}$$

Low-degree odd polynomials $g(W)$0 (overall degree $g(W)$1, for $g(W)$2) are fitted offline to $g(W)$3 on the normalized domain $g(W)$4 via weighted least-squares, yielding a one-parameter family of increasingly aggressive spectrum shaping transforms.

Figure 2: Piecewise-linear spectrum shaping targets and their least-squares polynomial approximations for varying aggressiveness levels.

Integration into Transformer Architectures

PC replaces select weight matrices in the transformer (primarily attention output $g(W)$5 and FFN matrices) by $g(W)$6, where $g(W)$7 is normalized by a streaming power-iteration estimate. After spectrum shaping, a rescaling step (with optional learnable scalar $g(W)$8) restores the norm, ensuring compatibility with existing training pipelines. The PC layer is formulated as a reparameterization active only during training—computation is absorbed into the static weight at inference.

Empirical Results

Training Efficiency and Downstream Performance

Deploying the PC layer in Llama-271M and Llama-1B models (on FineWeb, using AdamW and Muon optimizers) yields consistent improvements in token-efficiency. For AdamW, the PC-equipped model reaches baseline loss with approximately 39–50% fewer training tokens (i.e., 1.63–2.00× speedup) and achieves up to 0.070 lower final validation loss. Gains do not diminish at higher scale. Downstream task evaluation across diverse benchmarks shows improvements in average zero-shot accuracy by 1–2 percentage points across both optimizers.

Figure 3: Validation loss as a function of training tokens for Llama-271M, demonstrating improved token-efficiency and lower minimum achieved by PC.

Figure 4: Validation loss trajectories for Llama-1B with the PC layer, exhibiting larger gains.

Weight Spectrum and Conditioning

The global modified condition number (GMCN), which robustly measures the ratio of the top singular value to the mean of the bottom 10%, is significantly reduced by PC across targeted blocks. For instance, the global aggregate drops from 42.4 (baseline) to 25.0 (PC-equipped) models—a 41% reduction in spectrum spread—demonstrating that PC narrows the spectrum without eliminating anisotropy, a key factor in preserving model expressivity.

Figure 5: Evolution of the GMCN during training; PC leads to lower condition numbers across both targeted and non-targeted blocks.

(Figure 6, Figure 7)

Figure 6: Histogram of singular values for $g(W)$9 (layer 2) at the end of training; the PC-equipped model lifts the lower tail and reduces spectrum spread.

Ablations show that:

• Norm recovery (post-preconditioning rescaling) is essential for gains.
• A learnable scalar $g$0 has limited impact on final loss but stabilizes training.
• The optimal polynomial degree for spectrum shaping depends on the optimizer—higher for AdamW (pc_level=4), lower for Muon (pc_level=2).
• Overly aggressive spectral flattening (as in Polar/strict-orthogonalization polynomials) harms optimization, reinforcing the principle that "soft spectrum shaping" yields the best trade-off between conditioning and expressivity.

(Figure 8, Figure 9)

Figure 8: Ablation of polynomial degree (pc_level); under AdamW, highest degree (4) is optimal.

Figure 9: Over-flattening the spectrum leads to degraded performance, supporting nuanced control over weight conditioning.

Theoretical Insights

The central theorem establishes, for deep linear networks, that if all weight matrices remain within a uniformly bounded condition number, then gradient descent on the supervised MSE loss achieves geometric convergence to the optimum at a rate that depends polynomially on the global conditioning constant $g$1. This provides a clear actionable motivation for maintaining healthy weight spectra in each layer during training—aligning empirical design with established convergence theory.

Computational Overhead and Practical Considerations

PC layers incur negligible computational cost in practice: with streaming power iteration and polynomial calculations (Horner's method), forward+backward overheads constitute at most 0.39% (AdamW) and 0.24% (Muon) of training FLOPs per affected matrix block, and no inference penalty.

Implications and Future Directions

This work establishes that active, train-time spectrum conditioning via polynomial parameterization yields quantifiable improvements in training speed and downstream LLM accuracy, supporting a paradigm shift from passive initialization and generic norm control toward spectrum-aware architectural modules. Future work may target:

• Layerwise and blockwise adaptive selection of polynomial preconditioners,
• Scaling to much larger model and data regimes,
• Integration with other advanced optimizer-level spectral control mechanisms,
• Application to architectures beyond Llama or transformers.

Conclusion

The PC layer provides a theoretically grounded, computationally efficient, and highly effective method for modulating the spectrum of LLM weight matrices during pre-training. By avoiding restrictive orthogonalization and facilitating soft spectrum shaping, it improves both convergence and generalization, with substantial empirical improvements demonstrated in both efficiency and accuracy—all without inference cost or architectural disruption.