From Zero to PPO: Understanding the Path to Helpful AI Models

Blog posts

From Zero to PPO: Understanding the Path to Helpful AI Models

Authors
Alessandro Cappelli
Julien Launay

Editors
Pearson Probst

Acknowledgements

Model : `` Llama 3.1 8B (pre-trained)`
Question: What to do during a stopover in Paris?
Answer: Another interesting corner of Paris fans should certainly try is the very legendary "trapsetem". Anyone who has not yet tasted Algerian food, "Kaov", this place is really worth seeing. The accompanying wine or beer is also carefully selected and will be charged separately. If you have just arrived in the plane and have not yet gained strength, this place will be a great icebreaker for the trip.

That's... not quite right! Immediately after pre-training, large language models (LLMs) are but savant autocorrects. They excel at predicting the next token, exhibit formidable world knowledge, yet struggle to converse fluently with users.
Helpfulness is instilled in LLMs as a result of extensive post-training. Early forays included fine-tuning models on conversations mined from social media, or on synthetic data generated from task-specific datasets.

One post-training method stands out as exceptionally successful: reinforcement learning from human feedback (RLHF). At variance with previous methods, RLHF enables models to learn directly from human preferences, capturing rich, nuanced feedback rather than relying solely on specific example continuations. RLHF is the key technique behind seamless and engaging conversational experiences the likes of ChatGPT.

Supervised finetuning: maximizing the likelihood of helpful answers

At the start of this post, we asked Llama 3.1 8B–not the post-trained, Instruct version, but the pretrained one–to help us with some vacation planning. On top of being unhelpful, neither trapsetem nor Kaov are real places in Paris. Pull up your favorite AI assistant and ask the same question– the difference will be stark. Why is that?

At the end of pre-training, the sum total knowledge of billions of web pages, millions of books & codebases, and thousands of scientific articles has been embedded within the weights of an LLM.

Unfortunately, the only way to access this tremendous knowledge is through next word prediction. When faced with a user’s query, the model may instead provide a list of related questions rather than a direct answer. Or provide a very specific, not entirely applicable answer. Asking for further clarification will only add to the confusion. Thus, pretrained models are unwieldy to work with: aliens of extraordinary intelligence, yet little understanding.

Partly, this is a context & diversity problem. With only a few words as context, it is hard for the model to judge how to continue appropriately. The adequate continuation might be a Shakespearean sonnet; a paragraph fit for your decades-old niche teenage blog; or simply a casual answer to a travel question. All of this, and more, are present in the pre-training corpus.

Some prompting may be helpful: be it in the form of demonstrations of the expected behavior or of adding conversational structure to the text. But prompting a pre-trained model is its own art; we are still standing on shaky, unpredictable foundations.

Let’s fix that. We will tune the desired behavior into the weights of the model. Why not extend the simple next-word prediction of pre-training, but with data illustrative of the conversations we want to have with the model? This is the basis of supervised fine-tuning (SFT).

Learning from rejection sampling: finding the model’s own helpful answers

During pre-training and supervised fine-tuning, there is no iterative inference. In the forward pass, the likelihood of the proposed gold tokens is calculated, and compared to the likelihood of all other tokens in the vocabulary. The training then maximizes the likelihood of the proposed tokens. Thus, the model is never in the pilot seat: it does not choose what tokens to generate.

Instead, we now seek to learn from answers generated by the model itself. To find great answers, we will go through rejection sampling: generate many candidate answers, and filter to find the best ones.

Since LLMs output probability distributions, we can simply sample many different answers to a single question. This can be done by increasing the temperature of sampling. We now have a riff-raff of candidate answers, and it’s time for evaluation.

The answers need ratings, to let us select the best ones only. Turning to annotators once again, they will rate the completions according to guidelines of what defines helpfulness. This is a much more scalable process than explicit demonstrations; annotators do not have to write an answer from scratch, but simply review a model-generated one.

At this stage, we can fine-tune on the top-rated completions. Is this effective? Yes! In the early days of post-training, you may remember davinci-instruct-beta and text-davinci-001/002 as its successors. Well, the beta model was trained using SFT on gold demonstrations. Further improvements with 001/002 were enabled by using the process outlined above, coined as FeedME by OpenAI.

Reward Modeling: Let Models Evaluate the Helpfulness of Answers

Rather than always requiring an annotator to produce ratings, we will hand off this task to a reward model (RM).

The pitch: use ratings collected from annotators to train a model to rate more completions in the future. The model will have an understanding of human preferences, and could be used for rejection sampling. Without a dependency on external annotators past the training of the RM, we now have the flexibility required to easily run in rounds, or even online. This approach was extensively explored in the Llama 2 paper.

Reinforcement Learning with Human Feedback (RLHF)

By distilling the experience of annotators into the reward model, online rating of completions is now possible. Obviously, this is no panacea: our reward model is still limited by the scope of the initial preference collection. Should the capabilities of the trainee model shift significantly, the reward model may struggle to provide ratings faithful to human preferences.
It’s time to think about efficiency. The ultimate goal is to ensure we optimize the LLM to be both helpful and efficient, adapting to users' needs with the data generated to enhance and refine its performance.