Why they use KL divergence in Natural gradient?

Question

Natural gradient aims to do a steepest descent on the "function" space, a manifold that is independent from how the function is parameterized. It argues that the steepest descent on this function space is not the same as steepest descent on the parameter space. We should favor the former.

Since, for example in a regression task, a neural net could be interpreted as a probability function (Gaussian with the output as mean and some constant variance), it is "natural" to form a distance on the manifold under the KL-divergence (and a Fisher information matrix as its metric).

Now, if I want to be creative, I could use the same argument to use "square distance" between the outputs of the neural nets (distance of the means) which I think is not the same as the KL.

Am I wrong, or it is just another legit way? Perhaps, not as good?

Answer 1

The KL divergence has slightly different interpretations depending on the context. The related Wikipedia article contains a section dedicated to these interpretations. Independently of the interpretation, the KL divergence is always defined as a specific function of the cross-entropy (which you should be familiar with before attempting to understand the KL divergence) between two distributions (in this case, probability mass functions)

\begin{align} D_\text{KL}(P\parallel Q) &= -\sum_{x\in\mathcal{X}} p(x) \log q(x) + \sum_{x\in\mathcal{X}} p(x) \log p(x) \\ &= H(P, Q) - H(P) \end{align} where $H(P, Q)$ is the cross-entropy of the distribution $P$ and $Q$ and $H(P) = H(P, P)$.

The KL is not a metric, given that it does not obey the triangle inequality. In other words, in general, $D_\text{KL}(P\parallel Q) \neq D_\text{KL}(Q\parallel P) $.

Given that a neural network is trained to output the mean (which can be a scalar or a vector) and the variance (which can be a scalar, a vector or a matrix), why don't we use a metric like the MSE to compare means and variances? When you use the KL divergence, you don't want to compare just numbers (or matrices), but probability distributions (more precisely, probability densities or mass functions), so you will not compare just the mean and the variance of two different distributions, but you will actually compare the distributions. See the example of the application of the KL divergence in the related Wikipedia article.

Answer 2

Yes, Squared distances & KL Divergence are not the same. Squared distance between means is not a useful metric as it doesn't gauge the amount of similarity between 2 distributions.

When we compute \begin{align} D_\text{KL}(P\parallel Q) \end{align} We are computing the amount of information that is lost when we approximate P as Q. Ideally, we would want the KL divergence to be as low as possible. Here is an interesting article https://www.countbayesie.com/blog/2017/5/9/kullback-leibler-divergence-explained where the author has explained KL Divergence with a toy example.

I hope it helps :)

Answer 3

I have been reading a lot about Natural Gradient and its use to find a descent direction. I found that this post was the most clear.

Consider a model $p$ parameterized by some parameters $\theta$ and we want to maximize the likelihood of observing our data $x$ under this model: $p(x|\theta)$. To optimise this likelihood we can take steps in the the distribution space. Updating the parameters $\theta$ we need to measure how our likelihood changes and this is measure using the KLK divergence.

Even though the KL divergence is not a "proper" distance metric as it is not symmetric, it is still quite informative about the similarity between distributions. It's practical because it can capture differences between distributions that the Euclidean metric (parameter-dependent) could not (see the same post for a simple example).

So answering your question is essentially answering which is the best between Natural Gradient Descent and "Normal" Gradient Descent in the Euclidean space where your loss is measured with a L2 norm. You can train the same model using both methods and you will just find different descent directions.

Hopefully though, both will converge but in my opinion Natural Gradient descent should be superior in nature. It is just very expensive to actually compute because to find the direction in distribution space you need to compute the inverse Fisher matrix $F^{-1}$ or approximate it and that's quite costly as it is of size $n\times n$ where $n$ is the size of $\theta$ which is typically high in neural networks.