Reward Model Training: Architectures and Loss Functions

Human preference data is essential for training a reward model, $r_\theta(x, y)$, which is a core component of Reinforcement Learning from Human Feedback (RLHF). This data typically comes in the form of pairs $(x, y_w, y_l)$, where $y_w$ denotes the preferred ("winning") response and $y_l$ the rejected ("losing") response to a prompt $x$. The reward model's function is to learn a scalar representation of these human preferences, seeking to satisfy $r_\theta(x, y_w) > r_\theta(x, y_l)$ for each observed pair. The resulting learned reward function then serves to guide the fine-tuning of the large language model (LLM) policy.

Reward Model Architectures

The architecture of the reward model is a significant design choice. A prevalent and effective approach is to leverage the power of pre-trained language models themselves.

1. LLM-Based Architecture: The most common strategy involves initializing the reward model using the weights of the same pre-trained LLM that will eventually be fine-tuned (or a model from the same family, perhaps differing in size). The core idea is to adapt this pre-trained model to predict preference scores instead of generating text.

   • Input Processing: The prompt $x$ and a candidate response $y$ are typically concatenated into a single sequence, often separated by special tokens (e.g., prompt_text [SEP] response_text [EOS]).
   • Base Model: This concatenated sequence is fed through the transformer layers of the base LLM.
   • Reward Head: Instead of using the final layer's output logits for token prediction, a "head" is added on top of the LLM's final hidden state representation. Commonly, the hidden state corresponding to the last token (e.g., the [EOS] token) is extracted. This vector representation, which summarizes the entire input sequence (prompt and response), is then passed through a linear layer that projects it down to a single scalar value. This scalar output represents the predicted reward $r_\theta(x, y)$.

   A typical reward model architecture using a pre-trained LLM base, processing concatenated prompt and response, and outputting a scalar reward via a linear head attached to the final token's hidden state.

2. Initialization: Starting with pre-trained weights allows the reward model to benefit from the language understanding capabilities learned during the LLM's initial training. During reward model training, one might choose to fine-tune all the parameters of the base LLM or freeze the majority of the base layers and only train the final layers and the reward head. Fine-tuning the entire model can lead to better adaptation but is more computationally expensive and risks catastrophic forgetting. Freezing most layers is faster but might limit the model's ability to learn detailed preferences.

3. Alternative Architectures: While using the target LLM's architecture is standard, alternatives exist. One could train a smaller, independent transformer model or even use different model types. However, these approaches might struggle to capture the same level of detail required to accurately judge the quality of outputs generated by a large base model. Using a model from the same family generally ensures architectural compatibility and uses relevant pre-trained knowledge.

The Pairwise Ranking Loss Function

Given the preference data $(x, y_w, y_l)$, the reward model is trained to assign a higher score to $y_w$ than to $y_l$. The standard objective function for this is a pairwise ranking loss based on the Bradley-Terry model, which is commonly used for modeling preferences between pairs of items.

The loss function aims to maximize the probability that the difference in rewards between the chosen and rejected responses aligns with the human label. It is typically formulated as the negative log-likelihood of the preferences:

Let's break down this formula:

• $D$: The dataset of human preferences $\{(x^{(i)}, y_w^{(i)}, y_l^{(i)})\}_{i=1}^N$.
• $r_\theta(x, y)$: The scalar reward output by the model with parameters $\theta$ for prompt $x$ and response $y$.
• $r_\theta(x, y_w) - r_\theta(x, y_l)$: The difference between the predicted reward for the preferred response and the rejected response. We want this difference to be positive.
• $\sigma(z) = \frac{1}{1 + e^{-z}}$: The sigmoid function. It squashes the reward difference into the range $(0, 1)$. This value $\sigma(r_\theta(x, y_w) - r_\theta(x, y_l))$ can be interpreted as the model's estimated probability that $y_w$ is preferable to $y_l$.
• $\log(\cdot)$: Taking the logarithm transforms the probability into log-probability (or log-likelihood). Maximizing the log-probability is equivalent to maximizing the probability itself, but it's often more numerically stable and convenient.
• $-\mathbb{E}[\cdot]$: We average the negative log-likelihood over all preference tuples in the dataset $D$. Minimizing this negative log-likelihood is equivalent to maximizing the likelihood of the observed preferences according to the reward model.

During training, for each triplet $(x, y_w, y_l)$ in a batch, the reward model performs two forward passes: one for $(x, y_w)$ to get $r_\theta(x, y_w)$ and one for $(x, y_l)$ to get $r_\theta(x, y_l)$. The loss is computed based on the difference, and gradients are backpropagated to update the model parameters $\theta$.

Here's a sketch of the training step using PyTorch-like pseudocode:

import torch
import torch.nn.functional as F

# Assume:
# reward_model: The model taking tokenized input_ids and attention_mask, returning a scalar.
# tokenizer: The tokenizer corresponding to the reward_model base.
# optimizer: An optimizer like AdamW.
# dataloader: Provides batches of {'prompt': [...], 'chosen': [...], 'rejected': [...]}

def train_reward_model_step(batch, reward_model, tokenizer, optimizer, device):
    """Performs a single training step for the reward model."""

    prompts = batch['prompt']
    chosen_responses = batch['chosen']
    rejected_responses = batch['rejected']

    # Prepare inputs for chosen and rejected responses
    chosen_texts = [p + tokenizer.sep_token + r + tokenizer.eos_token for p, r in zip(prompts, chosen_responses)]
    rejected_texts = [p + tokenizer.sep_token + r + tokenizer.eos_token for p, r in zip(prompts, rejected_responses)]

    # Tokenize (handle padding and truncation appropriately)
    chosen_encodings = tokenizer(chosen_texts, padding=True, truncation=True, return_tensors="pt").to(device)
    rejected_encodings = tokenizer(rejected_texts, padding=True, truncation=True, return_tensors="pt").to(device)

    # Forward passes to get rewards
    # Assuming reward_model outputs a dict containing 'rewards' tensor
    rewards_chosen = reward_model(**chosen_encodings).rewards # Shape: (batch_size, 1)
    rewards_rejected = reward_model(**rejected_encodings).rewards # Shape: (batch_size, 1)

    # Ensure shapes are compatible if needed, e.g., squeeze last dimension
    rewards_chosen = rewards_chosen.squeeze(-1) # Shape: (batch_size,)
    rewards_rejected = rewards_rejected.squeeze(-1) # Shape: (batch_size,)

    # Calculate pairwise loss
    # loss = -log(sigmoid(chosen_reward - rejected_reward))
    loss = -F.logsigmoid(rewards_chosen - rewards_rejected).mean()

    # Optimization step
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss.item()

# --- Training Loop ---
# for epoch in range(num_epochs):
#     for batch in dataloader:
#         loss_val = train_reward_model_step(batch, reward_model, tokenizer, optimizer, device)
#         # Log loss_val, handle checkpoints, etc.

This training process yields a reward model $r_\theta(x, y)$ that ideally captures the nuances of human preference regarding the helpfulness, honesty, and harmlessness of LLM responses. Once sufficiently trained and evaluated, this model becomes an important component in the next stage: fine-tuning the LLM policy using reinforcement learning algorithms like PPO, where $r_\theta(x, y)$ serves as the reward signal.