Aligning LLMs with Direct Preference Optimization

Sam.

Hey everyone, My name is Dinah Chan Morgan and I'm part of the community team here@deeplearning AI. Today we have really special guests to talk to us about Direct Preference Optimization and really excited to dive in for everything. Just for everyone's information, the session will be recorded and the slides and notebooks will be available after the event. For any questions, please fill out the slido link. We've dropped it in the chat and you can vote on the questions you would like answered the most. For today's workshop we'll focus on the landscape of chatbots which has dramatically changed in the past year, which has been fueled by advancements of large language models like ChatGPT and Llama. While traditional supervised fine tuning has dominated the field, recent research highlights the potential of aligning language models with human preferences for increased helpfulness and safety. This workshop delves into Direct Preference Optimization dpo, a powerful technique used to train Zephyr, the current state of the art for 7B parameter chat models. For all of you attending, you'll gain practical knowledge on fine tuning LLMs for chat, understanding the theory and application of DPO with Hugging Face tools, and learn key metrics for evaluating chat model performance For Our event partner today is Hugging Face. Hugging Face is an AI company I'm sure all of you know that, specializing in natural language processing and machine learning and is known for its open source contributions and collaborative approach to AI research and development. The company is famous for developing the Transformers library which offers a wide range of pre trained models and tools for a variety of NLP tasks, making it easier for researchers and developers to implement state of the art AI solutions. Hugging Face also fosters a vibrant community for AI enthusiasts and professionals, providing a platform for sharing models, data sets and research which significantly contributes to the advancements of AI technology. So for today I want to introduce our first speaker. Hey Louis.

Hey, how's it going?

Happy to have you here. Lewis Tunstall is a machine learning engineer at Hugging Face. His work lies at the intersection of of the open source and research teams. He's the co author of the bestseller NLP with Transformers book and has previously built machine learning powered applications for startups and enterprises, the domains of natural language processing, topological data analysis and time series. He holds a PhD in theoretical physics, was a 2010 Fulbright Scholar and has held research positions in Australia, the US and Switzerland. His current work focuses on building tools and recipes to align language models with human and AI preferences through techniques into reinforcement learning. Happy to have you here Louis.

Thank you for having me.

Diana Absolutely. And for our second speaker we have Ed Beeching. Ed is.

Hi.

Has been a research scientist at Hugging Face for the last two years, focusing on building open tooling for LLM fine tuning, implementing state of the art LLM alignment algorithms, and building preference data sets. Ed is the creator of the open LLM leaderboard, co author of Zephyr and a maintainer of numerous open source ML libraries. His PhD was focused on deep reinforcement learning approaches for planning and navigation and robotics. His current research has predominantly embodied learning and LLM alignment. We're so happy to have you here today. Are you guys excited to dive in?

Yes, can't wait. Absolutely.

Well, take it away.

