Finetuning LLMs with RL

Chapter 1: Supervised Finetuning Loss

SFT works exactly like pre-training, but on curated data. Given a sequence of tokens x₁, x₂, ..., x_T, the model predicts each next token given all previous tokens. The loss measures how surprised the model is by the correct answer.

Let's derive the loss from scratch. At each position t, the model outputs a probability distribution over the entire vocabulary. We want the model to assign high probability to the actual next token x_t+1. The natural measure of "how wrong was the prediction" is cross-entropy.

Step 1: The per-token loss

At position t, the model sees tokens x₁, x₂, ..., x_t and outputs logits z ∈ ℝ^V where V is vocabulary size (e.g. 32,000 for LLaMA). We convert logits to probabilities with softmax:

The per-token loss is the negative log of this probability:

When the model assigns probability 1.0 to the correct token, ℓ_t = 0. When it assigns probability 0.01, ℓ_t = 4.6. The loss is always non-negative and is zero only when the model is perfectly confident and correct.

Step 2: The full SFT loss

We average the per-token loss over all T positions in the sequence:

What actually flows through the network

python
import torch
import torch.nn.functional as F

# logits: [B, T, V] from the model's final linear layer
# targets: [B, T] — the input tokens shifted left by 1
logits = model(input_ids)            # [B, T, V]
shift_logits = logits[:, :-1, :]    # [B, T-1, V]
shift_labels = input_ids[:, 1:]     # [B, T-1]

loss = F.cross_entropy(
    shift_logits.reshape(-1, V),     # [B*(T-1), V]
    shift_labels.reshape(-1),        # [B*(T-1)]
)  # scalar — mean cross-entropy over all tokens

Chapter 2: RL Fundamentals

Why do we need reinforcement learning at all? Because SFT requires explicit demonstrations of correct behavior. With RL, we only need to say whether an output is good or bad (a scalar reward), and the model figures out how to produce good outputs on its own.

RL has five core concepts. Let's build them one by one.

States, Actions, and Rewards

An agent lives in an environment. At each timestep t, the agent observes a state s_t, takes an action a_t, receives a reward r_t, and transitions to a new state s_t+1. This cycle repeats until the episode ends.

The Policy π

The policy π(a | s) is the agent's strategy: a probability distribution over actions given the current state. "When I see state s, what should I do?" A good policy assigns high probability to actions that lead to high reward.

Return: The Total Payoff

We don't just want high reward now — we want high reward over the entire episode. The return G_t is the sum of future rewards, with a discount factor γ ∈ [0, 1] that makes future rewards worth slightly less than immediate ones:

When γ = 0, the agent is completely myopic — only the immediate reward matters. When γ = 1, all future rewards are equally important. Typical values: γ = 0.99.

Value Functions

The value function V^π(s) is the expected return from state s when following policy π:

The action-value function Q^π(s, a) is the expected return from state s after taking action a, then following π:

Chapter 3: The Policy Gradient — REINFORCE

We have a policy π_θ(a | s) parameterized by θ (neural network weights). Our goal is to find the θ that maximizes the expected return:

where τ = (s₀, a₀, r₀, s₁, a₁, r₁, ...) is a trajectory and R(τ) = ∑_t r_t is the total return. We want ∇_θ J(θ) so we can do gradient ascent. But the expectation is over trajectories sampled from the policy, so how do we differentiate through sampling?

Derivation: The Log-Probability Trick

This is the key insight. Let's derive it step by step.

Step 1. Write the expected return as an integral over trajectories:

Step 2. Take the gradient with respect to θ. The reward R(τ) doesn't depend on θ, only the trajectory probability does:

Step 3. Here's the trick. We multiply and divide by P(τ; θ):

Step 4. Recognize that ∇f / f = ∇ log f (the log-derivative trick):

Step 5. That sum weighted by P(τ; θ) is just an expectation! So:

Step 6. The trajectory probability P(τ; θ) factors as p(s₀) ∏_t π_θ(a_t|s_t) p(s_t+1|s_t,a_t). Taking the log turns the product into a sum. The environment dynamics p(s_t+1|s_t,a_t) don't depend on θ, so their gradients vanish:

The REINFORCE gradient:

The variance problem

REINFORCE is unbiased but has very high variance. Why? Because R(τ) is the total return of the entire trajectory. If a trajectory got high reward, we reinforce every action — even the bad ones that happened to be on a lucky trajectory. Two standard fixes:

1. Baseline subtraction. Replace R(τ) with R(τ) − b, where b is a baseline (often V(s)). This doesn't change the expected gradient but reduces variance dramatically.

2. Use the advantage. Replace R(τ) with A_t = Q(s_t, a_t) − V(s_t). This says "how much better was this action than average?" rather than "was the whole episode good?"

Chapter 4: PPO — Proximal Policy Optimization

