Reward Model Training

Reinforcement Learning from Human Feedback (RLHF) relies on aligning a language model's behavior with human preferences. A primary step in this process involves training a separate model, the Reward Model (RM), denoted as $r_\phi(x, y)$. This model's objective is to learn a function that takes a prompt $x$ and a generated response $y$ as input and outputs a scalar value representing how much a human would likely prefer that response. Essentially, the reward model acts as a learned proxy for human judgment.

The Purpose of the Reward Model

Instead of directly using human feedback during the computationally intensive LLM fine-tuning phase (which would be slow and impractical), we first distill human preferences into the reward model. This RM can then provide dense feedback signals during the subsequent policy optimization stage (using algorithms like PPO), guiding the LLM $\pi_\theta(y|x)$ towards generating outputs that score highly according to the learned preference function.

Data Collection for Reward Modeling

Training the reward model requires a specialized dataset consisting of human preferences. While asking humans for absolute quality scores (e.g., rating a response from 1 to 10) is possible, it often suffers from inconsistency and poor calibration across different annotators and prompts.

A more common and often more reliable approach is to collect comparison data. In this setup, for a given prompt $x$, multiple responses ($y_1, y_2, ..., y_k$) are generated by one or more versions of the language model. Human annotators are then asked to rank these responses from best to worst, or more simply, to choose the single best response among a pair.

This comparison process yields data points typically structured as tuples: $(x, y_w, y_l)$, where $y_w$ is the preferred ("winning") response and $y_l$ is the less preferred ("losing") response for the prompt $x$. Compiling a large dataset of these comparisons ($D = \{(x^{(i)}, y_w^{(i)}, y_l^{(i)})\}$) forms the foundation for training the reward model.

Diagram illustrating the typical workflow for generating human preference data used in reward model training.

Reward Model Architecture

The architecture for the reward model often mirrors the base language model being fine-tuned. A common practice is to start with the pre-trained weights of the LLM (or a smaller version for efficiency) and replace or append a final linear layer. This new layer is trained to output a single scalar value (the reward score) instead of predicting the next token probabilities.

Initializing the RM from the pre-trained LLM is advantageous because the model already possesses a strong understanding of language structure, semantics, and context captured in the prompt $x$ and response $y$. The training process then focuses on adapting this understanding to predict the specific human preference signal represented in the comparison data.

Training the Reward Model

The core idea is to train the RM parameters $\phi$ such that the preferred response $y_w$ consistently receives a higher score than the rejected response $y_l$ for the same prompt $x$. This is typically framed as a classification or ranking problem.

A widely used objective function is based on the Bradley-Terry model, which models the probability that $y_w$ is preferred over $y_l$:

$$P(y_w \succ y_l | x) = \sigma(r_\phi(x, y_w) - r_\phi(x, y_l))$$

Here, $\sigma$ is the sigmoid function. The training objective is to maximize the likelihood of the observed human preferences in the dataset $D$. This translates to minimizing the negative log-likelihood loss:

$$\mathcal{L}(\phi) = -\mathbb{E}_{(x, y_w, y_l) \sim D} \left[ \log \sigma \left( r_\phi(x, y_w) - r_\phi(x, y_l) \right) \right]$$

This loss function encourages the reward model $r_\phi$ to output a larger difference between the scores of the winning and losing responses. The training proceeds using standard gradient-based optimization methods like Adam.

Training and Evaluation

• Initialization: As mentioned, initializing from a pre-trained model is standard. Sometimes, initializing from the supervised fine-tuned (SFT) model used in the data collection phase can provide a better starting point.
• Data Quantity and Quality: The performance of the RLHF process is highly sensitive to the quality and quantity of the preference data. Biases in annotation, inconsistent judgments, or insufficient data can lead to poorly calibrated or ineffective reward models.
• Calibration: Ideally, the difference in reward scores $r_\phi(x, y_w) - r_\phi(x, y_l)$ should correlate with the strength of the human preference. However, standard training objectives don't explicitly enforce this, potentially leading to an overconfident RM.
• Evaluation: The primary metric for evaluating the RM is its accuracy on a held-out set of preference pairs. That is, given a pair $(x, y_w, y_l)$ from the test set, does the model correctly predict the preference, i.e., is $r_\phi(x, y_w) > r_\phi(x, y_l)$? High accuracy indicates that the RM has successfully captured the patterns in the human preference data. Qualitative analysis and correlation studies with human ratings on a continuous scale can also provide insights into the RM's behavior.

Once a sufficiently accurate reward model is trained, it serves as the objective function for the next stage: fine-tuning the language model's policy using reinforcement learning.