KL Divergence (Distillation, and the RLHF KL Penalty Term)

The problem, from first principles

math-kl-divergence already built KL divergence as a pure information-theory concept, measuring how different two probability distributions are. This question shows KL divergence isn't just abstract theory: it's a genuinely load-bearing LOSS TERM in two completely different real training setups: knowledge distillation (training a small model to imitate a large one) and the KL penalty inside RLHF (rl-alignment-dpo-direct-preference-optimization's reference-relative theme, made explicit here as its own loss term).

From theory to code

Implement distillation_loss (student imitating a teacher's output distribution) and rlhf_kl_penalty (keeping a policy close to a frozen reference), both directly reusing math-kl-divergence's existing kl_divergence function.

Constraints

• distillation_loss(teacher_probs, student_probs) returns mean_i(kl_divergence(teacher_probs[i], student_probs[i], base=e)): the average KL(teacher || student) across a batch, using natural log.
• rlhf_kl_penalty(policy_probs, reference_probs, beta) returns beta * mean_i(kl_divergence(policy_probs[i], reference_probs[i], base=e)).
• Both must reduce to exactly 0 whenever every row of the two input arrays is identical.
• KL divergence is NOT symmetric in general: distillation_loss(p, q) and distillation_loss(q, p) can (and usually do) differ.

Hints

Open one at a time. Each gives away a little more than the last.