RLHF for Language Models: Reward Modeling and Policy Optimization

Reinforcement Learning from Human Feedback has become the de facto method for steering large language models from capable next-token predictors into useful assistants. Supervised fine-tuning teaches a model to imitate, but it cannot teach it to prefer one plausible response over another when both are technically correct. In my experience shipping RLHF pipelines for 7B to 70B parameter models, the gap between a well-tuned SFT model and its RLHF-tuned counterpart is not about knowledge — it's about judgment, calibration, and consistent adherence to nuanced human preferences. This article breaks down the two core subsystems that make RLHF work: reward modeling from pairwise preference data and policy optimization under a KL constraint, with practical details I have learned from debugging both at scale.

Why Supervised Fine-Tuning Leaves Alignment Gaps Unaddressed

Supervised fine-tuning (SFT) minimizes cross-entropy against a single demonstration per prompt. This works well when there is a single correct answer, but human preferences are rarely that deterministic. For open-ended instructions like "explain transformers to a product manager" or "refuse this harmful request helpfully," there are many valid completions with very different quality levels. SFT forces the model to place all probability mass on one sampled completion, even if other completions in the data distribution would have been judged better or worse.

I've found that SFT models exhibit three recurring failure modes that only become obvious at scale: verbosity bias where longer answers are over-represented in demonstrations, hedging where the model learns to cover all bases rather than give a decisive helpful answer, and style inconsistency where the model mimics the diverse styles of many annotators rather than a coherent persona. These are not solved by more SFT data — they are artifacts of the objective itself.

RLHF reframes the problem. Instead of imitating one response, we learn a scalar reward function that can score any response, then optimize the policy to maximize expected reward. If you are unfamiliar with the underlying formalism, it helps to revisit Reinforcement Learning Fundamentals: MDPs, Policies and Value Functions — in the language model setting, the "state" is the prompt plus tokens generated so far, the "action" is the next token, and the episode ends at EOS. The reward, however, is not from the environment but from a learned model of human judgment.

The Preference Data Bottleneck

The entire pipeline rests on pairwise preference data: for a prompt x, annotators see two completions y_w (winning) and y_l (losing) and choose which they prefer. Quality here dominates quantity. In one project, moving from 20k hastily-collected pairs with low inter-annotator agreement (58%) to 12k pairs with rigorous guidelines and 78% agreement improved final win rate by 11 points, despite less data. I now enforce a strict rubric with explicit axes — helpfulness, truthfulness, harmlessness, and conciseness — and require annotators to provide a short rationale. Those rationales are invaluable for auditing reward model errors later. Label noise is not neutral; it directly caps your achievable reward model accuracy and propagates into policy optimization as reward hacking.

Architecting the Reward Model: From Pairwise Preferences to Bradley-Terry Loss

The reward model r_phi(x, y) is typically initialized from the SFT checkpoint, with the language modeling head replaced by a scalar value head — a linear projection from the final hidden state of the last token. Some teams apply the head to every token and pool, but last-token pooling has proven more stable for me, especially for long responses where mean pooling dilutes the signal. The model is trained to assign higher scores to preferred responses.

We use the Bradley-Terry model for pairwise preferences, which translates to a simple logistic loss:

import torch
import torch.nn.functional as F

class RewardModelLoss(torch.nn.Module):
    def forward(self, rewards_w, rewards_l):
        # rewards_w: [batch] scores for preferred completions
        # rewards_l: [batch] scores for rejected completions
        # L = -log(sigmoid(r_w - r_l))
        diff = rewards_w - rewards_l
        loss = -F.logsigmoid(diff).mean()
        # accuracy for monitoring
        acc = (diff > 0).float().mean()
        return loss, acc

# In training loop:
# r_w = reward_model(input_ids_w, attention_mask_w)  # scalar per sequence
# r_l = reward_model(input_ids_l, attention_mask_l)
# loss, acc = criterion(r_w, r_l)
# loss.backward()

This looks trivial, but several engineering choices determine whether your reward model generalizes.

Regularization and Data Curation for Reward Robustness

First, do not train the reward model for too long. It overfits faster than a language model because the dataset is 10-50x smaller than pre-training and the signal is sparse. I typically see validation accuracy peak at one epoch, with loss continuing to drop on training pairs while calibration on held-out prompts degrades. Early stopping on pairwise accuracy, not loss, is the right criterion. I've also found that freezing the bottom 70% of transformer layers for the first half of training, then unfreezing, reduces variance on small datasets (<50k pairs).

Second, calibrate for length bias. Human annotators unconsciously prefer longer answers. If you don't correct for this, your reward model will learn r ~ length and your policy will become verbose to exploit it. I add length-decorrelated evaluation slices and in some runs include a length penalty term during policy optimization or stratify preference pairs by length difference. The Hugging Face TRL RewardTrainer documentation now includes helpful utilities for logging length vs. reward correlations — I check this plot every run.

