TensorFlow Probability负二项回归搭建出现loss为nan的问题求解

损失NaN核心原因 #

• total_count参数合法性：负二项的试验次数参数必须为正数，直接传递全连接层原始输出会出现负值，触发计算异常
• probs计算溢出风险：probs参数要求取值范围严格在(0,1)之间，当前实现用1e-3 + tf.math.softplus(0.01 * t[...,1:])计算，softplus本身输出无上界，很容易出现结果大于1的情况，导致概率计算非法出现NaN
• 损失函数不匹配：如果没有显式指定负对数似然作为损失函数，用默认的损失计算逻辑也会出现数值异常

可运行实现模板 #

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

model = tf.keras.Sequential([
    # 输出2个参数，对应负二项分布的两个待估参数
    tf.keras.layers.Dense(2, activation=None),
    tfp.layers.DistributionLambda(
        lambda t: tfd.NegativeBinomial(
            # 对total_count做softplus变换保证输出恒为正，加极小值避免出现0
            total_count=tf.math.softplus(t[..., :1]) + 1e-6,
            # 直接使用logits参数，无需手动计算probs，大幅提升数值稳定性
            logits=t[..., 1:]
        )
    ),
])

# 显式定义负对数似然作为损失函数
def nll(y_true, y_pred):
    return -y_pred.log_prob(y_true)

# 编译时使用较小的学习率，避免梯度爆炸
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=3e-4), loss=nll)

训练注意事项 #

• 标签Y必须是非负整数，符合负二项分布的取值要求
• 如果训练初期仍有轻微数值波动，可以对损失值做裁剪：return tf.clip_by_value(-y_pred.log_prob(y_true), -1e6, 1e6)
• 如果数据离散程度很高，可以在全连接层和DistributionLambda之间添加激活层或Dropout层提升泛化能力

内容的提问来源于stack exchange，提问作者sports-modelling