TensorFlow中是否有内置的KL散度损失函数？

Question

我有两个张量，prob_a和prob_b，形状为[None, 1000]，我想计算从prob_a到prob_b的KL散度。TensorFlow中是否有内置函数可以完成这个功能？我尝试使用tf.contrib.distributions.kl(prob_a, prob_b)，但是它会返回如下错误信息：

NotImplementedError: No KL(dist_a || dist_b) registered for dist_a type Tensor and dist_b type Tensor

如果没有内置函数可用，有什么好的解决方法吗？

- Transcendental

Answer 1

假设您的输入张量 prob_a 和 prob_b 是概率张量，沿最后一个轴上求和为1，您可以这样做：

def kl(x, y):
    X = tf.distributions.Categorical(probs=x)
    Y = tf.distributions.Categorical(probs=y)
    return tf.distributions.kl_divergence(X, Y)

result = kl(prob_a, prob_b)

一个简单的例子：

import numpy as np
import tensorflow as tf
a = np.array([[0.25, 0.1, 0.65], [0.8, 0.15, 0.05]])
b = np.array([[0.7, 0.2, 0.1], [0.15, 0.8, 0.05]])
sess = tf.Session()
print(kl(a, b).eval(session=sess))  # [0.88995184 1.08808468]

您可以使用同样的方法得到相同的结果：

np.sum(a * np.log(a / b), axis=1) 

然而，这个实现有一些小bug（在Tensorflow 1.8.0中已验证）。如果在a数组中有零概率，例如尝试[0.8、0.2、0.0]而不是[0.8、0.15、0.05]，即使按照Kullback-Leibler定义0×log(0/b)应该为零，你将得到nan。为了缓解这种情况，应该添加一些小的数字常数。此外，最好使用tf.distributions.kl_divergence(X, Y, allow_nan_stats=False)在这种情况下导致运行时错误。另外，如果b数组中有一些零，则会得到inf值，这些值不会被allow_nan_stats=False选项捕获，因此也必须处理。

- meferne

---

你的数组 a 和 b 在最后一个轴上似乎相加为1，而不是在第一个轴上。 - Luca Di Liello

是的，最好说“沿着轴1”，甚至更好的是最后一个轴。当我写“沿着第一个轴”时，我的意思是轴1，因为还有轴0。我会编辑答案。谢谢！ - meferne

AttributeError: module 'tensorflow' has no attribute 'distributions' - jtlz2

Answer 2

如果使用softmax_cross_entropy_with_logits，就不需要在KL上进行优化。

KL(prob_a, prob_b)  
  = Sum(prob_a * log(prob_a/prob_b))  
  = Sum(prob_a * log(prob_a) - prob_a * log(prob_b))  
  = - Sum(prob_a * log(prob_b)) + Sum(prob_a * log(prob_a)) 
  = - Sum(prob_a * log(prob_b)) + const 
  = H(prob_a, prob_b) + const 

如果prob_a不是常数，您可以将其重写为两个熵的差。

KL(prob_a, prob_b)  
  = Sum(prob_a * log(prob_a/prob_b))  
  = Sum(prob_a * log(prob_a) - prob_a * log(prob_b))  
  = - Sum(prob_a * log(prob_b)) + Sum(prob_a * log(prob_a)) 
  = H(prob_a, prob_b) - H(prob_a, prob_a)  

- Jiecheng Zhao

---

有时在优化过程中，目标概率prob_a会发生变化。这时它就不再是一个常数了。 - CyberPlayerOne

Answer 3

我不确定为什么它没有被实现，但也许有一个解决方法。KL散度的定义如下：

KL(prob_a, prob_b) = Sum(prob_a * log(prob_a/prob_b))

另一方面，交叉熵H的定义如下：

H(prob_a, prob_b) = -Sum(prob_a * log(prob_b))

因此，如果你创建一个变量y = prob_a/prob_b，你可以通过调用负数H(proba_a, y)来获得KL散度。在Tensorflow符号中，可以这样表示：

KL = tf.reduce_mean(-tf.nn.softmax_cross_entropy_with_logits(prob_a, y))

- E.J. White

---

当 prob_a = prob_b 时，KL散度必须为0。但是最后一行并没有得到0。 - Transcendental

是的。当 prob_a = prob_b 时，我们得到 y = 1。然后，从 log(y) 中得到 H(prob_a, y) 等于零。您是说您使用 Tensorflow 的 softmax_cross_entropy_with_logits(prob_a, y) 进行了检查，结果不为零吗？ - E.J. White

没错。TensorFlow的实现可能与实际公式略有不同。 - Transcendental

值得指出的是，softmax_cross_entropy_with_logits(prob_a,y)实际上并没有实现H(prob_a,y)，而是实现了H(softmax(a),y)。因此，只有在尝试计算softmax函数（prob_a）的激活和未缩放的logits（a）之间的KL散度时，使用softmax_cross_entropy_with_logits才有效。 - shapecatcher

Answer 4

tf.contrib.distributions.kl 接受 tf.distribution 的实例而不是 Tensor。

示例：

  ds = tf.contrib.distributions
  p = ds.Normal(loc=0., scale=1.)
  q = ds.Normal(loc=1., scale=2.)
  kl = ds.kl_divergence(p, q)
  # ==> 0.44314718

- jvdillon

Answer 5

假设您可以访问logits a和b：

prob_a = tf.nn.softmax(a)
cr_aa = tf.nn.softmax_cross_entropy_with_logits(prob_a, a)
cr_ab = tf.nn.softmax_cross_entropy_with_logits(prob_a, b)
kl_ab = tf.reduce_sum(cr_ab - cr_aa)

- Sara

---

不会工作！根据文档：“警告：此操作期望未经缩放的logits，因为它在内部对logits执行softmax以提高效率。不要使用softmax的输出调用此操作，否则将产生不正确的结果”（我强调）。 - mikkola

假设您可以访问logits a和b。这不是在prob_a和prob_b上调用它，而是在a和b上调用它。 - Sara

Answer 6

我使用了这段代码（来自this Medium 文章）中的函数来计算任意给定张量与正态高斯分布之间的 KL 散度，其中 sd 是标准差，mn 是张量。

latent_loss = -0.5 * tf.reduce_sum(1.0 + 2.0 * sd - tf.square(mn) - tf.exp(2.0 * sd), 1)

- generic_stackoverflow_user

Answer 7

我认为这可能有效：

tf.reduce_sum(p * tf.log(p/q))

其中p是我的实际概率分布，q是我的近似概率分布。

- Akshaya Natarajan