我正在尝试为我在论文中发现的分段连续函数构建自定义渐变

对数，线性和指数函数之间的连续体，及其在改善神经网络泛化性方面的潜力-Godfrey

当前问题：我的网络的自定义层中可学习参数的梯度爆炸，导致生成nans。该自定义层是上面论文中定义的可学习的激活功能。

尝试过的解决方案

• 剪切优化器的渐变->当我使用keras.optimizer.Adam(clipnorm=0.5)编译模型时，这没有影响
• 在自定义层中剪切可学习参数的值->这产生了一些我无法回忆的错误
• 创建自定义梯度计算以控制可学习参数的梯度->不确定应如何计算梯度。

我当前的梯度计算功能就是这样

@tf.custom_gradient
def call_lt0(alpha_actv,x):
    def grad(dx):
        grad_x = 1/(1-alpha_actv * (alpha_actv + x))
        grad_alpha = (tf.math.log(1-alpha_actv * (alpha_actv + x)) - ((2 * alpha_actv ** 2 + alpha_actv * x)/(alpha_actv ** 2 + alpha_actv * x - 1))) / alpha_actv ** 2
        return grad_x * dx,grad_alpha * dx  # denk ik..?
    return (tf.math.exp(alpha_actv * x) - 1) / alpha_actv + alpha_actv,grad

我给出x变量和可学习参数alpha的梯度。但是，当我用tf.GradientTape检查这些输出时，由于我对grad的定义，我只能得到一个输出..我期望两个输出。 另外，当我让keras / tensorflow计算所有东西，并再次用tf.GradientTape检查输出时，我只得到一个值..这是我对单个可学习参数的期望。 所以在某处我做错了事，或者误解了grad函数。

有人可以指点我吗？