怎样在 Eager 模式下利用 tf.GradientTape 实现复杂的二阶导数与海森矩阵计算

在深度学习优化和不确定性估计中，二阶导数（曲率信息）扮演着至关重要的角色，尤其是在牛顿法或拟牛顿法（如BFGS）中。TensorFlow 2.x 的 Eager 模式提供了灵活的自动微分机制 tf.GradientTape。虽然它主要用于计算一阶导数，但通过嵌套（Nesting）使用 tf.GradientTape，我们可以轻松实现二阶导数乃至海森矩阵（Hessian Matrix）的计算。

1. 原理：嵌套的 GradientTape

计算二阶导数的本质是计算“导数的导数”。

1. 内层 Tape (Inner Tape): 计算函数 $f(x)$ 的一阶导数 $f'(x)$。
2. 外层 Tape (Outer Tape): 计算 $f'(x)$ 关于 $x$ 的导数，即 $f”(x)$。

由于我们需要在内层导数计算完成后，再次对其结果进行求导，因此内层和外层 tf.GradientTape 都必须设置为 persistent=True，或者确保外部 Tape 能够追踪内层计算的结果。

2. 实战：计算海森矩阵（Hessian Matrix）

海森矩阵 $H$ 是损失函数 $L(oldsymbol{w})$ 关于变量向量 $oldsymbol{w}$ 的二阶偏导数矩阵。它实际上是梯度向量 $
abla L(oldsymbol{w})$ 关于 $oldsymbol{w}$ 的雅可比矩阵。

$$\boldsymbol{H} = \frac{\partial^2 L}{\partial \boldsymbol{w}^2} = \frac{\partial}{\partial \boldsymbol{w}} \left( \frac{\partial L}{\partial \boldsymbol{w}} \right)$$

代码示例

我们以一个简单的多元函数作为损失函数 $L(oldsymbol{w}) = w_0^2 + \sin(w_1)$，并计算其在特定点 $oldsymbol{w}=[2.0, 3.0]$ 处的海森矩阵。

import tensorflow as tf
import numpy as np

# 确保使用Eager模式
tf.executing_eagerly()

# 1. 定义变量向量
w = tf.Variable([2.0, 3.0], dtype=tf.float32)

# 2. 定义损失函数
def compute_loss(w):
    # L(w) = w[0]^2 + sin(w[1])
    return w[0]**2 + tf.sin(w[1])

print(f"--- 开始计算 ---")
print(f"变量 w 初始值: {w.numpy()}")

# 3. 嵌套 Tape 计算海森矩阵
# 外层 Tape 用于计算二阶导数 (梯度向量的导数)
with tf.GradientTape(persistent=True) as outer_tape:
    outer_tape.watch(w) 

    # 内层 Tape 用于计算一阶导数 (损失的导数)
    with tf.GradientTape(persistent=True) as inner_tape:
        inner_tape.watch(w)
        loss = compute_loss(w)

    # 第一次求导：计算梯度向量 (dL/dw)
    # 结果是一个向量 [g0, g1]
    grad = inner_tape.gradient(loss, w)

# 第二次求导：计算梯度向量 (grad) 关于变量 w 的雅可比矩阵，即海森矩阵
# 使用 outer_tape.jacobian(target=grad, sources=w) 
# 结果是一个矩阵 H = d(grad)/d(w)
Hessian = outer_tape.jacobian(grad, w)

# 4. 打印结果
print("\n--- 计算结果 ---")
print(f"损失 L(w): {loss.numpy():.4f}")
print(f"梯度向量 grad (dL/dw):
{grad.numpy()}")
print(f"海森矩阵 Hessian (d^2L/dw^2):
{Hessian.numpy().round(4)}")

# 清理 Tape 资源
del inner_tape
del outer_tape

# 理论验证 (w=[2.0, 3.0]):
# L = w0^2 + sin(w1)
# H = [[2, 0], [0, -sin(w1)]]
# H(2, 3) = [[2, 0], [0, -sin(3)]]
# -sin(3) ≈ -0.1411

结果分析

运行上述代码，输出的海森矩阵应为：

变量 w 初始值: [2. 3.]
...
海森矩阵 Hessian (d^2L/dw^2):
[[ 2.      0.    ]
 [ 0.     -0.1411]]

结果与理论计算完全一致，证明了通过嵌套 tf.GradientTape 和使用 tape.jacobian() 方法，可以在 Eager 模式下高效准确地计算复杂的二阶导数和海森矩阵。

3. 注意事项

1. 追踪变量 (watch): 确保所有需要求导的变量都被 watch 或声明为 tf.Variable，并且被相关的 Tape 追踪。
2. 持久化 (persistent=True): 当我们需要在一个 Tape 范围内多次调用 tape.gradient() 或像这里一样，将内层 Tape 的输出作为外层 Tape 的输入时，必须设置 persistent=True。
3. Jacobian vs. Gradient: 计算海森矩阵时，外层 Tape 的输出是向量（梯度）对向量（变量）的导数，因此我们必须使用 outer_tape.jacobian(grad, w)，而不是 outer_tape.gradient(grad, w)（这将导致错误或只计算雅可比向量积）。