KL Divergence in LLM Training — Preventing Policy Collapse

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: maximise 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 (regularises 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 ──
# D_KL(P||Q): "forward" — zero-avoiding (Q must cover all of P's support)
# D_KL(Q||P): "reverse" / inclusive — mass-covering (Q spreads over P's support)
# RLHF uses forward KL from current to reference policy
# Variational inference uses reverse KL (ELBO maximisation)
print("KL is asymmetric:")
print(f"  KL(P||Q) = {kl_manual:.4f}")
print(f"  KL(Q||P) = {(Q * (Q.log() - P.log())).sum():.4f}")