XPO Trainer

Overview

Exploratory Preference Optimization (XPO) was proposed in the paper Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF by Tengyang Xie, Dylan J. Foster, Akshay Krishnamurthy, Corby Rosset, Ahmed Awadallah, and Alexander Rakhlin. It is a simple online preference tuning method based on the DPO loss together with a reward model (RM). XPO augments the DPO objective with an exploration bonus allowing the method to explore outside the support of the intitial model and human feedback data.

The abstract from the paper is the following:

Reinforcement learning from human feedback (RLHF) has emerged as a central tool for language model alignment. We consider online exploration in RLHF, which exploits interactive access to human or AI feedback by deliberately encouraging the model to produce diverse, maximally informative responses. By allowing RLHF to confidently stray from the pre-trained model, online exploration offers the possibility of novel, potentially super-human capabilities, but its full potential as a paradigm for language model training has yet to be realized, owing to computational and statistical bottlenecks in directly adapting existing reinforcement learning techniques. We propose a new algorithm for online exploration in RLHF, Exploratory Preference Optimization (XPO), which is simple and practical — a one-line change to (online) Direct Preference Optimization (DPO; Rafailov et al., 2023) — yet enjoys the strongest known provable guarantees and promising empirical performance. XPO augments the DPO objective with a novel and principled exploration bonus, empowering the algorithm to explore outside the support of the initial model and human feedback data. In theory, we show that XPO is provably sample-efficient and converges to a near-optimal language model policy under natural exploration conditions, irrespective of whether the initial model has good coverage. Our analysis, which builds on the observation that DPO implicitly performs a form of Q*-approximation (or, Bellman error minimization), combines previously disparate techniques from language modeling and theoretical reinforcement learning in a serendipitous fashion through the perspective of KL-regularized Markov decision processes. Empirically, we find that XPO is more sample-efficient than non-exploratory DPO variants in a preliminary evaluation.

This post-training method was contributed by Kashif Rasul, Quentin Gallouédec and Lewis Tunstall.

Quick start

This example demonstrates how to train a model using the XPO method. We use the Qwen 0.5B model as the base model and the Qwen 0.5B reward model as the reward model. We use the prompts from the UltraFeedback dataset. You can view the prompts in the dataset here:

Below is the script to train the model:

# train_xpo.py
from datasets import load_dataset
from trl import XPOConfig, XPOTrainer
from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
reward_model = AutoModelForSequenceClassification.from_pretrained("trl-lib/Qwen2-0.5B-Reward", num_labels=1)
train_dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

args = XPOConfig(output_dir="nash-md-qwen2", logging_steps=10)
trainer = XPOTrainer(
    model=model,
    reward_model=reward_model,
    args=args,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
)
trainer.train()

Execute the script using the following command:

accelerate launch train_xpo.py

Expected dataset format

XPO requires a prompt-only dataset. The XPOTrainer supports both conversational and standard dataset format. When provided with a conversational dataset, the trainer will automatically apply the chat template to the dataset.

Usage tips

⚠️ Use the same chat template

Make sure that the SFT model and reward model use the same chat template. Otherwise, you may find the model completions are scored incorrectly during training.

Encourage EOS token generation

We can want the model to generate completion within a given length. During the learning, the model will generate completion up to the maximum completion length specified in the max_new_tokens argument of XPOConfig. I you want to penalize for not generating an EOS token before the maximum completion length, you can use the missing_eos_penalty argument of XPOConfig:

args = XPOConfig(..., max_new_tokens=128, missing_eos_penalty=1.0)

Logging Completions

To better understand your model’s behavior during training, you can log sample completions periodically using the LogCompletionsCallback.

trainer = XPOTrainer(..., eval_dataset=eval_dataset)
completions_callback = LogCompletionsCallback(trainer, num_prompts=8)
trainer.add_callback(completions_callback)

This callback logs the model’s generated completions directly to Weights & Biases.

Example script

We provide an example script to train a model using the XPO method. The script is available in examples/scripts/nash_md.py

To test the XPO script with the Pythia 14M model on the TL;DR summarization task, run the following command:

python examples/scripts/xpo.py \
    --model_name_or_path EleutherAI/pythia-14m  \
    --reward_model_path EleutherAI/pythia-14m \
    --dataset_name trl-lib/tldr \
    --learning_rate 5.0e-7 \
    --output_dir pythia-14m-tldr-xpo \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 32 \
    --num_train_epochs 3 \
    --max_new_tokens 64 \
    --warmup_ratio 0.1 \
    --missing_eos_penalty 1.0 \
    --push_to_hub

Logged metrics

The logged metrics are as follows:

• loss/xpo: The mean xpo part of the full loss.
• loss/dpo: The mean dpo part of the full loss.
• objective/kl: The mean KL divergence between the model and reference data.
• objective/entropy: The mean entropy of the model and reference data.
• objective/model_scores: The mean scores (according to the reward model) of the model completions.
• objective/ref_scores: The mean scores (according to the reward model) of the reference completions.
• objective/scores_margin: The mean score margin (according to the external reward model) between the chosen and rejected completions.
• rewards/chosen: The mean reward (according to XPO’s DPO implicit reward model) of the chosen completions.
• rewards/rejected: The mean reward (according to XPO’s DPO implicit reward model) of the rejected completions.
• rewards/accuracies: The accuracies of the XPO’s implicit reward model.
• rewards/margins: The mean reward margin (according to online DPO’s implicit reward model) between the chosen and rejected completions.
• logps/chosen: The mean log probabilities of the chosen completions.
• logps/rejected: The mean log probabilities of the rejected completions.
• val/model_contain_eos_token: The amount of times the model’s output contains the eos token.
• val/ref_contain_eos_token: The amount of times the reference’s output contains the eos token.
• alpha: The weight of the XPO loss term. Typically fixed, but can be made dynamic by passing a list to XPOConfig.
• beta: The parameter that controls the weight of the loss term representing the deviation from the reference model. Typically fixed, but can be made dynamic by passing a list to XPOConfig.

XPOTrainer

XPOConfig