Outcome School

What is a Reward Model and How Is It Trained?

Authors
What is a Reward Model and How Is It Trained?

A Reward Model is a model that reads a prompt and an answer, and gives back one number that tells how good that answer is as per human preference. It is trained on pairs of answers where humans picked the better one.

In this blog, we will learn about what a Reward Model is and how it is trained. We will also see why we need it, how its architecture turns text into one score, how human preference data is collected, how the pairwise loss trains it step by step, how it is used in RLHF and for Best-of-N, and where it fails.

We will cover the following:

  • Why do we need a Reward Model?
  • What is a Reward Model?
  • Collecting human preference data
  • The pairwise loss
  • Worked example of the pairwise loss
  • How is a Reward Model trained?
  • How is a Reward Model used?
  • Limitations of a Reward Model

I am Amit Shekhar, Founder @ Outcome School, I have taught and mentored many developers, and their efforts landed them high-paying tech jobs, helped many tech companies in solving their unique problems, and created many open-source libraries being used by top companies. I am passionate about sharing knowledge through open-source, blogs, and videos.

I teach AI and Machine Learning at Outcome School.

Let's get started.

Why do we need a Reward Model?

An LLM (Large Language Model) learns first by predicting the next token on a huge amount of text. After that, it can write fluent text. But fluent text is not always a helpful, honest, and safe answer.

So, we want to train the LLM further to give answers that humans like. For that, we can use reinforcement learning. In reinforcement learning, the LLM writes an answer, gets a reward (a number that says how good the answer was), and, over many such tries, learns to write more answers that get a high reward.

Now, the question is: who gives this reward?

For a math problem, it is easy. We can check if the final answer is correct. But, for a prompt like "Explain gravity to a 10-year-old", there is no simple formula that gives a score for a "good answer". Being clear, correct, polite, and the right length all matter together.

Another option is to ask humans to rate every answer. But, during training, the LLM generates millions of answers. Humans cannot read and rate millions of answers live, while the training is running. It is too slow and too costly.

So, here comes the Reward Model to the rescue. We collect a limited amount of human judgment once, and we train a model to copy that judgment. After that, the model can score millions of answers quickly, without a human in the loop.

What is a Reward Model?

A Reward Model is a model that takes a prompt and a response as input and outputs a single number, called the reward, that represents how much a human would prefer that response.

A higher number means a better answer. A lower number means a worse answer.

For example:

  • Prompt: "What is the capital of France?"
  • Response A: "The capital of France is Paris." gets a reward of 2.1
  • Response B: "I am not sure, maybe Lyon." gets a reward of -1.3

Here, we can see that the Reward Model prefers Response A, which matches what a human would pick.

Now, let's understand how it is built. A Reward Model is usually an LLM itself. Often, it starts as a copy of the LLM that we want to improve, or of a smaller LLM. A normal LLM ends with a layer that gives a score for every token in its vocabulary, so that it can predict the next token. In a Reward Model, we remove that last layer and put a small new layer in its place, called the scalar head. A scalar is a single number. So, the scalar head turns the LLM's understanding of the full text into one number.

Let's see this as a diagram:

┌───────────────────────────────────┐
│ Prompt + Response                 │
│ "What is the capital of France?   │
│  The capital of France is Paris." │
└───────────────────────────────────┘
                  ↓
┌───────────────────────────────────┐
│ LLM body (all its layers)         │
│ reads every token of the text     │
└───────────────────────────────────┘
                  ↓
┌───────────────────────────────────┐
│ Scalar head (new small layer)     │
│ looks at the last token           │
└───────────────────────────────────┘
                  ↓
┌───────────────────────────────────┐
│ Reward = 2.1                      │
└───────────────────────────────────┘

Here, we can see that the prompt and the response go in together as one piece of text. The LLM body reads all of it. The scalar head takes the LLM body's output for the last token, which has already seen the whole text, and turns it into one number.

We start from an LLM because it already understands language well. We need to teach it one new skill: judging.

If we want to go deep into LLM Fundamentals and LLM Internals and build a Large Language Model (LLM) from scratch, we can check out our AI and Machine Learning Program at Outcome School.

Collecting human preference data

To train the Reward Model, we need data that shows what humans prefer.

One way is to ask humans to give each answer a score from 1 to 10. But, these scores are not stable. One person gives a 7 to an answer, and another person gives a 5 to the same answer. Even the same person can give a different score on a different day. So, absolute scores are noisy.

