Illustrating Reinforcement Learning from Human Feedback (RLHF)

Reinforcement Learning from Human Feedback (RLHF) is a method used to align large language models (LLMs) with human values by using reinforcement learning to directly optimize a model based on human preferences. This process allows models to move beyond simple next-token prediction to generate text that is more helpful, truthful, and safe.

The Three-Step RLHF Training Process

RLHF is a multi-stage training pipeline that transforms a general-purpose pretrained model into an aligned assistant. The process is broken down into three core steps:

1. Pretraining the Language Model

The process begins with a language model pretrained using classical objectives, such as cross-entropy loss for next-token prediction. While the initial model can be fine-tuned on augmented data—such as human-generated "preferable" text (OpenAI) or distilled context clues for "helpful, honest, and harmless" criteria (Anthropic)—the core requirement is a model that responds well to diverse instructions.

2. Reward Model Training

Because defining a mathematical loss function for "good" text is intractable, RLHF uses a reward model (RM) to numerically represent human preference.

  • Objective: The RM takes a sequence of text and returns a scalar reward.
  • Data Collection: Prompts are sampled from a dataset and passed through the initial LM to generate multiple outputs. Human annotators then rank these outputs.
  • Ranking over Scoring: Humans rank outputs in head-to-head matchups (often using an Elo system) rather than providing absolute scalar scores, as rankings are less noisy and more calibrated across different annotators.
  • Model Capacity: Reward models vary in size. For example, 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 scores.

  • The RL Formulation:
    • Policy: The LM 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 scalar reward from the preference model ($r_{\theta}$) and a penalty based on the Kullback–Leibler (KL) divergence ($r_{\text{KL}}$) between the RL policy and the initial pretrained model.
  • The KL Penalty: The KL divergence term prevents the model from moving too far from the original pretrained model. Without this penalty, the model might generate gibberish that "fools" the reward model into giving a high score.
  • Update Rule: PPO is used as a trust region optimization algorithm to ensure that gradient updates do not destabilize the learning process.

Open-Source Tools for RLHF

Several PyTorch-based repositories provide the infrastructure necessary to implement RLHF:

  • TRL (Transformers Reinforcement Learning): Designed for fine-tuning pretrained LMs in the Hugging Face ecosystem using PPO.
  • TRLX: An expanded fork of TRL built by CarperAI to handle larger models (up to 33B parameters, with future support for 200B) for online and offline training, supporting PPO and Implicit Language Q-Learning (ILQL).
  • RL4LMs: A library from AllenAI that provides building blocks for various RL algorithms (PPO, NLPO, A2C, TRPO) and reward functions, benchmarked across 2,000 experiments.

Current Limitations and Future Directions

Despite its success in models like ChatGPT, RLHF faces several systemic challenges:

  • Data Cost and Quality: Gathering human preference data is expensive. High-quality human-generated text requires specialized staff, and human annotators often disagree, introducing variance into the training data.
  • Model Imperfections: RLHF-tuned models can still produce harmful or factually inaccurate text without expressing uncertainty.
  • Optimization Opportunities: PPO is an older algorithm, and there is significant unexplored design space. Potential improvements include using offline RL (like ILQL) to avoid costly forward passes of the reward model during training and exploring different exploration-exploitation balances.

Sources