Vol. XII · Mid-Training Plate 5 of 10continued pretraining, domain adaptation, DAPT, Code Llama, language adaptation, re-warming, learning rateLD

Continued Pretraining.

Take a trained model, keep training it on the new domain — and watch both the gain and the cost.

Training a model from scratch for every specialised domain is wasteful when a general model already knows most of what the domain needs. Continued pretraining takes a pretrained checkpoint and trains on more text from the target domain with the same objective. Gururangan and colleagues (2020) showed that this domain-adaptive pretraining improved results on biomedical, computer-science, news and review tasks; Code Llama was Llama 2 continued on 500 billion tokens of code; and models are adapted to new languages the same way, sometimes with an extended vocabulary.

Two choices dominate. The learning rate must be raised again — re-warmed — after a schedule that ended near zero, and the higher it goes, the faster the model moves towards the new domain and away from the old one. And some of the original data can be mixed back in (plate XII.6). The workshop trains a small bigram model on computing text for 2 000 steps, then continues it on mathematics and physics text.

Below: held-out loss on both domains through both phases, for the chosen continuation learning rate and length.

What you will learn
Workshop
Softmax bigram model on characters, Adam, batches of 256: 2 000 steps on the computing plates (learning rate 0.05), then the chosen number of steps on the mathematics and physics plates at the chosen learning rate. Held-out loss on 10% of each domain.

Gain where the data is, cost where it is not

After 2 000 steps on computing text, the model scores 2.567 nats per character on held-out computing text and 2.773 on mathematics and physics, a domain it has never seen but whose characters and common words it shares. Continuing for 1 000 steps on the new domain at learning rate 0.05 brings science to 2.699, a gain of 0.074, and lets computing slip to 2.624, a cost of 0.057. At a re-warmed rate of 0.01 the move is gentler — science 2.706, computing 2.601 — and at 0.1 both end worse, 2.716 and 2.646: too high a rate overshoots.

That trade-off is the core of continued pretraining at any scale. The domains here are close, so the cost is small; adapting to a new language, or to code from prose, moves the model further and costs more. Ibrahim and colleagues found that re-warming and re-decaying the learning rate, together with replaying a small fraction of the original data, let continued pretraining match the quality of retraining from scratch on all the data at a fraction of the compute — which is why it is the standard way to specialise a model.

Continued pretraining keeps the original tokeniser unless it is extended. For a new language that the tokeniser splits into many short pieces, adding vocabulary — and training the new embedding rows — can cut the number of tokens per sentence substantially (plate II.7).

Terms

Continued pretraining
Training a pretrained model further on text from a target domain or language.
Domain-adaptive pretraining
Continued pretraining on a domain’s unlabelled text before fine-tuning on its tasks.
Re-warming
Raising the learning rate again at the start of a continuation.
Stability gap
The temporary rise in loss on old data when a continuation begins.
Continued pretraining moves a model towards new data; the re-warmed learning rate sets how far, and how much of the old domain is given up.

Sources

  1. Gururangan et al., “Don’t Stop Pretraining: Adapt Language Models to Domains and Tasks”, ACL 2020.
  2. Rozière et al., “Code Llama: Open Foundation Models for Code”, arXiv:2308.12950, 2023 (Llama 2 continued on 500 billion tokens of code).
  3. Ibrahim et al., “Simple and Scalable Strategies to Continually Pre-train Large Language Models”, TMLR 2024.