高阶梯度计算

PyTorch 的自动求导引擎动态构建计算图并向后遍历以计算梯度，例如 $\frac{\partial L}{\partial w}$。这是训练大多数神经网络 (neural network)的基础。然而，某些高级方法需要计算梯度的梯度，这被称为高阶梯度。

考虑一个函数 $f(x)$。它的一阶导数是 $f'(x) = \frac{df}{dx}$。二阶导数是 $f''(x) = \frac{d^2f}{dx^2}$，它简单来说就是一阶导数的导数。同样地，我们可以计算三阶、四阶及更高阶导数。在多变量函数（例如带有参数 (parameter) $\theta$ 的神经网络损失函数 (loss function) $L(\theta)$）的上下文 (context)中，我们通常处理偏导数。一阶梯度构成梯度向量 (vector) $\nabla L$。高阶导数涉及 Hessian 矩阵（二阶偏导数矩阵，$\nabla^2 L$）或更高阶张量等结构。

PyTorch 的自动求导引擎能够处理这些计算。虽然标准的 .backward() 方法主要用于一阶梯度，但 torch.autograd.grad 函数式接口提供了高阶微分所需的灵活性。

为何计算高阶梯度？

计算高阶梯度对一些高级应用非常重要：

1. 优化算法： 像牛顿法或信任区域法这样的方法使用二阶信息（Hessian 矩阵）来可能实现比 SGD 或 Adam 等一阶方法更快的收敛。虽然对于大型网络来说，计算完整的 Hessian 矩阵通常不可行，但通过高阶自动微分可以高效地计算 Hessian-向量 (vector)积 ($\nabla^2 L v$)，并被用于一些优化策略中。
2. 曲率分析： 二阶导数（Hessian 矩阵）描述了损失函数 (loss function)的曲率。分析这种曲率可以提供关于优化过程、泛化特性以及局部最小值或鞍点存在情况的理解。
3. 元学习： 像模型无关元学习 (MAML) 这样的算法涉及根据模型在特定任务上经过一次或多次梯度更新后的表现来优化其参数 (parameter)。这需要对梯度更新步骤本身进行微分，因此需要计算梯度的梯度。
4. 正则化 (regularization)技术： 某些正则化项明确依赖于梯度范数或二阶导数。例如，带梯度惩罚的 Wasserstein GAN (WGAN-GP) 中的梯度惩罚需要计算判别器输出相对于其输入的梯度的范数。
5. 物理信息神经网络 (neural network) (PINNs)： PINNs 将物理定律（常以偏微分方程 (PDEs) 形式表示）融入损失函数。这些 PDE 常涉及网络输出相对于其输入坐标（例如，时间和空间）的二阶或更高阶导数。

使用 torch.autograd.grad 计算高阶梯度

在 PyTorch 中计算高阶梯度的主要工具是 torch.autograd.grad。与 tensor.backward() 方法隐式计算所有需要梯度的叶节点梯度不同，torch.autograd.grad 更显式。

其基本签名如下：

torch.autograd.grad(
    outputs,        # 要进行微分的标量或张量
    inputs,         # 计算梯度时所依据的张量
    grad_outputs=None, # 损失函数对 'outputs' 的梯度（用于向量-雅可比积）
    retain_graph=None, # 如果为 True，保留图；否则释放。
    create_graph=False, # 如果为 True，为梯度计算本身构建图
    allow_unused=False
)

对于高阶梯度，重要的参数 (parameter)是 create_graph=True。当您使用 torch.autograd.grad 并设置 create_graph=True 来计算一阶梯度时，PyTorch 不仅计算梯度，还会构建必要的图结构，以便您稍后可以对这次梯度计算进行再次微分。如果 create_graph=False（默认值），梯度计算被视为一个终端操作；生成的梯度只是张量，没有将它们通过微分过程连接回原始参数的任何历史记录。

我们来看一个简单的例子。假设我们有 $y = x^3$。我们想计算 $\frac{dy}{dx} = 3x^2$ 和 $\frac{d^2y}{dx^2} = 6x$。

import torch

# 输入张量需要梯度
x = torch.tensor([2.0], requires_grad=True)

# 第一次计算: y = x^3
y = x**3
print(f"y = {y.item()}")

# 计算一阶导数: dy/dx
# 使用 create_graph=True 以允许计算高阶梯度
grad_y_x = torch.autograd.grad(outputs=y, inputs=x, create_graph=True)[0]
print(f"x={x.item()} 处的 dy/dx: {grad_y_x.item()}") # 应该是 3 * (2^2) = 12

# grad_y_x 现在是一个带有自身计算图的张量
print(f"梯度张量 requires_grad: {grad_y_x.requires_grad}")

# 计算二阶导数: d^2y/dx^2 = d/dx (dy/dx)
# 我们对*一阶梯度* (grad_y_x) 相对于 x 进行微分
# 除非我们想要三阶梯度，否则这里不需要 create_graph=True
grad2_y_x2 = torch.autograd.grad(outputs=grad_y_x, inputs=x)[0]
print(f"x={x.item()} 处的 d^2y/dx^2: {grad2_y_x2.item()}") # 应该是 6 * 2 = 12

