Overview

This repository contains replication code for the following paper:

Aioli: A unified optimization framework for language model data mixing
Mayee F. Chen*, Michael Y. Hu*, Nicholas Lourie, Kyunghyun Cho, Christopher Ré
arxiv

Data mixing is an important step in data curation where practitioners must identify the optimal mixture of data groups to train on (i.e., code, law, math). While a brute-force search over the mixture proportions is a common technique in practice, this approach requires many training runs. Recent alternatives propose to algorithmically learn mixture proportions more efficiently; however, they are not well-understood and can sometimes even underperform a simple stratified sampling baseline.

In our work, we find that many mixing methods can be written as an optimization problem subject to a linear mixing law that describes the assumed relationship between loss and mixture proportion. We find that while methods are accurate in the parameterization of their mixing laws, they are incorrect in how they estimate the mixing law parameters. Notably, these inaccuracies in their parameters are correlated with when these methods do worse than stratified sampling, providing a more complete explanation of when and why data mixing methods fail.

Given these insights, the Aioli algorithm estimates the mixing law parameters from the current training trajectory and dynamically updates the mixture proportions. Aioli is able to consistently outperform stratified sampling and can enhance existing mixing methods by adjusting their learned proportions throughout the full training run.

Dependencies

To install dependencies:

git clone https://github.com/HazyResearch/aioli.git
cd aioli/
pip install -r requirements.txt

Usage

To run Aioli, refer to sample scripts scripts/arxiv_stackexchange/aioli/run.sh, scripts/arxiv_books_stackexchange/aioli/run.sh, and scripts/full/aioli/run.sh:

SLIMPAJAMA_DIR=your_directory
python3 main.py \
      --task_name slimpj \
      --train_data_dir $SLIMPAJAMA_DIR \
      --val_data_dir $SLIMPAJAMA_DIR \
      --selection_seed 0 \
      --max_steps 5000 \
      --sample_rule mixture \
      --slice_list arxiv stackexchange \
      --model EleutherAI/pythia-160m \
      --num_ckpts 20 \
      --batch_size 8 \
      --context_length 2048 \
      --lr 0.0005 \
      --aioli \
      --lp_rounds 4 \
      --lp_steps 4 \
      --eta 0.2 \
      --ema 0.1 \
      --update_steps 250 \
      --one_hot_factor 0.625 \
      --aioli_normalize_A \
      --checkpoint 0 \
      --lr_scheduler linear_warmup_cosine \
      --warmup_steps 500 \
      --use_flash_attention

Aioli parameters:

• --aioli: always set to use the algorithm.
• --lp_rounds: the number of sweeps for each mixture in LearnParams subroutine (Algorithm 2 in paper). Higher=more accurate