【Bug已解决】How to compute Hessian of the loss w.r.t. the parameters in PyTorch using autograd.grad 解决方案

【Bug已解决】How to compute Hessian of the loss w.r.t. the parameters in PyTorch using autograd.grad 解决方案

问题标签：PyTorch、Hessian矩阵、autograd.grad、二阶导数、二阶优化
适用版本：PyTorch 1.8+
难度等级：⭐⭐⭐⭐⭐（高）

一、问题描述

在深度学习中，大多数优化算法（SGD、Adam）只使用一阶梯度信息。但在某些场景下，我们需要计算二阶导数（Hessian 矩阵）：

• 二阶优化方法（Newton's Method, Trust Region）
• 模型不确定性估计（Laplace Approximation）
• 影响函数（Influence Functions）
• 剪枝与压缩（Optimal Brain Damage）
• 损失曲面分析

核心问题：

如何使用 PyTorch 的autograd.grad计算损失函数对模型参数的 Hessian 矩阵？

Hessian 矩阵的定义

对于损失函数 ( L(\theta) )，Hessian 矩阵 ( H ) 是一个 ( n \times n \ 的矩阵（( n ) 为参数数量）：

[ H_{ij} = \frac{\partial^2 L}{\partial \theta_i \partial \theta_j} ]

二、错误复现

场景：尝试用autograd.grad直接计算 Hessian

import torch # 创建参数和损失 x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) y = (x ** 2).sum() # 计算一阶梯度 grad = torch.autograd.grad(y, x, create_graph=True)[0] print(f"一阶梯度: {grad}") # tensor([2., 4., 6.]) # 尝试直接计算 Hessian（错误方式） try: hessian = torch.autograd.grad(grad, x)[0] # 看似合理 print(f"Hessian: {hessian}") except Exception as e: print(f"错误: {e}")

问题：上述代码对于简单函数可以工作，但对于复杂模型（如神经网络），直接计算 Hessian 会遇到以下问题：

1. 内存爆炸：Hessian 是 ( n \times n ) 矩阵，参数量百万级时无法存储
2. 计算效率低：需要 ( n ) 次反向传播
3. API 使用不当：autograd.grad的参数配置不正确

三、根因分析

3.1 autograd.grad 的工作原理

torch.autograd.grad(outputs, inputs, ...)计算outputs对inputs的梯度。

一阶导数: autograd.grad(loss, params) → ∂L/∂θ 二阶导数: autograd.grad(∂L/∂θ, params) → ∂²L/∂θ²

关键：计算一阶导数时必须设置create_graph=True，这样一阶导数本身也参与计算图，才能对其求导。

3.2 Hessian 矩阵的计算挑战

3.3 PyTorch 中的 Hessian 计算 API

PyTorch 提供了多种计算 Hessian 的方式：

四、解决方案

方案一：使用torch.autograd.functional.hessian（小规模推荐）

import torch from torch.autograd.functional import hessian # 定义函数 def loss_function(params): """示例: f(x, y) = x² + 3xy + 2y²""" x, y = params[0], params[1] return x**2 + 3*x*y + 2*y**2 # 计算点 point = torch.tensor([1.0, 2.0]) # 计算 Hessian H = hessian(loss_function, point) print(f"Hessian 矩阵:\n{H}") # 理论值: [[2, 3], [3, 4]] # ∂²f/∂x² = 2, ∂²f/∂x∂y = 3, ∂²f/∂y² = 4

方案二：手动使用autograd.grad计算 Hessian

import torch def compute_hessian_manual(loss, params): """ 手动计算 Hessian 矩阵 原理: 1. 计算一阶梯度（create_graph=True） 2. 对一阶梯度的每个分量求导 3. 拼接成 Hessian 矩阵 参数: loss: 标量损失 params: 参数张量（扁平化为一维） """ # 计算一阶梯度 grad = torch.autograd.grad( loss, params, create_graph=True, retain_graph=True )[0] n = grad.numel() hessian = torch.zeros(n, n) # 逐行计算 Hessian for i in range(n): # 对第 i 个梯度分量求导 grad2 = torch.autograd.grad( grad[i], params, retain_graph=True, create_graph=True )[0] hessian[i] = grad2 return hessian # 测试 x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) loss = (x ** 2).sum() + x[0] * x[1] # x0² + x1² + x2² + x0*x1 H = compute_hessian_manual(loss, x) print(f"参数: {x.data}") print(f"损失: {loss.item():.4f}") print(f"Hessian:\n{H}") # 理论值: [[2, 1, 0], [1, 2, 0], [0, 0, 2]]

方案三：计算神经网络参数的 Hessian

import torch import torch.nn as nn def compute_nn_hessian(model, loss_fn, inputs, targets): """ 计算神经网络损失对参数的 Hessian 注意: 仅适用于小模型（参数量 < 1000） """ # 前向传播 outputs = model(inputs) loss = loss_fn(outputs, targets) # 收集所有参数 params = [p for p in model.parameters() if p.requires_grad] # 扁平化参数 flat_params = torch.cat([p.view(-1) for p in params]) # 计算一阶梯度 grad = torch.autograd.grad( loss, params, create_graph=True, retain_graph=True ) flat_grad = torch.cat([g.view(-1) for g in grad]) n = flat_grad.numel() hessian = torch.zeros(n, n) # 逐行计算 for i in range(n): grad2 = torch.autograd.grad( flat_grad[i], params, retain_graph=True ) flat_grad2 = torch.cat([g.view(-1) for g in grad2]) hessian[i] = flat_grad2 return hessian # 测试小模型 model = nn.Linear(3, 2, bias=False) # 6 个参数 inputs = torch.randn(4, 3) targets = torch.tensor([0, 1, 0, 1]) loss_fn = nn.CrossEntropyLoss() H = compute_nn_hessian(model, loss_fn, inputs, targets) print(f"模型参数量: {sum(p.numel() for p in model.parameters())}") print(f"Hessian 形状: {H.shape}") print(f"Hessian:\n{H}")

方案四：Hessian-Vector Product (HVP) —— 大规模推荐

对于大模型，不需要显式计算 Hessian 矩阵，只需计算 Hessian 与向量的乘积。

import torch import torch.nn as nn def hvp(loss, params, vector): """ 计算 Hessian-Vector Product: H × v 原理（Pearlmutter, 1992）: H × v = ∇(∇L · v) 即: 先计算梯度和向量的点积，再对参数求导 优点: 只需 2 次反向传播，不需要存储 Hessian 矩阵 复杂度: O(n) 时间和内存（而非 O(n²)） """ # 计算一阶梯度 grad = torch.autograd.grad( loss, params, create_graph=True, retain_graph=True ) # 梯度与向量点积 grad_v_product = sum( (g * v).sum() for g, v in zip(grad, vector) ) # 对点积求导，得到 HVP hvp_result = torch.autograd.grad( grad_v_product, params, retain_graph=True ) return hvp_result # 测试 HVP model = nn.Linear(10, 5) inputs = torch.randn(8, 10) targets = torch.randn(8, 5) loss_fn = nn.MSELoss() outputs = model(inputs) loss = loss_fn(outputs, targets) params = [p for p in model.parameters() if p.requires_grad] vector = [torch.randn_like(p) for p in params] # 随机向量 # 计算 HVP hvp_result = hvp(loss, params, vector) for i, (p, h) in enumerate(zip(params, hvp_result)): print(f"参数 {i} HVP: shape={h.shape}, norm={h.norm():.4f}")

方案五：使用 HVP 近似 Hessian 的对角线

import torch import torch.nn as nn def hessian_diagonal(loss, params, num_iterations=1): """ 使用 HVP 计算 Hessian 的对角线元素 原理: H × e_i = H 的第 i 列 对角线: H_ii = (H × e_i)_i 对于每个参数，使用单位向量计算 HVP """ diag = [] for p in params: # 对每个参数张量，逐元素计算对角线 p_diag = torch.zeros_like(p) n = p.numel() for i in range(n): # 创建单位向量 e_i = torch.zeros(n, device=p.device) e_i[i] = 1.0 # 重塑为参数形状 e_i = e_i.view_as(p) # 构建完整向量（其他参数为零） vector = [torch.zeros_like(pp) for pp in params] vector[params.index(p)] = e_i # 计算 HVP hvp_result = hvp(loss, params, vector) # 取对应位置的对角线值 p_diag.view(-1)[i] = hvp_result[params.index(p)].view(-1)[i] diag.append(p_diag) return diag # 更高效的对角线估计方法 def hessian_diagonal_efficient(model, loss_fn, inputs, targets): """高效计算 Hessian 对角线（一次反向传播）""" outputs = model(inputs) loss = loss_fn(outputs, targets) params = [p for p in model.parameters() if p.requires_grad] # 计算一阶梯度 grad = torch.autograd.grad( loss, params, create_graph=True, retain_graph=True ) # 对每个参数计算对角线 diag = [] for g, p in zip(grad, params): # 对梯度的每个元素求导 g_diag = torch.zeros_like(p) n = g.numel() for i in range(n): if g.view(-1)[i].requires_grad: second_deriv = torch.autograd.grad( g.view(-1)[i], p, retain_graph=True )[0] g_diag.view(-1)[i] = second_deriv.view(-1)[i] diag.append(g_diag) return diag

五、完整修复代码

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset from torch.autograd.functional import hessian as functional_hessian # ==================== Hessian 计算工具类 ==================== class HessianCalculator: """Hessian 矩阵计算工具类""" @staticmethod def compute_full_hessian(loss, params): """ 计算完整的 Hessian 矩阵（仅适用于小模型） 复杂度: O(n²) 内存, O(n) 次反向传播 """ # 扁平化参数 if isinstance(params, torch.Tensor): params = [params] # 计算一阶梯度 grad = torch.autograd.grad( loss, params, create_graph=True, retain_graph=True ) flat_grad = torch.cat([g.view(-1) for g in grad]) n = flat_grad.numel() hessian = torch.zeros(n, n) for i in range(n): grad2 = torch.autograd.grad( flat_grad[i], params, retain_graph=True ) flat_grad2 = torch.cat([g.view(-1) for g in grad2]) hessian[i] = flat_grad2 return hessian @staticmethod def hvp(loss, params, vector): """ 计算 Hessian-Vector Product 复杂度: O(n) 内存, 2 次反向传播 """ if isinstance(params, torch.Tensor): params = [params] grad = torch.autograd.grad( loss, params, create_graph=True, retain_graph=True ) grad_v_product = sum( (g * v).sum() for g, v in zip(grad, vector) ) hvp_result = torch.autograd.grad( grad_v_product, params, retain_graph=True ) return hvp_result @staticmethod def compute_hessian_diagonal(model, loss_fn, inputs, targets): """计算 Hessian 对角线""" outputs = model(inputs) loss = loss_fn(outputs, targets) params = [p for p in model.parameters() if p.requires_grad] grad = torch.autograd.grad( loss, params, create_graph=True, retain_graph=True ) diag = [] for g, p in zip(grad, params): g_diag = torch.zeros_like(p) n = g.numel() for i in range(n): second_deriv = torch.autograd.grad( g.view(-1)[i], p, retain_graph=True, create_graph=False )[0] g_diag.view(-1)[i] = second_deriv.view(-1)[i] diag.append(g_diag) return diag @staticmethod def compute_hessian_eigenvalues(hessian_matrix, top_k=5): """计算 Hessian 矩阵的特征值""" eigenvalues = torch.linalg.eigvalsh(hessian_matrix) # 排序（绝对值从大到小） sorted_vals, indices = eigenvalues.abs().sort(descending=True) return eigenvalues[indices[:top_k]] @staticmethod def lanczos_iteration(loss, params, num_iterations=10): """ 使用 Lanczos 算法近似 Hessian 的最大/最小特征值 适用于大模型，只需要 HVP 操作 """ n = sum(p.numel() for p in params) # 初始化 v_prev = torch.zeros(n) v_curr = torch.randn(n) v_curr = v_curr / v_curr.norm() alphas = [] betas = [] for i in range(num_iterations): # 重塑向量为参数形状 vector = [] offset = 0 for p in params: size = p.numel() vector.append(v_curr[offset:offset+size].view_as(p)) offset += size # 计算 HVP hvp_result = HessianCalculator.hvp(loss, params, vector) hv_flat = torch.cat([h.view(-1) for h in hvp_result]) # Lanczos 递推 alpha = v_curr @ hv_flat alphas.append(alpha.item()) hv_flat = hv_flat - alpha * v_curr - (betas[-1] if betas else 0) * v_prev beta = hv_flat.norm() if beta < 1e-10: break betas.append(beta.item()) v_prev = v_curr v_curr = hv_flat / beta # 从三对角矩阵特征值近似 Hessian 特征值 T = torch.diag(torch.tensor(alphas)) for i in range(len(betas)): T[i, i+1] = betas[i] T[i+1, i] = betas[i] eigenvalues = torch.linalg.eigvalsh(T) return eigenvalues # ==================== 应用示例 ==================== class HessianAnalysis: """Hessian 分析应用""" @staticmethod def analyze_loss_landscape(model, loss_fn, inputs, targets): """分析损失曲面""" print("=" * 60) print("Hessian 分析") print("=" * 60) # 前向传播 outputs = model(inputs) loss = loss_fn(outputs, targets) params = [p for p in model.parameters() if p.requires_grad] n_params = sum(p.numel() for p in params) print(f"参数数量: {n_params}") if n_params <= 100: # 小模型: 计算完整 Hessian H = HessianCalculator.compute_full_hessian(loss, params) print(f"Hessian 形状: {H.shape}") print(f"Hessian 范数: {H.norm():.4f}") print(f"Hessian 迹: {H.trace():.4f}") # 特征值分析 eigenvalues = torch.linalg.eigvalsh(H) print(f"最大特征值: {eigenvalues.max():.4f}") print(f"最小特征值: {eigenvalues.min():.4f}") print(f"条件数: {eigenvalues.max() / eigenvalues.min():.4f}") # 正定性检查 positive = (eigenvalues > 0).sum().item() negative = (eigenvalues < 0).sum().item() zero = (eigenvalues.abs() < 1e-6).sum().item() print(f"正特征值: {positive}, 负特征值: {negative}, 零特征值: {zero}") if negative == 0: print("→ Hessian 正定（凸区域）") elif positive == 0: print("→ Hessian 负定（凹区域）") else: print("→ Hessian 不定（鞍点区域）") else: # 大模型: 使用 HVP 和 Lanczos print("模型较大，使用 HVP 近似...") # HVP 测试 vector = [torch.randn_like(p) for p in params] hvp_result = HessianCalculator.hvp(loss, params, vector) hvp_norm = sum(h.norm() for h in hvp_result) print(f"HVP 范数: {hvp_norm:.4f}") # Lanczos 近似特征值 eigenvalues = HessianCalculator.lanczos_iteration( loss, params, num_iterations=min(10, n_params) ) print(f"Lanczos 近似特征值: {eigenvalues}") @staticmethod def optimal_brain_damage_pruning(model, loss_fn, inputs, targets, prune_ratio=0.1): """使用 Hessian 对角线进行 Optimal Brain Damage 剪枝""" print("\n" + "=" * 60) print("Optimal Brain Damage 剪枝") print("=" * 60) # 计算 Hessian 对角线 diag = HessianCalculator.compute_hessian_diagonal( model, loss_fn, inputs, targets ) # 计算 saliency: (weight)² × hessian_diag saliencies = [] for p, h in zip(model.parameters(), diag): s = (p.data ** 2) * h saliencies.append(s.view(-1)) all_saliencies = torch.cat(saliencies) # 确定剪枝阈值 threshold = torch.quantile( all_saliencies.abs(), prune_ratio ) # 执行剪枝 total_pruned = 0 total_params = 0 for p, h in zip(model.parameters(), diag): mask = ((p.data ** 2) * h).abs() > threshold p.data[~mask] = 0 total_pruned += (~mask).sum().item() total_params += mask.numel() print(f"剪枝比例: {100.*total_pruned/total_params:.1f}%") print(f"剪枝数量: {total_pruned}/{total_params}") # ==================== 完整示例 ==================== def run_hessian_examples(): """运行 Hessian 计算示例""" # === 示例 1: 简单函数的 Hessian === print("=" * 60) print("示例 1: 简单函数的 Hessian") print("=" * 60) def quadratic(x): return x[0]**2 + 3*x[0]*x[1] + 2*x[1]**2 point = torch.tensor([1.0, 2.0]) H = functional_hessian(quadratic, point) print(f"函数: f(x,y) = x² + 3xy + 2y²") print(f"在点 {point} 的 Hessian:\n{H}") print(f"理论值: [[2, 3], [3, 4]]") # === 示例 2: 神经网络的 Hessian === print("\n" + "=" * 60) print("示例 2: 神经网络的 Hessian 分析") print("=" * 60) # 创建小模型 model = nn.Sequential( nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3) ) inputs = torch.randn(16, 4) targets = torch.randint(0, 3, (16,)) loss_fn = nn.CrossEntropyLoss() # 分析 HessianAnalysis.analyze_loss_landscape(model, loss_fn, inputs, targets) # === 示例 3: HVP === print("\n" + "=" * 60) print("示例 3: Hessian-Vector Product") print("=" * 60) # 创建较大模型 large_model = nn.Sequential( nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 10) ) inputs = torch.randn(32, 100) targets = torch.randint(0, 10, (32,)) outputs = large_model(inputs) loss = nn.CrossEntropyLoss()(outputs, targets) params = [p for p in large_model.parameters() if p.requires_grad] n_params = sum(p.numel() for p in params) print(f"模型参数量: {n_params}") # HVP vector = [torch.randn_like(p) for p in params] hvp_result = HessianCalculator.hvp(loss, params, vector) print(f"HVP 计算成功") for i, h in enumerate(hvp_result): print(f" 参数 {i}: shape={h.shape}, norm={h.norm():.4f}") if __name__ == '__main__': run_hessian_examples()

六、常见陷阱与注意事项

陷阱 1：忘记create_graph=True

import torch x = torch.tensor([1.0, 2.0], requires_grad=True) y = (x ** 2).sum() # ❌ 没有 create_graph=True，一阶梯度不参与计算图 grad = torch.autograd.grad(y, x)[0] try: hessian = torch.autograd.grad(grad, x)[0] except RuntimeError as e: print(f"错误: {e}") # ✓ 设置 create_graph=True grad = torch.autograd.grad(y, x, create_graph=True)[0] hessian = torch.autograd.grad(grad.sum(), x)[0] print(f"Hessian: {hessian}")

陷阱 2：对非标量求导

# ❌ autograd.grad 的 outputs 必须是标量（或需要 grad_tensors） x = torch.tensor([1.0, 2.0], requires_grad=True) y = x ** 2 # 向量 try: grad = torch.autograd.grad(y, x)[0] except RuntimeError as e: print(f"错误: {e}") # ✓ 对标量求导 loss = y.sum() grad = torch.autograd.grad(loss, x)[0]

陷阱 3：内存爆炸

import torch import torch.nn as nn # ❌ 对大模型计算完整 Hessian # model = nn.Linear(10000, 10000) # 1 亿参数 # Hessian 是 1亿 × 1亿 矩阵 → 不可能存储 # ✓ 使用 HVP 或对角线近似 # hvp_result = hvp(loss, params, vector) # O(n) 内存

陷阱 4：混淆 Hessian 和梯度的外积

import torch x = torch.tensor([1.0, 2.0], requires_grad=True) y = (x ** 3).sum() # ❌ 梯度外积不是 Hessian grad = torch.autograd.grad(y, x, create_graph=True)[0] outer_product = torch.outer(grad, grad) # 这是 ∇L ⊗ ∇L，不是 Hessian # ✓ 正确计算 Hessian hessian = torch.zeros(2, 2) for i in range(2): hessian[i] = torch.autograd.grad(grad[i], x, retain_graph=True)[0] print(f"Hessian: {hessian}")

陷阱 5：数值精度问题

import torch # 二阶导数可能放大数值误差 x = torch.tensor([1e-8], requires_grad=True) y = x ** 2 grad = torch.autograd.grad(y, x, create_graph=True)[0] hessian = torch.autograd.grad(grad, x)[0] print(f"Hessian: {hessian}") # 可能不精确 # ✓ 使用 double 精度 x = torch.tensor([1e-8], requires_grad=True, dtype=torch.float64) y = x ** 2 grad = torch.autograd.grad(y, x, create_graph=True)[0] hessian = torch.autograd.grad(grad, x)[0] print(f"Hessian (double): {hessian}") # 更精确

七、总结

选择决策树

需要 Hessian 信息？ ├── 参数量 < 100？ │ └── 使用 functional.hessian 或手动循环 ├── 参数量 < 10000？ │ └── 计算 Hessian 对角线 ├── 只需要 H × v？ │ └── 使用 HVP └── 需要特征值？ └── 使用 Lanczos 迭代

核心原则

对于大规模神经网络，永远不要尝试计算完整的 Hessian 矩阵。使用 Hessian-Vector Product (HVP) 是处理二阶信息的标准做法，它只需要 O(n) 的内存和常数次反向传播。

API 速查

# 一阶梯度 grad = torch.autograd.grad(loss, params, create_graph=True) # Hessian-Vector Product hvp = hvp_function(loss, params, vector) # 完整 Hessian（小规模） H = torch.autograd.functional.hessian(func, inputs) # Hessian 对角线 diag = compute_hessian_diagonal(model, loss_fn, inputs, targets) # 特征值近似 eigenvalues = lanczos_iteration(loss, params, num_iterations=10)

通过掌握 Hessian 矩阵的计算方法，可以深入理解神经网络的损失曲面结构，为二阶优化、模型压缩、不确定性估计等高级应用奠定基础。