Continual Pre-Training of Large Language Models: How to (re)warm your model?

Kshitij Gupta,Benjamin Thérien,Adam Ibrahim,Mats L. Richter,Quentin Anthony,Eugene Belilovsky,Irina Rish,Timothée Lesort

2023-09-07

Abstract:Large language models (LLMs) are routinely pre-trained on billions of tokens, only to restart the process over again once new data becomes available. A much cheaper and more efficient solution would be to enable the continual pre-training of these models, i.e. updating pre-trained models with new data instead of re-training them from scratch. However, the distribution shift induced by novel data typically results in degraded performance on past data. Taking a step towards efficient continual pre-training, in this work, we examine the effect of different warm-up strategies. Our hypothesis is that the learning rate must be re-increased to improve compute efficiency when training on a new dataset. We study the warmup phase of models pre-trained on the Pile (upstream data, 300B tokens) as we continue to pre-train on SlimPajama (downstream data, 297B tokens), following a linear warmup and cosine decay schedule. We conduct all experiments on the Pythia 410M language model architecture and evaluate performance through validation perplexity. We experiment with different pre-training checkpoints, various maximum learning rates, and various warmup lengths. Our results show that while rewarming models first increases the loss on upstream and downstream data, in the longer run it improves the downstream performance, outperforming models trained from scratch$\unicode{x2013}$even for a large downstream dataset.

Computation and Language,Machine Learning

What problem does this paper attempt to address?

### Problems Addressed by the Paper

The paper primarily explores how to achieve efficient continual pre-training on large-scale language models (LLMs). Specifically, the paper focuses on the following points:

1. **Avoiding Complete Retraining**: Researchers aim to update existing pre-trained models through continual pre-training when new data becomes available, rather than retraining the model from scratch.
2. **Reheating the Learning Rate**: Researchers believe that when training on new datasets, it is necessary to re-increase the learning rate to improve computational efficiency.
3. **Mitigating Performance Degradation**: Since the distribution changes in new data usually lead to performance degradation on old data, researchers hope to mitigate this issue through different warm-up strategies.

Through experiments, researchers found that reheating the learning rate and decaying it on new datasets can improve the performance of downstream tasks and is superior to models trained from scratch. Additionally, the study explores the impact of different pre-training checkpoints, maximum learning rates, and warm-up lengths on performance.