Comparing two answers is much easier. If we show a person two answers side by side and ask "Which one is better?", people decide faster, and different people agree with each other more often. This is why we collect comparisons instead of scores.

The process is as below:

Step 1: We take a prompt, for example "Explain gravity to a 10-year-old."

Step 2: We let the LLM generate two different answers to it.

Step 3: A human reads both answers and picks the better one.

Step 4: We save the prompt, the better answer, and the worse answer together as one example.

The better answer is called the chosen response, and the worse answer is called the rejected response. One saved example looks like below:

example = {
    "prompt": "Explain gravity to a 10-year-old.",
    "chosen": "Gravity is an invisible pull. The Earth pulls everything toward it, which is why a ball falls down.",
    "rejected": "Gravity is the curvature of spacetime described by general relativity."
}

Here, the rejected answer is not wrong, but it is not good for a 10-year-old. The human knows this, and the pair captures this judgment.

Sometimes, the LLM generates 4 answers, and the human ranks all 4 from best to worst. We can turn that one ranking into 6 pairs, because there are 6 ways to pick 2 answers out of 4.

We repeat this for tens of thousands or even hundreds of thousands of prompts. This collection of pairs is called the preference dataset.

Planning your move into AI?

No matter which tech domain you work in today, with your effort, you can transition to these roles:

  • AI Engineer
  • Forward Deployed Engineer
  • AI Agent Engineer
  • Harness Engineer
  • Voice AI Engineer
  • AI Architect
  • LLM Engineer
  • LLM Inference Engineer
  • Gen AI Engineer
  • MLOps Engineer
  • ML Engineer
  • Deep Learning Engineer

Our AI and Machine Learning Program at Outcome School takes you there.

The pairwise loss

Now, we have pairs. But, the Reward Model must output a single score for a single answer. So, how can we train a model that gives scores, using data that only says "A is better than B"?

The answer is: we do not care about the exact score. We only want the score of the chosen answer to be higher than the score of the rejected answer.

Let's call the score of the chosen answer r_chosen and the score of the rejected answer r_rejected. The thing that matters is the gap between them:

diff = r_chosen - r_rejected

If diff is a big positive number, the Reward Model agrees with the human. If diff is negative, the Reward Model disagrees with the human.

Next, we pass this gap through the sigmoid function. Sigmoid takes any number and squeezes it into a value between 0 and 1:

sigmoid(x) = 1 / (1 + e^(-x))

Here, e is a fixed number, about 2.718.

So, we can read sigmoid(diff) as the probability, as per the Reward Model, that the chosen answer is better than the rejected answer. A big positive gap gives a value close to 1. A gap of 0 gives exactly 0.5, which means "I cannot tell". A negative gap gives a value close to 0.

This idea of turning a score gap into a probability comes from the Bradley-Terry model, a method from 1952 that ranks players using only the results of matches between them. A Reward Model treats two answers like two players in a match.

Finally, we want this probability to be as high as possible. So, we take its negative log (the natural log), which turns "make the probability high" into "make the loss small". The log also punishes a confident wrong ranking very heavily:

loss = -log(sigmoid(r_chosen - r_rejected))

This is the pairwise loss. The loss is small when the chosen answer gets a much higher score than the rejected answer. The loss is big when the Reward Model ranks them the wrong way.

Worked example of the pairwise loss

The best way to learn this is by taking an example. Let's say we have one pair, and we look at three different situations.

Case 1: The Reward Model agrees with the human.

  • r_chosen = 2.0, r_rejected = 0.5
  • diff = 2.0 - 0.5 = 1.5
  • sigmoid(1.5) = 0.8176
  • loss = -log(sigmoid(1.5)) = 0.2014

The loss is small. The Reward Model is already doing well on this pair.

Case 2: The Reward Model cannot tell.

  • r_chosen = 1.0, r_rejected = 1.0
  • diff = 0
  • sigmoid(0) = 0.5
  • loss = -log(0.5) = 0.6931

The loss is in the middle.

Case 3: The Reward Model disagrees with the human.

  • r_chosen = 0.5, r_rejected = 2.0
  • diff = 0.5 - 2.0 = -1.5
  • sigmoid(-1.5) = 0.1824
  • loss = -log(sigmoid(-1.5)) = 1.7014

The loss is big. The Reward Model has made a mistake, and the training will push hard to fix it.

Let's see the code for this as below:

import math

def sigmoid(x):
    return 1 / (1 + math.exp(-x))

