Reward Modeling › Reward Model Training

Michael BrenndoerferDecember 24, 20257 min read
Language AI HandbookAll chapters
Part 37 of 67Alignment and RLHF
Chapter 275Reward Modeling › Reward Model Training
Reading section 4 of 10 · Reward Model Training · 7 minSection 4 of 10 · 7 min

Book contents

Section 4 of 10 · 7 min read

Reward Modeling › Reward Model TrainingOne section · Sections in this chapter4 / 10 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Training a reward model involves several practical considerations beyond the loss function, including data preparation, optimization settings, and regularization strategies. Getting these details right matters as much as the architectural and loss function choices. A perfectly designed reward model architecture can still produce poor results if the training procedure is poorly configured, especially when working with the relatively small preference datasets typical in RLHF pipelines.

Think of reward model training as sitting in a particularly challenging part of the spectrum between pretraining and fine-tuning. Unlike pretraining, where we have abundant data and are learning from scratch, reward model training uses a much smaller dataset of human comparisons. Unlike standard fine-tuning, where we have labeled examples of the correct output, reward model training has only relative judgments with no absolute quality labels. This combination of data scarcity and label ambiguity requires careful handling to produce a model that generalizes rather than memorizes.

The main point governing most reward model training decisions is that we are trying to learn a smooth, generalizable preference function from a small number of noisy comparisons. Every design choice should be evaluated against this goal. Conservative learning rates prevent catastrophic forgetting of useful pretrained representations. Early stopping prevents memorization of specific training pairs. Regularization smooths the learned reward function. And careful data preprocessing ensures the comparisons we train on reflect human values rather than annotation artifacts.

Data Preparation

Each training example consists of a prompt and two responses with a preference label. The standard format is:

In[5]:
Code
{
    "prompt": "Explain quantum entanglement to a 10-year-old.",
    "chosen": "Imagine you have two magic coins...",
    "rejected": "Quantum entanglement is a phenomenon...",
}
Out[5]:
Console
{'prompt': 'Explain quantum entanglement to a 10-year-old.',
 'chosen': 'Imagine you have two magic coins...',
 'rejected': 'Quantum entanglement is a phenomenon...'}

During training, we need to compute rewards for both responses. A common approach processes both responses in a single forward pass by concatenating them:

[prompt] [chosen response] [EOS] [PAD] ... [prompt] [rejected response] [EOS]

The model computes hidden states for both sequences, extracts the final token representations for each, passes them through the value head, and computes the loss on their difference.

Data quality matters enormously at this stage. Before training, the dataset should be audited for several issues. First, check for duplicate or near-duplicate pairs, which can cause the model to overfit to specific examples. Second, examine the distribution of preference pairs across different prompt types, since a dataset dominated by one domain (say, coding questions) will produce a reward model that generalizes poorly to other domains. Third, filter out pairs where the quality difference between chosen and rejected is very small or where the human label may be unreliable. Pairs where annotators showed high disagreement are informative but should be handled carefully, either by upweighting them during training or by assigning softer labels.

The formatting of prompt-response pairs also deserves attention. In production RLHF systems, the prompt and response are typically formatted using the same chat template that the policy model was fine-tuned on. This ensures the reward model sees the same input format it will encounter during RL training. Mismatches between the format seen during reward model training and the format seen during RL training can cause systematic scoring errors that are difficult to diagnose.

Optimization Configuration

Reward model training typically uses conservative hyperparameters. These choices reflect the need to adapt a pretrained model to a new task while preserving its core language understanding capabilities.

The standard configuration includes:

  • Learning rate: 1×10−51 \times 10^{-5} to 5×10−55 \times 10^{-5}, lower than instruction tuning to preserve pretrained knowledge
  • Batch size: Large batches (32-128 pairs) help with gradient stability
  • Epochs: 1-3 epochs over the preference data; more risks overfitting
  • Optimizer: AdamW with weight decay 0.01-0.1
  • Learning rate schedule: Linear warmup for the first 5-10% of training steps, followed by linear decay

The learning rate is kept low because we're building on a pretrained model that already has strong language understanding. We want to learn the preference structure without disturbing the underlying representations too much. Think of the pretrained backbone as a carefully calibrated instrument: we need to tune it slightly for our specific measurement task, but we should be careful not to disturb its basic sensitivity. A large learning rate would be like recalibrating an instrument with a blunt screwdriver, potentially destroying the precision that makes it useful.

Different learning rates for the backbone and the value head can improve performance. The value head is initialized randomly and needs to learn from scratch, so a higher learning rate (e.g., 1×10−41 \times 10^{-4}) helps it converge quickly. The backbone is already well-trained and only needs minor adjustments, so a lower rate (e.g., 1×10−51 \times 10^{-5}) prevents destructive updates. This pattern, called "differential learning rates" or "layerwise learning rate decay," is common in fine-tuning practice and applies especially well to reward modeling.

Regularization Considerations

Reward models are susceptible to overfitting, especially when trained on limited preference data. Several techniques help prevent the model from memorizing training pairs rather than learning the underlying preference structure.

Early stopping based on validation accuracy is needed. The model should generalize to held-out preferences, not memorize training pairs. Hold out at least 10% of your preference data as a validation set and monitor performance after each epoch. Stop training when validation accuracy plateaus or begins to decline, even if training accuracy continues to improve. This is a clearer signal of overfitting for reward models than it is for standard classification tasks, because the training distribution of human comparisons is often narrower than the distribution of responses the model will encounter during RL training.

Dropout in the value head (though not typically in the pretrained layers) adds regularization without affecting the language model's representations. A dropout rate of 0.1 to 0.2 on the value head prevents the model from becoming overly reliant on specific features that happen to correlate with the training labels. The pretrained backbone layers should generally not have their dropout rates modified, as those were set during pretraining for different purposes.

Label smoothing can help with noisy labels. If some human preferences are inconsistent or reflect borderline cases, treating them as probabilistic rather than hard labels reduces sensitivity to noise. With label smoothing parameter ϵ\epsilon, a preference label of 1 (chosen is better) becomes 1−ϵ1 - \epsilon, and the model is trained to predict a probability of 1−ϵ1 - \epsilon rather than 1. This discourages overconfident predictions and makes the model more reliable to annotation noise. A typical value is ϵ=0.1\epsilon = 0.1.

Training Stability

The reward scale is arbitrary, unlike classification where outputs are bounded probabilities. This can cause optimization instability, particularly in the early stages of training when the value head is freshly initialized and may produce extreme reward values. Several techniques help maintain stable optimization throughout training.

Reward normalization after training scales rewards to have zero mean and unit variance over a reference set. This makes the reward signal easier to use in downstream RL training. The normalization is computed on a held-out reference set rather than the training data. This keeps it represents the typical distribution of responses the reward model will encounter during RL. Specifically, we collect a large set of prompt-response pairs (using the SFT model to generate responses to a diverse set of prompts), compute the rewards for all of them, and record the mean and standard deviation. During RL training, rewards are normalized using these statistics: r^=(r−μ)/σ\hat{r} = (r - \mu) / \sigma.

Gradient clipping prevents large updates from outlier examples. A max gradient norm of 1.0 is typical. Without clipping, a single training example with an unusually large loss can cause a large gradient update that disrupts the model's existing knowledge. Clipping ensures that no single example has an outsized influence on the parameters, making training more reliable to label noise and distribution outliers.

Learning rate warmup over the first 5-10% of training helps stabilize early optimization. During warmup, the learning rate gradually increases from a very small value (e.g., 0) to its target value. This prevents the early training steps from making large, potentially destructive updates when the value head is still randomly initialized and creating arbitrary rewards.