Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization
📄 Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization #大语言模型 #优化器 #高效推理 ✅ 7.5/10 | 前25% | #大语言模型 | #优化器 | #高效推理 | arxiv 学术质量 6.0/7 | 选题价值 1.5/2 | 复现加成 0.0 | 置信度 中 👥 作者与机构 第一作者：Ruotong Sun（论文中未提供机构信息） 通讯作者：未说明 作者列表：Ruotong Sun， Ermin Wei（均未说明所属机构） 💡 毒舌点评 该论文巧妙地将KL-Shampoo预条件器的理论特性（特征值谱的“尖峰-平坦”结构）与正交化操作相结合，提出了一种计算更高效、内存更友好的优化器变体，体现了理论指导工程优化的优雅思路。不过，其创新更多是针对现有框架的“精装修”，而非开辟新赛道；此外，在LLM预训练如此依赖算力和数据的时代，仅靠在100M-450M规模模型上的验证，能否稳定泛化到更大规模模型并说服工业界采用，仍需打上一个问号。 📌 核心摘要 要解决什么问题：现有利用梯度矩阵结构的LLM预训练优化器（如KL-Shampoo和Muon）各有优势但独立发展。KL-Shampoo通过KL散度最小化来估计克罗内克分解预条件器，效果显著但计算成本和内存占用较高。本文旨在设计一种新优化器，以更高效的方式获得与KL-Shampoo相当甚至更好的性能。 方法核心是什么：方法的核心是发现并利用KL-Shampoo预条件器的Kronecker因子具有“尖峰-平坦”的特征值谱结构。作者将其中一个因子限制在一个参数化族中：在一个被跟踪的低维子空间（维度为r）上保留完整的谱结构（完整的特征值和特征向量），而在剩余的（n-r）个方向上使用一个共享的特征值。对于后者，应用正交化操作，该操作在代数上能恢复完整KL-Shampoo的预条件器形式。 与已有方法相比新在哪里：新方法（Pro-KLShampoo）在两大前沿——Kronecker-factored preconditioning与gradient momentum orthogonalization——之间建立了结构性桥梁。它并非简单结合，而是基于对KL-Shampoo内在结构的深刻理解，通过参数化限制大幅降低了计算和内存开销，同时通过正交化保证了数学上的等价性，从而在实践中实现了更优的权衡。 主要实验结果如何：在GPT-2（124M， 350M）和LLaMA（134M， 450M）四个预训练规模上，Pro-KLShampoo在所有测试的子空间秩（r）下，在验证损失、峰值GPU内存使用和达到特定损失水平所需的墙钟时间这三个指标上均一致性优于基线KL-Shampoo。论文摘要中未提供具体数值，但强调了其全面优势。 实际意义是什么：该工作提供了一种理论上更清晰、实践中更高效的LLM预训练优化器。它能降低训练过程中的计算开销和内存需求，有望加速模型开发迭代并降低训练成本，对于资源受限的预训练场景尤其有价值。 主要局限性是什么：论文摘要未明确提及该方法的局限性。可能的局限包括：对“尖峰-平坦”特征值谱结构的假设在更多样化的任务或模型架构上的普适性；子空间维度r的选择需要调参；以及在更大规模（如数百亿至千亿参数）模型上的有效性和扩展性尚未验证。 🔗 开源详情 代码：论文中未提及代码链接 模型权重：论文中未提及 数据集：论文中未提及 Demo：论文中未提及 复现材料：论文中未提及 论文中引用的开源项目：未提及 🏗️ 方法概述和架构 整体流程概述：Pro-KLShampoo是一个用于神经网络预训练的优化器。其完整流程是：在每个训练步骤，接收模型参数和计算得到的梯度，利用当前优化器状态（历史梯度动量）以及一个对预条件器矩阵结构的参数化估计，来更新模型参数。它本质上是一个单阶段的、基于一阶梯度信息构建二阶预条件信息的自适应学习率优化器。 主要组件/模块详解： ...