# 检查二阶导数的 requires_grad 状态
print(f"二阶导数张量 requires_grad: {grad2_y_x2.requires_grad}")

请注意，grad_y_x 的 requires_grad=True，因为我们在其计算过程中指定了 create_graph=True。这允许我们再次以 grad_y_x 作为输出调用 torch.autograd.grad。最终的 grad2_y_x2 的 requires_grad=False，因为我们在第二次调用中未指定 create_graph=True。

图修改示意

当使用 create_graph=True 时，反向传播 (backpropagation)过程本身会将节点添加到计算图中。

考虑 $y=x^2$，所以 $\frac{dy}{dx} = 2x$。

1. 前向传播： x -> pow(2) -> y
2. 反向传播 (create_graph=False)： 计算梯度 ($2x$) 并将其作为与用于计算它的图分离的新张量返回。
3. 反向传播 (create_graph=True)： 计算梯度 ($2x$)，但会将表示该梯度如何计算的操作添加到图中：x -> pow(2) -> y；grad_y -> MulBackward (使用保存的 x) -> grad_x。输出 grad_x 被连接到这个扩展图上。

该图对比了 torch.autograd.grad 在 create_graph=False (中间) 和 create_graph=True (右侧) 时的结果。当 create_graph=True 时，计算出的梯度 grad_x 通过梯度计算操作 (PowBackward) 保持与图的连接，从而允许进一步微分。

示例：Hessian-向量 (vector)积

我们来计算一个简单函数 $f(w_1, w_2) = w_1^2 \sin(w_2)$ 的 Hessian-向量积 (HVP)。梯度为 $\nabla f = [\frac{\partial f}{\partial w_1}, \frac{\partial f}{\partial w_2}] = [2w_1 \sin(w_2), w_1^2 \cos(w_2)]$。Hessian 矩阵为 $\nabla^2 f = \begin{pmatrix} \frac{\partial^2 f}{\partial w_1^2} & \frac{\partial^2 f}{\partial w_1 \partial w_2} \\ \frac{\partial^2 f}{\partial w_2 \partial w_1} & \frac{\partial^2 f}{\partial w_2^2} \end{pmatrix} = \begin{pmatrix} 2\sin(w_2) & 2w_1 \cos(w_2) \\ 2w_1 \cos(w_2) & -w_1^2 \sin(w_2) \end{pmatrix}$。

我们希望计算 $(\nabla^2 f) v$，其中 $v$ 为某个向量，且无需显式构造 $\nabla^2 f$。通过两次 torch.autograd.grad 调用即可实现。主要思路是 $(\nabla^2 f) v = \nabla (\nabla f \cdot v)$，其中 $\nabla f \cdot v$ 是点积（一个标量）。

import torch

w = torch.tensor([1.0, torch.pi / 2.0], requires_grad=True) # w1=1，w2=pi/2
v = torch.tensor([0.5, 1.0]) # 一个任意向量

# 定义函数
f = w[0]**2 * torch.sin(w[1])

# 计算一阶梯度: grad_f = nabla f
grad_f = torch.autograd.grad(f, w, create_graph=True)[0]
# 预期 grad_f: [2*1*sin(pi/2), 1^2*cos(pi/2)] = [2, 0]
print(f"梯度 (nabla f): {grad_f}")

# 计算点积: grad_f_dot_v = (nabla f) . v
# 这个操作需要成为图的一部分，以便进行第二次微分
grad_f_dot_v = torch.dot(grad_f, v)
print(f"点积 (nabla f . v): {grad_f_dot_v}") # 预期: 2*0.5 + 0*1 = 1.0

# 计算点积相对于 w 的梯度: nabla (nabla f . v)
# 这得到 Hessian-向量积 (nabla^2 f) v
hvp = torch.autograd.grad(grad_f_dot_v, w)[0]

# 预期 Hessian: [[2*sin(pi/2), 2*1*cos(pi/2)], [2*1*cos(pi/2), -1^2*sin(pi/2)]]
# = [[2, 0], [0, -1]]
# 预期 HVP: [[2, 0], [0, -1]] @ [0.5, 1.0] = [2*0.5 + 0*1, 0*0.5 + (-1)*1] = [1.0, -1.0]
print(f"Hessian-向量积 (nabla^2 f) v: {hvp}")

这种方法避免了显式生成可能非常大的 Hessian 矩阵，仅需要向量积和梯度计算，对于大型模型来说内存效率更高。

注意事项

• 计算成本： 计算高阶梯度的开销比一阶梯度大。每次调用 torch.autograd.grad 并设置 create_graph=True，在后续反向传播 (backpropagation)中，图的遍历深度实质上会翻倍。
• 内存占用： 存储高阶导数所需的图会消耗更多内存。
• 二次反向传播： 计算二阶导数的过程有时被称为“二次反向传播”。

了解如何使用 torch.autograd.grad 和 create_graph=True 标志计算高阶梯度，实现了在优化、模型分析以及 PyTorch 框架内实现复杂算法（如元学习和物理信息建模）方面的一系列高级能力。