Rafailov 2023 - Direct Preference Optimization

Here we trace the derivations from the DPO paper. Denote the model after SFT as ${\pi }^{SFT}(y|x)$, i.e. the policy ${\pi }^{SFT}$ is a probability function for each pair of inputs and answers (x, y). Naturally, we can use this policy to generate tokens by choosing the response with the highest probability (or approximate it in a greedy token-by-token manner).

To perform RLHF, we first need to build a reward model. First, we prompt the SFT model to obtain pairs of answers $(y_1,y_2)\sim {\pi }^{SFT}(y|x)$. These samples are presented to human labellers who record their preferences in the form $y_w▹y_l|x$, where $y_w$ wins and $y_l$ loses. These preferences are assumed to be generated from an underlying reward model $r^{\ast }(x,y)$ which we do not have access to.

We wish to learn this reward model. Since we only have access to pairwise preferences instead of the reward score, a common approach is to model the pairwise preferences using the Bradley-Terry model. Specifically, we assume that the observed human preference decisions are related to the underlying human reward model in the following way:

$$\begin{array}{rlr}{p}^{\ast }(y_1▹y_2|x)& =\frac{e^{r^{\ast }(x,y_1)}}{e^{r^{\ast }(x,y_1)}+e^{r^{\ast }(x,y_2)}}& (1)\end{array}$$

Suppose we have a static dataset of comparisons $\mathcal{D}=\{x^{(i)},y_w^{(i)},y_l^{(i)}{\}}_{i=1}^N$ sampled from the human preference model of $p^{\ast }$. We can create a reward model $r_{\varphi }(x,y)$ and use the BT-model to express the negative log likelihood of $\mathcal{D}$. Note that are we using the expression of the BT-model as a sigmoid function. With this NLL expression, we can optimize for $r_{\varphi }$ using gradient descent to learn the reward model from $\mathcal{D}$. Notice that the heart of Equation (2) is essentially just a difference in reward between the winning answer $y_w$ and losing answer $y_l$.

$$\begin{array}{rlr}\mathcal{L}(r_{\varphi },\mathcal{D})& =-{\mathbb{E}}_{(x,y_w,y_l)\sim \mathcal{D}}\left[log\sigma (r_{\varphi }(x,y_w)-r_{\varphi }(x,y_l))\right]& (2)\end{array}$$

Note that $r_{\varphi }$ is usually initialized using the SFT model ${\pi }^{SFT}$, but with an additional linear layer over the final transformer layer to generate a scalar value as the reward.

Having learned the reward model, we then need to use the learned reward function to fine-tune ${\pi }^{SFT}$ using reinforcement learning. Specifically, we set ${\pi }_{ref}:={\pi }^{SFT}$ as the reference model, and initialize a new ${\pi }_{\theta }$ model that we wish to train. Usually, ${\pi }_{\theta }$ is also initialized as a copy of ${\pi }^{SFT}$. The objective function we wish to maximize is:

$$\begin{array}{rlr}& \underset{{\pi }_{\theta }}{\max }{\mathbb{E}}_{x\sim \mathcal{D},y\sim {\pi }_{\theta }(y|x)}\left[r_{\varphi }(x,y)\right]-\beta \cdot {\mathbb{D}}_{KL}\left[{\pi }_{\theta }(y|x)||{\pi }_{ref}(y|x)\right]& (3)\end{array}$$

Inspecting this objective, we see that we are trying to tune ${\pi }_{\theta }$ such that it generates answers that maximize the learned reward function $r_{\varphi }$, while at the same time ensuring that we do not deviate too far from the original reference model. $\beta$ is a hyperparameter controlling the degree of deviation. This penalty constraint serves to:

1. Ensure that we do not drift too far from the (x, y) distribution on which the reward model is accurate
2. Ensure that we maintain generational diversity and not just collapse into a single high-reward answer for a given prompt

Objective (3) may not be directly optimized, because we need to generate $y\sim {\pi }_{\theta }(y|x)$ at each step from the current policy (not sure if I fully understand this). Hence typically this is optimized using reinforcement learning using PPO.

Direct Preference Optimization

The RLHF process in general is unstable, requires more memory / computation and requires tricks to make it work. Hence the authors of DPO set out to create an optimization procedure that:

1. Avoids fitting an explicit, standalone reward model
2. Avoids using reinforcement learning

DPO starts off with the KL-constrained reward maximization objective from Equation (3) above. The first step is to show that the optimal policy for this objective is of the following form for an arbitrary reward model $r(x,y)$:

$$\begin{array}{rlr}{\pi }_r(y|x)& =\frac{1}{Z(x)}{\pi }_{ref}(y|x)exp\left(\frac{1}{\beta }r(x,y)\right)& (4)\\ \text{where partition fn}Z(x)& =\sum _y{\pi }_{ref}(y|x)exp\left(\frac{1}{\beta }r(x,y)\right)& \end{array}$$

The derivation for Equation (4) is as follows. For a given reward function $r(x,y)$:

