Studying how data ordering can improve pretraining

[kothasuhas]

Description

Interested in how to properly order pretraining data for fixed available data.

To start, we can consider how you would want to do curriculum learning for transfer: if you are given a few samples from D_1 and samples from less related D_2, what ordering of samples minimizes test loss for D_1? How does this vary with the number of samples from D_1 vs D_2?

Hypothesis or Goal

Hypothesis 1: To improve test loss on D_1, it will be beneficial to bias D_1 samples toward the end of training
Hypothesis 2: The benefit of keeping D_1 samples closer to the end of training will decrease as D_1 has more samples

Goal: The amount we care about domains is not proportional to their frequency. Suppose we have a utility function over these domains (induced by correlation to downstream benchmarks, or a non-linear function of their sample count). In that case, we should be able to design curricula that bias the appropriate domains toward the end of training. Along the way, we build a deeper understanding of what data needs to be throughout training and what data can be treated as fine-tuning data.

Links

• Project proposal document: https://docs.google.com/document/d/1UGm1ve92FrYnHB_Iufcsfbfq_Sq7bcvg_add6UoMPxM/edit?usp=sharing

Results

TODO

[kothasuhas]

High-level findings (refer to comment after this for more details)

General setup: Two domains to train on and we only care about test perplexity on one of the domains. You can only repeat each data point once. You are allowed to set one threshold for where the mixture changes, and both stages combined must match your budget constraint. Most of the below intuitions are built in a setting where 1 of domain is rare (Python, Flan instruction-tuning) and another is frequent (C4) for 150m parameter language models.

1. Data ordering matters more when there are fewer target tokens
2. Given two stages, there is never a need to keep the rare domain in stage 1 (across experiments, the best strategy across the space of all two-stage schedules always has all the rare data in stage 2)
   • Corollary: intuitions such as "baking the knowledge in during pretraining" might be inaccurate at describing optimal strategy for test perplexity
3. It is very important to keep the frequent data in stage 2: if you only keep rare data in stage 2 (i.e. fine-tuning), not only does it drastically hurt frequent data performance, it also drastically damages rare data performance.
4. Learning rate schedule - When using cosine lr schedule, the data ordering does not give you much difference in your downstream performance
   • When using a constant lr schedule with linear decay at the end of training, data ordering is a much more important problem
     • When you don't change data order at all, the best lr schedule is linear decay for 40% of training
     • When you are allowed to change data order, the optimal lr schedule is 5% decay, with a stage 2 that is around 10% of the end of training
     • If you care about the tradeoff between rare and frequent domain perplexity, the Pareto frontier involves jointly changing the lr schedule and duration of stage 2
5. Model scaling
   • The optimal data schedule changes as you increase parameter and token count (4xing rare and frequent data). It seems to shift slightly closer to "only fine-tune"
   • Surprisingly, the optimal data schedule does not change that much when you increase the parameter count! This means that the optimal data schedule at small scale could transfer to larger scale.

[kothasuhas]

Current project status with plots and more detail: https://docs.google.com/document/d/18YFR-JsX8G7cWo-YpDUsCurVkzORgme7fgurNpefAoc/edit?usp=sharing

[kothasuhas]

Have made some more progress! Attaching abstract here as a summary of the main findings

What is the best way to improve performance on a task with a small number of data points? For such applications, it is common to initialize training with a language model pre-trained on general web data instead of training from scratch. We find that on tasks such as math, coding, and instruction following, simply replaying pre-training data while fine-tuning can outperform standard fine-tuning, effectively training on up to 1.87x more samples. Though replay is often used to reduce catastrophic forgetting on previous domains, we surprisingly find that it is helpful even if you only care about loss on the new rare data. Next, we study how much additional rare data efficiency we can achieve by changing the data mixture across all of training. By searching over possible two-stage mixture schedules, we can improve data efficiency by up to 4.8x over the new baseline. Interestingly, we find cases where the optimal data schedule only requires replay, removing the need to allocate rare data to pre-training. We use the findings from our controlled experiments to improve performance for specific rare fine-tuning setups, increasing agentic web navigation accuracy by 4.5% and Basque question-answering accuracy by 2%.

Here is a (pretty preliminary) draft of the main findings. Project still under progress, so please ask questions!

Paper draft

[percyliang]

Can you link to the experiment page in the data browser for the experiments and any WandB report / runs?

[kothasuhas]

Here is (one of many) data browser links for reference:

wandb report best read alongside paper