自定义TensorFlow损失函数用one-hot编码梯度异常极小问题

我实现了两种自定义交叉熵损失函数：一种通过tf.one_hot处理y_true，另一种使用tf.gather。两者输出值完全相同，但前者训练时产生极小梯度导致模型无法正常收敛，后者训练效果良好。希望了解该问题的通用原因，以及如何识别TensorFlow中可能引发此类问题的函数。

两种损失函数实现 #

基于tf.one_hot的版本 #

def custom_cross_entropy_loss(y_true, y_pred):
    return tf.reduce_mean(
        tf.reduce_logsumexp(y_pred, axis=-1) - tf.reduce_sum(y_pred * tf.one_hot(y_true, depth=10), axis=-1)
    )

基于tf.gather的版本 #

def custom_cross_entropy_loss(y_true, y_pred):
    return tf.reduce_mean(
        tf.reduce_logsumexp(y_pred, axis=-1) - tf.gather(y_pred, tf.cast(y_true, tf.int32), batch_dims=1)
    )

训练验证显示，两种损失函数计算值一致，但使用tf.one_hot的模型训练时损失几乎无变化、准确率极低，而使用tf.gather的模型训练正常。

问题通用原因 #

• 稀疏操作稀释梯度信号：tf.one_hot生成的是大部分元素为0的稀疏向量，和y_pred做元素乘法再求和时，反向传播需要对所有维度计算梯度，但绝大多数维度的梯度都是0，有效梯度被严重稀释，甚至因为浮点精度限制直接被截断为0，模型根本得不到有效的更新信号。而tf.gather直接提取对应索引的元素，梯度只会集中在被选中的关键位置，信号清晰且强度足够。
• 自动微分路径的效率差异：tf.one_hot的操作会把整数索引扩展成全维度的浮点向量，反向传播时的梯度链要经过所有维度的乘法和求和步骤，中间容易累积数值误差；tf.gather是直接的索引操作，梯度计算路径更短更直接，避免了冗余步骤带来的精度损失。
• 数值稳定性差距：当y_pred的数值波动较大时，one-hot乘法加求和的操作会引入更多浮点运算，容易触发梯度下溢（变成极小值甚至0），而tf.gather直接取元素，运算步骤少，数值稳定性更高。

如何识别TensorFlow中此类风险函数 #

• 生成稀疏张量的函数：比如tf.one_hot、tf.sparse.to_dense等，当它们生成的稀疏张量参与乘法、求和等后续操作时，要警惕梯度稀释或精度问题。
• 全维度扩展类函数：如果函数会把低维度输入（比如单索引）扩展成全维度张量，而实际只有少数维度有效，反向传播时会产生大量无效零梯度，干扰有效梯度传递。
• 涉及频繁类型转换的函数：比如把整数索引转成浮点张量的操作，类型转换过程可能引入精度损失，进而影响梯度计算的准确性。
• 手动验证梯度：对自定义损失或层，用tf.GradientTape手动计算梯度，对比不同实现的梯度值大小和分布。如果某一种实现的梯度均值远小于另一种，或者梯度里大量是0，那大概率存在问题。

内容的提问来源于stack exchange，提问作者samarendra chandan bindu Dash