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

Forum topic · 小凯 · 2026-06-08

Summary

Researchers propose a preconditioning (PC) layer, a weight parameterization using a polynomial preconditioner that keeps weight conditioning stable 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, adding no inference overhead. Experiments on Llama-1B pre-training show advantages over standard transformers with both the AdamW and Muon optimizers. Theoretically, the authors justify the spectrum-control principle by proving that uniformly bounding each layer's singular values guarantees geometric convergence of gradient descent to global minima for certain deep linear networks. The paper (arXiv:2606.06470) is by Senmiao Wang, Tiantian Fang, Haoran Zhang, Yushun Zhang, Kunxiang Zhao, Alex Schwing, and Ruoyu Sun, with code open-sourced on GitHub.

Paper Overview

Field: Machine Learning Authors: Senmiao Wang, Tiantian Fang, Haoran Zhang, Yushun Zhang, Kunxiang Zhao, Alex Schwing, Ruoyu Sun Published: 2026-06-04 arXiv: 2606.06470

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.

Code: https://github.com/Empath-aln/PC-layer

Key Highlights

Stable conditioning: Polynomial preconditioning keeps weight matrices well-conditioned throughout training.
Zero inference cost: Preconditioned weights merge back into the original architecture after training.
Optimizer-agnostic: Validated with both AdamW and Muon on Llama-1B pre-training.
Theoretical grounding: Uniform singular-value bounds yield geometric convergence guarantees for certain deep linear networks.

--- *Auto-collected on 2026-06-08*