Aligning LLMs to Human Preference using DPO

Training LLMs from scratch is hard, really hard. LLM training is not only hard but also very expensive. For instance, it took over 100 million USD to train ChatGPT and around 20 million USD to train Llama. This begs the question of whether we can properly align the LLM to human preference or not because we can’t keep retraining stuff again and again.

Direct Preference Optimization (DPO) in AI refers to a method where an AI system learns to make decisions or predictions based on directly optimizing user preferences. This contrasts with traditional methods where a model is trained to minimize a predefined loss function, which may or may not align perfectly with actual user preferences.

So, without further ado, let’s delve deeper into how we align our LLMs to human preferences.

LLM Training Pipeline

Let’s quickly describe the LLM training pipeline and then we will understand how DPO helps us.

Step 1: Pre-training

Data Collection and Diversity: A vast dataset is gathered from the internet, ensuring a wide spectrum of text sources. The diversity of data is crucial for the model to understand various language styles, topics, and nuances.
Data Cleaning and Preprocessing: This involves removing irrelevant content, formatting issues, and noise. It’s a delicate process, as it needs to balance cleaning data while preserving meaningful context.
Tokenization: The text is broken down into manageable units (words, subwords) using techniques like Byte-Pair Encoding or WordPiece. This step is critical for handling the nuances of language, including morphology and syntax.
Use of Transformer Architecture: Transformers are chosen for their effectiveness in handling sequential data and their ability to capture long-range dependencies in text.
Training for Next-Word Prediction: The model is trained to predict the next word in a sequence. This step is foundational for the model to learn language structure and generate coherent text. However, at this stage, the model lacks the ability to understand context or user intentions deeply.

Output After Step 1:

The model can generate text but may produce outputs that are contextually irrelevant or inappropriate, as it has yet to learn specific user interaction patterns or instructions.

Step 2: Supervised Fine-Tuning/Instruction Tuning

Learning from Targeted Inputs and Responses: The model is now trained with pairs of user inputs and appropriate responses. This training is crucial for the model to learn the context and relevance of responses.
Understanding Instructions and Knowledge Retrieval: The model starts to comprehend the meaning of instructions and learns how to retrieve and utilize its pre-trained knowledge effectively.

Output After Step 2:

The model now can understand and respond to questions more accurately. It begins to establish a relation between questions and contextually appropriate responses.

Step 3: Reinforcement Learning from Human Feedback (RHFL)

Aligning with Human Preferences (Helpful, Honest, Harmless): The focus here is on refining the model’s outputs to align with ethical and practical human standards, encompassing helpfulness, honesty, and harmlessness.
Training Reward Model Using Human Feedback: Multiple outputs are generated for the same prompt, and human labelers rank them. This ranking data trains another neural network, the reward model, which understands human preferences in content.
Replacing Humans with Reward Model for Scalability: The reward model, once trained, takes over the role of humans in providing feedback, allowing for large-scale, efficient fine-tuning of the LLM.

Purpose of RHFL:

This phase is critical for refining the model’s behavior. It ensures the model not only provides accurate and relevant responses but also does so in a manner that is ethically aligned with human values, avoiding misleading, harmful, or dangerous content.

In Summary:

Pre-training lays the foundation for understanding language.
Supervised Fine-Tuning/Instruction Tuning teaches the model to respond appropriately to specific instructions.
RHFL aligns the model’s responses with ethical and human-centric criteria, enhancing its reliability and safety in real-world applications.

End-to-end LLM training Pipeline (Img Src)

Need for DPO

Direct Preference Optimization (DPO) in AI is a crucial approach for several reasons, particularly in contexts where traditional models may not adequately align with complex and nuanced human preferences.

Predefined metrics (like accuracy, precision, recall) are proxies for what we think users want. They often fail to capture all aspects of user satisfaction, especially in complex decision-making scenarios. Different applications may require different types of optimizations. DPO enables customization of the model’s objectives to fit specific user needs and scenarios.

DPO, by directly incorporating user feedback, aligns the model’s outputs more closely with what users actually find valuable or preferable. Human preferences can change over time and vary across different contexts. DPO allows models to dynamically adapt to these evolving preferences, maintaining their relevance and effectiveness.

Applications like recommendation systems, personalized healthcare, or customer service, understanding and adapting to individual user preferences can significantly enhance the effectiveness of the AI system. And a method like DPO can potentially reduce biases in AI outputs. By learning from diverse user feedback, models can be more equitable and less likely to perpetuate existing biases present in the training data.

