RLHF: The Engine Tuning Human Values into Large Language Models

Author(s): Saurab

Originally published on Towards AI.

(Deep Dive into Reinforcement Learning from Human Feedback)

We stand in awe of Large Language Models (LLMs) like GPT-4, Gemini, Claude, and Llama. Trained on internet-scale data (> 10^12 tokens), they possess vast knowledge and linguistic fluency. Yet, transforming this raw potential into helpful, harmless, and honest AI assistants requires more than just massive datasets. It requires alignment — teaching these models to behave according to human preferences and values.

The key technology enabling this alignment for many state-of-the-art models is Reinforcement Learning from Human Feedback (RLHF). It’s a sophisticated process that fine-tunes behemoth models using human judgment as the ultimate guide.

This post dives deeper into RLHF, exploring the underlying mechanisms, mathematics, and computational steps involved.

The Alignment Problem: Why Pre-training Isn’t Enough

Standard LLM pre-training optimizes for predicting the next token in a sequence. Given a context c, the model learns parameters ψ\psi to maximize the likelihood of the true next token t across a massive dataset D:

This objective makes the model proficient at generating fluent and often factually knowledgeable text. However, it doesn’t inherently teach the model to:

• Follow Instructions Faithfully: Pre-trained models might ignore constraints or misunderstand the nuances of a user’s request.
• Maintain Truthfulness: They can “hallucinate” plausible-sounding but incorrect information.
• Avoid Harmful Outputs: They might reproduce biases or generate toxic content learned from the training data.
• Be Concise and Helpful: Responses can be verbose, evasive, or miss the user’s core need.
• Refuse Inappropriate Requests: They lack an inherent understanding of ethical boundaries.

We need a mechanism to explicitly steer the model towards desired behaviours, going beyond simple pattern completion.

RLHF: A Multi-Stage Alignment Process

RLHF tackles the alignment problem through a structured, three-stage process:

Stage 1: Supervised Fine-Tuning (SFT) — (Optional but Common)

Before diving into RLHF proper, models are often initially fine-tuned on a smaller, high-quality dataset of instruction-response pairs. This dataset contains examples of desired outputs for various prompts, often curated by humans.

• Goal: Adapt the pre-trained model to better follow instructions and respond in a specific style (e.g., like an assistant).
• Process: Fine-tune the model using the same next-token prediction objective as pre-training, but only on this curated dataset :

• Outcome: A model πϕ that is better at following instructions than the base pre-trained model πψ. This serves as a better starting point for the subsequent RLHF stages. Let’s call this starting policy πSFT.

Stage 2: Training a Reward Model (RM) — Learning Human Preferences

This is the core of the “Human Feedback” part. We need to capture human judgment about what constitutes a “good” response.

• Goal: Train a model rθ(x,y) that takes a prompt x and a generated response y and outputs a scalar score representing how much a human would prefer that response.
• Data Collection:

1. Select a diverse set of prompts x.
2. Use the SFT model (or even multiple model versions) to generate several candidate responses {y1, y2, … , yk} for each prompt x.
3. Present pairs of responses (yi, yj) for the same prompt x to human labellers.
4. Ask labellers to choose which response they prefer, Yw (winner) or Yl(loser).
5. Collect a large dataset DRM of preference tuples (x,Yw,Yl).

• Modeling Preferences:
  The reward model rθ(x,y) is trained to predict this preference. A common approach uses the Bradley-Terry model, which assumes the probability of preferring yw over yl is related to the difference in their reward scores via a sigmoid function σ :

• Training Objective:
  Maximize the likelihood of the observed human preferences. This translates to minimizing a logistic loss function over the dataset :

• The reward model often starts with the weights of the SFT model (or the pre-trained model) and adds a linear layer on top of the final token representation to output the scalar reward.
• Conceptual Code (Reward Model Loss):

import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Example
# Assume RewardModel outputs a single scalar score
# reward_model = AutoModelForSequenceClassification.from_pretrained('sft_model_path', num_labels=1)
# tokenizer = AutoTokenizer.from_pretrained('sft_model_path')

