用逐元素裁剪实现梯度谱控制,稳定训练并减少约7%训练量。
Can Entry-Wise Clipping Give Spectral Control of Stochastic Gradients?
- 提出逐元素裁剪方法,利用梯度噪声局部化特性实现谱控制。
- 在柯西污染噪声下保证O(ε⁻⁴)收敛性,实测节省约7%训练令牌。
- 适用于大模型训练,尤其适合优化器如Adam和Muon的改进。
训练不稳定性(如损失突增)常由随机梯度噪声引发。由于语言数据中的罕见表达和多层组合,噪声呈现重尾分布且难以通过小批量平均消除。现有方法需在结构保留与计算成本间权衡:向量范数裁剪忽略权重更新的矩阵结构,而谱归一化(如Muon (Jordan et al., 2024))虽保留结构但开销更大。本文发现该权衡可被平衡:真实梯度噪声类似逐元素重尾扰动,一阶扰动分析揭示其局部化特性,使简单逐元素方法实现谱控制。基于高斯信号先验,推导出贝叶斯最优逐元素估计器的可计算近似。在柯西污染噪声下建立$O(ε^{-4})$收敛保证。实验表明,平滑收缩能提升Adam在NanoGPT预训练中的表现,节省约7%训练令牌;进一步在谱归一化前使用逐元素裁剪,相较Muon再节省约2%令牌。
原文摘要 · Abstract (English)
Training instabilities such as loss spikes are frequently the result of stochastic gradient noise. Because of rare expressions in language training data, and multiple layer composition, the noise impact is heavy-tailed and survives mini-batch averaging. Existing remedies trade off structure against cost: vector-norm clipping ignores the matrix structure of weight updates, while spectral normalization (e.g., Muon (Jordan et al., 2024)) respects it at additional cost. We show that this trade-off can be balanced. Real gradient noise appears to be similar to entry-wise heavy-tailed contamination, and a first-order perturbation analysis reveals a localization property of such noise, under which a simple entry-wise method achieves spectral control. Exploiting this, we derive a tractable surrogate for the Bayes-optimal entry-wise estimator under a Gaussian signal prior. We establish $O(ε^{-4})$ convergence guarantee under Cauchy-contaminated noise. Empirically, we find that smooth shrinkage improves Adam on NanoGPT pretraining, saving ${\sim}7\%$ of training tokens. We further find that applying the entry-wise clipping before spectral normalization yields a ${\sim}2\%$ token saving on top of Muon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。