Search code examples
tensorflowprobabilityloss

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

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?

Solution

  • 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.