LIONs: An Empirically Optimized Approach to Align Language Models

TL;DR

This work conducts a rigorous analysis over a three-stage training pipeline consisting of supervised fine-tuning, offline preference learning, and online preference learning and finds that using techniques like sequence packing, loss masking in SFT, increasing the preference dataset size in DPO, and online DPO training can significantly improve the performance of language models.

Abstract

Alignment is a crucial step to enhance the instruction-following and conversational abilities of language models. Despite many recent work proposing new algorithms, datasets, and training pipelines, there is a lack of comprehensive studies measuring the impact of various design choices throughout the whole training process. We first conduct a rigorous analysis over a three-stage training pipeline consisting of supervised fine-tuning, offline preference learning, and online preference learning. We have found that using techniques like sequence packing, loss masking in SFT, increasing the preference dataset size in DPO, and online DPO training can significantly improve the performance of language models. We then train from Gemma-2b-base and LLama-3-8b-base, and find that our best models exceed the performance of the official instruct models tuned with closed-source data and algorithms. Our code and models can be found at \url{https://github.com/Columbia-NLP-Lab/LionAlignment}.

Figures (11)

• Figure 1: Measuring the effect of dataset size ($|D|$) and training steps (FLOPs) on final performance. While performance can quickly saturate given a fixed $|D|$, increasing the dataset size increases the point of saturation. Dotted lines are our interpolation using a degree 2 polynomial.
• Figure 2: Varying KL-divergence strength ($\beta$) under different training data sizes. We find the best $\beta$ stays relatively consistent across different dataset sizes.
• Figure 3: Effect of training on 10k data selected using different filtering algorithms. We find that simply training on a larger dataset (100k) outperforms all methods.
• Figure 4: We track the changes in the probability margin $\pi_\theta(y_w) - \pi_\theta(y_l)$ under various training configurations, and find that the best-performing models exhibit a parabolic pattern. Arena-Hard-Auto* results from left to right are 10.7, 14.8, and 13.2. Test loss from left to right is 0.63, 0.65, and 1.33.
• Figure A1: Illustration of efficient DPO implementation. Traditional DPO training requires adding padding tokens to the batch. Our implementation can remove the need of paddding tokens, and thus improving the training efficiency.