Margin Matching Preference Optimization

PyTorch implementation for the paper:

Margin Matching Preference Optimization: Enhanced Model Alignment with Granular Feedback
Kyuyoung Kim*, Ah Jeong Seo*, Hao Liu, Jinwoo Shin, Kimin Lee
In EMNLP 2024 Findings

conda create -n mmpo python=3.10
# install pytorch
conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=12.1 -c pytorch -c nvidia

# check gpu
import torch
torch.cuda.is_available()

# install the remaining package dependencies
python -m pip install -e .
pip install -r requirements.txt

# install flash attention
MAX_JOBS=4 pip install flash-attn --no-build-isolation

# for deepspeed
conda install -c conda-forge mpi4py mpich

We evaluate MMPO using both human and AI feedback data to assess their performance on feedback of varying qualities:

• For more precise evaluation, we used the filtered version released by AllenAI. Filtering details can be found in the dataset card.
  

• To build this filtered version of SHP, we sample uniformly across score differences to evaluate the methods across diverse quality margins.
• Following KTO, we ensured that the same prompt appears no more than five times to prevent overfitting. More details can be found in scripts/run_dpo_shp.py.

Use the following commands to train models:

• Gemma 2b:

accelerate launch --num_processes=2 --config_file recipes/accelerate_configs/multi_gpu.yaml scripts/run_dpo.py recipes/gemma-2b/dpo.yaml --bt_beta=2.2

• Gemma 7b:

accelerate launch --num_processes=4 --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_dpo.py recipes/gemma-7b/dpo.yaml --bt_beta=0.3

To train on the SHP dataset:

accelerate launch --num_processes=4 --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_dpo_shp.py recipes/gemma-7b/dpo_shp.yaml --bt_beta=1.1

• Llama 8b:

accelerate launch --num_processes=4 --config_file recipes/accelerate_configs/deepspeed_zero3.yaml run_dpo.py recipes/llama-3-8b/dpo.yaml --bt_beta=2.2

To evaluate on MT-bench, follow the instructions in FastChat. Below is the MT-bench results for the models trained with MMPO and DPO on UltraFeedback:

To evaluate on RewardBench, follow the instructions in RewardBench repository. As of September 2024, MMPO achieves state-of-the-art performance compared to other models at the same scale (see leaderboard):

More experimental details can be found in the paper.

@inproceedings{kim2024margin,
  title={Margin Matching Preference Optimization: Enhanced Model Alignment with Granular Feedback},
  author={Kim, Kyuyoung and Seo, Ah and Liu, Hao and Shin, Jinwoo and Lee, Kimin},
  booktitle={Findings of the Association for Computational Linguistics: EMNLP 2024},
  pages={13554--13570},
  year={2024}
}