DPO is super important because it can be used to align models towards safe response. We can easily imagine that in the pretraining stage, the model might have learned about sexual behaviors, chemical warfare, nuclear technologies, etc. DPO can seriously help us to not let these very powerful models wander into dangerous territories.

Understanding RLHF (Reward Model + PPO)

To understand this fully, let’s talk about the total number of models required for end-to-end RLHF. For this, we end up copying our base model four times.

Pre-trained Base LLM
Supervised Fine Tuned (SFT) LLM
Reward Model (LLM, but modified to be a reward model)
PPO Optimized Language Model (Final LLM aligned to preferences)

All four models and what is the difference between them RLHF (Reward model+PPO) (Img Src)

Let’s understand the Third and Fourth models in more detail, I’m assuming the above explanation clearly explained at least base LLM and SFT.

The Reward Model is simply the SFT model copied, and modified so the un-embedding layer (the final layers un-embedding layers convert the internal representation of text back into human-readable text) is removed, and a scalar reward head is added. This model takes in a prompt and outputs a scalar reward. Instead of outputting text, this layer outputs a scalar value — a single number that represents the ‘reward’ or ‘score’ for a given input. The Reward Model is trained using human feedback. This feedback could be in the form of rankings or ratings of the AI-generated responses.

The Reward Model serves as a crucial component in RLHF. It provides a quantifiable measure of the quality of responses generated by the AI. During further training of the AI, this reward signal is used to reinforce good responses (those that get higher rewards) and discourage bad ones (those that get lower rewards).

The Proximal Policy Optimization (PPO) Model is also a copy of the SFT model, where the parameters will be updated to align with human preferences. PPO, a reinforcement learning algorithm, is used to fine-tune the AI model (originally the SFT model). It optimizes the ‘policy’, which in this context refers to the strategy of generating text responses in accordance with human values. The goal of PPO training is to maximize the rewards received from the Reward Model.

In the PPO training process, the AI model generates responses to various prompts, which are then evaluated by the Reward Model. This Reward Model assigns a scalar value to each response, reflecting how well it aligns with human preferences. PPO, focusing on stability, updates the AI model’s parameters to ensure gradual improvement without drastic policy changes. The AI model learns to refine its text generation strategy, aiming to produce responses that garner higher rewards. This iterative cycle of response generation, evaluation by the Reward Model, and parameter updating via PPO gradually aligns the AI model’s outputs with human preferences. Over time, this leads to a model that not only understands language but also resonates more closely with human values and expectations.

I hope now a few more things are clear, I know you might have a few more questions on how ChatGPT is trained, please put them in the comment, otherwise, this will go too big.

Decoding Direct Preference Optimization (DPO)

DPO simplifies this process by removing the reward model all together.

They instead feed a comparison dataset directly to the final model by modifying the loss function of the final step. This is a great example of removing dependencies, and how keeping it simple can lead to better results.

Let’s look at how it works in practice

Sentiment Generation

Suppose the task is to align our model to only give positive reviews. The prompt x is a prefix of a movie review from the IMDB dataset and the policy must generate a completion Y with positive sentiment.

x: Iron man was ...

y_w: the best movie of all time because ...
y_l: the worse marvel movie I think I have ever seen ...

First, we train an SFT that generates movies-like data on the trainset of IMDB movie review data. Then we use a pre-trained sentiment classifier to generate the labeled preference dataset.

Based on the classifier’s assessment, label each completion as ‘preferred’ if it has a positive sentiment and ‘non-preferred’ if it has a negative sentiment. For example, y_w (the worst movie) would be labeled as non-preferred, and y_l (the best movie) as preferred.

Then we Fine-tune the model’s parameters to maximize the likelihood of generating responses that are classified as positive by the sentiment classifier. To achieve this goal we use a ranking loss function during training, where the model is penalized more for preferring negative completions over positive ones.

We Continuously generate new completions, apply the sentiment classifier, and use this data to further fine-tune the model. This iterative process helps the model become more adept at generating positive reviews. The model learns to adjust its language generation in a way that is more likely to produce positively sentimented completions.

The above graph shows how the model generates a bunch of positive sentiment sentences (high reward) across a variety of KL divergence scores. They ran a suite of 22 different hyperparameter settings to generate this data for different types of models. (Section 6.1)

