PC Layer: Polynomial Weight Preconditioning for Improving LLM Pre-Training
This work addresses the problem of unstable weight conditioning in large language model pre-training, offering a method that improves training dynamics and can be merged back without inference cost.
The paper introduces a preconditioning (PC) layer that stabilizes weight conditioning during LLM pre-training by reshaping singular-value spectra via low-degree polynomial preconditioning, achieving improved performance over standard transformers in Llama-1B pre-training with AdamW and Muon optimizers without inference overhead.
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.