Great, thank you so much. All right, so we're going to talk today about this idea of alignment in LLMs with, with the specific lens of an algorithm called direct preference optimization. And so for those of you who are hearing these words for the first time, this was a kind of big breakthrough from a set of Stanford researchers last year who showed that a lot of the complexity that goes into the kind of conventional alignment techniques that were pioneered by OpenAI and anthropic can actually be addressed in a much simpler and memory efficient way. And as you'll see throughout this talk, there are some sort of technical details that are useful to understand and we're going to try and do our best to kind of demystify a lot of the sort of mathematics that goes into these types of algorithms and also show you some practical examples of how you actually train language models using various tools in the hunting pace ecosystem. So to get started, I thought it's maybe nice to ask ourselves like, you know, why would we want to align a language model in the first place and maybe ask a bit of a philosophical question, what is alignment? So I think most people here probably have some familiarity with language models. And you know, that the first step when training a language model is a process called pre training, where you essentially feed a transformer neural network, typically a large amount of data from the Internet, and you get the model to basically predict the next token in the sequence given the previous tokens it's seen. And so this ends up being like a very, very powerful kind of autocomplete model which you can do tricks like few shot prompting to generate certain answers. But as an autocompleter, if you ask it a question that is something like information seeking, like, you know, is pineapple on pizza a crime? Then this base language model will just autocomplete what it thinks are the next most probable words to that. And in this Particular example it might say, well you know, this is one of many questions that will be answered at the pizza party because that could completion is from some sort of text that it saw during the pre training. So these base language models, they have their own use, but they typically need some additional tuning in order to be able to be used as chat models. So the second step is typically called supervised fine tuning. And what you do in this step is you provide the model with a few thousand examples where you've got a user asking a question and either a human written response or some LLM generated response that provides an answer to question. So in this particular example, maybe the human annotator really hates pineapple on their pizza and they wrote in the annotations that, you know, this is really bad. And so the model when you ask the question says, well, you know, yes, adding pineapple and pizza is a crime under the Geneva Convention. And so this SFT model has essentially learned from the kind of fine tuning data and also from the pre training data that there are some biases present there. And those biases can kind of be expressed in this particular way. Now the thing is, you might be a member of a company like you know, Domino's or Pizza Hut and you actually, you know, maybe your, your, your, your customers, they actually do like pineapple on pizza. And so one way you can kind of sort of adjust your, your language model to encode those kind of preferences from your customers is, is to essentially provide your language model with a bunch of questions. You can generate a few answers, in this case maybe yes or no. And then you can get your like customers or your human labelers to annotate this. And then this will give you a way of encoding some alternative preferences that are not directly present in the pre training of the supervised data. And essentially what you're trying to do is you're trying to like learn or teach the model how you can model the probability that the kind of preferred response is more likely than the dispreferred one. And you know, if you do this process correctly, you'll end up now adjusting the probabilities that the model will emit an answer that is more aligned with your human preferences typically. So in this case, you know, maybe you say no, adding pineapple is not a crime, it's a matter of personal preference or taste. So the thing to mention here is that alignment nowadays I think is often associated with things like censorship. So most people who have interacted with ChatGPT at some point get a refusal. As a language model, I can't do xyz, but it's really a broader concept which is that as humans we typically have some preferences. These are things that are diverse and depend on our culture and upbringing. But those preferences can in principle be modeled through neural networks. And then we can adjust our language models to behave kind of more in accordance with what our end users may want. So before we dive into dpo, I think it's good to sort of show the original approach that OpenAI pioneered. So they called it RLHF or Reinforcement Learning from Human Feedback. And the basic idea they had was this very similar recipe. You do supervised fine tuning. In this case they had humans basically annotating and creating these examples, these high quality examples of questions and answers. And then they fine tuned a variety of kind of GPT models on that data. And last year when sort of chatgpt had just come out, the open source community didn't really have many high quality data sets that you could use for supervised fine tuning. But this has really changed. And so now there are some really good examples. You have one called Open Hermes from Technium or Dolphin from Eric Hartford. And these are data sets that you can use today to kind of convert any sort of base language model into a chat model. And these models, if you train on these data sets, you'll typically get a very, very good chat model that you can converse with in certain ways. So what OpenAI did next is they said, okay, we've got this chat model. Now we need a way to sort of model the sort of human preferences from annotators. And so the way they did that is they basically took a collection of prompts, they fed it to the chat model and the chat model generates a variety of answers. And then those answers are rated by labelers from best to worst. So this gives you a kind of ranking of which answers are preferred over over the others. And then essentially you train a new type of model called a reward model, which essentially learns how to predict, given some response from a model, what is the kind of score or the preference score associated with what the human labelers gave. And if you do this, you've now got essentially kind of like a regression model that you can use to basically rate any sort of outputs. And then this gives you a way of sort of mapping those human labels into the preferences model. So we've got two models, we've got a chat model and a reward model. And just to mention that we have also now many public data sets that can do this, but you've got the chat model and the Reward model. And the final step is to somehow combine those two to create a sort of new chat model that has learned how to encode the kind of signal from that reward model and the way that OpenAI did it. And also Anthropic has been using reinforcement learning. So the basic idea here is that you provide your chat model with a prompt, you will generate an output. You can then rate or score that output using your reward model, and then you can use that score to basically update the weights of your original chat model to be kind of more aligned or closer in those preferences. And so if your human annotators are rating things in certain ways, that kind of behavior, those preferences will then be essentially encoded into your updated chat model. And the result of this will be something that is like ChatGPT or Claude from Anthropic. So mathematically, I don't want to make this too scary, but the basic idea is that you're trying to sort of maximize the rewards that you, that the model can get. So you give it all these prompts, you score the outputs, and you want the model to kind of get better at producing outputs that get high rewards. But if you just only use a reward model, it turns out there's a. There's a thing called reward hacking where effectively the chat model can learn that if it puts certain words or tokens that have very high reward, it can just keep doing this kind of ad nauseam. And you just get a kind of garbage model that will just output backslashes or emojis or something, because it turns out that might be something that is rewarded highly. So what OpenAI introduced was an extra term which they call the KL penalty. And basically what this does is this measures kind of the distance of your original chat model versus the model that you're currently optimizing. And there's a parameter called beta, which controls how closely you want those two models to be next to each other. And if you take these two things together and you optimize it with a reinforcement learning algorithm like ppo, you get something which is typically a very performant and aligned model. Now, there are certain challenges associated with this. The first one is that anyone who's done reinforcement learning knows that this is typically very unstable. There are many hyper parameters you have to tune. And one of our colleagues at Hugging Face called Costa, he's done this heroic effort to kind of re implement the, the original RLHF pipeline from OpenAI from scratch. And there were crazy things like changing one of the parameters in the Atom optimizer. Had a big difference on the outcome. So that's like a very kind of, let's say, finicky beast to get right. And the other thing that's also kind of a challenge is that you've basically got two models that you have to deal with. So you've got the chat model and the model that you're trying to sort of optimize and become more aligned. And then you've got the reward model as well. So you've got three large language models that you've got to juggle. And nowadays the compute hardware is getting quite good. But if you want to do anything at very large scales of like 30 billion parameters more, you're going to need a lot of compute to be able to juggle this. So this makes the whole process quite challenging, both from an engineering perspective, but also from a stability perspective. So the kind of idea behind direct preference optimization is to say, well, maybe we don't need reinforcement learning. And the question is, okay, if we don't use reinforcement learning, what are the alternatives? So the basic observation they made was that instead of trying to kind of maximize this combination of a reward model or a signal from reward model with this kind of KO penalty, the idea was maybe we can just change the kind of objective in a certain way that directly encodes or implicitly models the kind of score that a reward model would learn. So you effectively don't have a reward model anymore. You just have your original language model. But your language model, through the process of this optimization, learns how to model those preferences directly. So just to go through this equation kind of quickly, um, the, the basic idea is we're going to take a prompt and then we're going to have preferences, binary preferences of a good and a bad answer. So in this case, you might have, your prompt is like, is pineapple on pizza a crime? And then the good response is no. The bad response is yes. And then what we're going to do is we're going to basically compute two things. We're going to try and compute basically the kind of log probabilities that the model is able to kind of correctly predict the chosen response. So here what we do is we take a ratio of two terms. The term on the top is the model that we're optimizing kind of throughout the optimization process. And then we have a reference model underneath it, which kind of normalizes the probabilities so that we don't drift too far away from our reference in the same way that this KL penalty worked previously. And then we have Another term which is very similar. But now, instead of trying to compute the log probabilities of predicting the desired response, we try to predict the log probabilities of the sort of undesired response. And so if you just take these two terms and you combine them together, you're basically looking at trying to kind of maximize essentially the difference of two terms, where one is like the chosen response, one is the rejected response. And you're trying to make the model get much, much better at predicting the chosen response, which should hopefully then align with your human preferences. And the cool thing about this algorithm is that it's differentiable, which means that we can just use backpropagation to optimize the model. And this is also why we don't need reinforcement learning, because the previous objective that we showed before isn't differentiable. So just to walk through the algorithm, you basically take your prompt, you have these two responses of good and bad, you feed them through two models, you have your reference model and the model that you're optimizing, and then you do backprop, and then that's basically it. And in terms of Pytorch code, you can see here that it's roughly 10 lines of code. So when I kind of first saw this, this was like, first surprising that something so simple would work. But when we actually tested it ourselves, we found that compared to our reinforcement learning pipeline, it was much, much faster to get models that would converge and be much better aligned than previously. So just to sort of linger a little bit on sort of maybe the most important kind of idea here, there's basically this kind of parameter called beta, and this is the main hyper parameter that you need to kind of tune when you're doing dpo. And Ed will tell us a little bit later about some other parameters that we found are important. But essentially, when we are computing, like an update or gradient update, this is essentially when we're trying to take a direction in the optimization of the model. There's effectively kind of like three terms at play here. So the first term is essentially there's a kind of weighting factor that goes over the loss. And this weighting factor is basically measuring the difference between the kind of reward or the implicit reward of getting the incorrect response versus the correct one. And so if your model is like, continually getting, like, the kind of incorrect answer was continually predicting the incorrect label, then this weighting factor will become large because it's a sigmoid. So it will go up towards one, and then that will kind of penalize the model. To basically have a higher penalty for getting the incorrect estimates. And then you've got these other two terms which are essentially the difference of just increasing the likelihood that your model predicts the kind of chosen answer, and you're trying to decrease the likelihood of the rejected answer. And you've got a plot here, I've shown of some training runs where you can see that effectively what happens as we increase the beta parameter is we're effectively kind of pushing the model to kind of trains in such a way that the difference between the chosen and the rejected answers becomes larger and larger. And so this can be important because you may have quite noisy data that has kind of fuzzy labels between chosen and rejected. Maybe the quality of them is a bit mixed. And by tuning this beta parameter, you can kind of control how much you want model to focus on the chosen responses versus the rejected ones. And, you know, typically that you have to run a couple of experiments at different values to try and find which is, which is the best one. But the basic idea is that beta fundamentally controls the kind of target margin which is the difference of these kind of chosen and rejected rewards. And by kind of increasing that, you can get quite different behavior. And, you know, if math isn't your thing, then Tom Goldstein, who's a kind of grand master of memes, he showed this is roughly what's going on in dpo. Basically, you know, you just kind of do gradient descent on the good stuff and you do gradient descent on the bad stuff. And this is why, to some extent, DPO has been so popular, because it's a very simple algorithm, and it means you can basically train models without actually having to train the reward model independently, because you're modeling everything implicitly. So here's a couple of examples I'll walk through. So one of the ones we trained at, hugging face, it's called Zephyr. And what we were interested in looking at with Zephyr is that in the DPO paper, the authors had explored comparisons of DPO and PPO using kind of small models, so small being like around 2 to 3 billion parameters. And they had done it using kind of like academic data sets. And in particular, they had used a data set for chat that, you know, wasn't really optimized for performance. So the idea that we thought was, like, well, we know there's this trend in the community to kind of create synthetic data sets that are around essentially using ChatGPT models, conversing with each other or rating each other's responses. If we use those data sets together, with dpo, maybe we can get rid of human feedback kind of altogether and still produce a good performing model. And so in this example, what we did is we took a popular data set for supervised fine tuning called ultrachat, and we use that to train a chat model, the initial SFT model. And then we took another data set of synthetic feedback where you essentially have GPT4 rating responses across a diverse number of different models. And then we use those two data sets to basically do DPO using the Mistral 7B base model, which had only just come out a week or two before. And kind of one of the surprising things we found is that if you apply this process, you get a kind of 7B model that is like, you know, relatively competitive on, you know, various chat benchmarks like MT Bench, compared to much larger models like Llama 70B. Now the usual caveat here is that these benchmarks, they have certain biases. So for example, if you train models on GPT4 data, because GPT4 is acting as the evaluator, it tends to kind of favor its own, its own kind of style. But nevertheless, if you actually converse with the Zeppelin model, it's actually quite capable. And so this was one of the first public demonstrations that you can apply this recipe to larger models with kind of more capable data sets. And since then, the community has really taken this idea and pushed it in much further directions. So for example, some folks at Argilla, they did DPO on top of the Mistral model or the mixed roll model, sorry, when it came out, there's a guy called John Durbin who creates lots of very, very interesting DPO data sets and he used this to fine tune the Yi 34B model. And then of course, Mistral have been using DPO a lot to align their models. And now NOOS Research and Technium have been doing this themselves. So there's essentially, I would say DPO has become nowadays the kind of canonical alignment algorithm, at least in the open source community. And there's two main libraries that you can use to do this. So the one we have a hugging face is called TRL or Transformer Reinforcement Learning. But there's also Axolotl, which is very popular alternative. And Ed will show you later a little bit about how TRL works. So just to sort of wrap up this kind of quick overview, generally speaking, what you see in like the academic literature and machine learning is once there's a very popular idea, many, many people try to find ways to extend it and improve it. And DPO is no exception. And One of the sort of first improvements that was made was by researchers at DeepMind. They developed something called IPO or identity preference Optimization. And the kind of observation they made was that when you do dpo, there's a risk that you end up overfitting basically to the chosen response in your data set. And so if you want to prevent that from happening, you can add a regularization term in a similar way that we would do this. If we're doing any kind of regression, you basically add a weight. And they showed in their paper that this gave better results than dpo. And we at Hugging Face also did some experiments kind of comparing the two methods and we see that they're roughly comparable. There's a bit of trade off depending on the parameters. And then at the other level, there's a lot of recent work going on around this idea of iterative DPO or let's say online dpo. And there's a nice example from Snorkel where essentially what they did was they took a reward model from Allen AI and they used essentially the prompts from ultrafeedback to generate responses from their chat model. They use the reward model to basically rate and rank those prompts and then they retrain another DPO model and then they generate again on new prompts, rate and rank. And then you do this a set of times and you get a model that gets progressively better and better. And what's exciting about this is it's very close to how Llama 2 was trained. Llama 2 was trained in a sequence of progressively improved models. And it's also what we, as far as we understand, is how some of the anthropic models are trained. They're training this kind of online fashion where you're continually improving the model with new data. So I think this is quite an exciting direction. And one of the sort of other more interesting branches that we've seen is a new algorithm called KTO or Kahneman Tversky Optimization. And this is from researchers at Stanford and Contextual AI where they make the important point that if you want to collect preference data, you need to have essentially a prompt and two answers, where one answer is the desired one and the other one is kind of the counterfactual or the one one that you don't want. And if you try to do this with human feedback, it's very expensive to collect because you need annotators to interact with your chat model and pick at every step what is the preferred response. But you can actually, if you're clever like they are, you can actually define the loss in such a way that you only need to collect basically something that is deemed as good and something that is deemed as bad, and they don't have to be related to the same prompt. And so at HuggingFex, we're quite excited about this algorithm. We think it has a lot of promise and it's probably going to be the thing that you can hope to see from us soon, showing how it works at various scales. The paper itself is worth reading. They do a lot of great experiments and what seems to be missing currently is having some kind of good label data sets to run this at scale. But this is very, very exciting. So maybe it's a good time to take some questions. Ed.

