用多项式预条件层稳定大模型训练,提升收敛性。
PC Layer: Polynomial Weight Preconditioning for Improving LLM Pre-Training

- 通过低次多项式重塑权重矩阵的奇异值谱
- 在Llama-1B上相比标准Transformer加速收敛,适配多种优化器
- 理论证明奇异值有界可保证梯度下降全局收敛
我们提出一种预条件(PC)层,通过多项式预条件对权重参数化,确保大语言模型训练全程的权重条件稳定。该模块利用低次多项式重塑权重矩阵的奇异值分布。训练完成后,预条件权重可无缝合并回原始结构,不增加推理开销。我们在Llama-1B预训练中验证了该方法相较于标准Transformer的优势,适用于AdamW和Muon两种优化器。理论上,我们证明了对每一层奇异值进行统一有界控制,可使某些深度线性网络的梯度下降实现几何收敛至全局最小值。代码已开源:https://github.com/Empath-aln/PC-layer。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。