def pairwise_loss(r_chosen, r_rejected):
    return -math.log(sigmoid(r_chosen - r_rejected))

print(round(pairwise_loss(2.0, 0.5), 4))
print(round(pairwise_loss(1.0, 1.0), 4))
print(round(pairwise_loss(0.5, 2.0), 4))

This will print the following:

0.2014
0.6931
1.7014

Here, we can see that the three outputs match our three cases.

One more thing to notice: the loss only uses the gap. If we add 10 to both scores, the gap stays the same, and the loss stays the same. So, the reward values do not have any fixed meaning on their own. A reward of 2.1 is not "good" by itself. It is only good compared to other answers for the same prompt.

A quick note for you

No matter which tech domain you work in, get familiar with these topics:

  • LLM
  • RAG
  • MCP
  • Agent
  • Fine-tuning
  • Quantization

We put it all together in one video:

AI Engineering Explained: LLM, RAG, MCP, Agent, Fine-Tuning, and Quantization

No need to stop reading - bookmark it and watch later when you get time. Future you will thank you.

Now, let's get back to the topic.

How is a Reward Model trained?

Now, let's put all the pieces together.

Step 1: We take one example from the preference dataset: a prompt, a chosen answer, and a rejected answer.

Step 2: We feed "prompt + chosen answer" into the Reward Model and get r_chosen.

Step 3: We feed "prompt + rejected answer" into the same Reward Model and get r_rejected.

Step 4: We compute the loss: -log(sigmoid(r_chosen - r_rejected)).

Step 5: We use backpropagation to find how each parameter affects the loss, and an optimizer updates the parameters a tiny bit to make the loss smaller. This pushes r_chosen up and r_rejected down.

Step 6: We repeat this for all the pairs in the dataset, in batches.

Let's see this as a diagram:

┌────────────────────┐    ┌────────────────────┐
│ Prompt + Chosen    │    │ Prompt + Rejected  │
└────────────────────┘    └────────────────────┘
          ↓                         ↓
┌────────────────────┐    ┌────────────────────┐
│ Reward Model       │    │ Reward Model       │
│ (same parameters)  │    │ (same parameters)  │
└────────────────────┘    └────────────────────┘
          ↓                         ↓
┌────────────────────┐    ┌────────────────────┐
│ r_chosen = 2.0     │    │ r_rejected = 0.5   │
└────────────────────┘    └────────────────────┘
          ↓                         ↓
┌──────────────────────────────────────────────┐
│ diff = 2.0 - 0.5 = 1.5                       │
│ loss = -log(sigmoid(1.5)) = 0.2014           │
└──────────────────────────────────────────────┘
                       ↓
┌──────────────────────────────────────────────┐
│ Update the parameters:                       │
│ push r_chosen up, r_rejected down            │
└──────────────────────────────────────────────┘

Here, we can see that both answers go through the same Reward Model with the same parameters. One model scores both answers, and the loss compares the two scores.

The size of the push depends on how wrong the model is. For this loss, the push on each score is 1 - sigmoid(diff). In Case 1, this is 1 - 0.8176 = 0.1824, a small push. In Case 3, this is 1 - 0.1824 = 0.8176, a big push. So, the Reward Model learns the most from the pairs it gets wrong.

Often, we train for only 1 epoch, which means one full pass over the dataset. If we train longer, the Reward Model starts to memorize the pairs instead of learning the general idea of a good answer.

To check how good the Reward Model is, we keep some pairs aside that it has never seen. We count how often it gives the chosen answer the higher score. This is called accuracy on held-out pairs. On real human data, Reward Models often reach only about 70 percent. This sounds low, but human labelers themselves often agree with each other only about 70 percent of the time on these comparisons. So, the data itself has this limit.

To master Reward Models, Loss Functions, and Backpropagation in depth, check out our AI and Machine Learning Program at Outcome School.

How is a Reward Model used?

Now that we have learned how to train a Reward Model, it's time to learn how it is used. There are two common ways: in RLHF with PPO, and as a reranker.

Using it in RLHF with PPO

RLHF stands for Reinforcement Learning from Human Feedback. Here, the LLM that we want to improve is called the policy. PPO (Proximal Policy Optimization) is the reinforcement learning algorithm that updates the policy in small, controlled steps.

The loop works like below:

  • The policy reads a prompt and writes an answer.
  • The Reward Model scores that answer.
  • PPO updates the policy so that answers with high rewards become more likely.
  • This repeats with new prompts.

Let's see the full journey as a diagram:

