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

Question

我想使用张量流以数值稳定的方式计算比率let sourceImage = UIImage(named: "TestImage") let image = AWSRekognitionImage() image!.bytes = UIImageJPEGRepresentation(sourceImage!, 0.7) guard let request = AWSRekognitionDetectLabelsRequest() else { puts("Unable to initialize AWSRekognitionDetectLabelsRequest.") return } request.image = image request.maxLabels = 3 request.minConfidence = 90 的导数，但是当f = - a / b和a很小（b时遇到问题使用32位浮点表示）。当然，<1e-20的导数是f，但由于运算符优先级，分母中的平方首先计算，下溢，并导致未定义的渐变。

如果导数计算为df_db = a / b ** 2，则不会发生下溢，并且梯度将被明确定义，如下图所示，其中梯度显示为df_db = (a / b) / b的函数。蓝线对应于张量流可以计算导数的域。橙色线对应于分母下溢产生无限梯度的域。绿线对应于分母溢出产生零梯度的域。在两个有问题的域中，可以使用上面修改的表达式计算梯度。

通过使用丑陋的黑客，我能够获得更加数字稳定的表达

a = b

等同于g = exp(log(a) - log(b)) 但产生不同的张量流图。但如果我想计算一个更高阶的导数，我会遇到同样的问题。可以找到重现问题的代码here。

是否有建议的方法来缓解此类问题？如果一个人不想依赖自动分化，是否有可能在张量流中明确定义表达式的导数？

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)

完整笔记本是here。 PR是here。