Continued Pretraining for Domain Shift

A base model can write a polite paragraph about authentication and still fail the next-token test that matters for adaptation. Give it the start of an internal runbook, If AUTH_TIMEOUT fires, the client should, and watch it wander into generic retry advice that never mentions your jitter budget, idempotency key, or on-call owner. The chat format looks fine. The domain text doesn't.

The JAX chapter treated a training step as a state transition you can replay: parameters, optimizer state, RNG keys, and metrics all cross an explicit boundary. Continued pretraining (CPT) needs that same contract. Pin the starting checkpoint, the domain mix, the schedule, and two eval lanes, then you can tell whether the run actually moved the model.

CPT keeps the causal next-token objective you used in the scratch GPT lab. What changes is the text distribution: more compute lands on your domain's language, document structure, and recurring entities.^{[1]Reference 1Don't Stop Pretraining: Adapt Language Models to Domains and Tasks.https://aclanthology.org/2020.acl-main.740/} The weights shift because the data shifted, not because you switched losses.

Three tools get mixed up here:

• retrieval-augmented generation (RAG) leaves the weights frozen and injects documents at inference time. Use it when facts change often or must be cited.
• Supervised Fine-Tuning (SFT) changes behavior, format, and tone with curated prompt-response pairs. Use it when the model already reads the domain but answers in the wrong shape.
• Continued pretraining changes the weights with the same next-token objective so the model better fits unlabeled domain text. Use it when raw docs and runbooks still confuse the base model.

