使用 TF2 最小化 KL 散度

如何解决使用 TF2 最小化 KL 散度

我正在尝试更加熟悉 TensorFlow 2，因此我正在尝试一些练习。我正在做一个以最小化两个高斯之间的 KL 散度并更新一个相对于另一个的参数。

下面是运行良好但梯度计算为 NaN 的均值和方差的代码，即使 KL_divergence 具有有限值：

kl_divergence:  1947.9133
grads: [<tf.Tensor: shape=(1,),dtype=float32,numpy=array([nan],dtype=float32)>,<tf.Tensor: shape=(1,dtype=float32)>]

你知道这个错误可能来自哪里吗？我虽然 tf.where 无法传播渐变，但看起来它实际上可以。

import numpy as np
from scipy.special import kl_div
from scipy.stats import norm
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow.keras as Keras

x = np.arange(-10,10,0.001)
p = norm.pdf(x,5)

mean = tf.Variable(initial_value=tf.zeros((1,)),trainable=True,dtype=tf.float32,name='mean')
var = tf.Variable(initial_value=tf.ones((1,name='var')
optimizer = Keras.optimizers.SGD(0.1)
EPS = 1e-5

for _ in range(50):
    with tf.GradientTape() as tape:
        normal = tf.exp(-tf.square(x - mean / (2 * var)))
        q = normal / tf.reduce_sum(normal)
        kl_divergence = tf.reduce_sum(
            tf.where(q < EPS,tf.zeros(p.shape,tf.float32),p * tf.math.log(p / q))
        )

    print('kl_divergence: ',kl_divergence.numpy())
    grads = tape.gradient(kl_divergence,[mean,var])
    print(grads)
    # optimizer.minimize(kl_divergence,var],tape=tape)
    # print('mean: ',mean.numpy())
    # print('var: ',var.numpy())