🎓 A1Learner
🔍

PyTorch Entropy and KL Divergence

Entropy is one number that measures how spread out a distribution is. KL divergence measures the extra cost of using one distribution in place of another.

Entropy grows with spread

Wider Normal distributions carry higher entropy. Each doubling of the standard deviation adds the same step of about 0.6931.

for std in (0.5, 1.0, 2.0):
    d = Normal(0.0, std)
    print(std, round(d.entropy().item(), 4))

Entropy by hand

A standard Normal has mean 0 and a standard deviation of 1. Its entropy by hand matches the library value.

std = 1.0
hand = 0.5 * math.log(2 * math.pi * math.e * std ** 2)
print("By hand:", round(hand, 4))
print("By library:", round(Normal(0.0, std).entropy().item(), 4))

KL depends on direction

KL divergence depends on the direction of the swap. Swapping p and q changes the number, so KL is not a distance.

p = Normal(0.0, 1.0)
q = Normal(1.0, 2.0)
print("KL p to q:", round(kl_divergence(p, q).item(), 4))
print("KL q to p:", round(kl_divergence(q, p).item(), 4))

Example

import math

import torch
from torch.distributions import Normal, kl_divergence

p = Normal(0.0, 1.0)
q = Normal(1.0, 2.0)

entropy_hand = 0.5 * math.log(2 * math.pi * math.e * p.scale.item() ** 2)
entropy_lib = p.entropy().item()

mu1, s1 = p.mean.item(), p.scale.item()
mu2, s2 = q.mean.item(), q.scale.item()
kl_hand = (math.log(s2 / s1)
           + (s1 ** 2 + (mu1 - mu2) ** 2) / (2 * s2 ** 2) - 0.5)
kl_lib = kl_divergence(p, q).item()

print("Entropy by hand:", round(entropy_hand, 4))
print("Entropy by library:", round(entropy_lib, 4))
print("KL by hand:", round(kl_hand, 4))
print("KL by library:", round(kl_lib, 4))

Key points

  • Entropy: one number for how spread out a distribution is
  • Use dist.entropy() for the library entropy of a distribution
  • KL divergence: the cost of using distribution q in place of distribution p
  • Use kl_divergence(p, q) for the cost of using q in place of p
  • Closed form: a direct formula that needs no sampling or looping
  • KL divergence is not symmetric, so the order of the pair matters

In this example, p is a Normal with mean 0 and a standard deviation of 1. q is a Normal with mean 1 and a standard deviation of 2. The entropy of p is worked out by hand from the closed-form formula. The library entropy prints the same 1.4189. The KL divergence from p to q is also worked out by hand, in closed form. The library value matches at 0.4431. Both pairs agree, so the hand formulas and the library measure the same things.

Note

These values use natural logs, and their unit is called a nat

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

KL divergence from p to q prints 0.4431. What does KL from q to p print?