No history yet

Reward Model Architectures

Engineering a Reward Model

Once an LLM has been fine-tuned with supervised data (SFT), the next step in RLHF is to teach it our preferences. We need to move from simple instruction-following to understanding nuanced concepts like helpfulness, harmlessness, and tone. To do this, we build a separate model, the Reward Model (RM), whose only job is to read a prompt and a response, and then output a single number: a scalar reward score.

This process starts with human feedback. For a given prompt, we generate two or more responses from the SFT model. A human labeler then ranks these responses from best to worst. The most common format is a simple pairwise comparison: given two responses, which one is better? This collection of preferences, {prompt, winning_response, losing_response}, becomes the training dataset for our Reward Model.

The Reward Model itself is typically a pre-trained LLM, often the same one used for the SFT base model, but with a crucial modification. We remove its final language modeling layer and replace it with a new, randomly initialized 'head'. This head is a simple linear layer that takes the final embedding from the LLM and maps it to a single number.

During training, the core LLM's weights are usually 'frozen' or not updated. We only train the new linear head. This is computationally efficient and prevents the model from forgetting its powerful language understanding capabilities. The goal isn't to change what the model knows about language, but to teach it a new skill: scoring text based on human preferences.

The Math of Preference

To train this model, we need a loss function that understands ranking. A standard classification loss won't work, because we aren't assigning fixed labels. Instead, we are teaching the model a relative preference. The foundation for this comes from the , a statistical method for determining the strength of participants in a paired comparison tournament. In our case, the 'participants' are the LLM's responses, and the 'tournament' is the human ranking process.

The model calculates a reward score, rr, for both the winning response (ywy_w) and the losing response (yly_l) for a given prompt (xx). The core idea is that the reward for the winning response, r(x,yw)r(x, y_w), should be higher than the reward for the losing response, r(x,yl)r(x, y_l).

The training objective is to maximize the difference between the scores of the winning and losing responses. We use a specific loss function to achieve this.

loss(θ)=1(K2)E(x,yw,yl)D[log(σ(rθ(x,yw)rθ(x,yl)))]\text{loss}(\theta) = -\frac{1}{\binom{K}{2}} \mathbb{E}_{(x, y_w, y_l) \sim D} \left[ \log(\sigma(r_\theta(x, y_w) - r_\theta(x, y_l))) \right]

This formulation is a form of contrastive learning. It doesn't care about the absolute value of the reward scores. It only cares about the relative difference between them. The model learns by being shown a 'positive' example (the winner) and a 'negative' example (the loser) and is trained to push their scores apart. This makes the learning signal clear and stable.

Avoiding Pitfalls

Building a good Reward Model is as much about the data as it is about the architecture. A high-quality preference dataset is diverse. It should include a wide range of prompts, from simple questions to complex creative tasks. The responses should also vary in quality, showcasing subtle differences in tone, factual accuracy, and helpfulness. If the data only contains obvious good/bad pairs, the RM won't learn to make the fine-grained distinctions needed to guide the LLM toward excellence.

A Reward Model trained only on easy examples will fail when faced with nuanced decisions.

Even with great data, a major challenge is (also known as overoptimization). This happens when the LLM, during the reinforcement learning phase, finds a way to get a high score from the Reward Model that doesn't align with actual human preference. For example, the RM might have a slight bias for longer, more verbose answers. The LLM will exploit this, producing unnecessarily long responses to maximize its reward, even if humans would prefer a concise answer.

One way to mitigate this is by adding a penalty term to the RL objective function. A common approach is to use a Kullback-Leibler (KL) divergence penalty. This term measures how much the LLM's response distribution has strayed from its original SFT distribution. It essentially keeps the model from becoming too different from its initial, more general self, reining in its attempts to exploit the RM.

Time to check your understanding of how Reward Models are built and trained.

Quiz Questions 1/5

What is the primary function of a Reward Model (RM) in the RLHF process?

Quiz Questions 2/5

When training a Reward Model, what part of the model is typically updated?

By carefully designing the Reward Model architecture, curating a diverse dataset, and being mindful of potential failure modes like reward hacking, we can create a robust guide for the final stage of RLHF.