TensorFlow批量计算模型权重Hessian矩阵的技术求助

问题背景 #

复现论文时，需基于Keras CNN完成MNIST分类任务，核心需求为：针对每个训练样本计算模型参数关于该样本损失的Hessian矩阵，在全量训练数据上求平均后计算影响分数。当前仅能逐个样本计算Hessian，速度过慢，尝试两种批量方案均存在问题：

尝试方案1：嵌套GradientTape计算梯度Jacobian #

x=tf.convert_to_tensor(x_train[0:13])
with tf.GradientTape() as t2:
    with tf.GradientTape() as t1:
      y=model(x)
      mce = tf.keras.losses.CategoricalCrossentropy()
      y_expanded=y_train[train_idx]
      loss=mce(y_expanded,y)
    g = t1.gradient(loss, model.weights[4])
h = t2.jacobian(g, model.weights[4])
print(h.shape)

• 问题：目标层参数维度为20×30，输入13样本批次后，期望输出(13,20,30,20,30)的样本级Hessian矩阵，但实际输出(20,30,20,30)，丢失了样本维度，无法实现批量向量化计算。

尝试方案2：保留样本级损失的Jacobian计算 #

x=tf.convert_to_tensor(x_train[0:13])
mce = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.Reduction.NONE)

with tf.GradientTape() as t2:
    with tf.GradientTape() as t1:
      t1.watch(model.weights[4])
      y_expanded=y_train[0:13]
      y=model(x)
      loss=mce(y_expanded,y)
    j1=t1.jacobian(loss, model.weights[4])
j3 = t2.jacobian(j1, model.weights[4])

print(j3.shape)

• 问题：返回维度符合预期，但批量计算的Hessian求和结果与循环计算单个样本Hessian的求和结果不一致，无法保证正确性。

正确的批量Hessian计算方案 #

核心是确保每个样本的损失独立，二阶导数严格对应单个样本的损失对参数的二阶偏导：

import tensorflow as tf

# 加载批量训练数据
batch_x = tf.convert_to_tensor(x_train[:13])
batch_y = tf.convert_to_tensor(y_train[:13])
target_weights = model.weights[4]  # 指定需要计算Hessian的目标层参数

# 定义损失函数，保留每个样本的独立损失
loss_fn = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.Reduction.NONE)

# 启用persistent=True，支持多次调用梯度计算
with tf.GradientTape(persistent=True) as tape_second:
    tape_second.watch(target_weights)
    with tf.GradientTape(persistent=True) as tape_first:
        tape_first.watch(target_weights)
        y_pred = model(batch_x)
        sample_losses = loss_fn(batch_y, y_pred)  # 形状：(batch_size,)
    
    # 逐个计算每个样本损失对参数的一阶梯度，堆叠后保留样本维度
    sample_grads = tf.stack([tape_first.gradient(loss, target_weights) for loss in sample_losses], axis=0)

# 逐个计算每个样本一阶梯度对参数的Jacobian，即样本级Hessian
sample_hessians = tf.stack([tape_second.jacobian(grad, target_weights) for grad in sample_grads], axis=0)

# 释放persistent tape资源
del tape_first, tape_second

# 验证：批量求和结果与循环计算单个样本的结果一致
batch_hessian_sum = tf.reduce_sum(sample_hessians, axis=0)

关键细节

1. Reduction.NONE的必要性：必须保留每个样本的独立损失值，不能使用默认的均值或求和模式，否则会丢失样本维度导致批量计算失效。
2. Persistent GradientTape：由于需要对多个样本的损失/梯度分别计算导数，必须设置persistent=True，使用完毕后手动删除tape释放显存。
3. 显式堆叠样本梯度：通过循环单个样本损失计算梯度并堆叠，确保每个梯度严格对应单个样本，避免批量计算时的梯度混淆。

显存优化方案：向量积技巧（Vector Product Trick） #

当参数维度较大时，直接计算完整Hessian会占用大量显存。若仅需Hessian的平均或与其他向量的乘积，可使用向量积技巧间接计算，大幅降低显存占用：

# 生成随机向量v，用于间接计算Hessian
v = tf.random.normal(shape=target_weights.shape)

with tf.GradientTape() as tape_second:
    with tf.GradientTape() as tape_first:
        y_pred = model(batch_x)
        sample_losses = loss_fn(batch_y, y_pred)
        # 计算每个样本梯度与v的点积
        grad_v_dot = tf.einsum('ij,ij->i', tape_first.gradient(sample_losses, target_weights), tf.broadcast_to(v, (13,) + v.shape))
    # 点积对参数的梯度等价于Hessian与v的乘积
    hessian_v_product = tape_second.gradient(grad_v_dot, target_weights)

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