TensorFlow中多变量场景下计算Hessian矩阵的问题

这问题我之前也碰到过，确实有点坑！先给你理清楚为什么会出现这种情况，再给你两种靠谱的解决办法。

问题出在哪？ #

1. 直接传入变量列表的问题：当你调用tf.hessians(f, [x, y])时，TensorFlow会分别计算损失函数f对x的二阶导（2x2矩阵）和对y的二阶导（1x1矩阵），返回的是两个独立的矩阵，而不是把x和y当作一个整体的完整3x3 Hessian矩阵。
2. 拼接张量报错的原因：tf.concat([x, y], axis=-1)得到的是一个普通张量，不是可训练的tf.Variable，而tf.hessians要求目标必须是可训练变量（或者能追溯到变量的可导张量），所以会触发ValueError。

解决方案一：手动组合分块Hessian矩阵 #

既然TensorFlow会返回各个变量的二阶导，我们还可以手动计算交叉项，然后把这些分块拼接成完整的Hessian矩阵：

import tensorflow as tf

x = tf.Variable([1., 1.], dtype=tf.float32, name="x")
y = tf.Variable([1.], dtype=tf.float32, name="y")
f = (x[0] + x[1] ** 2 + x[0] * x[1] + y) ** 2

# 计算各分块矩阵
H_xx = tf.hessians(f, x)[0]  # f对x的二阶导（2x2）
H_yy = tf.hessians(f, y)[0]  # f对y的二阶导（1x1）

# 计算交叉二阶导：先求f对x的一阶导，再对y求导
grad_f_x = tf.gradients(f, x)[0]
H_xy = tf.gradients(grad_f_x, y)[0]  # 形状是(2,)，需要转置成(2,1)

# 计算y对x的交叉二阶导（也可以直接转置H_xy，因为Hessian是对称矩阵）
grad_f_y = tf.gradients(f, y)[0]
H_yx = tf.gradients(grad_f_y, x)[0]  # 形状是(1,2)

# 拼接成完整的3x3 Hessian矩阵
top_row = tf.concat([H_xx, tf.expand_dims(H_xy, axis=1)], axis=1)
bottom_row = tf.concat([tf.expand_dims(H_yx, axis=0), H_yy], axis=1)
full_hessian = tf.concat([top_row, bottom_row], axis=0)

# 运行验证结果
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(full_hessian))

运行后会输出你期望的完整Hessian矩阵：

[[ 8. 20.  4.]
 [20. 34.  6.]
 [ 4.  6.  2.]]

解决方案二：用GradientTape自动计算（推荐，适合TensorFlow 2.x） #

如果你用的是TensorFlow 2.x，推荐用tf.GradientTape配合jacobian方法，它能自动处理所有交叉项，不需要手动拼接分块，扩展性更强：

import tensorflow as tf

x = tf.Variable([1., 1.], dtype=tf.float32, name="x")
y = tf.Variable([1.], dtype=tf.float32, name="y")

# 把变量合并成一个列表，方便统一处理
variables = [x, y]

with tf.GradientTape() as tape2:
    with tf.GradientTape() as tape1:
        # 确保tape追踪所有变量的梯度
        tape1.watch(variables)
        f_val = (x[0] + x[1]**2 + x[0]*x[1] + y)**2
    # 计算一阶导数，拼接成一个梯度向量
    grads = tape1.gradient(f_val, variables)
    grad_vec = tf.concat(grads, axis=0)
# 对梯度向量求jacobian，得到完整的Hessian矩阵
hessian_blocks = tape2.jacobian(grad_vec, variables)

# 把分块拼接成3x3矩阵
full_hessian = tf.concat([
    tf.concat(hessian_blocks[:len(variables)], axis=1),
], axis=0)

print(full_hessian.numpy())

这种方法当你有更多变量拆分时，只需要把变量加入variables列表即可，不需要修改其他代码，非常方便。

内容的提问来源于stack exchange，提问作者Ruggero Turra