RLHF — Reinforcement Learning from Human Feedback

Summary

RLHF (Reinforcement Learning from Human Feedback) is an alignment technique that fine-tunes a pre-trained language model so that it produces responses aligned with human preferences. This method rests on three fundamental pillars: collecting preference data from human annotators, training a reward model capable of scoring any response, and optimizing the language model via a reinforcement algorithm such as PPO (Proximal Policy Optimization), regularized by a KL divergence penalty to prevent model degeneration.

RLHF is at the core of the success of foundational models such as ChatGPT, Claude, and Gemini. Without this technique, these systems would produce responses that are fluent but potentially toxic, biased, or useless — because a language model pre-trained solely on internet text has no notion of what is “polite,” “useful,” or “safe.” RLHF precisely fills this gap.

Mathematical Principle of RLHF

RLHF unfolds in three distinct steps, each grounded in rigorous mathematical foundations.

Step 1 — Training the reward model r_φ(x, y)

The first step consists of collecting pairwise human comparisons. For the same question (or prompt) x, human annotators are asked to compare two responses y₁ and y₂ generated by the model, and to designate the preferred response y_w (winner) and the rejected response y_l (loser).

With these pairs (x, y_w, y_l), a reward model parameterized by φ, denoted r_φ(x, y), is trained. Its loss function is written as:

L_RM(φ) = -E_[x,y_w,y_l)[log(σ(r_φ(x, y_w) - r_φ(x, y_l)))]

where σ denotes the sigmoid function. This loss is a rewritten version of logistic regression with pairs (Bradley-Terry model). Intuitively, the reward model learns to assign higher scores to responses preferred by humans. The larger the difference r_φ(x, y_w) – r_φ(x, y_l), the closer the predicted probability is to 1, and the lower the loss.

The reward model is typically initialized from the same pre-trained language model, with the output head replaced by a regression scalar. The final language layer is removed and a linear classifier is added that takes the representation of the last token and produces a scalar score.

Step 2 — Freezing the reference model π_ref

Once the reward model is trained, the initial language model is frozen. This frozen model, denoted π_ref, will serve as a reference during the fine-tuning step. Its role is crucial: it prevents the fine-tuned model from straying too far from its initial behavior, which would avoid degeneration (the model could learn to maximize the reward in a pathological way, producing nonsensical but well-scored responses).

Step 3 — PPO fine-tuning with KL penalty

This is the central step of RLHF. The language model π_θ (parameterized by θ) is optimized using the PPO (Proximal Policy Optimization) algorithm. The objective to maximize is:

J(θ) = E_[x~D, y~π_θ(·|x)][r_φ(x, y) - β · KL(π_θ(·|x) || π_ref(·|x))]

Let’s break down this objective:

• r_φ(x, y): the reward assigned by the reward model to response y for prompt x. The language model is encouraged to generate responses that receive high scores.
• KL(π_θ(·|x) || π_ref(·|x)): the Kullback-Leibler divergence between the fine-tuned model distribution and the reference model distribution. This measure quantifies how much π_θ has drifted from π_ref. It acts as a regularization term: it penalizes too-abrupt changes in model behavior.
• β (beta): the hyperparameter that controls the trade-off between alignment and fidelity to the original model.
• High β (e.g.: 0.2–0.5): the model stays very close to the reference behavior. It is stable but may be less well-aligned with human preferences.
• Low β (e.g.: 0.01–0.05): the model is strongly pushed to maximize the reward. It may align better but risks degeneration: producing strange or repetitive responses that exploit the reward model (a phenomenon called reward hacking).

The choice of PPO over other reinforcement algorithms (like REINFORCE or SAC) is explained by its stability: PPO uses a truncated probability ratio (clipping) that prevents overly aggressive updates, which is essential when training a model with several billion parameters.

Analogical Intuition

Imagine a pre-trained LLM as a brilliant student who has read every book in the library but never learned good manners. They can recite facts, write essays, and even make jokes — but they don’t know when to be quiet, when to qualify their statements, or when a response might be offensive.

RLHF is a coach who will teach the student how to behave in society:

1. The reward model is the coach who observes the student and says: “This response is polite, that one is arrogant.” They learn to judge response quality by human criteria.
2. KL regularization is the coach reminding the student: “Don’t lose your personality. You can be more polite without becoming robotic.” It prevents the model from forgetting what it learned during pre-training.
3. PPO is guided practice: the student tries to respond, the coach corrects them immediately, and they gradually adjust their behavior. No abrupt changes — small iterative corrections, like a musician refining their performance measure by measure.

