Natalie Abreu Nikhil Vyas Sham Kakade Depen Morwani

摘要

近期加速大语言模型（LLM）预训练的研究主要聚焦于利用二阶结构的计算高效近似方法。这引发了一个大规模训练中的关键问题：这些近似方法会损失多少性能？为探究这一问题，我们通过在参数规模达1.5亿的Transformer模型上应用完整的高斯-牛顿（Gauss-Newton, GN）预条件化，建立了一个实际的迭代复杂度上界。实验结果表明，与SOAP和Muon等强基线方法相比，完整的GN更新可显著提升训练效率，将训练迭代次数减少5.4倍。此外，我们发现一种精确的分层GN预条件器（忽略层间信息）几乎可达到完整GN方法的性能水平。综合来看，我们的研究结果表明：（1）GN近似在预条件化中具有极高的有效性，暗示高阶损失项对收敛速度的影响可能并不关键；（2）分层海森矩阵结构已包含实现绝大部分性能提升所需的关键信息；（3）当前近似方法与理想化的分层“预言机”之间仍存在显著的性能差距。