KL Divergence for two probability distributions in PyTorch

Question

I have two probability distributions. How should I find the KL-divergence between them in PyTorch? The regular cross entropy only accepts integer labels.

Answer 1

Yes, PyTorch has a method named kl_div under torch.nn.functional to directly compute KL-devergence between tensors. Suppose you have tensor a and b of same shape. You can use the following code:

import torch.nn.functional as F
out = F.kl_div(a, b)

For more details, see the above method documentation.

Answer 2

function kl_div is not the same as wiki's explanation.

I use the following:

# this is the same example in wiki
P = torch.Tensor([0.36, 0.48, 0.16])
Q = torch.Tensor([0.333, 0.333, 0.333])

(P * (P / Q).log()).sum()
# tensor(0.0863), 10.2 µs ± 508

F.kl_div(Q.log(), P, None, None, 'sum')
# tensor(0.0863), 14.1 µs ± 408 ns

compare to kl_div, even faster

Answer 3

If you have two probability distribution in form of pytorch distribution object. Then you are better off using the function torch.distributions.kl.kl_divergence(p, q). For documentation follow the link

Answer 4

If working with Torch distributions

mu = torch.Tensor([0] * 100)
sd = torch.Tensor([1] * 100)

p = torch.distributions.Normal(mu,sd)
q = torch.distributions.Normal(mu,sd)

out = torch.distributions.kl_divergence(p, q).mean()
out.tolist() == 0
True

Answer 5

If you are using the normal distribution, then the following code will directly compare the two distributions themselves:

p = torch.distributions.normal.Normal(p_mu, p_std)
q = torch.distributions.normal.Normal(q_mu, q_std)
    
loss = torch.distributions.kl_divergence(p, q)

p and q are two tensor objects.

This code will work and won't give any NotImplementedError.