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

tensorflow probability loss

Alberto Merciai · Apr 8, 2017 · Viewed 7.3k times · Source

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?

Answer

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.