ODM - Online Data Mixing (Under development)

This is the official repository for the methods provided in Efficient Online Data Mixing For Language Model Pre-Training.

Abstract

The data used to pretrain large language models has a decisive impact on a model's downstream performance, which has led to a large body of work on data selection methods that aim to automatically determine the most suitable data to use for pretraining. Existing data selection methods suffer from slow and computationally expensive processes, a problem amplified by the increasing size of models and of pretraining datasets. Data mixing, on the other hand, reduces the complexity of data selection by grouping data points together and determining sampling probabilities across entire groups. However, data mixing proportions are typically fixed before training and therefore cannot adapt to changing training dynamics. To address these limitations, we develop an efficient algorithm for Online Data Mixing (ODM) that combines elements from both data selection and data mixing. Based on multi-armed bandit algorithms, our online approach optimizes the data mixing proportions during training. Remarkably, our method trains a model that reaches the final perplexity of the next best method with 19\% fewer training iterations, and improves performance on the 5-shot MMLU benchmark by 1.9% relative accuracy, while adding negligible wall-clock time during pretraining.

Overview

This repository contains code for developing efficient online data mixing methods.

The majority of the functionality is found in the datasamplingutils.py file. Specifically, the SmoothedMeanWeightUpdater class is the data mixing implementation described in the paper, used to store and update mixture weights.

Note: This repository is under active development as is subject to change.

Setup

This repository is not setup to be directly run as-is. This is in part because it was developed for the now defunct Pile dataset. However, the core functions can be replicated with minimal effort.

If you already have access to the Pile dataset, you can use the code from this repo directly. To setup the environment, use: bash bash alon_setup.sh

For more details on setup and usage, you can follow the original gpt-neox README in the section below titled "Original GPT-NeoX README".

To Do:

• [ ] Isolate the functions specific to ODM so that it can be packaged with pip

Acknowledgements

This work is built on the GPT-NeoX framework developed by EleutherAI and the original repository can be found at https://github.com/EleutherAI/gpt-neox

How to Cite

If you extend or use this work, please cite the paper where it was first introduced: @misc{albalak2023efficient, title={Efficient Online Data Mixing For Language Model Pre-Training}, author={Alon Albalak and Liangming Pan and Colin Raffel and William Yang Wang}, year={2023}, eprint={2312.02406}, archivePrefix={arXiv}, primaryClass={cs.CL} }