Training the Reward Model (RM)

Training a Reward Model (RM) is a process within Reinforcement Learning from Human Feedback (RLHF) that utilizes human preference data. This data is typically gathered in the form of pairwise comparisons, comprising a prompt, a chosen response, and a rejected response. The RM's objective is to learn a function that maps a prompt and a potential response to a scalar value. This value represents how well the response aligns with human preferences for the given prompt. The resulting learned reward function will then serve as the objective signal for fine-tuning a language model using reinforcement learning.

Reward Model Architecture

The architecture for the reward model is often closely related to the base language model being aligned. A common practice is to start with the Supervised Fine-Tuned (SFT) model itself and modify its final layer. Instead of predicting the next token distribution, the RM's head is adapted to output a single scalar value.

Specifically, the input to the RM is the concatenation of the prompt ($x$) and a candidate response ($y$). This combined sequence is processed by the transformer architecture. The hidden state corresponding to the final token of the sequence (often the end-of-sequence token) is then passed through a linear layer (the reward head) to produce the scalar reward score.

Using the SFT model as the base for the RM offers significant advantages:

1. Initialization: The RM starts with weights already adapted to the target domain and style through SFT.
2. Understanding: The model inherently understands the nuances of language generation and structure relevant to the task.
3. Efficiency: It avoids training a completely separate model from scratch, potentially saving computational resources.

Let $r_\theta(x, y)$ denote the scalar reward output by the RM with parameters $\theta$ for prompt $x$ and response $y$.

Training Objective and Loss Function

The RM is trained on the collected preference dataset $D = \{(x^{(i)}, y_c^{(i)}, y_r^{(i)})\}_{i=1}^N$, where $y_c$ is the response preferred by humans (chosen) and $y_r$ is the response deemed less preferable (rejected) for prompt $x$. The objective is to train the RM such that it assigns a higher score to the chosen response compared to the rejected one for the same prompt:

$$r_\theta(x, y_c) > r_\theta(x, y_r)$$

This is typically framed as a binary classification problem on pairs of responses. A common approach adapts the Bradley-Terry model, which models the probability that $y_c$ is preferred over $y_r$. This probability can be modeled using the difference in their reward scores passed through a logistic sigmoid function $\sigma(z) = 1 / (1 + e^{-z})$:

$$P(y_c \succ y_r | x) = \sigma(r_\theta(x, y_c) - r_\theta(x, y_r))$$

The RM is trained by minimizing the negative log-likelihood of the human preferences in the dataset $D$. The loss function becomes:

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

This loss encourages the difference $r_\theta(x, y_c) - r_\theta(x, y_r)$ to be large and positive, effectively maximizing the probability of correctly classifying the preferred response according to the human labels. Sometimes, a margin term can be added, but this basic form is widely used.

Training Process

During training, each data point $(x, y_c, y_r)$ requires two forward passes through the RM: one for the prompt concatenated with the chosen response ($x \oplus y_c$) and one for the prompt concatenated with the rejected response ($x \oplus y_r$).

1. Input Preparation: Concatenate the prompt $x$ with the chosen response $y_c$ and the rejected response $y_r$. Ensure proper tokenization and attention masking, treating the prompt and response as a single sequence for each pass.
2. Forward Pass: Compute the scalar rewards $r_c = r_\theta(x, y_c)$ and $r_r = r_\theta(x, y_r)$.
3. Loss Calculation: Calculate the pairwise ranking loss using the formula above.
4. Optimization: Compute gradients of the loss with respect to the RM parameters $\theta$ and update the parameters using an optimizer like AdamW.

Here's a simplified PyTorch snippet illustrating the loss calculation within a training step:

import torch
import torch.nn.functional as F

# Assume 'reward_model' is your RM instance (e.g., a Transformer with a
# scalar head)
# Assume 'tokenizer' is your tokenizer instance
# Assume 'batch' contains tuples of (prompt, chosen_response,
# rejected_response) strings

def compute_rm_loss(reward_model, tokenizer, batch):
    """Computes the pairwise ranking loss for a batch of preference data."""

    prompts, chosen_responses, rejected_responses = batch

    # Tokenize and prepare inputs for chosen responses
    chosen_inputs = tokenizer(
        [p + c for p, c in zip(prompts, chosen_responses)],
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=1024 # Example max length
    )
    # Move tensors to the correct device
    chosen_inputs = {
        k: v.to(reward_model.device) for k, v in chosen_inputs.items()
    }

    # Tokenize and prepare inputs for rejected responses
    rejected_inputs = tokenizer(
        [p + r for p, r in zip(prompts, rejected_responses)],
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=1024 # Example max length
    )
    rejected_inputs = {
        k: v.to(reward_model.device) for k, v in rejected_inputs.items()
    }

    # Get reward scores from the model
    # The reward_model forward pass should return a scalar score
    # per sequence
    chosen_rewards = reward_model(**chosen_inputs).rewards
    # Assuming model output has .rewards attribute
    rejected_rewards = reward_model(**rejected_inputs).rewards

    # Calculate the loss
    # loss = -log(sigmoid(chosen_rewards - rejected_rewards))
    loss = -F.logsigmoid(
        chosen_rewards - rejected_rewards
    ).mean()

    return loss

# --- In your training loop ---
# optimizer.zero_grad()
# loss = compute_rm_loss(reward_model,
#                        tokenizer, batch_data)
# loss.backward()
# optimizer.step()

Evaluation

Evaluating the RM's performance is important before using it in the RL phase. The primary metric is accuracy on a held-out set of preference pairs. This measures how often the RM correctly predicts the human-preferred response ($r_c > r_r$). Accuracies typically range from 65% to 80%, depending on the task difficulty, data quality, and model capacity.

Basic workflow for training the Reward Model using preference data and a pairwise ranking loss.

Accuracy aside, qualitative analysis is useful. Examining cases where the RM strongly agrees or disagrees with human judgments can reveal biases or weaknesses in the model. It's also helpful to check if the reward scores correlate with other intuitive metrics of quality, like response length, coherence, or helpfulness, although these correlations might be weak.

Challenges

• Calibration: The absolute scale of the rewards is often not directly meaningful or well-calibrated. The training objective focuses only on the difference $r_c - r_r$. Normalizing reward scores (e.g., whitening) during RL training is a common practice.
• Overfitting: The RM can overfit the specific preferences in the training data, potentially learning spurious correlations or biases present in the labeling process. Regularization and using a sufficiently diverse dataset are important.
• Data Quality Dependency: The RM's quality is fundamentally limited by the quality and consistency of the human preference data. Noisy or biased labels will lead to a suboptimal RM.
• Scalability: Training the RM requires significant computation, comparable to fine-tuning the base LLM itself.

Successfully training a reward model is a significant step in the RLHF process. A well-trained RM provides the important signal needed to guide the LLM during the subsequent reinforcement learning phase, steering it towards generating outputs that better align with desired human characteristics like helpfulness, honesty, and harmlessness.