Yes, yes, absolutely. So let's start off with the first one. Just going to how do you evaluate the quality of open source alignment data sets?

Yeah, that's a really good question. I would say in general it requires a bit of what we would call vibes testing. So you just have to look at the data yourself. And I think this is often one of the things that I always try to sort of tell people is like look at the data because there's often strange errors in the labels. And probably the most popular method today is to use another language model to do the ratings. So probably the most common example is using GPT4 to basically judge the quality of responses. And what typically we find is GPT4 is quite good on many kind of, let's say creative tasks or tasks that require maybe like high school level reasoning. But if you ask it to kind of judge the quality of responses involving maybe some in mathematics or some code, it can kind of hallucinate and make errors in judgment. And so this isn't to say GPT was bad, it's just that like any sort of hard task, even for humans, it's, it would be hard to come up with the right answer. But currently, today this seems to be the way that the community is rating quality.

Absolutely. For this question, how much data is ideal for this process? For the alignment?

Yeah, great question. So what we've seen in practice is if you're doing this kind of multi stage process of, of SFT and then DPO, you probably can get by with, you know, 10 to 100,000 examples of SFT training data and then maybe around 50k examples for the DPO step. But I have seen examples lately. There's an algorithm called Data where they show that they can actually use only 6,000 examples for the SFT step, which radically reduces it. So I would say roughly 10,000 examples for SFT and maybe, you know, 10 to 20 for DPO. That would probably get you a pretty good model. The open question is that as far as we know, like the other labs, like OpenAI and Anthropic, they're using many, many, many more examples. And presumably if you do this kind of iteratively, you get much better models.

