tensorflow - 用 tensorflow custom_gradient 定义 jacobian

Question

我正在尝试定义一个 custom gradient in TensorFlow对于返回非标量的函数。

问题的简化玩具版本是:

import tensorflow as tf

@tf.custom_gradient
def f(x):
    y = tf.Variable([x[1], 2.0*x[0]])

    def grad(upstream):
        dg = tf.Variable([[0.0, 2.0], [1.0, 0.0]])
        return dg * upstream # ? no idea here

    return y, grad

x = tf.Variable([1.0, 2.0])
with tf.GradientTape() as tape:
    y = f(x)
tape.jacobian(y, x)
# tape.gradient(y,x)

我不明白 grad 函数需要返回什么形状。在上面的示例中，tape.gradient 给出了问题的完整雅可比矩阵(我希望从 tape.jacobian 得到)和 tape.jacobian什么也没给出(尽管我在重新启动内核之前一直收到错误消息)。

这是预期的形状和行为吗？对于向量值或矩阵值函数的导数，自定义梯度应返回什么形状？

Answer 1

我同意这可能会造成混淆。让我尝试扩展 Tensorflow docs 中给出的声明。并在数学和代码之间建立明确的联系。

让我们考虑一个函数

.

它的雅可比矩阵是 n x m 矩阵:

.

给定一个输入向量 x 和一个“上游”向量 u，函数 grad_fn 在 @custom_gradient 函数应返回“向量-雅可比积”vjp，即:

.

请注意，在此公式中，u 是给定的，不依赖于 x。因此，我们可以将vjp改写如下:

.

如果我们定义 ，这可以推广到更高阶的张量作为 a 和 b 的元素乘积之和(即 tf.reduce_sum(a * b))，这是常规点积对于向量和 Frobenius inner product对于矩阵。然后，grad_fn 的输出应该是一个与 x 形状相同的张量，计算如下:

.

请注意，这实际上是标量值函数 tf.reduce_sum(u * f(x)) 的梯度(因此更容易思考和使用)。

---

在您的特定情况下，您需要采用 u[0]*x[1]+u[1]*2.0*x[0] 相对于 x 的梯度[0],x[1]，即类似于:

def grad(upstream):
    return tf.Variable([2.0*upstream[1], upstream[0]])

这会产生所需的梯度。

不幸的是 jacobian 不能用这个(因为它不能用你的原始代码)因为你的代码有第二个问题:通过使用变量它打破了输出和输入之间的链接(创建变量就像复制/粘贴值一样，不会传播梯度)。您可以按如下方式解决此问题:

@tf.custom_gradient
def f(x):
    y = tf.stack([x[1], 2.0*x[0]])

    def grad(upstream):
        return tf.stack([2.0*upstream[1], upstream[0]])

    return y, grad