这项研究解决了LLM训练中权重矩阵病态条件数导致的收敛不稳定问题,做预训练优化的团队可以直接在Llama等模型上尝试,无需额外推理成本。
研究人员提出了一种名为PC Layer(Preconditioning Layer)的权重参数化方法,通过在训练过程中对权重矩阵进行低阶多项式预处理,稳定其奇异值谱,从而改善大型语言模型(LLM)的预训练效果。该方法与AdamW和Muon优化器兼容,在Llama-1B模型上验证了其优于标准Transformer架构。训练后,预处理权重可合并回原始架构,不增加推理开销。理论证明,均匀限制每层奇异值能确保梯度下降在深层线性网络中收敛到全局最小值。代码已开源。
PC Layer: Polynomial Weight Preconditioning for Improving LLM Pre-Training
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.