KL Divergence (Gaussians) - Generative Models

KL Divergence Between Two Gaussians

What you'll build

The closed-form KL divergence KL(P || Q) between two univariate Gaussians.

Background

KL divergence measures how many extra nats you pay, on average, for encoding samples from P with a code optimised for Q. It appears all over generative modelling: the VAE regulariser is KL(q(z|x) || N(0, 1)), diffusion's ELBO is a sum of Gaussian KLs, and RLHF penalises KL(π || π_ref). For Gaussians it has a closed form, so no sampling is needed.

Formula

With P = N(μ_p, σ_p²) and Q = N(μ_q, σ_q²) (σ are standard deviations):

KL(P || Q) = log(σ_q / σ_p) + (σ_p² + (μ_p − μ_q)²) / (2 σ_q²) − 1/2

Setting μ_q = 0, σ_q = 1 and writing log_var = log σ_p² gives the VAE term −0.5·(1 + log_var − μ² − exp(log_var)) used in VAE Loss Function.

Properties

• KL ≥ 0, with equality iff P = Q
• Not symmetric: KL(P || Q) ≠ KL(Q || P) in general
• Units are nats because the log is natural

Inputs

mu_p, sigma_p, mu_q, sigma_q: Python numbers (or NumPy scalars), with sigma_p, sigma_q > 0.

Expected return format

kl: a Python float rounded to 4 decimals.

Constraints

Use the closed form (no sampling).

Follow-up

• Which direction does a VAE minimise, and why does reverse KL lead to mode-seeking behaviour?
• Generalise to diagonal multivariate Gaussians (sum the per-dimension KLs).

Input: mu_p = 1, sigma_p = 1, mu_q = 0, sigma_q = 1
Output: 0.5

Input: mu_p = 0, sigma_p = 2, mu_q = 0, sigma_q = 1
Output: 0.8069

import numpy as np

def kl_divergence_gaussian(mu_p, sigma_p, mu_q, sigma_q):
    """
    KL(P || Q) for P = N(mu_p, sigma_p^2) and Q = N(mu_q, sigma_q^2).

    Args:
        mu_p, sigma_p: Mean and standard deviation of P (sigma_p > 0)
        mu_q, sigma_q: Mean and standard deviation of Q (sigma_q > 0)

    Returns:
        kl: Python float rounded to 4 decimals
    """
    # Your code here
    pass