Without this coaching, the model remains an awkward scholar. With RLHF, it becomes a competent, pleasant interlocutor.

Python Implementation with PyTorch

Here is a simplified but complete implementation of the RLHF pipeline, simulating a reward model on text classification and a PPO loop with KL penalty.

"""
Simplified RLHF — Implementation with PyTorch
Reward Model + PPO Loop + KL Penalty
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import numpy as np


# ============================================================
# 1. Reward Model
# ============================================================
class RewardModel(nn.Module):
    """
    Reward model based on a text encoder.
    Takes a sequence of embeddings and returns a scalar score.
    """
    def __init__(self, vocab_size=30522, embed_dim=128, hidden_dim=256):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=embed_dim,
                nhead=4,
                dim_feedforward=hidden_dim,
                batch_first=True
            ),
            num_layers=2
        )
        self.regression_head = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, x, attention_mask=None):
        """
        x: tensor of tokens [batch_size, seq_len]
        Returns a scalar score per sample.
        """
        emb = self.embedding(x)  # [B, T, E]
        encoded = self.encoder(emb)  # [B, T, E]
        # We use the last token vector
        if attention_mask is not None:
            lengths = attention_mask.sum(dim=1)  # [B]
            batch_indices = torch.arange(x.size(0))
            last_tokens = encoded[batch_indices, lengths - 1]
        else:
            last_tokens = encoded[:, -1, :]  # [B, E]
        score = self.regression_head(last_tokens).squeeze(-1)  # [B]
        return score


def train_reward_model(
    reward_model,
    pairs_y_w, pairs_y_l,
    lr=1e-4, epochs=10, batch_size=32
):
    """
    Trains the reward model on comparison pairs.
    pairs_y_w: preferred responses (winners) [N, seq_len]
    pairs_y_l: rejected responses (losers) [N, seq_len]
    """
    optimizer = torch.optim.Adam(reward_model.parameters(), lr=lr)
    dataset = TensorDataset(pairs_y_w, pairs_y_l)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    reward_model.train()
    for epoch in range(epochs):
        total_loss = 0.0
        for y_w, y_l in dataloader:
            r_w = reward_model(y_w)  # scores of preferred responses
            r_l = reward_model(y_l)  # scores of rejected responses

            # Bradley-Terry loss
            diff = r_w - r_l
            loss = -F.logsigmoid(diff).mean()

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(reward_model.parameters(), max_norm=1.0)
            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        print(f"  Epoch {epoch+1}/{epochs} — Reward loss: {avg_loss:.4f}")

    return reward_model


# ============================================================
# 2. Policy Model (the LLM being fine-tuned)
# ============================================================
class PolicyModel(nn.Module):
    """
    Simplified policy model for text generation.
    Uses the same architecture as the reward model
    but with a generation head (token classification).
    """
    def __init__(self, vocab_size=30522, embed_dim=128, hidden_dim=256, max_len=32):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=embed_dim,
                nhead=4,
                dim_feedforward=hidden_dim,
                batch_first=True
            ),
            num_layers=2
        )
        self.lm_head = nn.Linear(embed_dim, vocab_size)
        self.max_len = max_len
        self.vocab_size = vocab_size

    def forward(self, x, attention_mask=None):
        emb = self.embedding(x)
        encoded = self.encoder(emb)
        logits = self.lm_head(encoded)
        return logits  # [B, T, V]

    def generate(self, prompt, max_len=16):
        """Autoregressive sampled generation."""
        self.eval()
        generated = prompt.clone()
        for _ in range(max_len):
            logits = self.forward(generated)
            next_token_logits = logits[:, -1, :]
            probs = F.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            generated = torch.cat([generated, next_token], dim=1)
        return generated

    def compute_log_probs(self, tokens):
        """Computes token-by-token log-probabilities."""
        logits = self.forward(tokens)  # [B, T, V]
        log_probs = F.log_softmax(logits, dim=-1)  # [B, T, V]
        return log_probs


# ============================================================
# 3. KL Divergence Computation
# ============================================================
def compute_kl_divergence(log_pi_theta, log_pi_ref):
    """
    Computes KL(π_θ || π_ref) token by token.
    log_pi_theta: [B, T, V] — log-probs of the fine-tuned model
    log_pi_ref: [B, T, V] — log-probs of the reference model

    KL = sum_y π_θ(y|x) * log(π_θ(y|x) / π_ref(y|x))
    Sampling approximation:
    KL ≈ log_pi_theta - log_pi_ref (averaged over generated tokens)
    """
    # We use the approximation: KL ≈ E_{y~pi_theta}[log pi_theta - log pi_ref]
    # For each generated token, we take the log-prob of the token under pi_theta
    # and subtract the log-prob of the same token under pi_ref
    pi_theta_probs = torch.exp(log_pi_theta)
    kl = (pi_theta_probs * (log_pi_theta - log_pi_ref)).sum(dim=-1)  # [B, T]
    return kl.sum(dim=1)  # [B] — sum over the sequence


# ============================================================
# 4. PPO Loop for RLHF
# ============================================================
class PPOTrainer:
    """
    Simplified PPO trainer for RLHF.
    Implements the optimization loop with clipping and KL penalty.
    """
    def __init__(
        self,
        policy,
        ref_policy,
        reward_model,
        beta=0.1,
        clip_epsilon=0.2,
        ppo_epochs=4,
        lr=1e-5,
        gamma=0.99,
        gae_lambda=0.95,
        max_grad_norm=1.0
    ):
        self.policy = policy
        self.ref_policy = ref_policy
        self.reward_model = reward_model
        self.beta = beta            # KL coefficient
        self.clip_epsilon = clip_epsilon  # PPO clipping threshold
        self.ppo_epochs = ppo_epochs
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.max_grad_norm = max_grad_norm

        self.optimizer = torch.optim.Adam(policy.parameters(), lr=lr)

    def compute_advantages(self, rewards, values):
        """Computes advantages via GAE (Generalized Advantage Estimation)."""
        advantages = []
        gae = 0.0
        for t in reversed(range(len(rewards))):
            if t == len(rewards) - 1:
                delta = rewards[t] - values[t]
            else:
                delta = rewards[t] + self.gamma * values[t+1] - values[t]
            gae = delta + self.gamma * self.gae_lambda * gae
            advantages.insert(0, gae)
        return torch.tensor(advantages, dtype=torch.float32)

    def ppo_update(
        self,
        prompts,
        generated_responses,
        rewards,
        old_log_probs
    ):
        """
        One PPO update step.

        prompts: [B, prompt_len]
        generated_responses: [B, gen_len]
        rewards: [B] — scalar rewards
        old_log_probs: [B] — old log-probs (for the PPO ratio)
        """
        for _ in range(self.ppo_epochs):
            self.policy.train()

            # Re-generate and compute log-probs
            with torch.no_grad():
                # Reuse the generated responses
                full_responses = torch.cat([prompts, generated_responses], dim=1)
                new_log_probs_full = self.policy.compute_log_probs(full_responses)

            # Extract log-probs for the generated part only
            prompt_len = prompts.size(1)
            gen_len = generated_responses.size(1)

            # Log-probs of generated tokens under the current policy
            new_log_probs = []
            for i in range(prompts.size(0)):
                token_log_probs_i = 0.0
                for t in range(gen_len):
                    token_idx = prompt_len + t
                    token_id = generated_responses[i, t].item()
                    if token_id < new_log_probs_full.size(2):
                        token_log_probs_i += new_log_probs_full[i, token_idx, token_id]
                new_log_probs.append(token_log_probs_i)
            new_log_probs = torch.stack(new_log_probs)

            # PPO ratio: π_θ / π_old
            ratio = torch.exp(new_log_probs - old_log_probs)

            # PPO objective with clipping
            surr1 = ratio * rewards
            surr2 = torch.clamp(ratio, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * rewards
            policy_loss = -torch.min(surr1, surr2).mean()

            # KL penalty
            with torch.no_grad():
                log_pi_theta = self.policy.compute_log_probs(full_responses)
                log_pi_ref = self.ref_policy.compute_log_probs(full_responses)
                kl_penalty = compute_kl_divergence(log_pi_theta, log_pi_ref).mean()

            # Total objective
            total_loss = policy_loss + self.beta * kl_penalty

            # Gradient update
            self.optimizer.zero_grad()
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(
                self.policy.parameters(),
                max_norm=self.max_grad_norm
            )
            self.optimizer.step()

            return total_loss.item(), policy_loss.item(), kl_penalty.item()

    def train_step(self, prompts, reward_model=None):
        """
        One complete RLHF step:
        1. Generate responses
        2. Score with the reward model
        3. Update the policy with PPO
        """
        rm = reward_model if reward_model is not None else self.reward_model
        self.ref_policy.eval()
        rm.eval()

        # Generation phase
        with torch.no_grad():
            responses = self.policy.generate(prompts, max_len=16)
            # Compute reference log-probs
            ref_full = torch.cat([prompts, responses], dim=1)
            ref_log_probs = self.ref_policy.compute_log_probs(ref_full)

            # Score with the reward model
            rewards = rm(responses)  # [B]

            # Apply KL penalty to the reward
            prompt_len = prompts.size(1)
            kl_per_sample = compute_kl_divergence(
                self.policy.compute_log_probs(ref_full),
                ref_log_probs
            )
            adjusted_rewards = rewards - self.beta * kl_per_sample

            # Old log-probs
            old_log_probs = adjusted_rewards.detach().clone()

        # PPO optimization phase
        total_loss, pol_loss, kl_val = self.ppo_update(
            prompts, responses, adjusted_rewards, old_log_probs
        )

        return {
            "total_loss": total_loss,
            "policy_loss": pol_loss,
            "kl_penalty": kl_val,
            "mean_reward": rewards.mean().item()
        }


# ============================================================
# 5. Complete Usage Example of the RLHF Pipeline
# ============================================================
if __name__ == "__main__":
    torch.manual_seed(42)
    np.random.seed(42)

    print("=" * 60)
    print("RLHF Pipeline — Complete Implementation")
    print("=" * 60)

    # Configuration
    VOCAB_SIZE = 1000
    EMBED_DIM = 64
    HIDDEN_DIM = 128
    SEQ_LEN = 20
    BATCH_SIZE = 8
    NUM_PAIRS = 200

    # --- Step 1: Create and train the reward model ---
    print("\n[Step 1] Training the reward model...")
    reward_model = RewardModel(
        vocab_size=VOCAB_SIZE,
        embed_dim=EMBED_DIM,
        hidden_dim=HIDDEN_DIM
    )

    # Simulated data: preferred vs rejected responses
    pairs_y_w = torch.randint(0, VOCAB_SIZE, (NUM_PAIRS, SEQ_LEN))
    pairs_y_l = torch.randint(0, VOCAB_SIZE, (NUM_PAIRS, SEQ_LEN))
    # Make the "winners" slightly different from the "losers"
    pairs_y_w[:, :5] = 1  # Distinctive pattern for preferred responses

    reward_model = train_reward_model(
        reward_model,
        pairs_y_w, pairs_y_l,
        lr=5e-4, epochs=5, batch_size=BATCH_SIZE
    )

    # --- Step 2: Create the policy and the frozen reference ---
    print("\n[Step 2] Creating the policy model and the reference...")
    policy = PolicyModel(
        vocab_size=VOCAB_SIZE,
        embed_dim=EMBED_DIM,
        hidden_dim=HIDDEN_DIM,
        max_len=SEQ_LEN
    )
    # Copy initial weights as reference
    ref_policy = PolicyModel(
        vocab_size=VOCAB_SIZE,
        embed_dim=EMBED_DIM,
        hidden_dim=HIDDEN_DIM,
        max_len=SEQ_LEN
    )
    ref_policy.load_state_dict(policy.state_dict())
    ref_policy.eval() — Frozen — no training

    # --- Step 3: PPO loop (RLHF) ---
    print("\n[Step 3] PPO fine-tuning with KL penalty...")
    trainer = PPOTrainer(
        policy=policy,
        ref_policy=ref_policy,
        reward_model=reward_model,
        beta=0.05,            # Moderate KL coefficient
        clip_epsilon=0.2,     # Standard PPO clipping
        ppo_epochs=3,         # PPO iterations per batch
        lr=1e-5,              # Low learning rate for stability
        gamma=0.99,
        gae_lambda=0.95,
        max_grad_norm=1.0
    )

    # Simulated prompts for training
    prompts = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, 8))

    print(f"\n  PPO Configuration:")
    print(f"    β (KL coeff)    = {trainer.beta}")
    print(f"    clip_epsilon    = {trainer.clip_epsilon}")
    print(f"    ppo_epochs      = {trainer.ppo_epochs}")
    print(f"    lr              = {trainer.optimizer.param_groups[0]['lr']}")
    print()

    # Training loop
    num_steps = 20
    for step in range(num_steps):
        metrics = trainer.train_step(prompts)
        if (step + 1) % 5 == 0:
            print(f"  Step {step+1}/{num_steps} — "
                  f"Total loss: {metrics['total_loss']:.4f} | "
                  f"Policy loss: {metrics['policy_loss']:.4f} | "
                  f"KL: {metrics['kl_penalty']:.4f} | "
                  f"Mean reward: {metrics['mean_reward']:.4f}")

    print("\n[SUCCESS] RLHF pipeline completed successfully!")
    print("=" * 60)

Key Hyperparameters

The success of RLHF depends heavily on tuning the following hyperparameters:

Practical recommendation: Start with β = 0.1, ppo_epochs = 3, clip_epsilon = 0.2. Monitor the KL divergence at each step: if it exceeds 5–10 nats, increase β. If it stays below 0.5, decrease β for more aggressive alignment.

Advantages and Limitations of RLHF

Advantages

• Superior alignment: RLHF produces models whose responses are significantly more useful, harmless, and honest than those of a model fine-tuned by instruction alone (SFT only). This is the reason ChatGPT surpassed its competitors in 2022.
• Capturing complex preferences: Human preferences are often subtle and difficult to formalize in a manual reward function. RLHF learns them implicitly through human comparisons, capturing nuances such as tone, structure, or contextual relevance.
• Flexibility: Multiple reward models can be trained for different criteria (utility, safety, conciseness) and combined. Artificial rewards (heuristic rules, automatic verifiers) can also be used to complement human annotations.

Limitations

• High cost: Collecting tens of thousands of human comparisons is time-consuming and expensive. The process requires qualified annotators, and annotation quality directly determines the quality of the final model.
• Annotator bias: The captured preferences reflect the cultural, linguistic, and subjective biases of the annotators. An RLHF model trained with Western annotators will behave differently from one trained with Asian annotators — and both will be biased in their own way.
• Reward hacking and degeneration: If β is poorly tuned, the model can learn to “game” the reward model rather than produce good responses. Recent work (Gao et al., 2022) has shown that RLHF models can produce repetitive responses or token loops to maximize the score without improving actual quality.
• Limited preference scaling: Pairwise comparisons do not capture multidimensional preferences well. A response can be more informative but less concise — how do annotators decide? The reward model will learn an average that satisfies no one.
• Training instability: PPO can be unstable at scale. Subtle bugs in the implementation (especially KL computation) can lead to silent training that appears to converge but produces a degraded model.

4 Concrete Use Cases of RLHF

1. Conversational Assistants (ChatGPT, Claude, Gemini)

This is the most well-known use case. RLHF transforms a raw language model — capable of generating coherent but potentially toxic or useless text — into a useful and safe assistant. Human annotators compare responses on criteria of utility, truthfulness, and harmlessness. The final model is able to politely refuse to generate dangerous content, rephrase ambiguous responses, and structure complex explanations in a pedagogical way.

2. Automated Content Moderation

Platforms like Reddit, Stack Overflow, or corporate forums use RLHF to train more nuanced moderation classifiers. Instead of binary rules (this word is banned / allowed), the model learns gradations of severity from human moderators’ decisions. It can thus distinguish a direct insult from a harsh but legitimate critique, spam from excessive enthusiasm.

3. Document Summarization with Format Preferences

An RLHF summarization model can learn not only what to summarize, but how to do it: desired length, level of detail, formal or informal tone, inclusion or exclusion of certain types of information. Annotators compare summaries based on stylistic and informational criteria, and the fine-tuned model produces summaries adapted to the context of use — whether for an executive report or a quick note.

4. Intelligent Code Agents (GitHub Copilot, Code Llama)

RLHF also applies to code generation. Programmers compare code suggestions based on readability, efficiency, security, and compliance with project conventions. The model thus learns to favor solutions that are not only correct but also well-structured and maintainable — a nuance that supervised learning alone does not capture.

See Also

• Create Your Own Python Game: Reproduce the Action of Tom and Jerry with Pygame
• Solving a Square Recurrence Relation in Python: Complete Guide and Practical Tips