Eric Campbell 5 min read Statistics

How to Calculate KL Divergence

Kullback-Leibler (KL) divergence is a fundamental measure in information theory that quantifies how one probability distribution differs from another. If you've worked with variational autoencoders,...

Key Insights

  • KL divergence measures how one probability distribution differs from a reference distribution, but it’s asymmetric—D_KL(P||Q) ≠ D_KL(Q||P)—making it unsuitable as a true distance metric
  • Numerical stability is critical when implementing KL divergence; always add epsilon smoothing to prevent log(0) errors that can crash your calculations
  • Modern ML frameworks provide optimized KL divergence implementations, but understanding the manual calculation helps you debug issues and choose the right tool for discrete vs continuous distributions

Introduction to KL Divergence

Kullback-Leibler (KL) divergence is a fundamental measure in information theory that quantifies how one probability distribution differs from another. If you’ve worked with variational autoencoders, probabilistic models, or Bayesian inference, you’ve encountered KL divergence—it’s the mathematical backbone of how we measure distributional differences.

KL divergence answers a specific question: “If I use distribution Q to approximate distribution P, how much information am I losing?” This makes it invaluable for model comparison, compression algorithms, and any scenario where you need to quantify distributional mismatch. In machine learning, it’s particularly common in variational inference, where we approximate complex posterior distributions with simpler ones.

Unlike metrics such as Euclidean distance, KL divergence operates in probability space. It’s not symmetric and doesn’t satisfy the triangle inequality, so technically it’s a “divergence” rather than a “distance.” This distinction matters when choosing the right tool for your problem.

Mathematical Foundation

The KL divergence from distribution Q to distribution P is defined as:

For discrete distributions: D_KL(P||Q) = Σ P(x) log(P(x)/Q(x))

For continuous distributions: D_KL(P||Q) = ∫ p(x) log(p(x)/q(x)) dx

Three critical properties define KL divergence:

  1. Non-negativity: D_KL(P||Q) ≥ 0, with equality only when P and Q are identical
  2. Asymmetry: D_KL(P||Q) ≠ D_KL(Q||P) in general
  3. Not a metric: It doesn’t satisfy the triangle inequality

The asymmetry has practical implications. D_KL(P||Q) penalizes cases where P has high probability but Q has low probability (we’re surprised by events that should be common). Conversely, D_KL(Q||P) penalizes the opposite scenario. Choose the direction based on which type of error matters more for your application.

Calculating KL Divergence for Discrete Distributions

For discrete distributions, the calculation is straightforward: iterate through all possible outcomes, multiply each probability from P by the log ratio, and sum everything up.

Let’s calculate KL divergence between two loaded dice distributions:

import numpy as np

def kl_divergence_discrete(p, q):
    """
    Calculate KL divergence for discrete distributions.
    
    Args:
        p: True distribution (numpy array)
        q: Approximating distribution (numpy array)
    
    Returns:
        KL divergence D_KL(P||Q)
    """
    # Ensure distributions are normalized
    p = np.asarray(p, dtype=np.float64)
    q = np.asarray(q, dtype=np.float64)
    p /= p.sum()
    q /= q.sum()
    
    # Filter out zero probabilities in p (they don't contribute)
    mask = p > 0
    
    # Calculate KL divergence
    return np.sum(p[mask] * np.log(p[mask] / q[mask]))

# Example: Two different loaded dice
# Fair die
p_fair = np.array([1/6, 1/6, 1/6, 1/6, 1/6, 1/6])

# Loaded die (favors high numbers)
p_loaded = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.3])

# Calculate divergence
div_pq = kl_divergence_discrete(p_fair, p_loaded)
div_qp = kl_divergence_discrete(p_loaded, p_fair)

print(f"D_KL(Fair||Loaded): {div_pq:.4f}")
print(f"D_KL(Loaded||Fair): {div_qp:.4f}")

Notice the asymmetry in the results. The divergence changes depending on which distribution you treat as the reference.

Calculating KL Divergence for Continuous Distributions

For continuous distributions, we need to evaluate integrals. In practice, you’ll rarely compute these analytically except for well-known distribution families. For normal distributions, there’s a closed-form solution:

For two normal distributions N(μ₁, σ₁²) and N(μ₂, σ₂²): D_KL(P||Q) = log(σ₂/σ₁) + (σ₁² + (μ₁ - μ₂)²)/(2σ₂²) - 1/2

import numpy as np
from scipy.stats import norm

def kl_divergence_normal(mu1, sigma1, mu2, sigma2):
    """
    Calculate KL divergence between two normal distributions.
    
    Args:
        mu1, sigma1: Mean and std of distribution P
        mu2, sigma2: Mean and std of distribution Q
    
    Returns:
        D_KL(P||Q)
    """
    return (np.log(sigma2/sigma1) + 
            (sigma1**2 + (mu1 - mu2)**2) / (2 * sigma2**2) - 0.5)

# Example: Two normal distributions
mu1, sigma1 = 0, 1  # Standard normal
mu2, sigma2 = 0.5, 1.5  # Shifted and wider

