noahgolmant/pytorch-hessian-eigenthings

它解决了什么问题

对于大型神经网络来说，计算完整的 Hessian 矩阵（损失函数的二阶导数）在计算上是不可能的，因为内存需求随参数数量呈平方级增长。本库提供了一种可扩展的方法来分析模型损失景观（loss landscape）的曲率，而无需显式构建完整的矩阵。

工作原理

该项目使用 Hessian-vector products (HVPs)，它仅需要线性内存。通过将 HVPs 与迭代算法相结合，它可以在没有内存瓶颈的情况下计算 Hessian 及其他曲率矩阵（如 Generalized Gauss-Newton 和 empirical Fisher）的特定属性。

包含的关键算法有：

• Lanczos and stochastic power iteration，用于寻找前几个特征值和特征向量。
• Hutch++，用于估计矩阵的迹（trace）。
• Stochastic Lanczos Quadrature，用于估计谱密度（spectral density）。

对于大语言模型，它包含优化后的内核（使用 Triton 或 torch.compile），以加速交叉熵 Hessian-vector products。

适用对象

致力于神经网络优化、泛化分析以及研究损失景观几何结构（例如分析“flat minima”）的研究人员和开发人员。

亮点

• 可扩展分析：为真实世界的模型计算特征分解和谱密度，包括来自 HuggingFace 和 TransformerLens 的模型。
• 灵活的算子：支持 Hessian、Generalized Gauss-Newton (GGN) 和 Empirical Fisher 算子。
• 内存高效：通过 HVPs 使用线性内存，而非平方级内存。
• 性能优化：具有针对 LM 规模工作的融合内核（fused kernels），以减少峰值内存并提高速度。
• 参数过滤：允许使用基于名称的过滤器对特定的参数子集（例如 transformer 中的特定块）进行分析。