KL divergence — the mathematics

D_{KL}(P \| Q) = \sum_x P(x) \log \frac{P(x)}{Q(x)} = \mathbb{E}_{x \sim P}\left[\log \frac{P(x)}{Q(x)}\right]

import torch
import torch.nn.functional as F
import numpy as np

# ── 1. Basic KL Divergence computation ──
P = torch.tensor([0.3, 0.5, 0.2])   # True distribution (e.g., student output)
Q = torch.tensor([0.35, 0.4, 0.25]) # Reference distribution (e.g., teacher output)

# Manual computation
kl_manual = (P * (P.log() - Q.log())).sum()
print(f"KL(P||Q) manual: {kl_manual:.4f}")

# PyTorch (input must be log-probabilities for kl_div)
kl_torch = F.kl_div(Q.log(), P, reduction='sum')
print(f"KL(P||Q) torch:  {kl_torch:.4f}")

# ── 2. KL penalty in RLHF/GRPO (prevents policy collapse) ──
# RLHF objective: maximize reward while staying close to reference policy
# Total loss = -reward + β × KL(π_current || π_reference)

def rlhf_loss(reward, logprobs_current, logprobs_reference, beta=0.05):
    """
    reward:           scalar reward from reward model
    logprobs_current: log P(response | π_θ)  [current fine-tuned policy]
    logprobs_reference: log P(response | π_ref) [frozen reference: pretrained model]
    beta:             KL penalty weight (0.01-0.1 typical)
    """
    kl_penalty = (logprobs_current - logprobs_reference).mean()
    # KL(π_current || π_ref) ≈ logP_current - logP_ref (token-level average)
    loss = -reward + beta * kl_penalty
    return loss, kl_penalty.item()

# Example: reward model gives 2.5, but KL is very high (policy drifted)
reward = torch.tensor(2.5)
logprobs_cur = torch.tensor([-0.5, -1.2, -0.8, -0.3])   # Fine-tuned model probs
logprobs_ref = torch.tensor([-0.3, -0.5, -0.4, -0.2])   # Reference model probs (pretrained)

loss, kl = rlhf_loss(reward, logprobs_cur, logprobs_ref, beta=0.05)
print(f"RLHF loss: {loss:.3f} (reward={reward}, KL={kl:.3f})")

# ── 3. KL in knowledge distillation ──
# Soft targets: temperature scaling softens the teacher's distribution
def distillation_loss(student_logits, teacher_logits, hard_labels, T=3.0, alpha=0.5):
    """
    T:     Temperature (higher = softer distributions)
    alpha: Mix of soft (KL) and hard (CE) loss
    """
    # Soft targets from teacher
    soft_teacher = F.softmax(teacher_logits / T, dim=-1)
    soft_student = F.log_softmax(student_logits / T, dim=-1)

    # KL divergence loss (soft): match teacher distribution shape
    kl_soft = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (T**2)

    # Cross-entropy loss (hard): match ground truth labels
    ce_hard  = F.cross_entropy(student_logits, hard_labels)

    return alpha * kl_soft + (1 - alpha) * ce_hard

# ── 4. KL in VAE (regularizes latent space) ──
def vae_kl_loss(mu, logvar):
    """
    KL(N(μ,σ²) || N(0,1)) — keeps latent distribution close to standard Gaussian
    = -0.5 × Σ(1 + log(σ²) - μ² - σ²)
    """
    return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

# ── 5. Forward KL vs Reverse KL (P = target/true, Q = learned) ──
# D_KL(P||Q): "forward" — mass-covering / zero-avoiding
#             (Q must put mass everywhere P does)
# D_KL(Q||P): "reverse" — mode-seeking / zero-forcing
#             (Q concentrates on P's dominant modes)
# Variational inference minimizes reverse KL (ELBO maximization)
# RLHF penalizes KL(pi_theta || pi_ref), estimated ON-POLICY
# (samples from pi_theta) — reverse-KL-style toward the reference
print("KL is asymmetric:")
print(f"  KL(P||Q) = {kl_manual:.4f}")
print(f"  KL(Q||P) = {(Q * (Q.log() - P.log())).sum():.4f}")

Why the KL penalty prevents reward hacking

Without KL penalty, RLHF training quickly degrades: the LLM learns to output repetitive phrases that happen to fool the reward model ('The answer is: ' repeated 50 times scores high because it looks like it is about to give a long answer). The KL penalty penalizes any output distribution that deviates significantly from the pretrained model's distribution — forcing the fine-tuned model to stay linguistically coherent while improving alignment.

Forward vs reverse KL — and why on-policy distillation chooses reverse

The direction of the KL matters enormously in practice. With P the target distribution and Q the one being learned: minimizing forward KL D_KL(P‖Q) forces Q to spread mass over everything P can produce (mass-covering); minimizing reverse KL D_KL(Q‖P) lets Q concentrate on P's dominant modes but severely punishes Q for generating anything P considers impossible (mode-seeking, zero-forcing).

The one-line intuition: Forward KL says "be able to do everything the teacher does." Reverse KL says "never do something the teacher wouldn't." For small students and for RLHF-constrained policies, the second demand is both more achievable and more useful — which is why on-policy, reverse-KL-style objectives took over LLM post-training.

Practice questions

1. D_KL(P||Q) = 0.7 vs D_KL(Q||P) = 1.2. What does this asymmetry mean? (Answer: KL divergence is not symmetric — the direction matters. D_KL(P||Q) measures the extra bits needed to encode P using Q's code. Different directions penalize different types of mismatch. In RLHF, we use D_KL(π_current || π_ref) — penalizes current policy deviating from reference.)
2. In RLHF, why is β (the KL penalty weight) a crucial hyperparameter? (Answer: Too small β: reward hacking (model exploits reward model weaknesses). Too large β: model barely changes from pretrained policy (reward signal overwhelmed). Typical β = 0.05-0.1. Often adapted during training: start small, increase if reward hacking is detected, decrease if model is not improving.)
3. In knowledge distillation with temperature T=5, how does the teacher distribution change? (Answer: Temperature T scales logits before softmax: softmax(logits/T). Higher T → softer distribution (probabilities more uniform, smaller gaps between classes). The student learns not just the most likely class but the relative similarity structure: "this looks a bit like class 3 even though it's class 1". This dark knowledge is why distillation outperforms training from hard labels alone.)
4. Forward KL D_KL(P||Q) vs Reverse KL D_KL(Q||P) — which is mode-seeking and which is mode-covering? (Answer: Forward KL D_KL(P||Q) is mode-covering: P is the true/target distribution, Q is learned. Minimising forward KL forces Q to cover all modes of P (zero-avoiding — Q must not be zero where P is nonzero). Reverse KL D_KL(Q||P) is mode-seeking: Q concentrates on P's dominant modes and may ignore smaller ones. The RLHF penalty KL(π_current || π_reference) is estimated on-policy — with samples from the current policy — making it reverse-KL-style toward the reference: the policy may narrow within the reference's support, but is heavily punished for putting mass where the reference has none.)
5. In a VAE, what does KL(q(z|x) || p(z)) penalize? (Answer: It penalizes the encoder's posterior distribution q(z|x) from deviating from the prior p(z) = N(0,I). Without KL: each input maps to a distinct region of latent space — no overlap, no continuity, cannot sample from latent space for generation. With KL: all inputs map to distributions near N(0,I) — overlapping, smooth latent space that enables generation by sampling z ~ N(0,I).)