How do I compute the KL divergence in Keras with TensorFlow backend?

Question

How can I compute kl diveregence in keras while using tensorflow as backend? I compute L1 loss as follows:

def l1_loss(y_true, y_pred):
    return K.sum(K.abs(y_pred - y_true), axis=-1)

Answer 1

Keras already has the KL-divergence implemented, as it can be seen here, the code is just:

def kullback_leibler_divergence(y_true, y_pred):
    y_true = K.clip(y_true, K.epsilon(), 1)
    y_pred = K.clip(y_pred, K.epsilon(), 1)
    return K.sum(y_true * K.log(y_true / y_pred), axis=-1)

So just use kld, KLD or kullback_leibler_divergence as loss.

Answer 2

You can simply use the tf.keras.losses.kullback_leibler_divergence function.

If you want to use it as an activity regularizer, you can create a simple regularization function:

import keras # if using keras
# from tensorflow import keras # if using tf.keras
kullback_leibler_divergence = keras.losses.kullback_leibler_divergence
K = keras.backend

def kl_divergence_regularizer(inputs):
    means = K.mean(inputs, axis=0)
    return 0.01 * (kullback_leibler_divergence(0.05, means)
                 + kullback_leibler_divergence(1 - 0.05, 1 - means))

In this example, 0.01 is the regularization weight, and 0.05 is the sparsity target. Then use it like this:

keras.layers.Dense(32, activation="sigmoid",
                   activity_regularizer=kl_divergence_regularizer)

For example, this would be the encoding layer of a sparse autoencoder.

Note that the kullback_leibler_divergence expects all the class probabilities, even in the case of binary classification (giving just the positive class probability is not enough). This is why we compute the KLD for both 0.05 and 1-0.05 in the function above.