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。

適用對象

從事神經網路優化、泛化分析，以及研究損失地形幾何結構（例如分析「平坦極小值」）的研究人員與開發者。

重點特性

• 可擴展分析：可針對真實世界的模型（包括來自 HuggingFace 和 TransformerLens 的模型）計算特徵分解與譜密度。
• 靈活的算子：支援 Hessian、Generalized Gauss-Newton (GGN) 和 Empirical Fisher 算子。
• 記憶體效率高：透過 HVPs 使用線性記憶體，而非平方級記憶體。
• 效能優化：針對語言模型規模的工作提供融合內核（fused kernels），以降低峰值記憶體並提升速度。
• 參數過濾：允許使用基於名稱的過濾器來分析特定的參數子集（例如 transformer 中的特定區塊）。