Absolutely. Perfect. Well, thank you so much, Louis. I think we'll dive into the second section with Ed, so I'll bring you guys up. Take it away, Ed.

Great, thanks. So this is going to be a more sort of practical dive into how to actually run SFT on DPO on your dialog data sets. And one of those things is you have these data sets of conversations and how do you actually format these in order to tokenize them and feed them to a model in order to fine tune? We've got some links in this presentation to a bunch of SFT and DPO data sets. So if you're interested in getting started, there's a good starting point. And then I'll go through some metrics, mainly for the DPO step, of how to sort of evaluate whether your DPO is training well, how to identify problems with your training. And then once you've actually trained your chatbot, you probably want to evaluate it. You might test it out yourself by chatting with it, but you might want some automated way to do that. And so I'll suggest some ways of how we do this and some tools available in the community for doing this. Before I get onto dpo, however, I'll first talk about the supervised fine tuning the SFT step. Because it's important to do this step before dpo. We've actually run some experiments where you kind of skip this step and you jump directly to dpo. And actually it's very, very hard to perform the DPO step. We find that the model the LLM fails to learn the dialog template is one thing that we observe. Right, so I'll get started. Yeah. So to accompany this presentation, we've created two annotated notebooks for SFT and dpo. These have been written in a way that they'll run on a Google Colab gpu, so they should be what they should run. Initially I was thinking we can actually go through this live, but actually these models take many hours to train, so it's better to sort of run them offline. We don't plan to update these notebooks, so we actually have a more up to date version of our examples, which is the hugging face alignment handbook. And so check that out. And this has examples where if you're like GPU poor, if you've got a consumer GPU with like 24 gig of RAM, you can run some low resource examples with Lora. And if you're GPU rich, there's multi GPU multi node examples where we do distributed training with Accelerate and deepspeed. And in this alignment handbook we provide configurations high parameters scripts for running these examples. I should mention that the notebooks we provide use lora, which is a technique for fine tuning on like low GPU memory, GPU.