DPO Derivation › Deriving the Optimal Policy

The key insight behind DPO is that this optimization problem has a closed-form solution. Unlike many optimization problems in machine learning that require iterative gradient descent, this particular formulation admits an analytical answer. For a fixed prompt $x$, we're optimizing over the distribution $\pi(y|x)$ for all possible responses $y$. This is a constrained optimization problem over probability distributions: we must ensure our solution is a valid probability distribution that sums to one and assigns non-negative probability to every possible response.

The special structure of the RLHF objective allows for this closed-form solution. The expectation under $\pi$ combined with the KL divergence creates what's known as a variational problem over distributions. Such problems often have elegant solutions when the constraints are simple probability simplex constraints.

Let's work through the derivation step by step. For a single prompt $x$, the objective is:

$$\max_{\pi(\cdot|x)} \sum_{y} \pi(y|x) \left[ r(x, y) - \beta \log \frac{\pi(y|x)}{\pi_{\text{ref}}(y|x)} \right]$$

where:

• $\sum_{y}$: sum over all possible responses in the vocabulary
• $\pi(\cdot|x)$: the entire probability distribution over responses for prompt $x$
• $\pi(y|x)$: probability of response $y$ given prompt $x$
• $r(x, y)$: the reward function
• $\beta$: the KL penalty coefficient
• $\log$: the natural logarithm
• $\pi_{\text{ref}}(y|x)$: the reference policy probability

We need to maximize this subject to the constraint that $\pi(\cdot|x)$ is a valid probability distribution: $\sum_y \pi(y|x) = 1$ and $\pi(y|x) \geq 0$ for all $y$.

To solve this constrained optimization problem, we use the method of Lagrange multipliers, a classical technique from calculus. The idea is to incorporate the constraint directly into the objective by introducing a new variable (the Lagrange multiplier) that penalizes violations of the constraint. At the optimal solution, the gradient of the objective with respect to the decision variables must be proportional to the gradient of the constraint.

The Lagrangian is:

$$\mathcal{L} = \sum_{y} \pi(y|x) \left[ r(x, y) - \beta \log \pi(y|x) + \beta \log \pi_{\text{ref}}(y|x) \right] + \lambda \left( 1 - \sum_y \pi(y|x) \right)$$

where:

• $\mathcal{L}$: the Lagrangian function
• $\sum_{y}$: sum over all responses
• $\pi(y|x)$: probability of response $y$
• $r(x, y)$: the reward function
• $\beta$: the KL penalty coefficient
• $\log$: the natural logarithm
• $\pi_{\text{ref}}(y|x)$: the reference policy probability
• $\lambda$: the Lagrange multiplier for the constraint that probabilities sum to 1

Notice that we've expanded the log ratio term from the original objective, separating it into $-\beta \log \pi(y|x)$ and $+\beta \log \pi_{\text{ref}}(y|x)$. This separation makes taking derivatives more straightforward. The term $\lambda(1 - \sum_y \pi(y|x))$ enforces our normalization constraint: if the probabilities don't sum to one, this term will be nonzero, and the Lagrange multiplier $\lambda$ will adjust to push us toward a valid distribution.

Taking the derivative with respect to $\pi(y|x)$ for a specific $y$:

$$\frac{\partial \mathcal{L}}{\partial \pi(y|x)} = r(x, y) - \beta \log \pi(y|x) - \beta + \beta \log \pi_{\text{ref}}(y|x) - \lambda = 0$$

where:

• $\frac{\partial \mathcal{L}}{\partial \pi(y|x)}$: the partial derivative of the Lagrangian with respect to the probability of response $y$
• $r(x, y)$: the reward function
• $\beta$: the KL penalty coefficient
• $\log$: the natural logarithm
• $\pi(y|x)$: the policy probability
• $\pi_{\text{ref}}(y|x)$: the reference policy probability
• $\lambda$: the Lagrange multiplier

