Illustrating Reinforcement Learning from Human Feedback (RLHF)

Reinforcement Learning from Human Feedback (RLHF) is a method used to align large language models (LLMs) trained on general text corpora with complex human values and preferences. This process allows models to move beyond simple next-token prediction to generate text that is helpful, truthful, or creative based on subjective human criteria.

The Three-Step RLHF Process

RLHF is a multi-stage training pipeline that involves multiple models and distinct deployment phases. The process is broken down into pretraining, reward modeling, and reinforcement learning fine-tuning.

1. Pretraining the Language Model

The process begins with a language model pretrained using classical objectives. This base model must be capable of responding to diverse instructions. Examples of models used as starting points include smaller versions of GPT-3 for InstructGPT, transformer models ranging from 10 million to 52 billion parameters for Anthropic's research, and DeepMind's 280 billion parameter Gopher model.

While not strictly required, some organizations use augmented data to further refine the base model. OpenAI fine-tuned on "preferable" human-generated text, while Anthropic distilled an original LM using context clues for "helpful, honest, and harmless" criteria.

2. Reward Model Training

The reward model (RM), or preference model, is designed to take a sequence of text and return a scalar reward representing human preference. This scalar output is essential for integration with reinforcement learning algorithms.

Data Collection and Ranking To train the RM, prompts are sampled from a dataset and passed through the initial LM to generate multiple text outputs. Human annotators then rank these outputs. Because direct scalar scoring by humans is often noisy and uncalibrated, rankings (such as head-to-head matchups using an Elo system) are used to create a more regularized dataset.

Model Architecture Reward models can be fine-tuned LMs or models trained from scratch. For example, Anthropic used Preference Model Pretraining (PMP) to improve sample efficiency. The size of the reward model varies by implementation; OpenAI used a 6B reward model for a 175B LM, while DeepMind used 70B Chinchilla models for both the LM and the reward model.

3. Fine-Tuning with Reinforcement Learning

The final stage uses a policy-gradient RL algorithm, typically Proximal Policy Optimization (PPO), to optimize a copy of the initial LM based on the reward model's feedback.

The RL Formulation

  • Policy: The language model that takes a prompt and returns a sequence of text.
  • Action Space: The vocabulary of the language model (typically ~50k tokens).
  • Observation Space: The distribution of possible input token sequences.
  • Reward Function: A combination of the preference model's scalar reward ($r_{\theta}$) and a penalty on policy shift ($r_{\text{KL}}$).

The KL Divergence Penalty To prevent the model from generating gibberish that "fools" the reward model (reward hacking), a penalty is applied based on the Kullback–Leibler (KL) divergence. This penalizes the RL policy if it moves too far from the initial pretrained model, ensuring the output remains coherent.

The final reward calculation is: $r = r_{\theta} - \lambda r_{\text{KL}}$.

Update Rule PPO is used as a trust region optimization algorithm to ensure gradient updates do not destabilize the learning process. While PPO is common, DeepMind has utilized synchronous advantage actor-critic (A2C) for similar setups.

Open-Source Tools for RLHF

Several PyTorch-based repositories facilitate RLHF implementation:

  • TRL (Transformers Reinforcement Learning): Designed for fine-tuning pretrained LMs in the Hugging Face ecosystem using PPO.
  • TRLX: A fork of TRL by CarperAI optimized for larger models (up to 33B parameters currently, with goals for 200B) and supporting Implicit Language Q-Learning (ILQL).
  • RL4LMs: A library from AllenAI providing building blocks for various RL algorithms (PPO, NLPO, A2C, TRPO) and customizable reward functions.

Limitations and Future Directions

Despite its success, RLHF faces several technical and operational challenges:

  • Data Cost: Gathering human preference data is expensive. High-quality human-generated text is particularly costly, though preference labels (~50k samples) are more manageable.
  • Human Variance: Annotators often disagree, introducing variance into the training data without a ground truth.
  • Model Imperfection: Models can still produce harmful or factually inaccurate text without expressing uncertainty.

Potential Improvements Future research is focusing on improving the RL optimizer. Because the reward model requires costly forward passes for every piece of generated text, offline RL and algorithms like Implicit Language Q-Learning (ILQL) are being explored as more efficient policy optimizers.

Sources