TensorFlow中标量与矩阵逐元素乘积的导数计算结果异常问题

问题背景 #

x为标量（或形状为(1,1)的张量），B是5×5的二维张量，f是x与B的逐元素乘积（即f[i,j] = x*B[i,j]），需要计算f对x的导数，但得到的结果形状与预期不符。

测试环境：TensorFlow 2.17.0、2.12.0

复现代码 #

import tensorflow as tf
import numpy as np

with tf.GradientTape(persistent=True) as g:
    x = tf.Variable(np.random.uniform(size=(1,1)), dtype=tf.float64)
    B = tf.convert_to_tensor(np.random.normal(size=(5,5)), dtype=tf.float64)
    g.watch(x)
    g.watch(B)

    f = tf.multiply(B, x)

    gr = g.gradient(f, x)

结果对比 #

• 预期结果：gr.shape为(5, 5)，每个元素对应f[i,j]对x的导数（即B[i,j]）
• 实际结果：gr.shape为(1, 1)（若x是纯标量则为()），且值等于B所有元素的总和
• 验证：f.shape为(5, 5)，符合逐元素乘积的预期

原因分析 #

TensorFlow的GradientTape.gradient()方法默认会对输出张量的所有元素求和，再计算该总和对自变量的导数。对于本场景：

• sum(f) = sum(x*B[i,j]) = x * sum(B)
• 该总和对x的导数就是sum(B)，因此结果的形状与x一致（(1,1)或标量），而非与f一致。

解决方案 #

如果需要得到输出张量每个元素对x的导数（即形状与f匹配的梯度），需要使用GradientTape.jacobian()方法——它会计算输出张量每个元素对输入变量的导数，生成雅可比矩阵：

import tensorflow as tf
import numpy as np

with tf.GradientTape(persistent=True) as g:
    x = tf.Variable(np.random.uniform(size=(1,1)), dtype=tf.float64)
    B = tf.convert_to_tensor(np.random.normal(size=(5,5)), dtype=tf.float64)
    g.watch(x)
    g.watch(B)

    f = tf.multiply(B, x)

    # 计算雅可比矩阵，默认形状为(5,5,1,1)
    gr = g.jacobian(f, x)
    # 压缩多余维度，得到(5,5)的结果
    gr = tf.squeeze(gr, axis=[2, 3])

此时gr.shape为(5,5)，每个元素gr[i,j]等于B[i,j]，完全符合数学推导的结果。

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