KL-divergence

• A type of statistical distance: a measure of how one probability distribution P is different from a second, reference probability distribution Q.

• A simple interpretation of the KL divergence of P from Q is the expected excess surprise from using Q as a model instead of P when the actual distribution is P.

  • i.e. for every $x\sim p(x)$, how far away is the ratio $\frac{p(x)}{q(x)}$ from 1 ?

• $D_{KL}(P||Q)={\mathbb{E}}_{p(x)}[log(\frac{p(x)}{q(x)})]=\int log(\frac{p(x)}{q(x)}p(x))\mathrm{d}x$

Properties

• $D_{KL}(P||Q)\ge 0$ , a result known as Gibbs’ inequality.

Gaussian distributions

Computing the KL

• (https://github.com/openai/improved-diffusion/blob/1bc7bbbdc414d83d4abf2ad8cc1446dc36c4e4d5/improved_diffusion/losses.py#L12)

def normal_kl(mean1, logvar1, mean2, logvar2):
    return 0.5 * (
        -1.0
        + logvar2
        - logvar1
        + th.exp(logvar1 - logvar2)
        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
    )

• Full equation: $D_{KL}(p||q)=\frac{1}{2}\left[\log \frac{|{\Sigma }_q|}{|{\Sigma }_p|}-k+({\mathbf{\mu }}_{\mathbf{p}}-{\mathbf{\mu }}_{\mathbf{q}}{)}^T{\Sigma }_q^{-1}({\mathbf{\mu }}_{\mathbf{p}}-{\mathbf{\mu }}_{\mathbf{q}})+tr\left\{{\Sigma }_q^{-1}{\Sigma }_p\right\}\right]$
• For single-variate: $D_{KL}(p||q)=\log \frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}-\frac{1}{2}$
• The above code makes the simplifying assumption that $\Sigma$ is diagonal
• Thus, it applies the single-variate formula in parallel to all dimensions.