KL Divergence Penalty

easyFREEPreference OptimizationLLM 2: Training & Alignment

Compute the KL divergence between two discrete probability distributions.

$D_{KL}(P \| Q) = \sum_i P(i) \log \frac{P(i)}{Q(i)}$

Here $\log$ is the natural logarithm. This is used in RLHF to penalize the policy from deviating too far from the reference model.

Input:

P: a 1D NumPy array of floats (distribution P)
Q: a 1D NumPy array of floats (distribution Q), same length as P

Output: the KL divergence as a single float. The tests round it to 4 decimal places.

Example

Input

0.5 0.5
0.5 0.5

Output

0.0

Reasoning

P and Q are two probability distributions over the same 2 outcomes: $P = (0.5, 0.5)$ is the policy and $Q = (0.5, 0.5)$ is the reference model. Each entry is the probability of one outcome, and each sums to 1.
Outcome $i$ contributes $P(i) \ln \frac{P(i)}{Q(i)}$ : the log of how much more likely $P$ makes that outcome than $Q$ does, weighted by how often $P$ actually picks it.
Here both ratios are $\frac{0.5}{0.5} = 1$ and $\ln 1 = 0$ , so each term is $0.5 \cdot 0 = 0$ and the total is $0$ , shown as 0.0.
Notice that the divergence is zero exactly when the two distributions agree on every outcome: a policy that has not moved from the reference pays no penalty.
For contrast, $P = (0.9, 0.1)$ against the same $Q$ gives $0.9 \ln 1.8 \approx 0.5290$ and $0.1 \ln 0.2 \approx -0.1609$ . One term is negative, yet the total $0.3681$ is positive; KL divergence is never below zero. Swapping the roles of $P$ and $Q$ gives $0.5108$ , so it is not symmetric.

Constraints

P and Q are valid probability distributions (sum to 1, all positive)
Same length
Round to 4 decimal places

Learn this in

Safety & EthicsLLM Study Plan Deployment & OptimizationLLM Study Plan RLHF & AlignmentLLM Study Plan

More from LLM 2: Training & Alignment

RLHF Reward Shapinghard Implicit Reward from DPOmedium PPO Clipped Objectivemedium DPO Lossmedium Best-of-N Samplinghard

Concept primer

Background Knowledge

The KL divergence, or Kullback-Leibler divergence, is a measure of the difference between two probability distributions. It is commonly used in machine learning and information theory to quantify the similarity between two distributions. In the context of preference optimization and RLHF (Reinforcement Learning from Human Feedback), the KL divergence is used to penalize a policy for deviating too far from a reference model. This helps to ensure that the policy stays close to the desired behavior.

The KL divergence is defined as $D_{KL}(P \| Q) = \sum_i P(i) \log \frac{P(i)}{Q(i)}$ , where $P$ and $Q$ are the two probability distributions. The KL divergence is not symmetric, meaning that $D_{KL}(P \| Q) \neq D_{KL}(Q \| P)$ in general. This is important to keep in mind when interpreting the results. The KL divergence is also not a distance metric, as it does not satisfy the triangle inequality.

In the context of this problem, we are given two discrete probability distributions $P$ and $Q$ , and we need to compute the KL divergence between them. This requires understanding the definition of the KL divergence and how to work with discrete probability distributions. We will also need to use basic mathematical operations such as summation and logarithms.

Algorithm/Approach

The general approach to solving this problem is to use the definition of the KL divergence and apply it to the given probability distributions. This involves iterating over the possible outcomes, computing the probability ratios, and summing up the contributions to the KL divergence. We will need to use a programming language to implement this calculation and round the result to 4 decimal places.

Step-by-Step Strategy

To implement the solution, we can follow these steps:

Read in the probability distributions $P$ and $Q$ from the input
Initialize a variable to store the KL divergence
Iterate over the possible outcomes, computing the probability ratio $\frac{P(i)}{Q(i)}$ and the contribution $P(i) \log \frac{P(i)}{Q(i)}$ to the KL divergence
Sum up the contributions to compute the total KL divergence
Round the result to 4 decimal places and output it

Common Pitfalls

Some things to watch out for when implementing the solution include:

Ensuring that the probability distributions are valid (i.e., they sum up to 1)
Handling cases where $Q(i) = 0$ , as this would result in a division by zero
Using a stable method for computing the logarithm to avoid numerical issues

Time & Space Complexity

The time complexity of the solution is $O(n)$ , where $n$ is the number of possible outcomes, since we need to iterate over all outcomes to compute the KL divergence. The space complexity is $O(1)$ , since we only need to store a few variables to compute the result. Note that the input size is $n$ , as we need to read in the probability distributions.

solution.py

Python 3.13

Test results