The takeaway is that DPO provides the highest alignment to preferences across completions that vary in probability distribution from the original reference model.

Summarization

In another task of Summarization, they take a forum post from Reddit as the input X, then the policy must generate a summary Y of the main points in the post.

They use the Reddit TL;DR summarization dataset for this experiment along with human preferences already gathered.

Again, first, perform supervised fine-tuning on the Reddit responses to learn to generate summaries. Then they used Anthropic’s Helpful and Harmless dialogue dataset containing 170k dialogues between a human and an assistant to align the model with human preferences.

Maths behind DPO

Let’s look at some scary-looking equation and try to understand what’s going on here. The paper states that we can replace RL with “a simple classification loss”. Let’s see what they mean by that.

Loss is the measure of how well our model is doing, given the data. In the end, if we can minimize the loss, we are winning, our model is training in the right way.

Let’s start with a simplified version of the equation that has a winner (W) and a loser (L).

Loss = Winner (W) — Loser (L)

In our case, the winner will be the text completion that we labeled as 👍positive and the loser will be the text completion we labeled as 👎negative.

y_w: the best movie of all time because ...👍
y_l: the worse marvel movie I think I have ever seen ..👎

Now because we are minimizing the loss function we put a negative in front of it.

Loss = — {Winner (W) — Loser (L)}

Now the above equation looks a bit different but this is what it is doing precisely. I know we have logs and beta’s and sigmas and pi’s but in the end, we are trying to have a high score for W and a low score for L.

At a high level, this is a simple binary cross-entropy objective, and similar to what you learn in machine learning 101 when learning about classification loss functions.

But what are these pi’s, y_s, and x_s doing in here?

In this expression x is the prompt, y_w is the completion, and pi_theta is the language model. It can be read as the probability of the winning completion, given a prompt.

x: Iron man was ...

👍y_w: the (0.7) best (0.9) movie (0.3) of (0.7) all (0.5) time (0.54) because (0.2) ...
👎y_l: the (0.7) worse (0.4) marvel (0.15) movie (0.4) I (0.3) think (0.26) I (0.13) ...

basically the 0.7 + 0.9 + 0.3 + 0.7 + 0.5 + 0.54 give higher probability

In practice we use log probabilities for numerical stability, but we can think of the probability of the entire sequence as the sum of all its parts.

On the other side, we have y_l and the probability of the losing completion given the same prompt.

Hopefully, you see where we are going, it’s basically the probabilities of the winners subtracted losers.

What about rest of the terms?

pi_theta is the simply the language model we are training. pi_ref is the pre-trained language model we have already learned.

The reason we divide by the “reference model” is to make sure your DPO-trained language model does not deviate too far from the knowledge of the original model. The pre-trained model has already learned a lot about the world, and we do not want to overfit to just the preference data.

We know that the drift between two distributions is called the KL Divergence. This equation uses the KL Divergence so that the probability of y_w (winner) given the input x is high from the original language model, we want it to also be high from the new language model.

Overall we have the probability of the winning completion, minus the probability of the losing completion. Then we divide by the probability of the pre-trained reference language model in the denominator as regularization. Beta & Sigmoid are just other scaling factors used for numerical stability.

Now that we understand the loss equation better, the full DPO pipeline is relatively straightforward.

Sample two completions from our reference language model given a prompt x.
Optimize our new language model through backprop to pi_theta to minimize our loss.

The model gets rewarded if the completion y_w has a higher probability than y_l .
The model gets rewarded if the completions of W and L are close to the pre-trained model pi_ref's completions.

Conclusion

DPO is notable for efficiently enhancing an SFT model to better align with specific preferences, presenting an alternative to RLHF as introduced in OpenAI’s InstructGPT. Its effectiveness in academia, especially with limited resources, is impressive. While initially tested on models up to 6B parameters, DPO has since been successfully applied in various open source models. However, its scalability to larger models and generalization to out-of-distribution data are still open questions. Current experiments are exploring DPO’s potential in self-play scenarios, where a language model evaluates and improves upon its outputs, but we still need to evaluate it further to truly understand its impact.

Writing such articles is very time-consuming; show some love and respect by clapping and sharing the article. Happy learning ❤

Please don’t forget to subscribe to AIGuys Digest Newsletter

References

https://arxiv.org/abs/2305.18290