FOAM以单边预条件削减SOAP优化器内存与分解开销

摘要：NAVER AI 提出轻量级二阶优化器 FOAM，通过仅保留输入侧 Kronecker 预条件，并以按行、列统计构造的 Adam 矩估计替代 SOAP 的稠密旋转二阶矩。论文称，该方法在 1.24 亿至 7.6 亿参数的语言模型预训练中，验证损失与 SOAP、Muon 基线相当；优化器状态占用降至 SOAP 的一半，墙钟耗时接近 Muon，并在三倍学习率范围内保持稳定。其关键前提是输出侧曲率近似对角，因此大模型规模、混合架构与非 Transformer 场景仍需独立验证。

---

核心思路：保留 SOAP 的预条件收益，拆掉最昂贵的两块开销

SOAP 与 Shampoo 一类二阶优化器通过 Kronecker 分解近似 Hessian，能够修正 AdamW 对不同参数方向的一视同仁。问题在于，完整 SOAP 优化器通常需要保存稠密二阶矩、执行周期性特征分解，并在每步更新中处理旋转后的参数与梯度状态。

FOAM（Factored One-sided Adam-Moment）提出两项削减：

1. 只保留输入侧预条件，移除输出侧 Kronecker 因子及其对应的特征分解；
2. 以因子化 Adam 矩替代稠密旋转二阶矩，利用逐行、逐列统计构造更新方向。

论文的关键理论结果是：当输出侧曲率为对角时，输入侧单边预条件没有额外近似损失；在输出侧曲率仅近似对角时，损失通常较小。作者进一步在 K-FAC Fisher 假设下推导：当输出侧曲率为对角时，所提出的矩估计可以精确恢复 SOAP 的稠密二阶矩。实验测得 K-FAC 假设偏离约 15%–20%，且对实际参数更新影响较小。

机制拆解：为什么能减少状态与分解成本

完整 SOAP 的典型优化器状态包含稠密旋转二阶矩及其逆矩阵或等价表示。若参数矩阵形状为 m × n，稠密二阶状态的主要内存量级为 O(m² + n²)，而不是 AdamW 一阶矩的 O(mn)；当 m 或 n 很大时，二阶状态可能成为训练瓶颈。

FOAM 删除输出侧因子后，保留的预条件状态更接近 O(m²)；同时，Adam 型矩估计只需维护可分解的低阶统计，避免长期保存和更新完整的稠密旋转矩阵。对角线化相关的特征分解也被减少，使优化器更新能更接近普通矩阵优化器的执行成本。

不过，O(m²) 仍可能远高于 O(mn)。对于宽而扁或窄而高的权重矩阵，实际内存收益取决于矩阵形状、参数共享方式以及分布式优化器实现；不能仅凭“单边预条件”推断所有层都获得相同比例的节省。

工程权衡：更快，但并非无条件等价

FOAM 的核心交换条件是：输出侧曲率必须足够接近对角。它把一个昂贵的二维二阶结构问题，转化为一个依赖局部曲率近似的单边预条件问题。由此获得三项收益：

• 减少优化器状态，降低训练时显存或主机内存压力；
• 减少特征分解频率及计算成本；
• 在大规模训练中维持接近 Muon 的实际运行速度。

代价是近似不再对任意 Hessian 都成立。当某些层具有显著输出侧耦合、权重矩阵形状极端，或训练过程中局部曲率快速偏离 K-FAC 假设时，FOAM 可能需要更多迭代才能达到同等验证损失。论文报告的三倍学习率范围稳定性说明其鲁棒性较好，但并不意味着学习率可以无条件跨模型、跨架构放大。

评测表现：验证损失持平是关键，规模覆盖仍有限

论文在 1.24 亿至 7.6 亿参数的语言模型预训练上比较 FOAM、SOAP、Muon 和 AdamW 等基线。核心结果是：FOAM 在验证损失上匹配较强的 SOAP 与 Muon 基线，同时只使用约一半 SOAP 优化器状态，并以接近 Muon 的墙钟时间运行。

这组结果的工程意义在于，优化器收益不再必须通过持续承担完整 SOAP 的状态与分解成本获得。不过，验证损失接近并不等于所有任务指标一致：在多模态、长上下文、混合专家模型以及严格受约束的微调任务中，二阶预条件可能改变数据效率、超参数敏感性和最终质量，仍需增加相应实验。

此外，当前结果主要来自小至中等规模语言模型预训练。更大模型可能改变优化器状态占比、矩阵形状与通信模式，因此无法直接外推“SOAP 收益与成本兼得”在大规模训练中的幅度。官方性能数据也需要结合固定随机种子、硬件配置、吞吐口径及第三方复现来判断。

开发者落地建议

FOAM 代码已通过 naver-ai/foam 开放。工程团队可将 AdamW、Muon、SOAP 与 FOAM 放在同一训练配置下比较，重点记录：

• 每步耗时与端到端训练吞吐；
• 优化器状态的峰值显存或主机内存；
• 特征分解频率及其对通信、重计算的影响；
• 学习率扫描下的验证损失稳定性；
• 达到目标验证损失所需的 token 数或训练步数。

上线时建议先从学习率扫描开始，而不是直接照搬 SOAP 的超参数。分布式环境下还应检查优化器状态是否被分片、是否产生额外 all-gather，以及矩阵参数和非矩阵参数是否采用了不同实现。若业务目标是大模型预训练成本控制，FOAM 的吸引力主要来自状态占用和吞吐；若目标是微调、持续预训练或超参数稳定性，则应以真实数据分布上的验证结果作为最终选型依据。

结论

FOAM 提供了一条有工程吸引力的折中路线：用输出侧对角曲率近似换掉完整 SOAP 的双向 Kronecker 预条件，再通过因子化矩估计重建更新方向。论文证据支持它在小中规模语言模型预训练中以约一半 SOAP 状态达到相近验证损失，但结论仍受 K-FAC 近似、评测规模和任务覆盖范围约束。它更像一项可直接试用的优化器基础设施改进，而不是已经完成普适性验证的 SOAP 替代方案。

---

消息来源：Hugging Face Daily Papers (2026-11-12)

—— 完 ——

主题标签： #FOAM 关注 #SOAP 关注 #二阶优化器 关注 #Kronecker预条件 关注 #大模型预训练 关注