Lecture 21: RL with Sequence Models & Language Models, Part 2

What this lecture covers

Having covered RL under partial observability, the lecture turns to the opposite direction: using RL to train better language models. It sets up the single-step formulation used by systems like ChatGPT, where a prompt is the state, a full completion is the action, and treats getting the reward function right as the central design problem, since ground-truth rewards for open-ended text rarely exist.

After watching, you can explain how a language generation task is cast as a one-step RL (bandit) problem, why importance-weighted policy gradient estimators are preferred over plain REINFORCE for language models, how a reward model is trained from human preference comparisons rather than raw scores, and how KL penalties guard against reward over-optimization.

Key ideas

• Language models as next-token predictors: a transformer models p(x_t | x_1...x_{t-1}) and generates text by repeatedly sampling the next token, making the process inherently non-Markovian in the token history.
• RL objective for language models: supervised training matches the training data's distribution, while RL training optimizes a language model to maximize a reward function, useful for satisfying preferences, using tools, or holding goal-directed dialogue.
• Single-step formulation: the prompt is the state, the full completion (a variable-length sequence of tokens) is the action, and generating an entire response counts as one RL time step even though it spans many language modeling time steps, making this a bandit problem.
• Importance-weighted policy gradients: because sampling from and evaluating a language model is expensive, most methods generate a batch of completions once, then take several importance-weighted gradient steps against that fixed batch before resampling, rather than resampling for every gradient step as plain REINFORCE would.
• Reward models from preferences: rather than asking humans to assign raw numerical scores (hard for subjective quality), reward models are trained on pairwise preference labels, since people can more reliably say which of two completions they prefer.
• Bradley-Terry-style preference model: the probability that completion A1 is preferred over A2 is modeled as the exponential of A1's reward divided by the sum of both exponentials, the same mathematical form used in maximum entropy inverse RL, and the reward model is trained by maximizing the likelihood of observed human preferences.
• KL penalty against over-optimization: because repeatedly optimizing against a fixed, imperfect reward model risks the policy exploiting it (distributional shift, sometimes called over-optimization), a penalty subtracting the KL divergence from the original supervised policy is added to the reward.

Walkthrough

Language models as sequence predictors, and why use RL (0:00)

The lecture recaps how transformer language models generate text token by token, then explains that supervised training only teaches a model to match its training distribution. RL is introduced as the tool for going beyond that: satisfying human preferences, learning to use tools such as databases or calculators, and holding effective dialogues, all framed as needing a reward function rather than direct imitation of training text.

Casting language generation as a one-step MDP (4:06)

Using "What is the capital of France?" as a running example, the prompt becomes the state, and the completion (for example, the tokens for "Paris" and an end-of-sequence marker) becomes the action. The policy's probability of an action is the product of per-token probabilities from the transformer. A key clarification is that there are two distinct notions of time step: the many language-generation steps inside a completion, and the single RL time step that corresponds to producing an entire completion, making this a bandit problem for now.

Policy gradient estimators for language models (8:09)

The policy gradient for this objective reduces to the same gradients computed during standard cross-entropy training, summed across the completion's tokens. A plain REINFORCE estimator needs fresh samples from the current policy for every gradient step, which is costly since sampling and scoring completions is slow. Instead, most methods sample a batch of completions, fix an old policy pi_bar to describe how they were generated, and then run several importance-weighted gradient steps (in mini-batches) on that fixed batch before generating new samples, a PPO-style loop.

Building a reward model from human preferences (12:11)

Because a language model can produce open-ended answers with no fixed ground truth, the reward function itself must be learned, typically as a neural network. Asking humans to assign a raw numeric score to each completion is unreliable for subjective quality, so the lecture proposes collecting pairwise preferences instead ("I prefer this completion") and modeling the preference probability as the exponential of one reward divided by the sum of both completions' exponentiated rewards, the same functional form used in maximum entropy inverse RL. The reward model is then trained by maximizing the likelihood of the observed human preference labels, and this generalizes to k-way comparisons by decomposing them into pairwise terms.

The full RLHF-style loop and over-optimization (17:13)

The lecture assembles the pieces into the method behind InstructGPT and ChatGPT: start from a supervised fine-tuned policy, sample several completions per prompt, collect human preference labels, train the reward model on those labels, then run RL (with nested loops of sampling and importance-weighted updates) against the reward model, repeating as needed. Because human labeling is slow and expensive, most preference data comes from an early round rather than a continuously updated loop, which effectively makes this a model-based, sometimes offline model-based, RL method with the reward model playing the role of the learned model. The main risk is that the policy overfits and exploits weaknesses in the fixed reward model. The lecture's fix is a KL-divergence penalty, implemented by subtracting the current policy's log probability and adding the supervised policy's log probability to the reward, which keeps the optimized policy close to the original supervised model.

Before you watch

• Review the policy gradient and importance sampling material from earlier in the course, since this lecture reuses those estimators directly.
• Recall the maximum entropy inverse RL formulation from the IRL lecture, since the preference model here uses the same mathematical form.

Check your understanding

1. Why is generating a full language model response treated as a single RL time step even though it involves many token-by-token predictions?
2. Why do most RL methods for language models use importance-weighted updates from a fixed batch of samples rather than resampling for every gradient step?
3. Why are pairwise human preferences generally easier to collect reliably than raw numerical reward scores?
4. What problem does the KL-divergence penalty address, and how is it incorporated into the reward?