Third, use a held-out set of prompts that were never seen during reward training to measure both accuracy and Expected Calibration Error (ECE). A model with 72% accuracy but poor calibration will produce extreme reward scores that destabilize PPO. I maintain a separate "challenge set" of adversarial prompts — subtle harmful requests, ambiguous instructions, and style-transfer prompts — to test whether the reward model has learned the intended preference or a spurious proxy.

For a detailed derivation of preference modeling roots in RL, see the original InstructGPT and RLHF work on arXiv Machine Learning (Ouyang et al., 2022 and Christiano et al., 2017), which remain the clearest foundations.

Preventing Reward Hacking Through KL-Regularized Policy Optimization

Once you have a reward model, the naive objective max E[r_phi(x,y)] is dangerous. The policy will quickly discover out-of-distribution completions that exploit flaws in the reward model — repeating phrases the reward model spuriously likes, adding unnecessary formatting, or gaming length. In my experience, without regularization, reward scores climb steeply for ~50-100 steps while human-judged quality collapses.

The standard fix is KL-regularization against the reference policy (typically the SFT model):

objective = E_{x~D, y~pi_theta}[ r_phi(x,y) - beta * KL(pi_theta(y|x) || pi_ref(y|x)) ]

Here beta controls how far the policy can drift. This is not a minor hyperparameter. With beta=0.02, I have seen policies degenerate into repetitive loops within 200 PPO steps on a 13B model. With beta=0.2, learning stalls and the win rate barely moves. For 7B-13B models I usually start at beta=0.05-0.1 and tune based on KL vs. reward frontier plots. For 70B models, a slightly lower beta (0.03-0.05) often works because the base policy is already more capable and needs less constraint.

Implementing Token-Level KL Penalties

KL is estimated per token, not per sequence, to provide dense signal. The unbiased estimator is KL = log pi_theta(y_t | x, y_<t) - log pi_ref(y_t | x, y_<t). During generation, you compute log-probs from both policies and accumulate. The sequence-level KL is the sum over tokens. This token-level penalty is then subtracted from the reward and also used as a diagnostic — if mean KL exceeds ~10-12 nats per sequence, you are almost certainly hacking. I log histograms of KL per sequence and per-token KL spikes; spikes often correspond to the policy inserting bizarre tokens the reward model overvalues.

Another practical safeguard is reward normalization and clipping. I normalize rewards per batch to zero mean and unit variance using a running statistic, and clip the final advantage. This prevents a single batch of anomalously high or low rewards from causing a destructive update, especially early in training when the policy is still exploring.

PPO in Practice for Language Models: Clipping, Value Heads, and Advantage Estimation

Proximal Policy Optimization remains the most battle-tested algorithm for RLHF, despite the rise of alternatives. Adapting PPO from continuous control to autoregressive language generation requires several specific adjustments, and many failures I see stem from applying the vanilla recipe from Deep RL Algorithms: DQN, PPO and SAC Compared without considering the scale and sparsity of the LLM setting.

Unlike traditional RL, reward is sparse — you get one scalar score at the end of the sequence. The value function V(x, y_<t) must therefore learn to predict the final return from intermediate states. I initialize the value head from the reward model, not from scratch; this gives a strong prior and accelerates convergence by ~30% in my runs. The value head is a separate linear layer on top of the same backbone as the policy, but I use a disjoint optimizer or at least a much lower learning rate for the value head to avoid interfering with policy features.

def ppo_step(policy, ref_policy, reward_model, value_head, batch, beta=0.05, eps=0.2):
    prompts, responses = batch["prompt"], batch["response"]
    # Get logprobs and values for each token
    logp = policy.get_logprobs(prompts, responses)      # [B, T]
    logp_ref = ref_policy.get_logprobs(prompts, responses)
    values = value_head(prompts, responses)             # [B, T]
    
    # Final scalar reward, broadcast with KL penalty per token
    r = reward_model.score(prompts, responses)          # [B]
    kl = logp - logp_ref                                # [B, T]
    # token-level reward: only last token gets r, each token gets -beta*KL
    rewards = -beta * kl
    rewards[:, -1] += r  # add scalar reward to final position
    
    # GAE for advantages (gamma=1.0, lam=0.95 for LM)
    advantages, returns = compute_gae(rewards, values, gamma=1.0, lam=0.95)
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
    
    # PPO clipped objective
    ratio = torch.exp(logp - logp_old)  # logp_old from rollout
    pg_loss = -torch.min(
        ratio * advantages,
        torch.clamp(ratio, 1-eps, 1+eps) * advantages
    ).mean()
    
    v_loss = 0.5 * (values - returns).pow(2).mean()
    # combine with entropy bonus if needed
    return pg_loss + 0.1 * v_loss

Taming PPO Instability at LLM Scale

