PyTorch中快速计算模型参数Hessian矩阵的优化方法咨询

我希望在PyTorch中计算损失关于模型参数的Hessian矩阵，但无法使用torch.autograd.functional.hessian，因为该函数会重新计算我已通过前置调用得到的模型输出与损失。我的当前实现如下：

import torch
import time

# 创建模型
model = torch.nn.Sequential(torch.nn.Linear(1, 100), torch.nn.Tanh(), torch.nn.Linear(100, 1))
num_param = sum(p.numel() for p in model.parameters())

# 在随机数据集上计算损失
x = torch.rand((1000,1))
y = torch.rand((1000,1))
y_hat = model(x)
loss = ((y_hat - y)**2).mean()

''' 计算Hessian矩阵 '''
start = time.time()

# 初始化Hessian矩阵
H = torch.zeros((num_param, num_param))

# 计算损失关于模型参数的Jacobian
J = torch.autograd.grad(loss, list(model.parameters()), create_graph=True)
J = torch.cat([e.flatten() for e in J]) # 展平为一维向量

# 逐行填充Hessian矩阵
for i in range(num_param):
    result = torch.autograd.grad(J[i], list(model.parameters()), retain_graph=True)
    H[i] = torch.cat([r.flatten() for r in result]) # 展平

print(time.time() - start)

请问是否存在更快的实现方式？比如避免使用循环，因为循环会为每个模型变量调用autograd.grad。

可以通过批量计算梯度避免循环调用autograd.grad，利用torch.autograd.grad的grad_outputs参数一次性计算整个Jacobian的梯度，大幅提升效率。

核心思路：构造与Jacobian同维度的单位矩阵，将Jacobian与单位矩阵的每一列做点积（等价于取出Jacobian的每个元素），然后一次性对所有点积结果求导，直接得到完整的Hessian矩阵。

优化后的代码如下：

import torch
import time

# 创建模型
model = torch.nn.Sequential(torch.nn.Linear(1, 100), torch.nn.Tanh(), torch.nn.Linear(100, 1))
num_param = sum(p.numel() for p in model.parameters())

# 在随机数据集上计算损失
x = torch.rand((1000,1))
y = torch.rand((1000,1))
y_hat = model(x)
loss = ((y_hat - y)**2).mean()

''' 快速计算Hessian矩阵 '''
start = time.time()

# 计算损失关于模型参数的Jacobian（保留计算图）
J = torch.autograd.grad(loss, list(model.parameters()), create_graph=True)
J = torch.cat([e.flatten() for e in J])

# 构造单位矩阵，用于批量计算每个Jacobian元素的梯度
eye = torch.eye(num_param, device=J.device)

# 一次性计算所有Jacobian元素的梯度，得到完整的Hessian矩阵
H = torch.autograd.grad(J, list(model.parameters()), grad_outputs=eye, retain_graph=False)
# 将结果展平并拼接成二维矩阵
H = torch.cat([h.flatten() for h in H]).reshape(num_param, num_param)

print(time.time() - start)

优化说明 #

1. 消除循环开销：原循环需要调用num_param次autograd.grad，优化后仅需1次调用，彻底消除循环带来的额外开销。
2. 利用批量自动微分：grad_outputs参数允许同时对多个目标（Jacobian的每个元素）求导，PyTorch会自动并行处理计算，充分利用硬件并行能力（如GPU的CUDA核心）。
3. 降低内存损耗：批量计算减少了中间张量的创建与销毁次数，内存利用效率更高。

额外优化建议 #

如果模型参数数量较大，Hessian矩阵会占用大量内存（例如10000个参数的单精度Hessian约占400MB内存），可考虑：

• 使用稀疏Hessian矩阵：若模型结构导致Hessian存在大量零元素，可借助torch.sparse相关API存储，节省内存。
• 分块计算Hessian：将参数划分为若干块，逐块计算对应的Hessian子矩阵，降低单批次内存占用。

内容的提问来源于stack exchange，提问作者Thomas Wagenaar