Tensorflow GradientTape 在使用 MultivariateNormalTriL 分布时返回 NaN

Question

我正在使用网络输出实例化来自 tensorflow_probability 和 GradientTape 的 MultivariateNormalTriL 分布，以在计算损失后记录梯度。

经过随机数量的训练步骤后，GradientTape 返回网络某些层的 NaN 值。网络输出和分布中的所有值看起来都不错。它们是一些随机数。不小，也不大。所有计算的损失值也可以。除了 GradienTape 的渐变之外，其他任何地方都没有 NaN。

此外，使用正态分布时一切正常。

据我了解，GradientTape 只应在函数不可微分时返回 NaN。所以事实证明，MultivariateNormalTriL 对于某些特定值是不可微的。

我错过了什么吗？你知道如何解决这个问题或者至少去哪里看吗？

Answer 1

似乎问题出在标准差矩阵上。我使用tfp.bijectors.FillScaleTriL 来转换矩阵，它现在似乎正在工作。