┌──────────────────────────────────────┐
│ Stage 1: Collect preference data     │
│ 50,000 prompts x 2 answers each      │
│ a human picks the better answer      │
└──────────────────────────────────────┘
                   ↓
┌──────────────────────────────────────┐
│ Stage 2: Train the Reward Model      │
│ pairwise loss on 50,000 pairs        │
│ chosen score up, rejected down       │
└──────────────────────────────────────┘
                   ↓
┌──────────────────────────────────────┐
│ Stage 3: RLHF with PPO               │
│ policy writes an answer              │
│ Reward Model scores it (no human)    │
│ PPO updates the policy               │
└──────────────────────────────────────┘

Here, we can see that humans are needed only in the first stage. After that, the Reward Model gives the reward for every answer during RLHF. We have a detailed blog on InstructGPT that explains how this full pipeline taught GPT-3 to follow instructions.

Now, we also add a penalty based on KL divergence, which measures how far the policy's token probabilities have moved away from the original LLM. The final reward becomes:

final_reward = reward_model_score - beta * KL(policy || original_model)

Here, beta is a small number that controls how strong the penalty is. This penalty stops the policy from drifting too far from the original LLM in order to chase a higher score. We will see why this is needed in the limitations section.

Using it as a reranker (Best-of-N)

We can also use a Reward Model without any further training of the LLM. We let the LLM generate N answers for the same prompt, for example N = 4. The Reward Model scores all 4, and we show the user only the one with the highest score.

For example:

  • Answer 1: reward 0.8
  • Answer 2: reward 2.3
  • Answer 3: reward -0.4
  • Answer 4: reward 1.1

We pick Answer 2. This is called Best-of-N sampling or reranking. It is simple, and it often improves quality, but it costs N times more generation work for every prompt.

Stay updated: Subscribe to our newsletter to get our latest AI and Machine Learning blogs straight to your inbox.

Limitations of a Reward Model

A Reward Model copies human judgment from a limited set of pairs. So, it has the following limitations:

  • Bias of the labelers: If labelers often pick longer answers, the Reward Model learns that "longer is better", even when a short answer is clearer. The same happens with a confident tone and with lists, even when the content is not better.

  • Reward hacking: The policy can find answers that get a high score from the Reward Model but are not actually good for humans.

  • Overoptimization: At the beginning of RLHF, a higher reward usually means a better answer. But, if we push the policy too hard on the Reward Model's score, the score keeps going up while the true quality, as judged by humans, starts going down. This happens because the Reward Model was trained only on answers like the ones in the preference dataset. When the policy moves far away from those answers, the Reward Model's scores are no longer reliable. This is why we use the KL penalty and also stop the training early.

  • Cost: Collecting hundreds of thousands of human comparisons takes a lot of time and money. And, the Reward Model can be as big as the LLM itself. During RLHF, the policy, the original LLM, and the Reward Model often all sit in GPU memory together.

We have a detailed blog on Direct Preference Optimization (DPO) that covers how to train an LLM directly on preference pairs, without a separate Reward Model.

Now we must have understood what a Reward Model is, how it is trained with human preference pairs and the pairwise loss, and how it is used in RLHF and for Best-of-N.

Frequently Asked Questions

Why do we collect comparisons instead of scores to train a Reward Model?

Comparisons are more reliable than scores. When people give an answer a score from 1 to 10, different people give different scores, and even the same person can change the score on a different day. When people only pick the better of two answers, they decide faster and agree with each other more often.

Does a reward value mean anything on its own?

No. A reward value has no fixed meaning on its own. The pairwise loss only uses the gap between two scores, so adding 10 to both scores does not change the loss. A reward of 2.1 is only good when compared to the rewards of other answers for the same prompt.

How accurate is a Reward Model?

On real human data, a Reward Model often reaches only about 70 percent accuracy on held-out pairs. This sounds low, but human labelers themselves often agree with each other only about 70 percent of the time on these comparisons. So, the limit comes from the data itself.

Why is a KL penalty added to the reward in RLHF?

The KL penalty stops the policy from drifting too far from the original LLM in order to chase a higher score. The Reward Model was trained only on answers like the ones in the preference dataset, so its scores are no longer reliable when the policy moves far away from those answers. Without this penalty, the score can keep going up while the true quality goes down.

That's it for now.

Thanks

Amit Shekhar
Founder @ Outcome School

You can connect with me on:

Follow Outcome School on:

Read all of our high-quality blogs here.