kl_div = kl_divergence_normal(mu1, sigma1, mu2, sigma2)
print(f"D_KL(N(0,1)||N(0.5,1.5)): {kl_div:.4f}")

# Verify with numerical integration
x = np.linspace(-5, 5, 1000)
p = norm.pdf(x, mu1, sigma1)
q = norm.pdf(x, mu2, sigma2)
dx = x[1] - x[0]

# Numerical approximation
kl_numerical = np.sum(p * np.log(p / q) * dx)
print(f"Numerical approximation: {kl_numerical:.4f}")

Using Built-in Library Functions

Don’t reinvent the wheel. Modern libraries provide optimized, numerically stable implementations:

import numpy as np
from scipy.stats import entropy
import torch
import torch.nn.functional as F
import tensorflow as tf

# Sample distributions
p = np.array([0.1, 0.4, 0.5])
q = np.array([0.2, 0.3, 0.5])

# SciPy: entropy(p, q) calculates D_KL(P||Q)
kl_scipy = entropy(p, q)
print(f"SciPy KL divergence: {kl_scipy:.4f}")

# PyTorch: expects log probabilities for q
p_torch = torch.tensor(p)
q_torch = torch.tensor(q)
kl_pytorch = F.kl_div(q_torch.log(), p_torch, reduction='sum')
print(f"PyTorch KL divergence: {kl_pytorch.item():.4f}")

# TensorFlow: KLDivergence loss
kl_loss = tf.keras.losses.KLDivergence()
kl_tf = kl_loss(p, q)
print(f"TensorFlow KL divergence: {kl_tf.numpy():.4f}")

Important: PyTorch’s kl_div expects the second argument to be log probabilities and has reversed argument order compared to the mathematical convention. Always check the documentation.

Common Pitfalls and Best Practices

The biggest issue with KL divergence is numerical instability. When Q(x) is zero but P(x) is non-zero, you get log(∞), which crashes your calculation. When both are zero, you get 0 * log(0/0), which is undefined.

Always implement epsilon smoothing:

def kl_divergence_stable(p, q, epsilon=1e-10):
    """
    Numerically stable KL divergence calculation.
    
    Args:
        p: True distribution
        q: Approximating distribution
        epsilon: Small constant for numerical stability
    
    Returns:
        D_KL(P||Q)
    """
    p = np.asarray(p, dtype=np.float64)
    q = np.asarray(q, dtype=np.float64)
    
    # Normalize
    p = p / p.sum()
    q = q / q.sum()
    
    # Add epsilon to prevent log(0)
    p = p + epsilon
    q = q + epsilon
    
    # Renormalize after adding epsilon
    p = p / p.sum()
    q = q / q.sum()
    
    return np.sum(p * np.log(p / q))

# Test with edge case
p_edge = np.array([0.5, 0.5, 0.0])
q_edge = np.array([0.3, 0.3, 0.4])

# Without epsilon, this would fail
kl_stable = kl_divergence_stable(p_edge, q_edge)
print(f"Stable KL divergence: {kl_stable:.4f}")

When NOT to use KL divergence:

  • When you need symmetry: use Jensen-Shannon divergence instead
  • When distributions have different supports: use Wasserstein distance
  • When comparing high-dimensional distributions: consider Maximum Mean Discrepancy (MMD)

Practical Applications

KL divergence shines in real-world machine learning scenarios. Here’s how to use it for model comparison:

import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from scipy.stats import entropy

# Generate classification dataset
X, y = make_classification(n_samples=1000, n_features=20, 
                           n_informative=15, n_classes=3, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# Train two models
rf = RandomForestClassifier(n_estimators=100, random_state=42)
lr = LogisticRegression(max_iter=1000, random_state=42)

rf.fit(X_train, y_train)
lr.fit(X_train, y_train)

# Get predicted probability distributions
rf_probs = rf.predict_proba(X_test)
lr_probs = lr.predict_proba(X_test)

# Compare how differently the models see each test sample
kl_divergences = []
for i in range(len(X_test)):
    kl = entropy(rf_probs[i], lr_probs[i])
    kl_divergences.append(kl)

print(f"Mean KL divergence between models: {np.mean(kl_divergences):.4f}")
print(f"Max KL divergence: {np.max(kl_divergences):.4f}")

# Find samples where models disagree most
most_different_idx = np.argmax(kl_divergences)
print(f"\nMost disagreement on sample {most_different_idx}:")
print(f"Random Forest: {rf_probs[most_different_idx]}")
print(f"Logistic Regression: {lr_probs[most_different_idx]}")

This approach helps identify samples where models fundamentally disagree, which is valuable for ensemble methods, model debugging, and understanding prediction confidence.

KL divergence is also central to Variational Autoencoders (VAEs), where it regularizes the learned latent distribution to match a prior (typically standard normal). The VAE loss function explicitly includes a KL divergence term that prevents the encoder from learning arbitrary distributions.

Master KL divergence and you’ll have a powerful tool for comparing distributions, debugging probabilistic models, and implementing advanced machine learning architectures. Just remember: always handle numerical stability, respect the asymmetry, and choose your reference distribution carefully.

Liked this? There's more.

Every week: one practical technique, explained simply, with code you can use immediately.