开发者重写 PyTorch 曲率分析库：引入 Triton 内核，支持大规模模型优化

开发者 Noah Golmant 近日发布了 PyTorch 曲率分析开源库 pytorch-hessian-eigenthings 的 1.0 大版本重写。该库最初发布于 2018 年，旨在解决神经网络损失景观的曲率指标计算问题，这些指标与神经网络的泛化能力密切相关。原有的完整海森矩阵计算内存消耗随参数量呈二次方增长，在大规模模型中难以实施。新版库通过 Hessian-vector 积结合兰乔斯算法、幂迭代等迭代方法，成功将特征值分解的内存复杂度优化至线性级别。此次重写引入了多项重大更新：新增了广义高斯-牛顿（GGN）和经验 Fisher 矩阵等曲率算子；支持 Hutchinson 和 Hutch++ 等迹估计算法以及谱密度估计；针对拥有超大词表的基础模型，实现了基于 Triton 和 torch.compile 的融合交叉熵 Hessian-vector 内核，有效解决了标准实现在此场景下的内存爆炸问题。此外，新版本加强了数值分析验证，在线性和逻辑回归模型上进行了闭合形式正确性测试，并与 curvlinops 库进行了交叉验证以防止回归。

事件分析

💡 核心观点：二阶优化算法的工程化落地正逐步突破算力瓶颈，高性能开源算子库将成为加速 AI 科研与工业应用的重要基石。