Don't rank those tools from labels alone. In Ovadia et al.'s knowledge-injection experiments, RAG beat unsupervised fine-tuning on MMLU and current-events questions, while repeated paraphrases of the same fact helped fine-tuning on a new-fact task.^{[2]Reference 2Fine-Tuning or Retrieval? Comparing Knowledge Injection in LLMshttps://aclanthology.org/2024.emnlp-main.15/} That's evidence for factual-update failures, not a universal ranking. CPT earns an experiment when the model poorly fits the domain text distribution, not when you only need fresher facts or a different answer format.

The architecture stays put. The loss stays causal language modeling (mean per-token negative log-likelihood) over a sequence of tokens $w_1,\ldots,w_n$:

$$\mathcal{L}_{\mathrm{CLM}}(\theta) = -\frac{1}{n}\sum_{i=1}^{n} \log P(w_i \mid w_{<i}; \theta)$$

Here $w_{<i} = (w_1, \ldots, w_{i-1})$ is the prefix context, and $P(w_i \mid w_{<i}; \theta)$ is the model's predicted probability for token $w_i$. Exponentiating this token-averaged loss yields the sequence perplexity $\mathrm{PPL} = \exp(\mathcal{L}_{\mathrm{CLM}})$.

$\theta$ starts from the base checkpoint, not from random initialization. Later SFT still uses a next-token loss, but the examples are prompt-response pairs and the loss is usually applied to the response tokens. That's a different supervision shape, even when the math looks related.

If more than one failure shows up, start with the earliest one in the stack. A model that can't continue AUTH_TIMEOUT runbooks is not ready for assistant-style SFT on those same codes.

Diagnose the shift before spending training compute

A useful direct signal is held-out raw-text loss and its exponentiated form, perplexity. Record the base model's value on domain documents, then test whether CPT lowers it while a general-text control stays inside budget. A base model scoring worse on domain than general text is only a screening clue, because corpora can have different inherent predictability. It doesn't prove CPT will improve product tasks.

Fragmentation during tokenization is a weaker diagnostic. A fixed tokenizer may spend more tokens on unfamiliar identifiers, which raises context cost, but CPT doesn't change that tokenizer unless you redesign embeddings and retrain compatible weights. Use fertility as a corpus inspection signal, not a promise that continued pretraining will shorten tokenized documents.

The tiny longest-match tokenizer below is not a production BPE. It makes the inspection concrete: timeout is in the vocabulary, AUTH_TIMEOUT is not, so the error code falls apart into characters.

Vocabulary extension vs. fixed tokenizer

High token fertility raises context consumption and inference latency. If domain identifiers consistently fragment into character-level pieces, you face a structural choice:

1. Keep the fixed tokenizer (standard practice): Retain the base vocabulary $V$ and embedding matrix $E \in \mathbb{R}^{V \times d}$. The model adapts its attention weights and representations to domain n-grams without disrupting existing token coordinates. This is the default in models like Code Llama because it avoids initialization instability.^{[3]Reference 3Code Llama: Open Foundation Models for Codehttps://arxiv.org/abs/2308.12950}
2. Extend the vocabulary: Add $K$ domain-specific tokens (such as frequent API methods or specialized identifiers), expanding the embedding matrix $E \to \mathbb{R}^{(V+K) \times d}$ and the output language-modeling head $W_{\mathrm{head}} \to \mathbb{R}^{(V+K) \times d}$.

If you do extend the vocabulary, don't initialize new rows with random Gaussian noise. Random vectors produce large gradient spikes that can destabilize pretrained attention layers during early steps. Instead, initialize each new token embedding with the mean of the embeddings of its constituent subwords from the base tokenizer. Every new token also needs enough occurrences in the adaptation corpus for those rows to move; a handful of AUTH_TIMEOUT hits won't do it. That's why small domain runs usually keep the base tokenizer intact.

Don't split a validation corpus by shuffled token chunks. Near-duplicates, revisions of the same OpenAPI page, or two runbook copies from the same source can land in both training and validation and make CPT look stronger than those splits justify. Assign a provenance or deduplication group to one split before tokenization.

Those two measurements still leave an open question: if you do start training, how hard should you push the weights?

The two failure dynamics: forgetting and underfitting

Resuming on a new distribution pulls the weights in two directions, and a good run balances them.

Catastrophic forgetting is loss of previously learned ability as parameters shift to absorb new data. Push too hard on AUTH_TIMEOUT runbooks and broad validation quality can regress.

Underfitting is the opposite failure: train too gently and the domain leaves no real impression. The model still can't continue the runbook.

Two major controls for this balance are the learning-rate schedule and the data mixture. Run length and corpus quality matter too.

Learning rate re-warming and re-decaying

A base checkpoint often finished its original cosine schedule at a very small learning rate. If you resume at that floor, adaptation may be inefficient. If you resume too aggressively, general-text loss may regress.

Ibrahim et al. (2024) study a related decoder-only continual-pretraining setting: updating a model with large new datasets after its original cosine schedule ended.^{[4]Reference 4Simple and Scalable Strategies to Continually Pre-train Large Language Modelshttps://openreview.net/forum?id=DimPeeCxKO} For 405M models under English-to-English and English-to-German shifts, and a 10B-parameter model under the English-to-English shift, learning-rate re-warming, re-decaying, and replay matched retraining baselines on their reported losses and evaluation averages while spending less compute. Their experiment is evidence for testing this recipe, not permission to copy one peak learning rate into every domain run.

Code Llama is a useful second data point. Those models start from Llama 2 and continue pretraining on 500B code-heavy tokens (1T for the 70B). The authors kept the original Llama 2 peak learning rates (for example $3\times 10^{-4}$ at 7B and 13B) rather than shrinking them the way people often do for fine-tuning.^{[3]Reference 3Code Llama: Open Foundation Models for Codehttps://arxiv.org/abs/2308.12950} That agrees with Ibrahim's warning not to sit at the cosine floor. It still isn't a universal peak to paste into your schedule.

One subtlety from Ibrahim et al.: re-warming can itself increase loss on old data. Sweep the peak and measure both lanes instead of assuming adaptation is free. The paper also explores schedules that aren't tied to one fixed token budget.

The snippet below only shows the shape: climb from a small floor to a peak, then cosine-decay back. The peak is an order-of-magnitude fixture from Ibrahim's published cosine max, not a hyperparameter you should inherit.

Replay: keep prior-data signal in the mix

The second knob is replay: mix a fraction of previous or representative general-purpose data back into the incoming domain corpus. It provides training signal on broad text while the domain stream shifts the model, so it's a practical candidate for limiting regression.

The effective training objective becomes a weighted expectation across data distributions:

$$\mathcal{L}_{\mathrm{mix}}(\theta) = (1 - \alpha)\,\mathbb{E}_{x \sim \mathcal{D}_{\mathrm{domain}}}[\mathcal{L}_{\mathrm{CLM}}(x; \theta)] + \alpha\,\mathbb{E}_{y \sim \mathcal{D}_{\mathrm{replay}}}[\mathcal{L}_{\mathrm{CLM}}(y; \theta)]$$

where $\alpha \in [0, 1]$ is the replay ratio.

How much replay? Treat it as a sweep, not a standard percentage. In Ibrahim et al.'s headline comparison, the chosen mixes use 5% replay for the SlimPajama update and 25% replay for the larger English-to-German shift.^{[4]Reference 4Simple and Scalable Strategies to Continually Pre-train Large Language Modelshttps://openreview.net/forum?id=DimPeeCxKO} Code Llama's 500B code stage is 85% code, 8% natural language related to code, and 7% general natural language.^{[3]Reference 3Code Llama: Open Foundation Models for Codehttps://arxiv.org/abs/2308.12950} Llama 3 later branched a code expert from the main pretraining run and continued on a 1T-token mix that was mostly (>85%) code, following a similar recipe.^{[5]Reference 5The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783} Those values belong to those datasets and compute budgets. With a fixed token budget, replay also replaces some new-domain tokens, so it can reduce adaptation opportunity while controlling general regression.

Your continued-pretraining run resumes from the base checkpoint at its final tiny learning rate and uses 100% domain text. Domain perplexity barely moves. After you raise the re-warm peak, domain perplexity improves but general-text loss regresses. Which two sweeps should you run?

Answer

Sweep a moderate re-warm peak followed by re-decay: inheriting the old floor can underfit, while an aggressive peak can damage broad behavior. Sweep replay ratios too, then select against held-out domain and general metrics. Replay is a regression-control candidate, not the explanation for the original underfitting.

When continued pretraining is the right tool

Reach for continued pretraining when the domain has its own language that the base model under-serves:

• internal API docs and error catalogs (AUTH_TIMEOUT, AUTH_REJECTED)
• on-call runbooks and incident notes
• SDK guides with domain-specific method names
• long compliance or protocol documents
• codebases whose APIs and identifiers barely appeared in public pretraining

The trigger isn't "the product team wants custom behavior." The model needs more exposure to the domain's text distribution before post-training behavior shaping makes sense.

Good signals

Bad signals

A practical split for the AUTH_TIMEOUT assistant: test raw domain-text continuation and prompt-response behavior separately. If the model can't continue the runbook or OpenAPI description, that points to continued pretraining. If raw continuation is competent but assistant behavior is weak, that points more directly to SFT.

The 2020 "Don't Stop Pretraining" paper made a related distinction in masked-language-model experiments with RoBERTa:^{[1]Reference 1Don't Stop Pretraining: Adapt Language Models to Domains and Tasks.https://aclanthology.org/2020.acl-main.740/}

• DAPT (domain-adaptive pretraining): keep training on large unlabeled domain text such as API docs or runbooks
• TAPT (task-adaptive pretraining): continue on the task's own unlabeled inputs, even when the corpus is smaller

The decision remains useful for decoder-only LLM projects, but don't silently transfer RoBERTa's quantitative gains to a causal base model. You still have to measure whether more exposure to the target text distribution improves your model and downstream task.

Data for continued pretraining

The same discipline from large-scale pretraining still applies:

• filter low-quality text
• deduplicate aggressively
• remove benchmarks and eval leakage
• scrub PII and secrets from runbooks and traces
• keep provenance and usage rights for every corpus slice

The corpus can be narrower and more targeted. Domain data can also be more sensitive than public pretraining text, so provenance, access control, and removal procedures are product requirements, not cleanup tasks.

Gate the corpus before tokenization

Keep a manifest that records whether a source may be trained on, whether it contains unresolved sensitive content, and whether it's reserved for evaluation. A high-quality domain document that fails one of these gates doesn't belong in the training stream.

oncall-notes is licensed but still contains unresolved PII, while heldout-probes is clean and licensed. Which source can enter continued-pretraining blocks?

Answer

Neither. Unresolved PII blocks oncall-notes, and the evaluation-only flag keeps heldout-probes out of training. Corpus quality never overrides rights, privacy, or evaluation-isolation gates.

Keep evaluation text out of training

For a small exact-overlap gate, normalize text and hash it before building token blocks. Production pipelines also need near-duplicate detection, because formatting changes and partial copies will evade exact hashes.

Mixing strategy

Don't train on 100% domain text just because you have it. Teams usually mix:

• a high-quality domain slice
• a smaller replay slice of general text

That replay is one guardrail against forgetting. The exact ratio is empirical: define candidate ratios, hold total training tokens fixed, and select with domain-gain and broad-regression metrics. If the model forgets too much general language while specializing, the run overshot.

BloombergGPT is a useful contrast, not replay evidence: it was trained from scratch on 51.27% financial and 48.73% public tokens, and reports strong financial performance while remaining competitive on general-purpose benchmarks.^{[6]Reference 6BloombergGPT: A Large Language Model for Financehttps://arxiv.org/abs/2303.17564} It shows that corpus composition should be explicit. It doesn't identify the right CPT replay ratio for your checkpoint.

Pack blocks and preserve the mixture

CPT uses the same causal objective as base pretraining. A common loader recipe joins document token sequences with end-of-document markers and emits full blocks. The separator marks a boundary, but it doesn't prevent cross-document attention by itself. As the data-pipeline chapter explained, choose explicitly between an ordinary causal mask and a document-isolated block-diagonal mask. Small integer token sequences make separator placement inspectable.

This small example drops an incomplete final block instead of padding it. Production loaders need an explicit remainder policy.

Once domain and replay streams are packed, make mixture selection explicit and auditable. Here each twenty-block training window uses a seeded shuffle with the requested replay count.

Evaluation: domain gain without lying to yourself

Continued pretraining needs two evaluation lanes at the same time.

Lane 1: domain gain

Measure:

• domain validation perplexity
• retrieval or classification tasks in the domain
• generation quality on held-out domain documents
• downstream task lift after later SFT

Lane 2: general regression

Measure:

• a small broad-language validation slice
• a lightweight general benchmark set
• free-form generations outside the target domain

If you only watch domain gain, you can accidentally produce a model that sounds like one AUTH_TIMEOUT runbook and forgot how to write broadly coherent language.

Evaluate loss in comparable token units. Perplexity is exp(mean negative log-likelihood), so aggregate token-level loss before exponentiating; don't average document perplexities and call the result a corpus metric.

Runnable checkpoint ledger

The simplest useful artifact is a checkpoint ledger. It doesn't train a model; it shows how to choose between checkpoints after a continued-pretraining sweep. Domain perplexity can improve while general text gets worse, so the chosen checkpoint needs to pass both lanes. Use general regression as a hard gate. Among survivors, rank downstream probe accuracy first and use domain perplexity as a tie-breaker. That keeps the policy visible instead of hiding trade-offs inside an arbitrary weighted score.

cpt-12k has the best domain perplexity but exceeds the allowed general-regression budget. cpt-4k passes the gate and has the best downstream probe accuracy among survivors. Which checkpoint wins?

Answer

Choose cpt-4k. General regression is a hard gate, so cpt-12k is ineligible. Among eligible checkpoints, downstream probe accuracy ranks first and domain perplexity only breaks ties.

Stopping rules

Because continued pretraining keeps the same objective, it can feel deceptively safe. It isn't safe by default.

Good stopping cues:

• domain validation loss flattens
• downstream gains after a probe SFT stop improving
• general regressions start to outweigh domain benefits

Bad stopping cues:

• "we still have more domain text"
• "loss is still going down a little"

More steps aren't a free lunch once the domain shift is already absorbed.

Where it fits relative to LoRA and SFT

Compare the choices by asking what you want to change.

LoRA and QLoRA are parameter-efficient implementation choices; QLoRA also stores the frozen base model in quantized form.^{[7]Reference 7QLoRA: Efficient Finetuning of Quantized Language Models.https://arxiv.org/abs/2305.14314} They don't determine what supervision teaches. An adapter can be trained with a next-token domain-text objective or with prompt-response SFT. First choose objective from the failure mode, then choose full-weight or parameter-efficient training from budget and deployment constraints.

A strong training stack often looks like:

1. base model
2. continued pretraining on domain corpus
3. SFT on curated prompt-response data
4. preference optimization if needed

Not every product needs every stage. Choose the stage that matches the failure you observe. Code Llama is one shipped version of that stack: continued pretraining on code, then instruction data.^{[3]Reference 3Code Llama: Open Foundation Models for Codehttps://arxiv.org/abs/2308.12950} Llama 3's code expert followed the same idea, then used the expert to collect better annotations rather than treating CPT as the product surface.^{[5]Reference 5The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783}

Common pitfalls

Using continued pretraining to fix assistant tone

• Symptom: the model still formats answers badly after a long domain-text run.

• Cause: the issue was interface behavior, not domain language exposure.

• Fix: move to SFT sooner.

Over-specializing on one corpus

• Symptom: domain completions improve, but the model becomes narrow or brittle elsewhere.

• Cause: no replay mixture, or too many adaptation steps.

• Fix: keep a general-text regression lane and stop earlier.

Skipping the downstream check

• Symptom: domain perplexity improves, but the final task model barely benefits.

• Cause: the adaptation run optimized text fit that did not transfer to the product task.

• Fix: probe the adapted checkpoint with a small downstream SFT instead of judging only by perplexity.