在深度学习二阶分析领域，Hessian 特征值长期以来是 flat minima 假说、低秩结构推断和曲率感知优化的核心工具。然而，完整 Hessian 的存储开销与参数量呈二次关系，在十亿级参数模型上完全不可行。eigenthings 库（pytorch-hessian-eigenthings）在沉寂多年后，于 2025 年发布了 v1.0 重写版本，将 HVP（ Hessian-Vector Product ）的计算图抽象、曲率算子设计和 Lanczos 族算法整合为一套生产可用的工具链。本文聚焦于该重写的核心工程决策，重点拆解 CurvatureOperator 计算图重构、线性内存反向模式设计，以及生产级特征值工具链的参数化实践。
1. 核心问题：为何需要 HVP 而不是显式 Hessian
在进入 API 设计之前，有必要明确 HVP 的数学意义与内存收益。对于一个具有 $n$ 个参数的模型，完整 Hessian 矩阵 $H \in \mathbb {R}^{n \times n}$ 需要 $O (n^2)$ 存储空间，即使参数量仅为一百万，所需显存也超过 8TB（ float64 ）。而 HVP 操作定义为 $Hv$，即 Hessian 矩阵与任意向量 $v$ 的乘积，计算时无需显式构造 $H$，仅需两次反向传播（backward-on-backward），内存开销线性降为 $O (n)$。
具体实现路径为：首先对损失函数 $L$ 执行标准前向 - 反向传播，得到参数空间梯度 $g = \nabla_\theta L$；随后以向量 $v$ 作为梯度输出的权重，执行第二次反向传播得到 $\nabla_\theta (g^\top v)$，结果即为 $Hv$。PyTorch 的 autograd 引擎天然支持这一链式操作：第一次 autograd.grad(loss, params, create_graph=True) 保留计算图，第二次以 v 作为 grad_outputs 调用 autograd.grad，即可得到曲率信息而不存储完整 Hessian。
2. CurvatureOperator 抽象：计算图重构的核心
v2 版本最重要的架构决策是将曲率算子抽象为统一的 CurvatureOperator 接口层。在此之前，v0.x 版本的实现散落在各个独立函数中，缺乏统一的算子规范，导致扩展新曲率矩阵（如 GGN 或经验 Fisher）时需要重复实现 Lanczos 集成逻辑。
重写后的设计将所有曲率矩阵组织为一个可调用对象，接口契约简化为：输入任意形状与 $v$ 兼容的向量 $u$，返回 $Cu$。三个核心算子均实现了这一接口：
HessianOperator 对应真实二阶导数 Hessian，适用于回归损失，在分类损失下可能非半正定；GGNOperator（ Generalized Gauss-Newton ）始终保持半正定，其数学近似精度通常足以支撑分析需求，在分类场景下更符合直觉；EmpiricalFisherOperator 基于样本梯度外积的平均，对应自然梯度中的 Fisher 信息矩阵估计。
这种统一抽象的核心价值在于下游算法的完全解耦：Lanczos 特征提取、Hutch++ 追踪估计和谱密度计算均可作用于任意 CurvatureOperator，无需为每个算子单独适配求解器逻辑。在工程实践中，这意味着新增曲率矩阵类型或算法变体时，只需实现或调用标准接口，而不必改动求解器核心。
3. 线性内存反向模式：梯度回收与检查点策略
在 HVP 的两次反向传播中，中间激活值需要被保留至第二次反向传播完成，这导致峰值显存接近前向传播的两倍。v2 重写通过两类策略控制内存开销。
第一类是梯度图的精细管理。通过 retain_graph 参数显式控制计算图的生命周期：在 Lanczos 迭代的连续向量 - 矩阵乘积中，同一参数化的 HVP 运算可以共享同一个反向图，避免每次 HVP 触发重新前向传播。但需要注意图保留与 PyTorch 版本间的数值稳定性，部分场景下连续多次 retain_graph=True 可能累积浮点误差，此时应考虑显式重新构造计算图。
第二类是激活检查点（ checkpointing ）。PyTorch 的 torch.utils.checkpoint.checkpoint 系列 API 可在计算图的关键位置插入反向重计算点，将前向传播的内存占用从 $O (\text {层数})$ 压低至 $O (\text {层数}/\text {检查点数})$。在 HessianFlow 原始论文的实现中，这一策略被广泛用于大 batch 训练的内存优化。eigenthings v2 支持用户传入自定义的 checkpoint 函数，使得在超大规模模型上运行 Lanczos 迭代成为可能。
此外，HessianOperator 还提供了有限差分（ finite_difference ）路径，通过数值微分而非 autograd 计算 HVP，适用于 double_backward 存在困难的场景，例如 FSDP（ Fully Sharded Data Parallel ）训练框架或含有不透明 C++ 算子的模型。该路径牺牲了部分精度换取兼容性，在标准反向模式可用的绝大多数场景中不推荐使用。
4. 特征值与特征向量：Lanczos 参数化实践
Lanczos 算法是 eigenthings v2 求解特征值分解的核心工具。其基本思想是构建 Krylov 子空间 ${v, HV, H^2V, \ldots}$ 的正交基，将无限维算子在低维子空间上投影为三对角矩阵 $T$，其特征值作为原始 Hessian 特征值的近似。
eigenthings 提供的 lanczos(H, k, seed) 函数中，k 参数控制待求解的顶部特征值数量。实践中，k 的选择需要在精度与计算开销间权衡：较大的 k 值（如 20–50）能够捕捉更丰富的谱结构，但每次迭代的 orthogonalization 开销从 $O (k)$ 上升至 $O (k^2)$；较小的 k（如 5–10）运行速度快，但谱密度估计的分辨率下降。默认实践中，对于分析模型收敛性，$k=10$ 是一个合理的初始值；对于研究谱间隙（ spectral gap ）结构，建议 $k \geq 20$。
seed 参数用于控制随机向量初始化的可重复性。在谱分析的严谨实验中，固定种子是确保结果可复现的基本要求，但需注意不同种子可能收敛至不同的局部特征向量集合，特别是在 Hessian 存在退化特征值时。建议对每个 k 配置运行 3–5 个随机种子取平均，以降低采样方差的影响。
Lanczos 迭代的终止条件默认基于 Rayleigh 商的变化率：当连续迭代间最大特征值的相对变化低于阈值时，认为已达到收敛。对于大多数实验配置，默认阈值（$10^{-5}$）足够使用；在需要极致精度时，可通过 lanczos_steps 参数显式控制迭代步数（通常设为 $k \times 3$ 至 $k \times 5$）。
5. 追踪估计：Hutch++ 与随机方差缩减
对于完整谱分析，仅知道顶部特征值是不够的。Hessian 的谱追踪（ trace，即对角线元素之和）可用于估计条件数、计算有效参数维度，以及验证低秩假设的成立程度。精确追踪需要完整 Hessian 的对角线，在大规模模型上同样不可行。
eigenthings v2 实现了 Hutch++ 追踪估计器，这是一种基于随机投影的方差缩减方法。相比原始 Hutchinson 估计器（使用单个随机向量的蒙特卡洛方法），Hutch++ 通过两次随机采样将方差从 $O (1/\epsilon^2)$ 降至 $O (1/\epsilon)$，显著减少了所需的矩阵 - 向量乘积次数。
trace(H, num_matvecs, seed) 函数中，num_matvecs 是核心调参点。该参数控制随机投影数量，直接影响追踪估计的方差与置信区间。实践中，$99$ 次随机投影可提供较为稳定的估计；若需 $95%$ 置信区间宽度小于 $5%$ 的相对误差，$199$ 次投影通常是安全阈值。需要注意的是，Hutch++ 的收敛行为与 Hessian 谱分布相关：高度病态（ ill-conditioned ）的矩阵需要更多采样才能平滑极端特征值的贡献。
6. 谱密度：随机 Lanczos 二次型（SLQ）
谱密度（ spectral density ）描述 Hessian 特征值在实数轴上的分布密度，即 $\rho (\lambda) = \sum_i \delta (\lambda - \lambda_i)$。它比单个特征值或追踪提供了更丰富的谱结构信息，可用于识别特征值聚集区间、检测低秩近似精度，以及可视化训练动态中的曲率演化。
eigenthings v2 通过随机 Lanczos 二次型（ Stochastic Lanczos Quadrature，SLQ ）计算谱密度。其基本原理是：对多个随机初始向量执行 Lanczos 迭代，将各子空间三对角矩阵的特征值 - 特征向量通过高斯加权求和近似原始谱密度。
spectral_density(H, num_runs, lanczos_steps, seed) 函数中，num_runs 控制随机向量的采样数量（默认为 8），lanczos_steps 控制每个 Lanczos 过程的迭代步数（默认为 40）。这两个参数的乘积决定了总计算量与分辨率的权衡：更多的 num_runs 平滑随机采样方差，更多的 lanczos_steps 提高单次 Lanczos 的谱分辨率。对于参数规模在 $10^6$ 至 $10^8$ 的模型，num_runs=8 与 lanczos_steps=40 的组合提供良好的工程精度；对于更大模型，可适当降低步数换取速度，因为 SLQ 的分辨率在高维空间天然受限。
7. 大词表交叉熵 HVP：Triton 融合内核
在语言模型分析中，输出层的交叉熵（ Cross-Entropy ）损失是常见的损失函数。当词表规模达到数万甚至数十万时，标准的 HVP 计算路径会产生两个瓶颈：对数 Softmax 的反向传播涉及大矩阵转置与归一化操作，内存占用与词表大小呈线性关系；梯度累积需要 $O (V)$ 大小的临时缓冲区（ $V$ 为词表大小）。
eigenthings v2 针对这一场景实现了融合交叉熵 Hessian-Vector 内核，在 CUDA 设备上使用 Triton 编写融合 kernel，在其他设备上使用 torch.compile 编译优化路径。根据基准测试，Triton 版本相比 eager 模式实现约 3.4 倍加速，峰值显存降低约 50%；torch.compile 路径提供约 2.6 倍加速和相近的显存缩减。
这一优化通过将 log-softmax 反向与 CE 梯度融合为单个 kernel，减少了中间结果的全局内存访问次数。对于使用 HuggingFace 或 TransformerLens 的因果语言模型，可通过 hf_lm_loss_of_output() 辅助函数自动调用该优化路径，用户无需感知底层实现细节。若需强制使用非融合路径进行调试，可传入 fused="eager" 参数。
8. 参数子集过滤：param_filter 与分块分析
在实际分析中，研究者经常只关注模型特定部分（如 Transformer 编码器的注意力层、或是 ResNet 的残差分支）的曲率结构，而非整个模型。eigenthings v2 通过 param_filter 参数支持参数子集过滤，该参数接受一个可调用对象或名称模式字符串，仅将匹配的参数纳入 Hessian 构造。
match_names("blocks.*.attn.*") 等模式支持通配符匹配，可快速定位特定子模块。在分块分析的场景中，一个典型工作流是：首先对整个模型运行 lanczos(k=5) 获取全局谱概况，然后针对每个注意力块单独构造 HessianOperator 并运行 Lanczos，最后比较不同块间特征值分布的差异。参数子集过滤不仅降低了算子构造的计算量，还使得分块曲率比较成为可能，这对诊断特定子模块的训练动态尤为有用。
9. 数值正确性验证体系
v2 重写的一个关键工程改进是建立了系统的数值正确性验证体系。该体系包含三个层次的验证。
第一层是解析闭式验证。在线性回归（ MSE 损失）和逻辑回归（交叉熵损失）中，Hessian 具有已知的解析形式：线性回归的 Hessian 是 $X^\top X$（加正则项），逻辑回归的 Hessian 是 $X^\top \text {diag}(\sigma (x) x (1-\sigma (x))) X$。eigenthings 通过在这些解析可验证的模型上运行算子，逐一比对数值结果与解析结果，确保 autograd 路径的正确性。
第二层是跨库回归测试。库中集成了与 curvlinops 库的交叉测试套件，在相同的模型配置与数据 batch 上对比两库的计算结果。任何数值偏差超过浮点精度容忍度的结果都会触发 CI 失败，这一机制有效防止了 API 变更或算法重构中的回归。
第三层是随机矩阵对照。对于无法构造解析 Hessian 的非线性网络，通过与随机初始化的网络在小参数量（ $n < 1000$ ）下的显式 Hessian 构造进行对比，验证 HVP 计算结果的正确性。该验证覆盖了 HessianOperator、GGNOperator 和 EmpiricalFisherOperator 三种算子。
10. 工程实践清单与选型建议
基于上述技术解析，以下给出面向生产使用的选型决策清单：
在算子选型上，默认推荐 GGNOperator 作为通用选择，其半正定性保证了优化方向的数值稳定性；若需理论上的真实 Hessian 且损失为回归型，HessianOperator 是精确选择；若关注自然梯度或 K-FAC 类近似，EmpiricalFisherOperator 对应 Fisher 信息矩阵。
在特征值求解上，对于仅需顶部 1–5 个特征值的场景，幂迭代（ power iteration ）通常比 Lanczos 更快收敛；对于需要估计特征值分布或追踪的场景，Lanczos 是唯一可行选项，此时建议 k=10 作为初始配置、k=20 以上用于高分辨率分析。
在内存约束场景下，首先启用梯度检查点以削减激活内存，然后考虑使用 method="finite_difference" 绕过 double-backward；最后可通过 torch.float16 混合精度在内存与精度间取折中（需注意半精度下特征值的数值误差）。
在词表规模超过 10,000 的语言模型上，优先使用 hf_lm_loss_of_output() 触发 Triton 融合内核，除非需要单步调试精度。
在结果复现性上，每次实验固定 seed，对关键结果运行至少 3 个不同种子取平均，并将 torch.backends.cudnn.deterministic 设为 True 以消除 cuDNN 非确定性带来的小幅波动。
eigenthings v2 的重写体现了二阶分析工具从研究原型向生产级库演进的标准路径：以 HVP 为核心计算基元，通过统一的算子抽象解耦算法与曲率矩阵定义，借助 Triton 融合内核和梯度检查点策略突破内存瓶颈，最终通过系统化的数值验证体系保障工程可靠性。对于需要在大规模模型上开展曲率感知优化或损失景观分析的研究者和工程师，这套工具链提供了从单卡 toy 模型到多卡 Foundation Model 的完整技术路径。
资料来源：GitHub - noahgolmant/pytorch-hessian-eigenthings（ https://github.com/noahgolmant/pytorch-hessian-eigenthings ）；作者 Hacker News 声明帖（ https://news.ycombinator.com/item?id=48132232 ）；官方文档（ https://noahgolmant.github.io/pytorch-hessian-eigenthings/ ）
内容声明：本文无广告投放、无付费植入。
如有事实性问题，欢迎发送勘误至 i@hotdrydog.com。