def compute_reward_model_loss(reward_model, tokenizer, prompt, response_winner, response_loser):
 """Conceptual function for RM loss calculation"""
 # Tokenize inputs (handle truncation, padding etc.)
 winner_inputs = tokenizer(prompt + response_winner, return_tensors='pt', truncation=True, padding=True)
 loser_inputs = tokenizer(prompt + response_loser, return_tensors='pt', truncation=True, padding=True)

 # Get reward scores from the model
 # The actual way to get the scalar might differ based on model head architecture
 # Assuming model outputs logits, and we take the first logit as the score
 score_winner = reward_model(**winner_inputs).logits.squeeze() # Shape: [batch_size]
 score_loser = reward_model(**loser_inputs).logits.squeeze() # Shape: [batch_size]

 # Calculate loss based on pairwise comparison
 # log(sigmoid(score_winner - score_loser))
 loss = -torch.log(torch.sigmoid(score_winner - score_loser)).mean()
 return loss

# --- Training Loop Sketch ---
# optimizer = optim.Adam(reward_model.parameters(), lr=1e-5)
# for batch in preference_dataloader: # Batches of (prompt, y_w, y_l)
# optimizer.zero_grad()
# loss = compute_reward_model_loss(reward_model, tokenizer,
# batch['prompt'], batch['response_winner'], batch['response_loser'])
# loss.backward()
# optimizer.step()

• Disclaimer: The code above is highly conceptual and omits many details like data loading, precise model architecture for scoring, device handling, and distributed training.

Stage 3: Fine-tuning with Reinforcement Learning (RL) — Optimizing the Policy

Now, we use the trained reward model to improve the SFT language model (our initial policy) using RL.

• Goal: Train a policy (which is the LLM we are tuning) that generates responses y to prompts x which maximize the expected reward predicted by the reward model.
• RL Formulation:
• State: The prompt x received by the agent.
• Action: Generating the response :

• Reward: The score assigned by the reward model.
• The Challenge:
  Naively maximizing can lead the LLM to deviate drastically from realistic language generation (finding exploits in the RM) or forget knowledge from pre-training.
• The Solution (PPO):
  Proximal Policy Optimization (PPO) is commonly used. It aims to maximize the reward while constraining how much the policy changes from a reference policy, which is typically the initial SFT policy . This constraint is enforced using a KL divergence penalty.
• PPO Objective Function:
  The objective for tuning the LLM policy is approximately:

• Where:
• E(x,y)∼πϕ denotes the expectation when prompts x are drawn from a dataset and responses y are generated by the current policy πϕ.
• rθ(x,y) is the reward from the trained reward model.
• KL is the Kullback–Leibler divergence between the current policy and the initial SFT policy. It measures how much the current policy has “drifted” from the reference.
• β is a hyperparameter controlling the strength of the KL penalty. A higher β keeps πϕ closer to πSFT.
• (Note: PPO uses a slightly more complex clipped surrogate objective in practice for stability, but the core idea is captured above.)
• Conceptual Code (RL Update Step):

import torch
# Assume policy_llm is the model being trained (e.g., AutoModelForCausalLM)
# Assume reference_llm is the frozen SFT model (same architecture)
# Assume reward_model is the trained RM from Stage 2
# Assume tokenizer is shared

