<p>Kullback-Leibler divergence is a measure of similarity between two probability distributions. The KL divergence implemented in Keras assumes two discrete probability distributions (hence the sum).</p>

<p>The exact format of your KL loss function depends on the underlying probability distributions. A common usecase is that the neural network models the parameters of a probability distribution P (eg a Gaussian), and the KL divergence is then used in the loss function to determine the similarity between the modelled distribution and some other, known distribution (potentially Gaussian as well). E.g. a network outputs two vectors mu and sigma^2. Mu forms the mean of a Gaussian distribution P while sigma^2 is the diagonal of the covariance matrix Sigma. A possible loss function is then the KL divergence between the Gaussian P described by mu and Sigma, and a unit Gaussian N(0, I). The exact format of the KL divergence in that case can be derived analytically, yielding a <em>custom</em> keras loss function that is not at all equal to the KL divergence implemented in Keras.</p>

<p>In the original paper that introduces Variational Auto-Encoders, the loss function is summed over the samples in the minibatch and then multiplied by a factor (N/M), where N is the size of the entire dataset, and M is the size of the minibatch. See equations 8 and 10 in <a href="https://arxiv.org/abs/1312.6114" rel="nofollow noreferrer">https://arxiv.org/abs/1312.6114</a>.</p>
What is the meaning of the implementation of the KL divergence in Keras?
<p>I'm a little confused about how the KL divergence is applied, specifically in Keras, but I think the question is general to deep learning applications. In keras, the KL loss function is defined like this:</p>

<pre><code>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)
</code></pre>

<p>In my model, <code>y_true</code> and <code>y_pred</code> are matrices; each row of <code>y_true</code> a one-hot encoding for one training example, and each row of <code>y_pred</code> the output of the model (a probability distribution) for that example. </p>

<p>I can run this KL divergence calculation on any given pair of rows from <code>y_true</code> and <code>y_pred</code> and get the expected result. The mean of these KL divergence results over the rows matches the loss reported by Keras in the training history. But that aggregation - running KL divergence on each row and taking the mean - doesn't happen within loss function. In contrast, I understand MAE or MSE to aggregate across the examples: </p>

<pre><code>def mean_squared_error(y_true, y_pred):
    return K.mean(K.square(y_pred - y_true), axis=-1)
</code></pre>

<p>For the KL divergence, it's not totally obvious to me that taking the mean across the examples is the right thing to do. I guess the idea is that the examples are random samples from the true distribution, so they should appear in proportion to their probability. But that seems to make a pretty strong assumption about how the training data was collected. I haven't really seen this aspect (aggregating across samples from a dataset) addressed in the online treatments of the KL divergence; I just see a lot of redefinition of the basic formula.</p>

<p>So my questions are: </p>

<ol>
<li><p>Is this interpretation of what Keras is doing to come up with the KL divergence loss (i.e. averaging over the KL divergence of the rows) correct?</p></li>
<li><p>Why is this the right thing to do?</p></li>
<li><p>From an implementation perspective, why doesn't the definition of the loss function in Keras do the aggregation over the rows the way MAE or MSE does?</p></li>
</ol>