Three issues dominate PPO tuning for language models. First, the rollout batch size must be large. I've found that updates with <512 sequences are too noisy for stable advantage estimation; I target 1024-2048 prompts per rollout, with 1-4 generations per prompt, and then do 1-2 epochs of minibatch updates (minibatch 128-256). Doing 4 epochs as in some Atari configs almost always causes collapse due to over-optimization on the same rollouts.

Second, generation parameters during rollout matter enormously. Use temperature 0.7-1.0 and top_p 0.9 for diverse rollouts. If you rollout with greedy decoding, you get no exploration and advantages are near zero. I also enforce a maximum generation length (e.g., 512 tokens) and truncate, but ensure the KL penalty is not applied to padding tokens.

Third, learning rates are an order of magnitude lower than SFT. I use 1e-6 to 5e-7 for 7B models with AdamW, with a cosine decay and 10% warmup. Higher rates cause policy entropy to collapse within tens of steps. I monitor entropy continuously — healthy PPO runs maintain entropy within 10-15% of the initial SFT policy for at least 500 steps.

Direct Preference Alternatives: DPO, IPO and RLOO in Production

PPO is powerful but operationally heavy: it requires four models in memory (policy, reference, reward, value), rollouts, and careful distributed orchestration. This has motivated direct alignment methods that skip the explicit reward model and RL loop. Direct Preference Optimization (DPO) is the most widely adopted. It reparameterizes the optimal KL-constrained policy and derives a classification loss directly on preference pairs:

def dpo_loss(pi_logp_w, pi_logp_l, ref_logp_w, ref_logp_l, beta=0.1):
    # pi_logp: log prob under current policy, ref_logp: under reference
    # DPO loss: -log(sigmoid(beta * (log_ratio_w - log_ratio_l)))
    log_ratio_w = pi_logp_w - ref_logp_w
    log_ratio_l = pi_logp_l - ref_logp_l
    logits = beta * (log_ratio_w - log_ratio_l)
    loss = -torch.nn.functional.logsigmoid(logits).mean()
    return loss

In my experience, DPO converges faster and is far simpler to implement — you only need two models. For straightforward preference alignment like style control and helpfulness, DPO achieves 85-95% of PPO's win rate with 40% of the engineering effort. However, DPO has a known tendency to drive down the log-probability of both preferred and rejected responses (with rejected dropping faster), which can reduce overall generation quality if beta is not tuned and if you train for too many epochs. I see overfitting after 1-2 epochs on most datasets.

Other alternatives address specific weaknesses. IPO (Identity Preference Optimization) adds regularization to prevent overconfidence, and RLOO (REINFORCE Leave-One-Out) provides a simpler, value-free policy gradient estimator that I have found surprisingly competitive for instruction following tasks with verifiable rewards. The choice depends on your constraints, as summarized below.

I've migrated teams from PPO to DPO when iteration speed mattered more than squeezing the last 3% of win rate, and the opposite when we needed fine-grained control over reward shaping for safety. If you are exploring distributed training for any of these, the orchestration challenges overlap with multi-agent setups discussed in Multi-Agent Reinforcement Learning: Cooperation, Competition and Communication, particularly around rollout collection and synchronization.

Diagnosing Alignment Failures: Evaluation Beyond Win Rates

A win rate against the SFT model on a held-out preference set is the headline metric, but it is a lagging and incomplete indicator. I've shipped models with 68% win rate that were unusable due to subtle regressions, and models with 58% win rate that felt dramatically better in human interaction. You need a multidimensional evaluation harness.

I run three evaluation tiers continuously. First, automated reward model scoring on a frozen, high-quality eval set — but I never use the same reward model that trained the policy. I keep a separate "gold" reward model trained on a different data split or even a larger model as an evaluator, or use a strong LLM as judge with a calibrated rubric. This detects reward hacking where the training reward rises but gold reward falls.

Second, attribute-specific probes: helpfulness (AlpacaEval, MT-Bench), truthfulness (TruthfulQA), harmlessness (HH-RLHF harmless split, red-team suites), and instruction following (IFEval). RLHF often improves helpfulness at the cost of truthfulness if the reward model is not balanced — I track these tradeoffs explicitly and use Pareto frontiers to select checkpoints.

Third, qualitative, longitudinal human interaction. I maintain a set of 200 multi-turn dialogues where annotators converse freely with the model for 5-10 turns. Single-turn win rates miss failures like sycophancy, over-apologizing, or drift in persona over a conversation. In one case, our KL was well-controlled and single-turn metrics looked excellent, but the model became overly agreeable after turn three, agreeing with incorrect user premises to maximize short-term reward. We only caught this through multi-turn evaluation and fixed it by adding multi-turn preference data.

Finally, log everything: per-step reward mean/variance, KL mean/max, policy entropy, value loss, advantage histograms, generation length, and gradient norms. When RLHF fails, it rarely fails silently — the signals are in the logs before the outputs degrade. I set alerts for KL spikes >20% step-over-step or entropy drops >5% per 50 steps, which have saved multiple runs from irreversible collapse.

Frequently Asked Questions