$$\begin{array}{rl}& \underset{\pi }{\max }{\mathbb{E}}_{x\sim \mathcal{D},y\sim \pi (y|x)}\left[r(x,y)\right]-\beta \cdot {\mathbb{D}}_{KL}\left[\pi (y|x)||{\pi }_{ref}(y|x)\right]\\ & =\underset{\pi }{\max }{\mathbb{E}}_{x\sim \mathcal{D},y\sim \pi (y|x)}\left[r(x,y)-\beta \log \frac{\pi (y|x)}{{\pi }_{ref}(y|x)}\right]\\ & =\underset{\pi }{\min }{\mathbb{E}}_{x\sim \mathcal{D},y\sim \pi (y|x)}\left[\log \frac{\pi (y|x)}{{\pi }_{ref}(y|x)}-\frac{1}{\beta }r(x,y)\right]\\ & =\underset{\pi }{\min }{\mathbb{E}}_{x\sim \mathcal{D},y\sim \pi (y|x)}\left[\log \frac{\pi (y|x)}{\frac{1}{Z(x)}{\pi }_{ref}(y|x)\exp \left(\frac{1}{\beta }r(x,y)\right)}-\log Z(x)\right]\\ \text{where }Z(x)& =\sum _y{\pi }_{ref}(y|x)exp\left(\frac{1}{\beta }r(x,y)\right)\end{array}$$

For line 2 above, recall that ${\mathbb{D}}_{KL}(P||Q)$ is the expected value of $\log P(Y)-\log Q(Y)$ if the random variable $Y$ is drawn from $P$. Since the outer expectation is over draws from $\pi$, we can breakdown the KL-divergence by bringing the log difference into the expectation. Line 3 simply divides by $-\beta$ and flips max to min. Line 4 uses $\log \exp$ to bring the reward term into the denominator of the left term, then introduces an arbitrary $Z(x)$. Note that the two $Z(x)$ can be cancelled out if we brought them together, but we will be using them later on.

Now let us define the optimal policy ${\pi }^{\ast }(y|x)=\frac{1}{Z(x)}{\pi }_{ref}(y|x)\exp \left(\frac{1}{\beta }r(x,y)\right)$. We will need to prove that ${\pi }^{\ast }$ is indeed optimal. Note that ${\pi }^{\ast }$ is a valid probability distribution as:

• ${\pi }^{\ast }(y|x)\ge 0\forall y$; and
• ${\sum }_y{\pi }^{\ast }(y|x)=1$, since the denominator is just the sum over $y$ of the numerator

Since $Z(x)$ is not a function of $y$, we can sub in ${\pi }^{\ast }$ and take $\log Z(x)$ out. The left term becomes a KL-divergence between $\pi$ which we are optimizing over and the optimal policy ${\pi }^{\ast }$.

$$\begin{array}{rl}& \underset{\pi }{\min }{\mathbb{E}}_{x\sim \mathcal{D}}\left[{\mathbb{E}}_{y\sim \pi (y|x)}\left[\log \frac{\pi (y|x)}{{\pi }^{\ast }(y|x)}\right]-\log Z(x)\right]\\ & =\underset{\pi }{\min }{\mathbb{E}}_{x\sim \mathcal{D}}\left[{\mathbb{D}}_{KL}\left(\pi (y|x)||{\pi }^{\ast }(y|x)\right)-\log Z(x)\right]\end{array}$$

Finally, note that $Z(x)$ does not depend on $\pi$, so we only need to consider the KL-divergence term. Gibb's inequality tells us that KL-divergence is minimized at $0$ if and only if the two distributions are identical. This completes our derivation of (4) by showing that ${\pi }^{\ast }$ is indeed the optimal policy.

Now that we have completed the derivation, let's consider what Equation (4) is saying. It tells us that we have an analytical solution for the policy ${\pi }_r$ that optimizes (3), and that it can be expressed in terms of ${\pi }_{ref}$ (which we already have) and a given reward function $r(x,y)$.

Since we previously learned a reward model $r_{\varphi }$, we could simply plug that into (4) to get our optimal policy. Specifically, for a new input prompt $x$, we can compute ${\pi }_{ref}(y|x)\exp \left(\frac{1}{\beta }{r}_{\varphi }(x,y)\right)$ for all possible values of $y$ and pick the best $y$. We can ignore $Z(x)$ since it is fixed for a given prompt $x$. Intuitively, the new model scales the probability of high reward answers up with an exponential multiplier, and the degree of scaling is controlled by $\beta$. However, this is not computationally practical as we need to evaluate over a very large space of $y|x$ (i.e. all possible answers for a given prompt $x$).

Hence, we want to find a form of the optimal policy which does not involve the partition function $Z(x)$ nor the reward model $r(x,y)$. We start by re-arranging Equation (4), taking log on both sides and re-arranging:

$$\begin{array}{rlr}& r(x,y)=\beta \log \frac{{\pi }_r(y|x)}{{\pi }_{ref}(y|x)}+\beta \log Z(x)& (5)\end{array}$$