def compute_rl_loss(policy_llm, reference_llm, reward_model, tokenizer, prompts, responses, beta):
 """Conceptual function for RL loss/objective calculation"""
 # --- Get Rewards ---
 # Need to tokenize prompts + responses for the reward model
 reward_inputs = tokenizer(prompts, responses, return_tensors='pt', padding=True, truncation=True)
 with torch.no_grad(): # RM is frozen during RL update
 rewards = reward_model(**reward_inputs).logits.squeeze() # Shape: [batch_size]

 # --- Get Log Probabilities ---
 # Tokenize for the policy LLM (inputs = prompts, labels = responses)
 policy_inputs = tokenizer(prompts, return_tensors='pt', padding=True, truncation=True)
 policy_labels = tokenizer(responses, return_tensors='pt', padding=True, truncation=True).input_ids
 # Calculate log probability of generating 'responses' given 'prompts'
 # under the current policy_llm
 outputs_policy = policy_llm(**policy_inputs, labels=policy_labels)
 log_probs_policy = outputs_policy.loss * -policy_labels.size(1) # Reconstruct sum log P(token|context)
 # Calculate log probability under the reference model (frozen)
 with torch.no_grad():
 outputs_ref = reference_llm(**policy_inputs, labels=policy_labels)
 log_probs_ref = outputs_ref.loss * -policy_labels.size(1)

 # --- Calculate KL Divergence (Approximation) ---
 # KL(policy || reference) approx = log_probs_policy - log_probs_ref
 kl_div = log_probs_policy - log_probs_ref

 # --- Combine into Objective ---
 # Maximize: E[rewards - beta * KL] <=> Minimize: E[-rewards + beta * KL]
 # We usually average over the batch
 loss = (-rewards + beta * kl_div).mean()

 # PPO might involve importance sampling ratios and clipping, omitted here for clarity
 return loss, rewards.mean(), kl_div.mean() # Return loss and stats

# --- RL Training Loop Sketch ---
# rl_optimizer = optim.Adam(policy_llm.parameters(), lr=1e-6)
# for prompts_batch in rl_prompt_dataloader:
# # 1. Generate responses using current policy_llm (autoregressive sampling)
# responses_batch = policy_llm.generate(tokenizer(prompts_batch, return_tensors='pt').input_ids, ...)
# responses_text = tokenizer.batch_decode(responses_batch) # Decode to text
# # 2. Calculate loss and update
# rl_optimizer.zero_grad()
# loss, avg_reward, avg_kl = compute_rl_loss(policy_llm, reference_llm, reward_model, tokenizer,
# prompts_batch, responses_text, beta=0.1)
# loss.backward()
# rl_optimizer.step()

• Disclaimer: This RL code is extremely simplified. Real implementations involve complex handling of generation, tokenization, batching, KL estimation, PPO-specific details like value functions, advantage estimation, and distributed training.

Benefits Revisited

This complex RLHF pipeline delivers significant improvements:

• Superior Instruction Following: Models adhere more closely to complex user requests.
• Reduced Harmful Outputs: Explicitly penalized by the RM trained on human safety preferences.
• Improved Truthfulness: Preference for factual accuracy can be incorporated into human labeling guidelines.
• Enhanced Controllability: Better alignment makes the model’s behaviour more predictable and steerable.

Challenges Magnified

The complexity introduces challenges:

• Data Cost & Quality: Obtaining consistent, high-quality human preference data at scale is a major bottleneck.
• Reward Model Limitations: The RM is an imperfect proxy for true human preference and can be exploited (reward hacking). Its accuracy limits the final policy.
• Alignment Tax: The KL constraint can sometimes limit the model’s capabilities or creativity if set too high. Finding the right balance is key.
• Complexity & Stability: RL training can be unstable and sensitive to hyperparameters (β, learning rates, PPO parameters).
• Bias Propagation: Biases held by human labellers can be encoded into the RM and amplified by the RL process.

The Evolving Landscape: Beyond Standard RLHF

RLHF is a rapidly evolving field. Techniques like Direct Preference Optimization (DPO) aim to bypass the explicit reward modeling step, directly optimizing the language model on preference data using a derived loss function related to the RL objective. This can simplify the pipeline but comes with its own set of trade-offs. Constitutional AI uses AI feedback based on predefined principles to scale alignment efforts.

Conclusion: The Human Element in the Loop

RLHF, despite its complexity and challenges, represents a landmark achievement in aligning powerful AI models with human intent. By incorporating human judgment directly into the training loop via preference modeling and reinforcement learning, we can build LLMs that are not just knowledgeable but also significantly more helpful, harmless, and trustworthy. It underscores a critical theme in modern AI development: the most capable systems often arise from a sophisticated synergy between machine learning algorithms and human guidance. As of April 16, 2025, RLHF and its derivatives remain central to pushing the frontiers of safe and useful AI.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI