使用自定义softplus激活函数时Keras训练出现NaN损失是什么原因？

• 核心原因是你自定义的softplus实现没有做数值稳定性优化，存在上溢出风险：
  你当前的实现逻辑为log(exp(z)+1)，当z为较大的正值时（float32精度下z>88时，exp(z)就会超过float32的最大值上限），tf.exp(z)会直接计算为无穷大（inf），后续log(inf + 1)也会输出inf，前向传播就会产生异常值，反向传播计算梯度时自然会得到NaN。
  你小范围测试时输入的数值都在[-3,2]的小区间内，exp计算不会溢出，所以结果和官方实现接近，不会暴露问题。而模型训练过程中，全连接层的输出值会随着权重更新动态变化，很容易出现较大的正值，触发溢出问题。
• TensorFlow官方的softplus激活函数内置了数值稳定优化，等价实现逻辑为：当z>0时返回z + log(1 + exp(-z))，当z<=0时返回log(1 + exp(z))，全程不会触发exp的上溢出，因此训练过程完全正常。

你可以修改自定义softplus为数值稳定版本，逻辑和官方实现对齐即可解决NaN问题：

def my_softplus(z):
    z = tf.cast(z, tf.float32)
    return tf.math.maximum(z, 0) + tf.math.log(1 + tf.math.exp(-tf.math.abs(z)))

内容的提问来源于stack exchange，提问作者An old man in the sea.