难以理解 tf.custom_gradient() 的高阶导数

Question

基于 tensorflow 网站中引用的示例：https://www.tensorflow.org/api_docs/python/tf/custom_gradient

@tf.custom_gradient
def op_with_fused_backprop(x):
     y, x_grad = fused_op(x)

     def first_order_gradient(dy):
         @tf.custom_gradient
         def first_order_custom(unused_x):
             def second_order_and_transpose(ddy):
                 return second_order_for_x(...), gradient_wrt_dy(...)
             return x_grad, second_order_and_transpose
         return dy * first_order_custom(x)
     return y, first_order_gradient

没有关于second_order_and_transpose(ddy) 返回两个对象的详细信息。根据 tf.custom_gradient 的文档，grad_fn（即second_order_and_transpose()）应该返回一个张量列表，这些张量是 dy w.r.t 的导数。 unused_x。甚至不清楚他们为什么将其命名为unused_x。有人对这个示例有任何想法，或者通常为高阶导数创建自定义渐变吗？

Answer 1

1. 没有关于为什么 second_order_and_transpose(ddy) 返回两个对象的详细信息。

根据我使用的一些示例，我相信您是正确的。官方文档在某种程度上模棱两可（或不正确）。 second_order_and_transpose(ddy) 应该只返回一个对象，即计算出的二阶梯度。

1. 也不清楚他们为什么将其命名为used_x。

这是棘手的部分。 unused_x 解释了他们为什么命名它（因为你永远不会使用它......）。这里的目标是将您的二阶计算函数包装在一个名为first_order_custom 的函数中。您从fused_op 计算 x 的梯度，并将其用作返回值，而不是 unused_x。

为了更清楚地说明这一点，我通过了一个从官方文档扩展的示例来定义log1pexp的二阶梯度：

注意：二阶梯度在数值上是不稳定的，所以我们用 (1 - tf.exp(x)) 来表示它，只是为了方便我们的生活。

@tf.custom_gradient
def log1pexp2(x):
    e = tf.exp(x)
    y = tf.math.log(1 + e)
    x_grad = 1 - 1 / (1 + e)
    def first_order_gradient(dy):
        @tf.custom_gradient
        def first_order_custom(unused_x):
            def second_order_gradient(ddy):
                # Let's define the second-order graidne to be (1 - e)
                return ddy * (1 - e) 
            return x_grad, second_order_gradient
        return dy * first_order_custom(x)
    return y, first_order_gradient

要测试脚本，只需运行：

import tensorflow as tf

@tf.custom_gradient
def log1pexp2(x):
    e = tf.exp(x)
    y = tf.math.log(1 + e)
    x_grad = 1 - 1 / (1 + e)
    def first_order_gradient(dy):
        @tf.custom_gradient
        def first_order_custom(unused_x):
            def second_order_gradient(ddy):
                # Let's define the second-order graidne to be (1 - e)
                return ddy * (1 - e) 
            return x_grad, second_order_gradient
        return dy * first_order_custom(x)
    return y, first_order_gradient

x1 = tf.constant(1.)
y1 = log1pexp2(x1)
dy1 = tf.gradients(y1, x1)
ddy1 = tf.gradients(dy1, x1)

x2 = tf.constant(100.)
y2 = log1pexp2(x2)
dy2 = tf.gradients(y2, x2)
ddy2 = tf.gradients(dy2, x2)

with tf.Session() as sess:
    print('x=1, dy1:', dy1[0].eval(session=sess))
    print('x=1, ddy1:', ddy1[0].eval(session=sess))
    print('x=100, dy2:', dy2[0].eval(session=sess))
    print('x=100, ddy2:', ddy2[0].eval(session=sess))

结果：

x=1, dy1: 0.7310586
x=1, ddy1: -1.7182817
x=100, dy2: 1.0
x=100, ddy2: -inf