The $-\beta$ term comes from differentiating $\pi \log \pi$, which gives $\log \pi + 1$. This is a standard result from calculus: when we differentiate $\pi \log \pi$ with respect to $\pi$, we apply the product rule, obtaining $\log \pi + \pi \cdot (1/\pi) = \log \pi + 1$. Setting this derivative equal to zero gives us the first-order optimality conditions.

Solving for $\log \pi(y|x)$:

$$\begin{aligned} \beta \log \pi(y|x) &= r(x, y) + \beta \log \pi_{\text{ref}}(y|x) - \beta - \lambda \\ \log \pi(y|x) &= \frac{1}{\beta} r(x, y) + \log \pi_{\text{ref}}(y|x) - 1 - \frac{\lambda}{\beta} \end{aligned}$$

where:

• $\log \pi(y|x)$: the log-probability of response $y$
• $r(x, y)$: the reward function
• $\beta$: the KL penalty coefficient
• $\pi_{\text{ref}}(y|x)$: the reference policy probability
• $\lambda/\beta$: constant terms related to the normalization constraint

This expression for the log-probability has an illuminating structure. The log-probability under our optimal policy equals the log-probability under the reference, adjusted by the scaled reward $r(x,y)/\beta$, plus some constants that don't depend on $y$. Higher reward responses get higher log-probabilities, with $\beta$ controlling how much the reward matters relative to staying close to the reference.

Exponentiating both sides:

$$\pi(y|x) = \pi_{\text{ref}}(y|x) \cdot \exp\left( \frac{r(x, y)}{\beta} \right) \cdot \exp\left( -1 - \frac{\lambda}{\beta} \right)$$

where:

• $\pi(y|x)$: the policy probability
• $\pi_{\text{ref}}(y|x)$: the reference policy probability
• $\exp(\cdot)$: the exponential function
• $\exp(r(x, y)/\beta)$: the reward-scaling term
• $r(x, y)$: the reward function
• $\beta$: the KL penalty coefficient
• $\exp(-1 - \lambda/\beta)$: the normalization term constant across all $y$

The term $\exp(-1 - \lambda/\beta)$ is just a normalizing constant that ensures the distribution sums to 1. This term doesn't depend on $y$; it only depends on $x$ through the constraint that all probabilities must sum to one. We can absorb it into a partition function $Z(x)$:

$$\pi^*(y|x) = \frac{1}{Z(x)} \pi_{\text{ref}}(y|x) \exp\left( \frac{r(x, y)}{\beta} \right)$$

where:

• $\pi^*(y|x)$: the optimal policy distribution
• $Z(x) = \sum_{y'} \pi_{\text{ref}}(y'|x) \exp\left( \frac{r(x, y')}{\beta} \right)$: the partition function (normalizing constant) for prompt $x$
• $\pi_{\text{ref}}(y|x)$: the reference policy distribution
• $r(x, y)$: the reward function
• $\beta$: the temperature parameter

This is the optimal policy for the RLHF objective. It has a clear interpretation: the optimal policy takes the reference distribution and reweights each response by the exponentiated reward, then normalizes. The exponential function ensures all probabilities are positive, naturally satisfying the $\pi(y|x) \ge 0$ constraint. Responses with higher reward get exponentially more probability mass, with $\beta$ controlling how aggressively we reweight.

Consider what happens at extreme values of beta to see how this reweighting works. When $\beta$ is very large, the reward term $r(x,y)/\beta$ becomes small, and $\exp(r(x,y)/\beta) \approx 1$ for all responses. In this limit, the optimal policy stays very close to the reference distribution, barely adjusting for reward at all. Conversely, when $\beta$ is very small, the exponential amplifies small reward differences into massive probability shifts, putting all probability mass on the highest-reward response. The choice of $\beta$ therefore controls the exploration-exploitation trade-off: larger values favor staying close to the reference (exploration), while smaller values favor chasing high reward (exploitation).