Reinforcement learning with human feedback (RLHF)

Why RLHF?

In classical RL, the reward function is known. But for many tasks, the reward is hard to specify:

• It’s easy to judge which poem is better, but hard to write a masterpiece: the discriminator-generator gap
• It’s easy to spot a helpful answer, but hard to specify what “helpfulness” means as a formula
• Pre-training optimizes next-token prediction, the most likely continuation isn’t necessarily the most useful

RLHF converts human judgments, often pairwise comparisons or rankings, into a learned proxy reward model. This allows optimization in subjective and non-verifiable domains, although the resulting reward model is noisy and can be exploited.

Note

Classical RL only requires an accessible reward signal, RLHF learns that signal from human preference data.

Relation to post-training

RLHF is a subset of post-training, which is a set of methods and best practices to make language models more useful for downstream tasks. Post-training has three primary methods:

1. Instruction / Supervised Fine-tuning (IFT/SFT), where we teach formatting and form the base of instruction-following abilities. This is largely about learning features in language.
2. Preference Fine-tuning (PreFT), where we align to human preferences via RLHF and related methods (and get a smaller bump in capabilities at the same time). This is largely about style of language and subtle human preferences that are hard to quantify.
3. Reinforcement Learning with Verifiable Rewards (RLVR), the newest type of post-training that boosts performance on verifiable domains with more RL training.

RLHF lives in the second subset. RLHF made the breakout success with the release of InstructGPT and ChatGPT early in 2023, so it encompassed all of post-training back then. Now in modern times, post-training is much more complicated with many stages and complimenting methods to RLHF.

RLHF lineage

Classical RL vs. RLHF

Classical RL:

• Agent takes actions in an environment with states
• Reward is a known function from the environment per step
• Optimize cumulative return over a trajectory (total steps )

RLHF:

• No environment, prompts sampled from a dataset, agent generates completions
• Reward is learned from human preferences (a proxy)
• Response-level reward (bandit-style, not per-token)
• Regularized with KL penalty to stay close to the base model

The 3-stage pipeline: SFT, RM, RLHF

Popularized by InstructGPT:

Instruction fine-tuning

The goal is to format the next-word continuation responses into a helpful assistant format.

Start from a pretrained model checkpoint. Collection desired human-assistant response pairs. Train with standard next-token prediction loss (same in pre-training) on these pairs with different learning rate and batch size.

Prompt masking: One important distinction with pre-training is that the loss on the prompt is masked out since that is given. The language modeling loss is only calculated on the response tokens conditioned on the prompt. This is so that the model don’t waste compute recreating the prompt.

Model can now answer questions. Easy to use IFT to quickly adapt base models to many domains.

Reward modeling

Same as GPT-1 task-specific adapters or foundation vision models, chop off the final linear projection layer (project LM’s hidden state back to token space) and pass final non-padding hidden state through a new linear/MLP layer to get a scalar reward (pseudo-regression problem). This would be the reward model. Outputs a scalar reward given a prompt and completion.

1. Collects comparisons between different outputs generated by the model for the same prompt.
2. Have a human (or in today’s world, another AI) to score them. Use these scores to produce winner-loser pairs of the responses.
3. Train a reward model to rank winner completions higher.

The probability model says a response should win when it gets a higher reward score. We do this because the scalar reward is unbounded, sigmoid squishes the relative difference into a preference probability:

Training then minimizes the negative log-likelihood of the preferred response beating the rejected one, essentially binary cross-entropy:

This is a contrastive loss that pulls apart the winners from the losers in the model’s representation.

The reward modeling loss punishes cases where the model rewards losing answers more than winning answers.

Note

The reward used in RLHF is the reward model predicting the probability of a given piece of text being the winner in a pair/batch!

Note

The reward model recipe can also be applied to train a proxy for data quality, to be used in filtering data!

Reinforcement learning with a reward model

1. Take prompts representative of tasks the model should be good at
2. Generate a bunch of completions
3. Have the reward model rank them
4. Use RL to figure out how to change the model and make it better

The basic primitive is that reinforcement learning is given a signal of which actions are good, in the form of tokens that a language model generates, and derives update rules that attribute different actions to different parameters in the model. The final RLHF stage shifts parameters to make good tokens more likely, and does so iteratively to maintain the general capabilities of the initial model.

Once RL is complete, and performance has saturated, this is often the final model served to the user.

• Expected value of logprob is just entropy of the distribution.
• The RL policy distribution: . This is a dynamic, on-policy distribution because you sample from your training prompt dataset, then you feed that into your current, actively updating RL model . The model generates a response . As the model updates ( changes), this distribution changes.
• The pretraining distribution: . This is the static, offline dataset used during pre-training of the base model. It is used as an unchanging anchor so the model doesn’t forget general knowledge or grammar.

Note

We cannot use back-propagation to update policy parameters because back-propagation cannot pass through discrete token sampling, so policy gradients estimate how token probabilities should change from the scalar reward.

Step-by-step walkthrough of the calculation (Gemini)

To calculate this objective during a training step, the system processes a batch of data through three distinct components:

Step 1: Generate Data and Gather Policy Logs

First, the system samples a batch of prompts from the RL training prompt set.

1. The prompts are passed through the active RL Policy () to generate completions . This gives us our pairs.
2. For every token generated in , the system computes and saves the log-probability of the active model generating it: .

Step 2: Compute the Reference Model Logs

Next, the exact same pairs generated in Step 1 are fed into a frozen copy of the model—the SFT Reference Model ().

1. The system calculates the log-probability of the reference model generating that exact same response: .
2. Note: The reference model does not generate text here; it just evaluates the mathematical probability of the text the RL model generated.

Step 3: Score with the Reward Model

The pairs are also fed into the trained Reward Model ().

1. The Reward Model outputs a single scalar score representing how good, safe, or helpful the response is.

Step 4: Calculate the First Expectation (Reward + KL)

Now, the system computes the value inside the first bracket for every sample in the batch:

1. Calculate the Log Ratio: Subtract the SFT log-prob from the RL log-prob:

1. Apply Scaling: Multiply this value by the hyperparameter .
2. Combine: Subtract this penalized log-ratio from the scalar reward:

1. Average: Average these scores across the entire batch to get the empirical mean ().

Step 5: Calculate the Pretraining Term

In parallel, the system takes a separate batch of text sequences from the raw pretraining dataset ().

1. It passes them through the active RL Policy ().
2. It calculates the log-probabilities of those sequences: .
3. It averages these log-probabilities across the batch and multiplies the result by .

Step 6: Total Objective and Gradient Update

Finally, the system adds the average from Step 4 to the average from Step 5.

This final scalar value is the total objective(ϕ). Because this is a maximization problem, the system computes the gradient of this objective with respect to the active model’s weights () and performs a gradient ascent update:

This pushes the active model parameters () to generate text that gets higher rewards, stays close to the SFT model, and maintains its original pretraining capabilities.

Elicitation theory of post-training

The idea is that there is a lot of intelligence and ability within base models, but because they can only answer in next-token prediction and not question-answering format, it takes a lot of work building around them, through post-training, in order to make excellent final models.

A related idea is from LIMA: Less is More for Alignment paper, the Superficial Alignment Hypothesis

A model’s knowledge and capabilities are learnt almost entirely during pretraining, while alignment teaches it which subdistribution of formats should be used when interacting with users. If this hypothesis is correct, and alignment is largely about learning style, then a corollary of the Superficial Alignment Hypothesis is that one could sufficiently tune a pretrained language model with a rather small set of examples.

Takeaway

How do we map the complexities of human values and objectives into systems we use on a regular basis?