如何以数值稳定的方式计算梯度

Question

我想使用tensorflow以数值稳定的方式计算比率f = - a / b的导数，但当a和b很小时（使用32位浮点表示时小于<1e-20），会遇到问题。当然，f的导数是df_db = a / b ** 2，但由于运算符优先级，分母中的平方首先被计算，导致下溢并导致梯度未定义。

如果导数计算为df_db = (a / b) / b，则不会发生下溢，梯度将被定义，如下图所示，它显示了梯度作为a = b函数的情况。蓝线对应于tensorflow可以计算导数的域。橙色线对应于分母下溢产生无限梯度的域。绿线对应于分母上溢产生零梯度的域。在两个有问题的域中，可以使用上述修改后的表达式计算梯度。

我通过使用丑陋的技巧，得到了一个更加数值稳定的表达式。

g = exp(log(a) - log(b))

这相当于f，但会产生不同的tensorflow图。但如果我想计算更高阶导数，我会遇到同样的问题。可以在此处找到重现问题的代码。

有推荐的方法来缓解这些问题吗？如果不想依赖于自动微分，是否可能明确定义表达式在tensorflow中的导数？

- Till Hoffmann

---

也许这是一个愚蠢的问题，但你不能把所有的值乘以1000或更多来摆脱这个小范围吗？ - fafl

如果您想为各种操作插入自己的反向 AD 运算符实现，可以使用 gradient_override_map。 - Yaroslav Bulatov

你也可以通过将表达式包装到 TensorFlow 函数中来定义表达式的梯度，一个例子在 function_test.py 中。 - Yaroslav Bulatov

太好了，谢谢你的指引。我会去看看的。 - Till Hoffmann

Answer 1

感谢 Yaroslav Bulatov 的指针，我能够实现一个带有所需渐变的自定义函数。

# Define the division function and its gradient
@function.Defun(tf.float32, tf.float32, tf.float32)
def newDivGrad(x, y, grad):
    return tf.reciprocal(y) * grad, - tf.div(tf.div(x, y), y) * grad


@function.Defun(tf.float32, tf.float32, grad_func=newDivGrad)
def newDiv(x, y):
    return tf.div(x, y)

完整的笔记本在此链接中，有关PR的信息在这里。

- Till Hoffmann