Towards Efficient Exact Optimization of Language Model Alignment

Overview

This is the official pytorch implementation of the EXO algorithm for efficient exact optimization of aligning language models (LMs) with human preferences, as described in our ICML 2024 paper Towards Efficient Exact Optimization of Language Model Alignment.

EXO essentially minimizes the reverse KL between the empirical distributions defined by the policy and the reward. As a comparison, DPO corresponds to minimizing the forward KL. The above figure illustrates the distinct behavior of policies obtained by minimizing (a) the reverse KL and (b) the forward KL.

Dependencies

python >= 3.9.0
transformers >= 4.34.1
deepspeed >= 0.11.2

It is recommended to precomplie the required extensions when installing deepspeed (For more details refer to the installation guideline of deepspeed):

DS_BUILD_FUSED_ADAM=1 DS_BUILD_TRANSFORMER=1 DS_BUILD_TRANSFORMER_INFERENCE=1 pip install deepspeed 

General Guideline

EXO supports the settings of (i) training directly on the preference data and (ii) training with the supervision provided by a learned reward model. The pipeline is comprised of the following stages:

• Supervised fine-tuning (SFT) stage:

  • Train the SFT policy: train_sft.sh.
  • (Optional) Sample from the SFT policy: inference_sft.sh.

• Reward Model (RM) stage:

  • (Optional) Train the RM: train_rm.sh.
  • Score the prompt-generation pairs with reward: inference_rm.sh.

• Alignment stage:

  • Train the policy to align with human preferences using EXO: train_exo.sh.
  • Sample from the learned policy: inference_sft.sh.

To train on a custom dataset, one should create a dataset class by inheriting from the base class PromptRawDataset located in src/utils/data/raw_datasets.py. Then, simply add a few lines of code to the get_raw_dataset method in src/utils/data/data_utils.py to utilize the custom dataset.

SFT Stage

Training

In the SFT stage, the LM is fine-tuned with supervised MLE on the data that is supposed to be obtained from the same distribution as the preference data. One can simply finetune on the chosen texts of the preference data if no such data is available.

Inference

(Optional but recommended) To utilize supervision of the reward model for alignment, one need to sample from the SFT model and later use the reward model to score the inference results.

RM stage

Training

(Optional but recommended) In order to utilize the continuous preference signal, one can train a reward model on the preference data to predict the human preference.

Inference

(Optional but recommended) Then use the reward model to score the SFT generated data with continuous reward.

Alignment stage

EXO & DPO

In the alignment stage, the SFT model is fine-tuned to align with human preferences by training on either the preference data or the RM labeled data.

Before training, the preference dataset should be converted to the same format as the RM labeled data:

python src/utils/data/pref_to_rw.py /local/path/to/preference/data

Reproducing experiments in the paper

We also provide guidelines to reproduce the experiments on the three public datasets: IMDB, TL;DR and Anthropic-HH to facilitate future study:

• Reproducing the IMDB experiment
• Reproducing the TL;DR experiment
• Reproducing the Anthropic-HH experiment

Citing

@article{Ji2024TowardsExact,
  title={Towards Efficient Exact Optimization of Language Model Alignment},
  author={Haozhe Ji, Cheng Lu, Yilin Niu, Pei Ke, Hongning Wang, Jun Zhu, Jie Tang, Minlie Huang},
  year={2024},
  journal={The Forty-first International Conference on Machine Learning},
  url={https://arxiv.org/abs/2402.00856}
}

Please kindly cite our work if you find the paper or this repository useful :)