Since Equation (5) holds for any arbitrary reward model, we can use the optimal (unknown) human reward model back in Equation (1), $r:=r^{\ast }$. Also, note that ${\pi }_r$ in Equation (5) refers to the optimal policy under $r$, so since we are using the optimal reward $r^{\ast }$, we can call this optimal policy ${\pi }^{\ast }$ as well. Now we plug Equation (5) back into Equation (1).

$$\begin{array}{rlr}{p}^{\ast }(y_1▹y_2|x)& =\frac{\exp \left(r^{\ast }(x,y_1)\right)}{\exp \left(r^{\ast }(x,y_1)\right)+\exp \left(r^{\ast }(x,y_2)\right)}& \\ & =\frac{\exp \left(\beta \log \frac{{\pi }^{\ast }(y_1|x)}{{\pi }_{ref}(y_1|x)}+\beta \log Z(x)\right)}{\exp \left(\beta \log \frac{{\pi }^{\ast }(y_1|x)}{{\pi }_{ref}(y_1|x)}+\beta \log Z(x)\right)+\exp \left(\beta \log \frac{{\pi }^{\ast }(y_2|x)}{{\pi }_{ref}(y_2|x)}+\beta \log Z(x)\right)}& \\ & =\frac{1}{1+\exp \left(\beta \log \frac{{\pi }^{\ast }(y_2|x)}{{\pi }_{ref}(y_2|x)}-\beta \log \frac{{\pi }^{\ast }(y_1|x)}{{\pi }_{ref}(y_1|x)}\right)}& \\ & =\sigma \left(\beta \log \frac{{\pi }^{\ast }(y_2|x)}{{\pi }_{ref}(y_2|x)}-\beta \log \frac{{\pi }^{\ast }(y_1|x)}{{\pi }_{ref}(y_1|x)}\right)& (6)\end{array}$$

The derivation of the above is quite simple, we just need to note that it is of the form $exp(x)/(\exp (x)+\exp (y))$, and use the expression of the BT-model as a sigmoid function. The great thing is that the pesky partition function $Z(x)$ cancels out because the BT-model simply ends up with the difference between two scores / rewards.

Equation (6) looks simple but is quite remarkable when we compare it to Equation (1), because we now have the probability of the human preference data in terms of just the reference policy ${\pi }_{ref}$ and the optimal policy ${\pi }^{\ast }$, without the reward model at all! Of course, the reward model is implicitly embedded in the equation. If we stare at Equation (6), we see the implicit reward function is $r(x,y)=\beta \log \frac{{\pi }^{\ast }(y|x)}{{\pi }_{ref}(y|x)}$.

The benefit of this new formulation is that we can write the maximum likelihood objective in terms of the optimal policy. We do not know the optimal policy, but we can parametrize it as ${\pi }_{\theta }$ and use the maximum likelihood objective to train the model. Our new policy objective becomes:

$$\begin{array}{rlr}& {\mathcal{L}}_{DPO}({\pi }_{\theta };{\pi }_{ref})={\mathbb{E}}_{(x,y_w,y_l)\sim \mathcal{D}}\left[\log \sigma \left(\beta \log \frac{{\pi }_{\theta }(y_w|x)}{{\pi }_{ref}(y_w|x)}-\beta \log \frac{{\pi }_{\theta }(y_l|x)}{{\pi }_{ref}(y_l|x)}\right)\right]& (7)\end{array}$$

Application to Dual Encoder Retrieval

The DPO framework offers a way to fine-tune a policy according to human preferences, whilst ensuring stability against a reference model. In the original formulation, $\pi (y|x)$ represents the probability of generative model generating $y$ given an input prompt $x$.

A related problem is that of fine-tuning embedding models for search retrieval, in what is known as the dual encoder framework. In this case, we also have preference data in the form of triplets (query, positive_passage, negative_passage), and we wish to fine-tune embeddings such that the dot-product or cosine similarity between (query, positive_passage) is high whilst that of (query, negative_passage) is low. In this formulation, we could let the policy $\pi (q,p)$ represent the normalized probability of a relevant match between query and passage. We can then borrow the framework of DPO to fine-tune our embeddings. The benefit of DPO compared to typical optimization objectives for dual encoder is the stability of the policy against the reference policy, which hopefully is a good form of regularization even when preference data is limited.

Specifically, we define a dual encoder policy $\pi (q,p)$, where $q$ represents a query and $p$ represents a passage, like so:

1. Encode $q$ into a vector $\in {\mathbb{R}}^d$ using a BERT model
2. Encode $p$ into a vector $\in {\mathbb{R}}^d$ using a BERT model (can be same model as the query encoder)
3. Take the dot product
4. Run the dot product through a sigmoid layer to convert it into a probability

We can then apply this method to a reference encoder model and call it ${\pi }_{ref}(q,p)$, and optimize another encoder model against this reference model. Will need to conduct experiments to see if this method offers gains against the typical dual encoder objectives, such as Karpuhkhin 2020.

References

• Ouyang 2022 - InstructGPT
• Rafailov 2023 - Direct Preference Optimization
• Huggingface RLHF Blog