IT资源栈开发者 Noah Golmant 近日发布了 PyTorch 曲率分析开源库 pytorch-hessian-eigenthings 的 1.0 大版本重写。该库最初发布于 2018 年，旨在解决神经网络损失景观的曲率指标计算问题，这些指标与神经网络的泛化能力密切相关。原有的完整海森矩阵计算内存消耗随参数量呈二次方增长，在大规模模型中难以实施。新版库通过 Hessian-vector 积结合兰乔斯算法、幂迭代等迭代方法，成功将特征值分解的内存复杂度优化至线性级别。此次重写引入了多项重大更新：新增了广义高斯-牛顿（GGN）和经验 Fisher 矩阵等曲率算子；支持 Hutchinson 和 Hutch++ 等迹估计算法以及谱密度估计；针对拥有超大词表的基础模型，实现了基于 Triton 和 torch.compile 的融合交叉熵 Hessian-vector 内核，有效解决了标准实现在此场景下的内存爆炸问题。此外，新版本加强了数值分析验证，在线性和逻辑回归模型上进行了闭合形式正确性测试，并与 curvlinops 库进行了交叉验证以防止回归。
事件分析
此次开源库的重写体现了 AI 基础设施向高性能计算与深度学习框架深度融合的趋势。在 AI 领域，二阶优化方法（如 K-FAC、自然梯度）理论上收敛速度快，但因计算复杂度和显存占用过高长期受限。该项目通过 Hessian-vector 积这一数学技巧规避了显存墙，并引入 Triton 语言编写融合内核，展示了在现有 PyTorch 生态下通过自定义算子突破性能瓶颈的路径。特别是针对大模型词表的优化，直接解决了现代 LLM 训练中的特定算子瓶颈。这不仅为科研人员提供了研究损失景观和优化动力学的利器，也侧面验证了高性能算子定制对于推动 AI 算法落地的关键作用。
💡 核心观点：二阶优化算法的工程化落地正逐步突破算力瓶颈，高性能开源算子库将成为加速 AI 科研与工业应用的重要基石。
原文链接：Hacker News
评论前必须登录！
立即登录 注册