How can I implement the Kullback-Leibler loss in TensorFlow?

Question

I need to minimize KL loss in tensorflow.

I tried this function tf.contrib.distributions.kl(dist_a, dist_b, allow_nan=False, name=None), but I failed.

I tried to implement it manually:

def kl_divergence(p,q):
    return p* tf.log(p/q)+(1-p)*tf.log((1-p)/(1-q))

Is it correct?

Possible duplicate of [KL Divergence in TensorFlow](http://stackoverflow.com/questions/41863814/kl-divergence-in-tensorflow) — Transcendental, Apr 09 '17 at 00:11

score 10 · Accepted Answer · edited May 17 '19 at 18:57

10

What you have there is the cross entropy, KL divergence should be something like:

def kl_divergence(p, q): 
    return tf.reduce_sum(p * tf.log(p/q))

This assumes that p and q are both 1-D tensors of float, of the same shape and for each, their values sum to 1.

It should also work if p and q are equally sized mini-batches of 1-D tensors that obey the above constraints.

edited May 17 '19 at 18:57

Anubhav Singh

answered Apr 08 '17 at 17:56

Daniel Slater

Thanks a lot, and in case that p and q are multidimensional? – Alberto Merciai Apr 08 '17 at 18:16
Do you mean in the case that p and q are mini-batches of distributions that you want to optimize? I think what I have should be fine for that case as well. If it's not that then I would need more context. – Daniel Slater Apr 08 '17 at 18:52
i get nan when i try to compute the division 0/0 – Alberto Merciai Apr 10 '17 at 18:53
i add a **costant** = 0.00001 to each p and q to avoid nan, is correct? – Alberto Merciai Apr 11 '17 at 11:44
@srabb that is one option, you can also do tf.max so it doesn't modify other results, I wonder if there is a better way though – Daniel Slater Apr 11 '17 at 15:34
adding a small const is the version I see in tensorflow examples from Google, I would guess this is marginally faster than doing a max? – Daniel Slater Oct 23 '17 at 16:35
Instead of adding a constant to handle zeros smooth it with an uniform distribution. – kauDaOtha Oct 25 '18 at 19:46
tf.log no longer exists - tf.math.log rather – jtlz2 Apr 08 '20 at 14:43

1 Answers1