使用tf.math.pow(x, 0.5)作为激活函数时训练损失为NaN问题咨询

• 边界处的梯度计算差异
  两个函数的正向计算逻辑近似，但反向传播的梯度处理逻辑完全不同：tf.math.sqrt是针对平方根场景做了边界优化的专用算子，输入x趋近于0时会隐性截断梯度，避免出现无穷大；而tf.math.pow(x, 0.5)是通用幂运算，严格按照幂函数求导公式计算梯度：0.5 * x^(-0.5)，当x等于0时梯度直接为无穷大，会触发梯度爆炸，参数更新后直接出现异常值，最终损失变为NaN。
• tf.where的隐式分支计算逻辑
  TensorFlow计算tf.where的梯度时，会同时计算两个分支的前向值和梯度，再根据mask筛选对应位置的梯度回传。如果输入x存在极小的负值，哪怕cond判断为False会走负分支，正分支的tf.math.pow(x, 0.5)仍然会先对负x做0.5次幂运算，直接得到NaN，该NaN会被带入后续梯度计算，扩散到整个计算图。
• 运算实现的数值稳定性差异
  通用幂运算tf.math.pow的内部实现需要做对数、乘法、指数三次转换：x^0.5 = exp(0.5 * ln(x))，当x非常接近0时，ln(x)趋近于负无穷，中间步骤的数值精度损失极大，很容易出现下溢、溢出或者异常值；而tf.math.sqrt直接调用硬件优化的原生平方根指令，没有中间转换步骤，数值稳定性远高于通用幂运算。

如果一定要用幂运算实现该激活函数，可以在分支中加入极小的epsilon避免边界问题：

def pwsqrt(x, epsilon=1e-12):
    cond = tf.greater_equal(x, 0)
    return tf.where(
        cond, 
        tf.math.pow(tf.maximum(x, epsilon), 0.5), 
        -tf.math.pow(tf.maximum(-x, epsilon), 0.5)
    )

内容的提问来源于stack exchange，提问作者ag2718