REINFORCE has a fatal flaw: the gradient can be huge. One lucky trajectory can massively change the policy, destroying what was learned before. We need controlled updates — change the policy, but not too much at once.

The probability ratio

Define the ratio between the new and old policy:

When r_t = 1, the new policy is identical to the old one. When r_t > 1, the action became more likely. When r_t < 1, it became less likely.

The surrogate objective

Instead of REINFORCE, we optimize a surrogate objective that uses the ratio:

This is called the Conservative Policy Iteration (CPI) objective. It's equivalent to the policy gradient but lets us take multiple gradient steps on the same batch of data. The problem: if A_t > 0 (good action) and we keep increasing r_t, we get an ever-larger objective — but the policy might change so much it collapses.

PPO's clipping trick

PPO clips the ratio to prevent it from straying too far from 1. The clipped surrogate is:

where ε is a small number, typically 0.2. Let's unpack what happens in each case:

The full PPO loss

PPO also adds a value function loss (to train the baseline) and an entropy bonus (to encourage exploration):

where L^VF = (V_θ(s) − V_target)² and H is the entropy of the policy.

Chapter 5: RL for Language Models

How do you apply RL to text generation? The mapping is surprisingly direct:

The RLHF Objective

We want to maximize reward, but with a crucial constraint: the finetuned model shouldn't drift too far from the SFT model. Why? Because the reward model is imperfect — if we optimize it too aggressively, the LLM finds degenerate outputs that "hack" the reward (high reward, low quality). The solution is a KL penalty:

Where:

• r_φ(x, y) is the reward model's score for prompt x, response y
• π_ref is the frozen SFT model (our "anchor")
• β controls the strength of the penalty (β ≈ 0.01–0.2)
• KL divergence measures how different π_θ is from π_ref

The KL Divergence Penalty

KL divergence measures the difference between two probability distributions. For two distributions p and q over vocabulary V:

KL is always ≥ 0, and equals 0 only when p = q. In the RLHF context, it penalizes the model for deviating from the reference — the more different the token distribution, the larger the penalty. This prevents reward hacking: the model can't just memorize one high-scoring response and output it for every prompt.

Chapter 6: Reward Modeling

RLHF needs a reward signal, but we can't have a human rate every single model output during training. Instead, we train a reward model to predict human preferences. Then we use this learned reward model as the scoring function for RL.

Collecting Preference Data

The process is straightforward:

The Bradley-Terry Model

We need a mathematical model of human preferences. The Bradley-Terry model (1952) assumes there exists a latent reward r*(x, y) such that the probability of preferring y₁ over y₂ is:

This is just a sigmoid! Let Δ = r*(x, y₁) − r*(x, y₂). Then:

When Δ > 0 (y₁ has higher reward), the probability of preferring y₁ is above 0.5. When Δ = 0, it's a coin flip. When Δ < 0, y₂ is preferred.

Deriving the Reward Modeling Loss

We parameterize the reward with a neural network r_φ(x, y) — typically the SFT model with a linear head that outputs a scalar. The loss maximizes the log-likelihood of the observed preferences:

Let's unpack this:

• r_φ(x_i, y_i^w) is the reward model's score for the preferred response
• r_φ(x_i, y_i^l) is the score for the rejected response
• σ(·) is the sigmoid function
• We want the winner's score to be higher than the loser's, so Δ = r_w − r_l > 0
• −log σ(Δ) is small when Δ is large positive (model agrees with human), large when Δ is negative (model disagrees)

python
# Reward model training step
def reward_loss(r_phi, x, y_w, y_l):
    # r_phi: reward model (LLM + linear head)
    # x: prompts [B]
    # y_w: preferred responses [B]
    # y_l: rejected responses [B]
    r_w = r_phi(x, y_w)    # [B] — scalar reward for winners
    r_l = r_phi(x, y_l)    # [B] — scalar reward for losers
    loss = -torch.log(
        torch.sigmoid(r_w - r_l)
    ).mean()                # scalar
    return loss

Chapter 7: The Full RLHF Loop

Now we have all the pieces. Let's assemble the complete RLHF pipeline, step by step, with concrete data shapes at every stage.

Stage 1: Supervised Finetuning

Stage 2: Reward Model Training

Stage 3: PPO Training

This is where the magic happens. Four models are active simultaneously:

One PPO Iteration

Advantage Estimation for LLMs

In the LLM setting, a common simplification is to set the discount factor γ = 0. The reward is given once at the end of the full response. This means the advantage for the entire sequence simplifies to:

That is: "How much better was the actual reward than what the value model predicted?" If the response was surprisingly good (A > 0), increase its probability. If surprisingly bad (A < 0), decrease it.

Chapter 8: Connections

Cheat Sheet

Method Comparison

What to explore next

• Reward & Alignment — deeper dive into DPO, RLAIF, Constitutional AI
• GPT — Generative Pre-trained Transformer — how the base model works
• The Transformer — the architecture underlying everything