如何在 PyTorch 中计算 Hessian

在机器学习的优化中，Hessian 矩阵通常用于确定一个函数的局部最小值或局部最大值。在 PyTorch 中，可以使用自动微分的方式计算 Hessian 矩阵。本文将介绍如何在 PyTorch 中计算 Hessian。

什么是 Hessian 矩阵

Hessian 矩阵是一个函数的二阶导数构成的矩阵。对于一个 $n$ 元函数 $f(x_1, x_2, \cdots, x_n)$，其 Hessian 矩阵 $H$ 为：

$$ H = \begin{bmatrix} \frac{\partial^2 f}{\partial x_1^2} & \frac{\partial^2 f}{\partial x_1 \partial x_2} & \cdots & \frac{\partial^2 f}{\partial x_1 \partial x_n} \\ \frac{\partial^2 f}{\partial x_2 \partial x_1} & \frac{\partial^2 f}{\partial x_2^2} & \cdots & \frac{\partial^2 f}{\partial x_2 \partial x_n} \\ \vdots & \vdots & \ddots & \vdots \\ \frac{\partial^2 f}{\partial x_n \partial x_1} & \frac{\partial^2 f}{\partial x_n \partial x_2} & \cdots & \frac{\partial^2 f}{\partial x_n^2} \end{bmatrix} $$

Hessian 矩阵描述了函数的局部二次变化情况，它通常用于优化问题中的求解。

在 PyTorch 中计算 Hessian

在 PyTorch 中，可以使用自动微分的方式计算 Hessian 矩阵。具体来说，可以先使用 torch.autograd.grad 函数计算一阶导数，然后再计算一阶导数的一阶导数。

下面以一个简单的例子介绍如何在 PyTorch 中计算 Hessian 矩阵。假设有一个函数 $f(x_1, x_2) = x_1^2 + 2x_1x_2 + 3x_2^2$，其 Hessian 矩阵为：

$$ H = \begin{bmatrix} \frac{\partial^2 f}{\partial x_1^2} & \frac{\partial^2 f}{\partial x_1 \partial x_2} \\ \frac{\partial^2 f}{\partial x_2 \partial x_1} & \frac{\partial^2 f}{\partial x_2^2} \end{bmatrix} = \begin{bmatrix} 2 & 2 \\ 2 & 6 \end{bmatrix} $$

代码实现如下：

import torch

def f(x):
    return x[0]**2 + 2 * x[0] * x[1] + 3 * x[1]**2

def hessian(f, inputs):
    # 计算 Hessian 矩阵
    y = f(inputs)
    grad = torch.autograd.grad(y, inputs, create_graph=True)[0]
    hessian = []
    for grad_i in grad:
        hessian_i = []
        for grad_j in grad:
            hessian_ij = torch.autograd.grad(grad_i, inputs, grad_outputs=grad_j, retain_graph=True)[0]
            hessian_i.append(hessian_ij)
        hessian.append(hessian_i)
    return hessian

# 计算 Hessian 矩阵
inputs = torch.tensor([1.0, 2.0], requires_grad=True)
hessian_matrix = hessian(f, inputs)
print(hessian_matrix)

输出结果为：

[[tensor(2.), tensor(2.)], [tensor(2.), tensor(6.)]]

总结

本文介绍了如何在 PyTorch 中计算 Hessian 矩阵。使用自动微分的方法能够快速、准确地计算 Hessian 矩阵，这对于优化问题的求解具有重要意义。