DPO Derivation › The RLHF Objective

Let's begin with the objective that RLHF aims to optimize. As we discussed in the chapters on PPO for Language Models and KL Divergence Penalty, the goal is to find a policy $\pi$ that maximizes expected reward while staying close to a reference policy $\pi_{\text{ref}}$ (typically the supervised fine-tuned model).

Before diving into the mathematics, it helps to understand the intuition behind this objective. We want our language model to generate responses that humans prefer, which is captured by the reward function. However, if we optimize the reward too aggressively, the model might find unexpected shortcuts or produce degenerate outputs that technically achieve high reward but don't represent helpful behavior. The reference policy is an anchor, representing the model's pre-trained knowledge and natural language capabilities. By penalizing deviations from this anchor, we encourage the model to improve its outputs while maintaining coherent, fluent generation.

Think of the reference policy as the starting point on a map, and the reward function as a compass pointing toward a better destination. The KL penalty acts like a tether that keeps you from wandering too far afield in a single step. You want to move in the direction of higher reward, but you don't want to teleport so far that you lose the path entirely. In language modeling, that "losing the path" corresponds to mode collapse or reward hacking: outputs that score well on a proxy reward metric but are nonsensical or harmful. The KL constraint is what prevents those degenerate solutions.

The constrained optimization problem is:

$$\max_{\pi} \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi(y|x)} \left[ r(x, y) \right] - \beta \cdot D_{\text{KL}}\left( \pi(y|x) \| \pi_{\text{ref}}(y|x) \right)$$

where:

• $\pi$: the policy being optimized
• $\mathbb{E}$: the expectation operator
• $x$: a prompt sampled from the data distribution $\mathcal{D}$
• $y$: a response sampled from the policy
• $r(x, y)$: the learned reward function
• $\beta$: a hyperparameter controlling the strength of the KL constraint
• $D_{\text{KL}}$: the Kullback-Leibler divergence
• $\pi_{\text{ref}}$: the reference policy

To work with this objective more directly, we need to expand the KL divergence term. Recall that the KL divergence measures how different one probability distribution is from another, expressed as the expected log ratio of the two distributions. Writing out the KL divergence explicitly, this becomes:

$$\max_{\pi} \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi(y|x)} \left[ r(x, y) - \beta \log \frac{\pi(y|x)}{\pi_{\text{ref}}(y|x)} \right]$$

where:

• $\mathbb{E}$: the expectation operator
• $\mathcal{D}$: the data distribution
• $r(x, y)$: the reward function
• $\beta$: the KL penalty coefficient
• $\log$: the natural logarithm
• $\pi(y|x)$: probability of response $y$ given prompt $x$ under the policy
• $\pi_{\text{ref}}(y|x)$: probability of response $y$ given prompt $x$ under the reference model

This reformulation reveals the objective's structure more clearly. The log ratio $\log \frac{\pi(y|x)}{\pi_{\text{ref}}(y|x)}$ measures how much the policy has shifted away from the reference for a particular response. When this ratio is positive (the policy assigns more probability than the reference), the KL term subtracts from the objective, penalizing the deviation. When the ratio is negative (the policy assigns less probability), the term adds to the objective, which might seem like a reward, but since we're sampling from $\pi$, we're unlikely to generate responses where $\pi$ assigns low probability.

The key insight is that this scalar objective, combining a reward signal and a distributional constraint, can be solved analytically. That analytical solution is the foundation of everything that follows. The reward-KL tradeoff creates a specific mathematical structure that admits an exact solution in closed form, rather than acting as an arbitrary heuristic regularizer.

This objective captures a basic tension in alignment: we want the model to produce high-reward outputs (as judged by human preferences), but we don't want it to deviate too far from its pre-trained behavior. Without the KL penalty, the model might find degenerate solutions that exploit flaws in the reward model, a phenomenon we explored in the chapter on Reward Hacking.

Before proceeding to the derivation, it is worth spending a moment on why the KL divergence is the natural choice of regularizer here. Many other distance measures between distributions exist, such as total variation distance, Renyi divergences, or the Wasserstein distance. Each would produce a different optimization problem and a different optimal policy. The KL divergence $D_{\text{KL}}(\pi \| \pi_{\text{ref}})$ is particularly convenient because it has exactly the right algebraic structure to yield a clean closed-form solution. Its definition as $\mathbb{E}_\pi[\log(\pi / \pi_{\text{ref}})]$ means that when we write the full objective as an expectation under $\pi$, everything becomes a single expectation of a per-token scalar quantity, which is tractable.

The KL divergence also has information-theoretic grounding. It measures the number of extra "bits" required to encode samples from $\pi$ using a code optimized for $\pi_{\text{ref}}$. Keeping this quantity small means the policy's outputs remain within the statistical neighborhood of the reference, which in practice means they retain the grammatical fluency, factual grounding, and stylistic coherence of the original supervised fine-tuned model.