ArXiv: 2308.04014
🎯 Pitch
Re-heating a language model’s learning rate after initial training initially hurts performance, but ultimately lets it adapt to a massive new dataset better than training from scratch—even when the new data is as large as the original.
1. Executive Summary
This paper studies how to efficiently continue pre-training large language models on new datasets rather than re-training from scratch, examining the effect of re-warming — the strategy of re-increasing the learning rate after pre-training has completed and the rate has decayed to a small value — on the Pythia 410M model when transitioning from the Pile (upstream, 300B tokens) to SlimPajama (downstream, 297B tokens) under a linear warmup with cosine decay schedule. The central finding is that rewarming then decaying the learning rate is necessary for effective downstream adaptation, and that varying the maximum learning rate during rewarming creates a controllable tradeoff: higher maximum rates (6×10⁻⁴) improve downstream SlimPajama performance at the cost of upstream Pile forgetting, while lower rates (1.5×10⁻⁴) preserve upstream performance at the cost of reduced downstream adaptation. Continual pre-training with rewarming ultimately outperforms models trained from scratch on SlimPajama alone — even when the downstream dataset is as large as the upstream one — establishing that positive transfer between datasets makes re-training unnecessary, though the length of the warmup phase itself has no significant effect and using earlier (less-converged) pre-training checkpoints does not improve downstream adaptation.
2. Context and Motivation
The Core Problem: Pre-Training From Scratch Is Unsustainable as Data Evolves
The fundamental problem this paper addresses is deceptively straightforward: when a new, improved pre-training dataset becomes available, should practitioners re-train their language models from scratch, or can they efficiently update existing pre-trained models? This question matters because the economics of large language model development are reaching a breaking point. Training frontier models from scratch costs millions of dollars in compute (Hoffmann et al., 2022; Brown et al., 2020) and produces substantial carbon emissions. If every new dataset release demands a ground-up re-training, the cumulative cost quickly becomes prohibitive — not just for the largest industrial labs, but also for academic and mid-scale research groups who want to keep their models current.
The paper frames this as a continual pre-training problem (Section 1, paragraph 2):
"Our approach circumvents the need for complete re-training by continuing to pre-train existing models on new data. We refer to this as 'continual pre-training' and the goal is to minimize the loss on new data while maintaining low loss on previous data."
This is distinct from standard fine-tuning. In typical fine-tuning, a pre-trained model is adapted to a relatively small, task-specific dataset. Here, the downstream dataset — SlimPajama, at 297B tokens — is roughly the same scale as the upstream dataset — the Pile, at 300B tokens. This is domain-adaptation-scale continual training, where the model faces a distribution shift between two massive corpora. The goal is not just task performance, but preserving general language modeling capability across both old and new data distributions.
Why This Problem Is Urgent: The Data Landscape Is Accelerating
The paper situates its motivation in a specific, observable trend (Section 1, paragraph 1):
"As the amount of data available for pre-training is ever-growing, new and improved datasets (e.g. RedPajama and SlimPajama (Together.xyz, 2023; Soboleva et al., 2023; Touvron et al., 2023)) will continue to become available."
This is not a hypothetical concern. In the span of a few years, the community has moved from the Pile (2020, 800GB) to RedPajama (2023, 5TB of raw data) to SlimPajama (2023, 627B cleaned and deduplicated tokens), with LLama-derived datasets representing entirely new curation strategies. Each release incorporates lessons from prior work — better deduplication, improved quality filtering, broader coverage — making them genuinely superior to their predecessors. A team that trained on the Pile in 2021 faces a concrete decision: retrain on SlimPajama at full cost, or find a way to continue training that leverages the compute already spent.
The paper also connects to a broader trend: the emergence of temporally evolving corpora (Jin et al., 2022; Han et al., 2021; Loureiro et al., 2022), where the data distribution itself shifts over time as language and knowledge evolve. In such settings, re-training from scratch is not merely expensive — it is structurally inadequate, because the model needs to track changing distributions continuously. While the paper does not itself study temporal dynamics (the focus is on domain shift between static datasets), the continual pre-training paradigm it investigates is necessary infrastructure for temporal adaptation.
The Central Challenge: Catastrophic Forgetting Meets Large-Scale Training
The fundamental tension in continual pre-training is the same one that pervades all of continual learning: catastrophic forgetting (French, 1999). When a neural network is trained on new data, gradient updates that reduce loss on the new distribution typically increase loss on the old distribution unless explicit countermeasures are taken. The paper explicitly names this as the challenge:
"Continual pre-training is a critical challenge since it can lead to catastrophic forgetting (French, 1999)."
However, the paper argues that standard continual learning remedies are poorly suited to this setting (Section 1, paragraph 2). Replay-based methods (Rebuffi et al., 2017; Ostapenko et al., 2022), which store and interleave samples from previous datasets, incur storage and compute overhead proportional to the number of previous training stages — a cost that compounds as more datasets are encountered. Regularization-based methods (Kirkpatrick et al., 2017; Farajtabar et al., 2020), which penalize parameter changes that harm performance on previous tasks, add computational overhead to each gradient step and scale poorly with model size. The paper frames these as:
"...not compute efficient enough"
for the long sequence of training stages that continual pre-training envisions. This is a critical insight: in continual learning research, methods are often evaluated on relatively small models and short task sequences. Scaling these approaches to billion-token pre-training runs with hundreds of millions of parameters is an entirely different regime where per-step overhead is unacceptable.
Prior Approaches to Learning Rate Management in Continual Training
Given that standard CL methods are too expensive, the paper turns to a simpler, more scalable lever: the learning rate schedule. This is where the paper situates its primary contribution.
The dominant paradigm in large language model training is a linear warmup followed by cosine decay to approximately 10% of the maximum learning rate (Brown et al., 2020; Hoffmann et al., 2022; Touvron et al., 2023; Scao et al., 2022). After pre-training completes, the learning rate is at its minimum — often three orders of magnitude below the peak. If you simply resume training on new data at this tiny learning rate, weight updates will be too small to meaningfully adapt to the new distribution. This is the core practical problem: the learning rate is too low to learn, but raising it risks forgetting.
Prior work has explored several approaches to this dilemma, which the paper reviews in Section 3:
Constant learning rate. Some studies (Ke et al., 2023a; Scialom et al., 2022) maintain a fixed, moderate learning rate throughout continual training, effectively treating the downstream phase as an extension of pre-training. This ensures the model continues to learn, but provides no mechanism to control the forgetting rate or to stabilize training after the new data is absorbed.
Progressive decrease. Winata et al. (2023) advocate for progressively lowering the learning rate as new data stages are encountered. The intuition is that each new dataset requires smaller and smaller adjustments from the model, eventually converging to a stable parameter configuration. However, this approach has a structural limitation that the paper identifies (Section 1, paragraph 2):
"...repeatedly decreasing the learning rate would cause it to eventually become too small if the number of training stages becomes high."
In other words, a monotonically decreasing schedule paints itself into a corner: after enough stages, the learning rate approaches zero and the model loses all capacity to adapt to genuinely new data.
Warmup then decrease. Caccia et al. (2021) apply a warmup-then-decay schedule for each new stage of an online language learning setting, re-increasing the learning rate before decaying it again. This is the closest precursor to the paper's approach, but the focus of prior work was on task-level performance in classical continual learning benchmarks, not on the specific dynamics of large-scale, distribution-shifted pre-training where both upstream and downstream datasets are massive.
The gap the paper identifies. Critically, the paper states (Section 3, final paragraph):
"to the best of our knowledge, no existing work studies specifically the influence of the warm-up phase in the context of continual pre-training for large language models."
Prior work that used warmup in continual settings treated it as a standard training hyperparameter — something you include because it helps optimization — without systematically studying how the warmup length, the maximum learning rate chosen after rewarming, and the training convergence state interact with the dual objectives of downstream adaptation and upstream retention. The paper positions itself as filling this specific gap: understanding whether, how, and how much to rewarm.
The Paper's Framing: Recasting Warmup as a Strategic Control Knob
The paper reframes the warmup phase — typically treated as a detail of optimizer configuration — as a strategic lever for managing the upstream-downstream tradeoff. The hypothesis is stated in the abstract:
"Our hypothesis is that the learning rate must be re-increased to improve compute efficiency when training on a new dataset."
This is more subtle than it sounds. The claim is not merely that rewarming helps (which is intuitive — a tiny learning rate cannot produce meaningful adaptation), but that the parameters of rewarming — specifically the peak learning rate — can be tuned to balance forgetting and adaptation. Higher peak rates maximize the model's capacity to absorb new data but destabilize previously acquired knowledge. Lower peak rates preserve stability but limit adaptation. There is no free lunch: the paper expects a Pareto tradeoff frontier, and the experiments are designed to map it.
The paper also explicitly frames its experimental setup as analogous to warm-starting for image classification studied by Ash & Adams (2020):
"Our experimental setup is comparable to the setup of (Ash & Adams, 2020), where they train a classifier on half of the samples of a dataset first, and fine-tune it later on all samples. They show that warm starting for image classification is challenging."
This comparison is instructive. Ash & Adams found that simply resuming training on the full dataset from a model pre-trained on half the data is surprisingly difficult — the model can perform worse than training from scratch due to optimization challenges. By drawing this parallel, the paper signals that its investigation is not about whether continual pre-training is a good idea in principle (it clearly saves compute), but about understanding and overcoming the optimization barriers that make it fail in practice when not handled carefully.
The Underlying Tension: Data Similarity Makes Continual Training Both Easier and Harder
A subtle but important point that the paper grapples with is the nature of the shift between Pile and SlimPajama. The two datasets are drawn from similar sources (Common Crawl, Wikipedia, GitHub, books, academic papers — see Table 1) but differ in scale, curation strategy, and cleaning procedures. SlimPajama is an extensively deduplicated version of the LLama training corpus, while Pile is a more diverse, less aggressively filtered dataset.
This creates an interesting dynamic. On one hand, the high similarity means positive transfer is likely — knowledge learned on Pile (grammar, factual knowledge, reasoning patterns) should be useful when training on SlimPajama. This is what the paper hypothesizes will make continual pre-training outperform training from scratch on SlimPajama alone. On the other hand, the similarity means data overlap is present — some content from Pile likely appears (possibly in slightly different form) in SlimPajama. This complicates the interpretation of forgetting measurements: a model that "remembers" Pile content may be benefiting from having seen it again in SlimPajama, rather than genuinely retaining it.
The paper acknowledges this in Section 5 (Discussion):
"Since in continual learning, different types of shifts can lead to variations in performance (Lesort et al., 2021), our results may not generalize to setups with different distribution shifts, such as language domain adaptation pre-training setups."
This is both a limitation and a deliberate choice. By starting with highly similar datasets, the paper establishes an upper bound on transferability — if rewarming is challenging even under favorable conditions (similar data, same architecture, same optimizer), then more dramatic domain shifts (e.g., general text to biomedical literature, English to multilingual) are likely to be even harder. The paper is exploring the easy case first to establish baselines and understand fundamental dynamics before tackling harder distribution shifts.
Summary: What the Paper Aims to Establish
The paper's motivational argument can be distilled into four claims that motivate the experimental design:
-
The practical problem is real and growing: New pre-training datasets are released regularly, and re-training from scratch is economically unsustainable, creating an urgent need for efficient continual pre-training methods.
-
Standard continual learning methods don't scale: Replay and regularization incur overhead prohibitive for large-scale pre-training, motivating a focus on lightweight, schedule-based interventions.
-
Learning rate management is the obvious but underexplored lever: Prior work uses warmup in continual settings but has not systematically studied how warmup parameters — length, peak rate, and checkpoint selection — govern the adaptation-forgetting tradeoff in the specific context of billion-token LLM pre-training.
-
Data similarity sets the stage for isolating optimization dynamics: By choosing upstream and downstream datasets with high overlap, the paper can focus on optimization phenomena (the stability gap, the tradeoff curve, the effect of convergence state) rather than confounds from extreme distribution shift, establishing baselines that can inform future work on more challenging domain adaptation scenarios.
The experiments that follow are designed to map the functional relationship between rewarming hyperparameters (warmup length, maximum learning rate, checkpoint age) and the dual performance metrics (upstream Pile perplexity, downstream SlimPajama perplexity), ultimately providing practitioners with actionable guidance on how to continue pre-training efficiently.
3. Technical Approach
3.1 Reader Orientation
This is an empirical analysis paper that systematically studies what happens when you resume pre-training a large language model on a new dataset — specifically, whether and how to re-increase ("rewarm") the learning rate after it has decayed to a small value, and how the parameters of that rewarming (peak rate, warmup duration, checkpoint selection) govern the tradeoff between learning new data and forgetting old data. The system being analyzed is a pre-trained language model checkpoint on which training is resumed with a modified learning rate schedule; the paper does not propose a novel architecture or algorithm, but rather provides the first careful characterization of how optimizer schedule design choices determine success or failure in the continual pre-training regime.
3.2 Big-Picture Architecture (Diagram in Words)
The experimental system has four major components:
-
Upstream Pre-Training Pipeline — the Pythia 410M model (GPT-NeoX architecture) is pre-trained on the Pile dataset (300B tokens) following standard practice: linear warmup then cosine decay to 10% of peak learning rate. This produces multiple checkpoints at different stages of convergence.
-
Downstream Continual Pre-Training Pipeline — a selected upstream checkpoint is loaded and training resumes on SlimPajama (297B tokens), but with a new learning rate schedule: the rate is re-warmed from its decayed minimum to a new peak value (which may differ from the upstream peak), then cosine-decayed again. The warmup length, peak rate, and choice of checkpoint are the experimental variables.
-
Dual Validation Evaluation — throughout downstream training, the model's perplexity is evaluated on held-out validation sets from both SlimPajama (measuring adaptation) and the Pile (measuring forgetting). These two metrics are tracked jointly to reveal the Pareto tradeoff.
-
Baseline Comparators — (a) A model trained from scratch on SlimPajama with the standard warmup-cosine schedule, (b) a model continually trained with a constant (non-decayed) learning rate, and (c) in one experiment, a model re-warmed and fine-tuned on the same data (Pile) to isolate the effect of distribution shift from the effect of rewarming itself.
Information flows: a pre-trained checkpoint is loaded → the optimizer state (moments, parameter values) is inherited → a new learning rate schedule is imposed (linear warmup from near-zero to MaxLr, then cosine decay to 0.1 × MaxLr over 240B tokens, flat thereafter) → training proceeds on SlimPajama → perplexity is logged on both upstream and downstream validation sets at regular intervals → the co-evolution of these two losses reveals the adaptation-forgetting dynamics.
3.3 Roadmap for the Deep Dive
- First, the experimental design: datasets, model, and the standard training recipe inherited from Pythia, establishing the baseline against which rewarming strategies are compared.
- Second, the learning rate schedule parameterization — the central mechanism being varied — including the formal definition of the warmup-cosine-decay schedule, the three control knobs (warmup length, MaxLr, checkpoint selection), and the computational budget.
- Third, the evaluation framework: how upstream and downstream performance are measured simultaneously, why perplexity is the chosen metric, and how the tradeoff is visualized (the joint perplexity plots).
- Fourth, the baseline training runs that establish reference points — from-scratch training on SlimPajama and constant-LR continual training — needed to interpret whether rewarming "works."
- Fifth, the same-data ablation (Section 4.4) that isolates whether observed performance degradation is caused by distribution shift or by the rewarming operation itself.
- Sixth, the checkpoint selection experiment (Section 4.5) that varies the convergence state at which rewarming begins.
3.4 Detailed, Sentence-Based Technical Breakdown
This is an empirical analysis paper whose core idea is that the learning rate must be re-increased (rewarmed) when continuing pre-training on new data, and that the peak value chosen for this rewarming constitutes a direct control knob for the tension between downstream adaptation and upstream forgetting—no single setting is universally optimal, but the rewarming-then-decaying paradigm always outperforms training from scratch on the new data.
3.4.1 Experimental Design: What Is Held Fixed and What Is Varied
The paper fixes essentially everything except the learning rate schedule parameters, so that the observed effects can be confidently attributed to schedule choices rather than model architecture, dataset scale, or optimizer configuration. All experiments use the Pythia 410M parameter model (Biderman et al., 2023), which implements the GPT-NeoX architecture (Black et al., 2022). This is a decoder-only transformer with 24 layers, a hidden dimension of 1024, and 16 attention heads, trained with the same tokenizer used for the Pile — a BPE tokenizer trained specifically on the Pile corpus, producing a vocabulary of 50,304 tokens. The model uses learned positional embeddings and standard causal attention masking. Flash attention (Dao et al., 2022) is explicitly not used (Section 2, Model paragraph), which is a relevant detail because memory-efficient attention implementations can affect the feasible batch size and training throughput, potentially interacting with schedule design.
The optimizer is AdamW (Loshchilov & Hutter, 2018) with hyperparameters fixed exactly as in the original Pythia pre-training: $\beta_1 = 0.9$, $\beta_2 = 0.95$, $\epsilon = 10^{-8}$, and a weight decay of 0.1. Weight decay is applied directly to the weights (decoupled from the adaptive learning rate scaling, as per the AdamW formulation), meaning that the effective regularization strength is independent of the learning rate schedule. Gradient clipping is set to a global norm maximum of 1.0. Training is conducted in half-precision (FP16) without dropout. These optimizer settings are important because the interaction between learning rate magnitude and Adam's per-parameter scaling (via the second-moment estimate $v_t$) means that simply increasing the learning rate does not linearly increase the effective step size — the adaptive scaling partially compensates. Understanding this interaction requires knowing that AdamW updates take the form:
where $\theta_t$ is the parameter vector at step $t$, $\eta$ is the learning rate, $\hat{m}_t$ is the bias-corrected first moment estimate (exponential moving average of gradients), $\hat{v}_t$ is the bias-corrected second moment estimate (exponential moving average of squared gradients), $\epsilon = 10^{-8}$ prevents division by zero, and $\lambda = 0.1$ is the weight decay coefficient. The final term $-\eta \lambda \theta_t$ is the weight decay contribution.
What it computes: the update direction for each parameter $\theta$ at step $t$, combining three terms: scaled momentum (moving in the direction of recent gradients, normalized by recent gradient magnitude), numerical stability (the $\epsilon$ prevents blow-up when $\hat{v}_t$ is near zero), and weight decay (shrinking all parameters toward zero, scaled by both the weight decay coefficient and the learning rate).
Why this form: the decoupled weight decay in AdamW means that the regularization strength ($\eta \lambda$) is directly proportional to the learning rate. This is crucial for interpreting the paper's results: when the learning rate is increased during rewarming, the effective weight decay also increases proportionally. This could confound the interpretation that larger learning rates cause more forgetting solely through larger gradient steps — part of the forgetting may be due to stronger regularization pushing parameters toward zero. The paper does not discuss this interaction, but it is a property of the chosen optimizer that affects the mechanism.
The upstream dataset is the Pile (Gao et al., 2020), an 800GB curated collection of 22 diverse text sources including academic papers (ArXiv, PubMed), code (GitHub), books, web text (Common Crawl via Pile-CC, OpenWebText2), Wikipedia, and domain-specific corpora (legal, medical, patents). The Pile is used with the same sampling weights as Black et al. (2022) for validation, meaning that during evaluation, different sub-domains are reweighted to avoid Common Crawl dominating the perplexity signal. However, the paper does not detail how these weights are applied to the upstream perplexity computation — the validation loss on the Pile is reported as a single aggregated number.
The downstream dataset is SlimPajama (Soboleva et al., 2023), an extensively deduplicated and quality-filtered version of RedPajama (Together.xyz, 2023), which itself is an open reproduction of the LLama training corpus (Touvron et al., 2023). The paper subsamples SlimPajama to form a training set of approximately 297B tokens and a validation set of approximately 316M tokens. Table 1 provides the exact composition:
| Subset | Sampling % | Training Tokens | Validation Tokens |
|---|---|---|---|
| CommonCrawl | 67.0% | 153.25B | 147.28M |
| C4 | 15.0% | 78.49B | 72.49M |
| GitHub | 4.5% | 15.41B | 22.42M |
| Book | 4.5% | 14.22B | 22.04M |
| Wikipedia | 4.5% | 11.78B | 15.79M |
| ArXiv | 2.5% | 13.77B | 22.73M |
| StackExchange | 2.0% | 9.95B | 13.08M |
| Total | 100% | 296.86B | 315.83M |
The sampling percentages represent the probability that any given token during training is drawn from that subset. CommonCrawl dominates at 67%, reflecting web text's prevalence in the LLama-style corpus design, while curated sources like Wikipedia and ArXiv each contribute only 2.5–4.5%. This is a substantially different mixture than the Pile, which has a more uniform distribution across sources and includes domains (patents, DM Mathematics, Freelaw) absent from SlimPajama. The paper explicitly notes that no replay buffer is used — the model never sees Pile data during downstream training.
Why this dataset pairing: the Pile and SlimPajama share similar high-level domains (web, code, academic text, Wikipedia, books) but differ in curation strategy, deduplication aggressiveness, and source weighting. This makes the setup a "realistic best case" for positive transfer — knowledge from Pile should help on SlimPajama — while the distribution shift is still genuine enough to stress-test the optimization dynamics of rewarming. The paper acknowledges that domain overlap is present (Section 5), and this is a deliberate choice to study the fundamental optimization behavior without the confounding factor of extreme domain mismatch.
3.4.2 The Learning Rate Schedule Parameterization
The central mechanism being studied is the learning rate schedule, specifically how its parameters are set when transitioning from upstream to downstream training. The schedule follows the standard two-phase structure used in almost all LLM pre-training (Brown et al., 2020; Hoffmann et al., 2022; Touvron et al., 2023): a linear warmup from near-zero to a maximum value, followed by cosine decay to a minimum value.
The schedule during downstream continual pre-training is parameterized by three variables:
-
MaxLr: the peak learning rate reached at the end of warmup. Tested at three values:$1.5 \times 10^{-4}$,$3 \times 10^{-4}$(the default used for original Pythia pre-training), and$6 \times 10^{-4}$. -
Warmup length: the fraction of the total downstream training budget allocated to linearly increasing the learning rate from its minimum to
MaxLr. Tested at four values: 0%, 0.5%, 1%, and 2% of the full 297B-token downstream dataset. A warmup of 1% on 297B tokens means the learning rate increases linearly over approximately 2.97B tokens. The 0% condition means the learning rate jumps instantaneously toMaxLrat step 1. -
Minimum learning rate: always set to
$0.1 \times \text{MaxLr}$(i.e., 10% of the chosen peak). This is the standard practice from prior work and is not varied. This means that across the threeMaxLrconditions, the final learning rate after cosine decay is$1.5 \times 10^{-5}$,$3 \times 10^{-5}$, or$6 \times 10^{-5}$respectively.
The full downstream training budget is 240B tokens of learning rate decay (the cosine schedule reaches the minimum at 240B, and the rate remains constant at the minimum for any remaining tokens up to the full 297B). The paper reports perplexity curves out to 240–250B tokens (the end of the decay period).
The upstream schedule (for context). During original pre-training on the Pile, the Pythia 410M model was trained for 143,000 iterations with MaxLr = 3 × 10^{-4}, linear warmup over approximately 1% of the total training steps, and cosine decay to $0.1 \times 3 \times 10^{-4} = 3 \times 10^{-5}$. After this pre-training, the learning rate is at its minimum of $3 \times 10^{-5}$. If one were to simply resume training on SlimPajama at this rate without rewarming, the effective step size would be 10× smaller than the peak value used during pre-training — potentially too small to escape the current loss basin or to make meaningful progress on the new data distribution. This is precisely the problem that rewarming is designed to solve.
Why these particular MaxLr values: the default of $3 \times 10^{-4}$ is the peak used for original Pythia pre-training and serves as the "matched" condition — rewarming to the same peak as pre-training. $1.5 \times 10^{-4}$ is half the original peak and represents a conservative rewarming (prioritizing stability over adaptation). $6 \times 10^{-4}$ is double the original peak and represents an aggressive rewarming (maximizing adaptation at the cost of potential instability). The factor-of-two increments allow the paper to test whether the relationship between MaxLr and the adaptation-forgetting tradeoff is monotonic.
Why cosine decay to 10%: this is the de facto standard in LLM pre-training, established by works like GPT-3, Chinchilla, and LLaMA. Cosine decay has the property that the learning rate decreases slowly at first (near the peak) and then accelerates in its decrease as it approaches the minimum. This provides a long period at relatively high learning rates for exploration, followed by a convergence phase. Holding this decay shape constant isolates the effect of the peak value and warmup length.
An important schedule detail: the paper notes that the cosine decay reaches the minimum learning rate at 240B tokens and is constant thereafter. This means the final phase of training (from 240B to 297B tokens) is effectively constant-LR fine-tuning at 10% of the peak MaxLr. This detail matters because it means the different MaxLr conditions experience different constant-LR plateaus: the MaxLr = 6 × 10^{-4} condition finishes training at $6 \times 10^{-5}$, which is double the final rate of the MaxLr = 3 × 10^{-4}$ condition ($3 \times 10^{-5}$). Any continued movement in the perplexity curves after 240B tokens would likely be attributable to this constant-LR phase.
3.4.3 The Evaluation Framework: Simultaneous Upstream and Downstream Tracking
The paper's evaluation framework is designed to measure two things simultaneously — how well the model learns the new data and how much it forgets the old data — and to visualize their co-evolution throughout training. This is achieved through dual validation perplexity tracking.
Perplexity as the metric. Perplexity (PPL) is the exponentiated average negative log-likelihood per token:
where $N$ is the total number of tokens in the validation set, $x_i$ is the $i$-th token, $x_{<i}$ are the preceding tokens, and $p_\theta(x_i | x_{<i})$ is the model's predicted probability for token $x_i$ given its context, under parameters $\theta$.
What it computes: the geometric mean of the model's inverse confidence — lower perplexity means the model assigns higher probability to the actual next tokens. A perplexity of 2.5 means the model is, on average, as uncertain as if choosing uniformly from 2.5 options at each step; a perplexity of 10 means uncertainty equivalent to 10 options.
Why perplexity: for causal language modeling on large validation sets, perplexity is a smooth, continuous signal that correlates well with downstream task performance and is sensitive to both distribution shift (the model encounters tokens or patterns it wasn't trained on) and optimization state (the model moves away from previously learned minima). Perplexity can be measured on any text without labels, making it ideal for tracking both upstream and downstream performance simultaneously. However, the paper acknowledges (Section 5) that perplexity is an imperfect proxy — improvements in perplexity do not always translate to improvements on downstream benchmarks — and lists evaluating on HELM or Harness as important future work.
The upstream validation set is a held-out portion of the Pile, evaluated using the same sampling weights as Black et al. (2022). The exact size of this validation set is not specified in the paper, but the Pile's standard validation split provides approximately 1–2M tokens per sub-domain, aggregated across the 22 domains.
The downstream validation set is the 316M-token SlimPajama validation split described in Table 1, with tokens distributed across subsets proportional to their training weights (CommonCrawl at 67%, C4 at 15%, etc.). The paper tracks this loss at regular intervals, producing curves like those in Figures 2 and 3 that show SlimPajama Val. Loss and Pile Val. Loss as functions of downstream tokens processed.
The joint perplexity plot (Figure 4, Figure 6). The paper introduces a visualization that plots downstream perplexity on the x-axis against upstream perplexity on the y-axis for each time step of training. Each condition traces a curve in this space:
- The starting point (red dot in Figure 4) is the model's performance before any downstream training begins — low perplexity on both upstream and downstream, since the model was trained on Pile and Pile and SlimPajama are similar.
- As training begins, the model moves up and right (both losses increase) during the warmup phase, a phenomenon the paper calls the "stability gap" or "chaotic phase."
- As training continues and the learning rate decays, the model moves down and left, generally settling at a point with higher upstream loss than the start but lower downstream loss.
- The green points mark the end of the warmup phase, showing the transient degradation before recovery.
Why this visualization matters: it reveals the Pareto frontier. No single MaxLr dominates — higher rates produce lower downstream loss but higher upstream loss (further right, higher up), while lower rates preserve upstream performance but adapt less. A practitioner can pick their desired tradeoff point on this frontier by choosing the appropriate MaxLr. Additionally, the shape of the trajectory (not just the endpoint) reveals the optimization dynamics — the initial loss spike, the rate of recovery, and whether the model converges back toward its starting point or settles at a new equilibrium.
3.4.4 Baseline Comparators: Establishing Reference Points
To determine whether rewarming "works," the paper needs reference points that establish what performance looks like under alternative strategies. Three baselines are defined:
Baseline 1: Training from scratch on SlimPajama. A randomly initialized Pythia 410M model is trained directly on SlimPajama using the same schedule as original Pile pre-training: linear warmup to $3 \times 10^{-4}$ over 1% of the 297B tokens, followed by cosine decay to $3 \times 10^{-5}$. This baseline answers the question: is it better to continue pre-training from a Pile checkpoint, or to simply start over? Since the total compute used in pre-training on Pile (300B tokens) is approximately equal to the compute saved by not pre-training on SlimPajama from scratch, this comparison is essentially a transfer efficiency test: does the knowledge in the pre-trained checkpoint provide enough benefit to justify the compute already spent on it?
The from-scratch model has one structural disadvantage: it only sees SlimPajama data and thus cannot benefit from the broader coverage of the Pile (which includes domains like patent text and DM Mathematics absent from SlimPajama). This is the positive transfer hypothesis — the continual pre-training model should outperform the from-scratch model on SlimPajama because it benefits from knowledge learned on Pile.
Baseline 2: Constant learning rate. Instead of rewarming and decaying, the downstream training uses a fixed learning rate of $3 \times 10^{-5}$ for the entire 240B tokens. This is the learning rate at which upstream pre-training ended — effectively, "just keep training at the final rate." This baseline answers: is the cosine decay after rewarming necessary, or can you just train at a moderate constant rate?
The constant-LR baseline is particularly interesting because it represents the simplest possible continual training strategy: no schedule engineering, no hyperparameter tuning, just continuing as if pre-training never ended. If this baseline performs competitively, the entire novelty of rewarming would be questionable. As the results show (Section 4.2), the constant-LR model achieves the best downstream perplexity early in training (within the first 100B tokens) but is eventually overtaken by rewarmed-then-decayed models as the decay improves convergence precision. This reveals that the advantage of decay is in the late phase: by reducing the learning rate, the model can settle into a deeper, narrower minimum on the loss landscape.
Baseline 3 (implicit): No rewarming (MaxLr = current minimum). While not explicitly labeled as a separate condition, the comparison to the starting point (the red dot in Figure 4) serves as the "do nothing" baseline — leaving the model as-is with no additional training on SlimPajama. This gives perfect upstream performance (no forgetting) and baseline downstream performance (whatever the model achieves by zero-shot transfer from Pile to SlimPajama). All other strategies must improve on at least one of these axes to be useful.
3.4.5 The Same-Data Ablation: Isolating Optimization from Distribution Shift
A critical confound in interpreting rewarming results is that two things change simultaneously when continuing pre-training on a new dataset: (1) the data distribution shifts, and (2) the learning rate is re-increased. Observed performance degradation (the initial spike in loss) could be caused by either factor, or by their interaction. Section 4.4 addresses this through a clean ablation: re-warm the model while continuing to train on the same data (the Pile).
Setup. The model pre-trained on 300B tokens of the Pile is loaded, and then training is resumed on the Pile (not SlimPajama) for 50B additional tokens using the same rewarming approach: linear warmup to various MaxLr values over 1% of these 50B tokens (so, 500M tokens of warmup), followed by cosine decay. The only difference from the main experiment is the training data distribution — everything else (model, optimizer, schedule parameterization) is identical.
What this tests. If the distribution shift between Pile and SlimPajama is the primary cause of the initial loss spike, then rewarming on the same data should produce no spike — the model should smoothly continue improving or remain at its converged loss. If the spike still appears, the cause is intrinsic to the rewarming operation itself (the sudden increase in learning rate disrupts the optimizer state and pushes the model out of its current minimum, regardless of data).
Result preview. The paper obtains the latter result (Figure 5): rewarming on the Pile does produce an initial loss spike, and the model does not recover to its pre-rewarming loss level (Figure 6 shows a linear relationship between upstream and downstream losses, unlike the tradeoff curve in Figure 4). This establishes that rewarming itself, not the distribution shift, is a primary driver of the forgetting-phenomenon observed. The distribution shift matters (the trajectories in Figures 4 and 6 are not identical), but the core optimization disruption from re-increasing the learning rate is present even without domain change.
Why this matters for interpreting all other results. If the loss spike were purely due to distribution shift, one might try to mitigate it through better data mixing or domain alignment. The fact that it's an optimization phenomenon suggests that solutions should focus on the schedule design — perhaps through smoother warmup transitions, momentum reset strategies, or alternative schedule shapes. The paper does not fully pursue these implications but the ablation is essential for correctly attributing causality.
3.4.6 The Checkpoint Selection Experiment: Does Convergence State Matter?
The final experimental variable is the age of the checkpoint from which downstream training begins. The paper's hypothesis is that earlier (less converged) checkpoints might be located at "more favorable points in the loss landscape" (Section 4.5, setup paragraph) that allow easier adaptation to new data. This hypothesis draws on the intuition from the loss landscape literature that early-stopped models reside in wider basins that generalize better and may be more amenable to further training.
Setup. Three checkpoints are selected from the full Pythia 410M pre-training trajectory (which spans 143,000 iterations):
-
Iteration 143,000: the final, fully converged checkpoint (Pythia 410M as released). Pile validation loss is at its minimum.
-
Iteration 27,000: a checkpoint whose Pile validation loss approximately matches the maximum Pile validation loss observed during re-warming experiments in Figure 1 (bottom) — about 2.5 PPL. This is roughly 19% through total pre-training.
-
Iteration 10,000: roughly halfway to the 27,000-iteration checkpoint in terms of training progress (about 7% through total pre-training), intended to provide a finer-grained view of the impact of convergence state.
The paper does not provide exact validation loss values for these checkpoints, but the relative ordering (iter 10,000 > iter 27,000 > iter 143,000 in loss) is clear. Each checkpoint is then loaded and subjected to the same rewarming procedure (linear warmup to $3 \times 10^{-4}$, then cosine decay) on SlimPajama for 50B tokens, with two warmup lengths tested: 0% (instant jump) and 1%.
What this tests. The experiment probes whether the model's plasticity — its capacity to learn from new gradients — degrades as it converges. If earlier checkpoints achieve lower downstream loss (or the same downstream loss with less upstream forgetting), that would suggest that pre-training to full convergence is detrimental for continual learning. This connects to the broader literature on loss of plasticity in continual learning and the "stability-plasticity dilemma": highly converged models may have sharp minima that are hard to escape, while earlier models occupy flatter regions from which new learning is easier.
Why cross with warmup length. The instant (0%) vs. gradual (1%) warmup comparison at each checkpoint age tests whether earlier checkpoints — which may have different optimizer moment statistics — are more or less sensitive to the aggressive schedule change. An earlier checkpoint with differently-scaling second-moment estimates in AdamW might experience a larger or smaller effective step size from the same MaxLr, which could interact with the warmup duration.
Result preview. The paper reports that earlier checkpoints do not improve downstream performance — the final checkpoint (iter 143,000) achieves the best downstream loss (Figure 7), and also the best upstream loss (Figure 8 in the appendix). This is taken as evidence that "pre-training did not lead the model into a loss of plasticity that would make the model difficult to re-warm" (Section 4.5, results paragraph). Whether this conclusion extends to larger models, longer training runs, or different dataset shifts remains an open question that the paper flags for future work.
3.5 Summary of Design Choices and Their Justifications
-
Pythia 410M with matched optimizer settings: inherits all hyperparameters from published Pythia pre-training, ensuring that any observed effects are attributable to the rewarming variables rather than to differences in model architecture or optimizer configuration. The specific size (410M parameters, roughly 1.2B training FLOPs per token) makes full training runs feasible on academic compute while being large enough to exhibit phenomena relevant to larger models.
-
~300B token upstream and downstream datasets of matched scale: eliminates the confound of dataset size — the downstream task is not "fine-tuning on a small dataset" but genuinely "continuing pre-training on an equally large corpus," testing whether transfer helps even when the new data is as abundant as the old.
-
Cosmetic schedule rather than alternative shapes: cosine decay is the standard in LLM pre-training; using it makes the results directly comparable to the broader training literature and ensures that any observed behavior is not an artifact of an unusual schedule shape.
-
Three
MaxLrvalues with factor-of-two spacing: provides sufficient resolution to detect monotonic trends in the adaptation-forgetting tradeoff while keeping the experimental budget manageable. A single value would miss the tradeoff; many values would be redundant. -
Simultaneous upstream/downstream tracking + joint perplexity plots: instead of reporting a single aggregate metric, the paper visualizes the full Pareto frontier, giving practitioners the information needed to choose their preferred operating point. This is essential because real-world deployments differ in their tolerance for forgetting.
-
From-scratch and constant-LR baselines: establish both a floor (training from scratch represents the lowest acceptable performance given that continual pre-training is trying to improve on it) and a simple alternative (constant LR represents the cheapest schedule intervention). Both must be beaten for rewarming to be justified.
-
Same-data ablation (Pile → Pile): the only clean way to separate distribution-shift effects from optimization-dynamics effects. Without this control, one might incorrectly attribute the loss spike to domain mismatch and pursue mitigation strategies (data mixing, domain classifiers) that miss the root cause.
-
Checkpoint age sweep with warmup-length crossing: tests a plausible hypothesis from the plasticity literature (earlier is better) and interacts it with the warmup variable. The negative result (later is better) is as informative as a positive result would have been — it tells practitioners to use their most converged checkpoint, and it tells researchers that Pythia-scale models do not lose plasticity over the course of standard pre-training.
4. Key Insights and Innovations
Innovation 1: Recasting the Learning Rate Warmup as a Strategic Control Knob for the Adaptation-Forgetting Tradeoff, Not a Stabilization Detail
The critical conceptual move in this paper is elevating the learning rate warmup phase from a routine optimizer stabilization detail — the way it is treated in essentially all large language model pre-training work (Brown et al., 2020; Hoffmann et al., 2022; Touvron et al., 2023) — into a first-class strategic lever for managing the tension between downstream adaptation and upstream forgetting. This reframing is what makes the paper's investigation distinctive, not the specific numeric findings about warmup length or peak learning rate values.
In standard pre-training, warmup exists for one reason: to prevent training instability in the first few thousand steps, when the randomly initialized weights produce gradients with deceptively large variance that can cause Adam's adaptive scaling to take destructively large steps. Warmup solves this by starting the learning rate at zero and linearly increasing it to the peak over approximately 1% of training, allowing the second-moment estimates in Adam to accumulate reliable statistics before the learning rate reaches its operational value. Once this stabilization is achieved, warmup's job is done — it is a transient phase you endure, not a phase you design.
This paper argues that when warmup is applied a second time — during continual pre-training — its function fundamentally changes. The model is no longer randomly initialized; it occupies a specific location in the loss landscape with accumulated optimizer statistics (first and second moment estimates) encoding the entire history of upstream training. Re-warming the learning rate now means deliberately destabilizing a converged model to enable new learning. The warmup parameters — peak learning rate, warmup duration, and the checkpoint from which rewarming begins — collectively determine how far from the upstream minimum the model is pushed and how effectively it can navigate the downstream loss landscape from that displaced position.
This reconceptualization has two implications that the paper demonstrates but does not fully articulate. First, the peak learning rate becomes a direct tradeoff dial: the evidence in Figures 2–4 shows that higher MaxLr values (6×10⁻⁴) consistently produce lower final downstream perplexity at the cost of higher upstream forgetting, while lower MaxLr values (1.5×10⁻⁴) preserve upstream performance but limit downstream adaptation. This monotonic relationship means practitioners can choose their desired operating point on the Pareto frontier by adjusting a single scalar — the peak learning rate — without changing architecture, data mixture, or regularization. This is a cleaner control mechanism than alternative continual learning interventions like replay ratio selection or regularization strength tuning, which interact non-monotonically with forgetting and often require per-task calibration.
Second, the warmup itself becomes a mechanism for intentional destabilization. The finding that 0% warmup (instantaneous jump to MaxLr) produces a brief "chaotic phase" with a spike in both upstream and downstream loss, but ultimately convergences to the same final performance as gradual warmup (Figure 1), reveals that the transient instability is not harmful in the long run. This is a negative result with practical value: warmup length can be set to zero or near-zero without long-term penalty, reducing the engineering complexity of schedule design. The dominant assumption in the field — that warmup is necessary for stability — does not hold in the continual pre-training context because the model's optimizer state already provides the statistical stabilization that random initialization lacks.
The significance of this reframing extends beyond the specific experiments. It suggests that as language model deployment cycles accelerate and the number of sequential training stages grows, schedule engineering will become as important as architecture engineering for managing model lifecycle. The paper's conceptual contribution is to establish that the standard pre-training schedule — linear warmup, cosine decay — can be profitably decomposed and its components assigned new strategic functions in the continual setting, a perspective absent from prior work that treated continual pre-training schedules as either simple extensions of pre-training (constant LR) or ad-hoc modifications (progressive decrease).
Innovation 2: Empirically Demonstrating That Positive Transfer Between Large-Scale Datasets Makes Re-Training from Scratch Unnecessary — and That This Holds Even When the Downstream Dataset Is as Large as the Upstream One
The paper's most practically impactful empirical finding is that a model continually pre-trained on SlimPajama outperforms a model trained from scratch on SlimPajama alone, despite the downstream dataset being approximately the same size as the upstream dataset (297B vs. 300B tokens). This result, visible in Figures 2 and 3 where the from-scratch baseline (blue curves) consistently underperforms all rewarmed models after sufficient downstream training, establishes a crucial economic fact: the compute already invested in pre-training on the Pile is not sunk cost — it provides transferable knowledge that continues to pay dividends even when training on a new, comparably-sized corpus.
This finding is more surprising than it might first appear. The intuition from standard transfer learning is that pre-training helps when the downstream dataset is small — the pre-trained model provides a feature extractor or initialization that generalizes beyond what limited downstream data could support. This is the classic fine-tuning paradigm: pre-train on ImageNet (millions of images), fine-tune on a medical imaging task (thousands of images). When the downstream dataset is as large as the upstream dataset, however, the from-scratch model has enough data to learn high-quality representations on its own — the transfer benefit should diminish or vanish as the downstream data volume approaches the upstream volume.
The paper's results contradict this expectation directly. After 200B tokens of downstream training, the from-scratch model on SlimPajama achieves worse SlimPajama validation perplexity (approximately 2.62 in Figure 2) than the MaxLr = 3 × 10^{-4} rewarmed model (approximately 2.55), and this gap emerges despite the from-scratch model having the advantage of being optimized solely for the downstream data from initialization. The paper diagnoses this as positive transfer: knowledge acquired from the Pile — which includes domains like patent text, DM Mathematics, and Freelaw absent from SlimPajama — generalizes to improve modeling of SlimPajama's distribution, even when the SlimPajama data is abundant enough to train a competitive model on its own.
This finding has direct implications for how organizations should think about their pre-training investments. If a 14 million runs on new datasets, and the resulting model outperforms a ground-up $14 million training run on the new dataset alone, then the cumulative value of pre-training increases with each dataset cycle. The Pile checkpoint is not merely a cost to be recovered; it is an asset that appreciates as new corpora become available. This flips the economic argument from "continual pre-training is cheaper because it avoids re-training" to "continual pre-training is better because transfer produces models that cannot be obtained through independent training, even at the same cost."
The paper also reports a subtle but telling secondary observation: the from-scratch model trained on SlimPajama actually improves its Pile validation perplexity over the course of training (Figure 3, blue curve), dropping from around 2.9 to below 2.65. This means that training exclusively on SlimPajama produces a model that generalizes backward to the Pile — further evidence of strong positive transfer between the two datasets. This cross-dataset generalization is bidirectional, which strengthens the case that the two corpora share a common underlying structure that makes continuing pre-training more synergistic than competing approaches would suggest.
It is worth distinguishing this finding from the standard continual learning narrative. In classical continual learning, the goal is typically mitigating forgetting — maintaining performance on old tasks while learning new ones. The paper's result goes beyond mitigation to demonstrate genuine cross-dataset synergy: the model trained on both datasets in sequence outperforms the model trained on either alone. This is a stronger claim than "we avoided catastrophic forgetting" — it is "we built a better model by training sequentially than we could have built by training on a single dataset, even with equivalent total compute."
Innovation 3: The Diagnostic Separation of Distribution Shift from Optimization Disruption as Causes of Forgetting, via the Same-Data Ablation
The paper makes a methodological contribution that sharpens how the continual learning community should think about performance degradation during task transitions. By running the same-data ablation — rewarming the learning rate while continuing to train on the Pile, with no distribution shift at all (Section 4.4) — the paper cleanly isolates whether the forgetting observed in the main experiments is caused by the new data distribution or by the act of rewarming itself.
The result (Figure 5) is unambiguous: rewarming on the same data produces a loss spike qualitatively similar to the spike observed when switching to SlimPajama. The Pile validation loss transiently increases after rewarming, then partially recovers but settles at a higher value than before rewarming. Figure 6 shows a tight linear relationship between upstream Pile perplexity and downstream SlimPajama perplexity during this same-data training — a fundamentally different trajectory shape from the Pareto curves in Figure 4, where the two losses move in opposing directions (downstream improves while upstream degrades). This visual difference is the signature of two distinct mechanisms: optimization disruption (the same-data case, where both losses move together) and distribution-shift-driven adaptation (the cross-dataset case, where the losses trade off).
This is a conceptual advance over the standard continual learning framing, which attributes forgetting almost entirely to distribution shift — the model overwrites old knowledge because the new gradients point in directions that conflict with previously learned functions. The paper's ablation demonstrates that a substantial component of the performance degradation is mechanically induced by the optimizer schedule, independent of what data the model sees. When the learning rate jumps from 3×10⁻⁵ (its decayed minimum) to 3×10⁻⁴ (its rewarmed peak), the model takes large steps that push parameters away from the converged minimum, regardless of whether those steps are driven by gradients from the same distribution or a new one. The loss increases because the optimizer has been deliberately destabilized, not because interference has occurred.
The practical implication of this diagnostic separation is significant for algorithm design. If forgetting were purely caused by distribution shift, the solution would involve better data mixing strategies (replay, interleaving), gradient projection methods, or regularization terms that penalize parameter changes in important directions. The paper's result suggests that these interventions, while potentially helpful, address only part of the problem. The optimization-disruption component of forgetting requires schedule-level interventions: smoother warmup transitions, momentum buffer resets, or learning rate schedules that avoid the initial overshoot entirely (e.g., warm-starting from zero with fresh optimizer state).
The paper does not pursue these schedule-level interventions in depth — that is left as future work — but the diagnostic framework it establishes provides the conceptual vocabulary for doing so. A researcher interested in reducing forgetting now has a clear experimental protocol: run the same-data ablation to measure the optimization-disruption component, then test whether a schedule modification reduces that component without needing to confound the measurement with distribution shift. This is a transferable methodological contribution that applies to any continual learning setting where learning rate re-increase is used, not just LLM pre-training.
Innovation 4: The Negative Result That Earlier (Less-Converged) Checkpoints Do Not Improve Continual Pre-Training, Resolving a Plausible Hypothesis from the Plasticity Literature
The paper includes a negative result that carries substantial weight for practitioners: using an earlier, less-converged pre-training checkpoint does not improve downstream adaptation (Section 4.5, Figures 7 and 8). This finding directly contradicts a plausible hypothesis grounded in the loss landscape and continual learning literatures — that models stopped before full convergence might reside in wider, flatter minima from which new learning is more effective. The hypothesis is intuitive: a sharply converged model might be "brittle," with parameters locked into a narrow basin that gradients from a new distribution cannot easily escape, whereas an early-stopped model retains the plasticity to adapt.
The paper tested this by selecting three checkpoints spanning the full range of pre-training convergence (7%, 19%, and 100% through Pile training) and applying identical rewarming procedures to each. The fully converged checkpoint (iteration 143,000) achieved the best downstream SlimPajama performance and the best upstream Pile retention — the exact opposite of what the plasticity hypothesis would predict. The earlier checkpoints performed worse on both axes.
This negative result is valuable for two reasons. First, it provides actionable guidance: practitioners should use their most trained checkpoint, not an intermediate one, when continuing pre-training. This simplifies operational decisions — there is no need to maintain and evaluate multiple candidate checkpoints to find the optimal restart point. Second, it suggests that loss of plasticity is not a practical concern at the Pythia 410M scale with approximately 300B tokens of upstream pre-training. The model has not entered a regime where continued training at a fixed learning rate would have been impossible, or where the convergence state impedes new learning. This is evidence that current pre-training practices (300B tokens, cosine decay to 10%) do not push models into the plasticity-loss regime that affects other continual learning settings (e.g., reinforcement learning with long training horizons, online class-incremental classification).
That said, this negative result comes with an important scope caveat that the paper acknowledges: the experiment was conducted at a single model scale (410M parameters) and a single upstream data volume (300B tokens). It remains possible that plasticity loss becomes significant at larger scales (billions of parameters) or with longer upstream pre-training (trillions of tokens), where the optimization trajectory is longer and the model may settle into sharper minima. The paper positions this as an open question rather than a settled fact, but the result establishes a baseline: at the scale studied, earlier is not better, and the simplest strategy (use the final checkpoint) is also the optimal one. This prevents the community from prematurely adopting checkpoint selection heuristics that are unnecessary in the current regime and may not generalize.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The upstream (pre-training) dataset is the Pile (Gao et al., 2020), an 800GB collection of 22 diverse text sources, used with the same sampling weights as Black et al. (2022) for validation. The downstream (continual pre-training) dataset is SlimPajama (Soboleva et al., 2023), a cleaned and extensively deduplicated version of RedPajama (Together.xyz, 2023) built based on the LLama dataset (Touvron et al., 2023); the paper subsamples SlimPajama to form a training set of approximately 297B tokens and a validation set of approximately 316M tokens, with composition detailed in Table 1 (67% CommonCrawl, 15% C4, 4.5% each GitHub/Book/Wikipedia, 2.5% ArXiv, 2.0% StackExchange). The authors explicitly note that SlimPajama is built from similar sources as the Pile but with higher data quantity, meaning some upstream data may be repeated during downstream pre-training; no replay buffer is used during downstream training — the model never explicitly sees Pile data again.
-
Base model. All experiments use the Pythia 410M model (Biderman et al., 2023), implementing the GPT-NeoX decoder-only transformer architecture (Black et al., 2022) with 24 layers, hidden dimension 1024, 16 attention heads, learned positional embeddings, and a BPE tokenizer trained specifically on the Pile (vocabulary size 50,304). Flash attention (Dao et al., 2022) is explicitly not used. The authors argue this model size is sufficient to observe phenomena relevant to larger models while keeping full training runs feasible on academic compute; the Pythia suite provides publicly available intermediate checkpoints at multiple training stages (10,000, 27,000, and 143,000 iterations), which is methodologically essential for the checkpoint selection experiments in Section 4.5.
-
Metrics. The primary metric throughout is validation perplexity (PPL), computed as the exponentiated average negative log-likelihood per token: PPL = exp(−(1/N) Σᵢ log p_θ(xᵢ | x_{<i})), where N is the total number of validation tokens, xᵢ is the i-th token, and p_θ(xᵢ | x_{<i}) is the model's predicted probability for that token given its context. Two perplexity values are tracked simultaneously for every experiment: downstream validation loss on the SlimPajama validation set (measuring adaptation to new data) and upstream validation loss on the Pile validation set (measuring forgetting of previously learned data). The paper also introduces a joint perplexity visualization (Figures 4 and 6) that plots downstream PPL on the x-axis against upstream PPL on the y-axis at each training step, revealing the Pareto tradeoff trajectory. The authors acknowledge (Section 5) that perplexity is an imperfect proxy and that future work should validate findings on benchmarks like HELM (Liang et al., 2022) or Harness (Gao et al., 2021), which can provide insight into the evolution of model capabilities beyond token-level prediction quality.
-
Baselines. Three baselines are established: (1) Training from scratch — a randomly initialized Pythia 410M model is trained directly on SlimPajama using the standard warmup-cosine schedule (same MaxLr = 3 × 10⁻⁴, warmup over 1% of 297B tokens, cosine decay to 10% of MaxLr); this tests whether the pre-trained Pile checkpoint provides transfer benefit that justifies the compute already spent on it, and is the primary comparator for the claim that continual pre-training outperforms re-training. (2) Constant learning rate — instead of rewarming and decaying, the downstream training uses a fixed learning rate of 3 × 10⁻⁵ (the rate at which upstream pre-training ended) for the entire 240B tokens; this represents the simplest possible continual training strategy (continue as if pre-training never ended) and tests whether the cosine decay after rewarming is necessary or whether a moderate constant rate suffices. (3) Implicit no-training baseline — the model's performance at the starting checkpoint (red dot in Figure 4), representing zero additional downstream training, which gives perfect upstream retention and baseline downstream transfer performance.
-
Generation budget / compute accounting. The paper does not use "generations" as a compute unit (since there is no sampling or search — this is pure training). Instead, compute is measured in tokens processed during training, a direct proxy for total FLOPs since the model architecture and sequence length are fixed. The upstream pre-training on the Pile consumed approximately 300B tokens (143,000 iterations at a fixed batch size inherited from Pythia's published recipe). The downstream continual pre-training budget is 240B–297B tokens depending on the experiment: Sections 4.1, 4.4, and 4.5 use 50B tokens (short runs to test transient dynamics), while Sections 4.2 and 4.3 use 240B tokens (the full cosine decay period). The paper does not perform a formal FLOPs accounting that would enable direct comparison between the cost of pre-training on Pile + continuing on SlimPajama versus the cost of training from scratch on SlimPajama alone; such an accounting is implicit — the from-scratch model on SlimPajama consumes approximately the same compute as the downstream phase of the continual model, while the continual model has the additional sunk cost of upstream pre-training. The paper's framing (Section 4.3) is that this sunk cost has already been paid and the relevant question is whether continuing from the existing checkpoint produces a better model than starting over.
-
Cross-validation / statistical protocol. The paper does not report any cross-validation, statistical significance testing, or confidence intervals. All perplexity curves are single-run trajectories; the test sets (Pile validation split, 316M-token SlimPajama validation split) are fixed. The authors do not discuss variance across random seeds, data orderings, or subset splits. This is a notable methodological gap — with a single 410M model run per condition, it is impossible to quantify whether the observed differences (e.g., the ~0.1 PPL gap between MaxLr = 3 × 10⁻⁴ and MaxLr = 6 × 10⁻⁴ in Figure 2) are statistically reliable or within the range of run-to-run variation. The paper implicitly relies on the size of the validation sets (hundreds of millions of tokens) to ensure that perplexity estimates are precise, but this does not address variance in the training trajectory itself.
Main Quantitative Results
4.1: Warmup Length Has No Significant Effect on Final Performance
The first experiment (Section 4.1, Figure 1) tests whether the fraction of downstream data allocated to linear warmup — 0%, 0.5%, 1%, or 2% of the 297B-token dataset — affects the model's ability to adapt to SlimPajama or retain Pile performance. All conditions use MaxLr = 3 × 10⁻⁴, the same peak rate as original Pythia pre-training, and are evaluated over the first 50B tokens of downstream training.
Headline result: The length of the warmup phase does not significantly influence perplexity on either the downstream SlimPajama validation set or the upstream Pile validation set. At 50B tokens, all four warmup lengths produce essentially identical SlimPajama validation loss (~2.70 PPL, reading from Figure 1 top) and Pile validation loss (~2.30 PPL, reading from Figure 1 bottom). The trajectories throughout training are nearly indistinguishable after the first few billion tokens.
Transient behavior—the stability gap: The 0% warmup condition (instantaneous jump to MaxLr = 3 × 10⁻⁴ at step 1) exhibits a brief but dramatic spike in both downstream and upstream loss during the first ~1B tokens of training. At its peak, the SlimPajama validation loss reaches approximately 2.875 PPL versus ~2.70 for the gradual warmup conditions, and the Pile validation loss spikes to approximately 2.65 versus ~2.25. This spike is transient — by roughly 5B tokens into training, the 0% warmup curve has converged to the same trajectory as the 1% and 2% warmup curves. The paper names this phenomenon a "chaotic phase" and connects it to the stability gap concept from Lange et al. (2023) and Caccia et al. (2022): re-increasing the learning rate abruptly displaces the model from its converged minimum, causing a temporary loss degradation before the optimizer settles into the new loss landscape.
Implication for practice: Since the length of warmup does not affect final performance, practitioners can set warmup to 0% or near-0% without long-term penalty. This is a counterintuitive finding: the standard rationale for warmup in pre-training (stabilizing early gradients when the optimizer state is uninitialized) does not apply in continual pre-training because the optimizer state (Adam's first and second moment estimates) is already well-populated from upstream training. The model can tolerate an instantaneous jump to the peak learning rate without diverging. The only cost is the transient stability gap, which has no lasting consequences on the final converged loss.
4.2: Maximum Learning Rate Controls the Adaptation-Forgetting Tradeoff
The core experiment (Section 4.2, Figures 2–4) varies the maximum learning rate reached after rewarming — 1.5 × 10⁻⁴, 3 × 10⁻⁴, or 6 × 10⁻⁴ — with warmup fixed at 1% of the 297B-token dataset and all conditions using cosine decay to 10% of their respective MaxLr (reaching the minimum at 240B tokens and held constant thereafter). The two baselines (from-scratch training on SlimPajama, constant learning rate at 3 × 10⁻⁵) are plotted alongside for comparison.
Downstream performance (Figure 2): At the end of training (~240B tokens), there is a clear monotonic relationship between MaxLr and downstream SlimPajama validation perplexity. The MaxLr = 6 × 10⁻⁴ condition achieves the best downstream performance at approximately 2.52 PPL, followed by MaxLr = 3 × 10⁻⁴ at approximately 2.55, and MaxLr = 1.5 × 10⁻⁴ at approximately 2.60. All three rewarmed models outperform the from-scratch baseline, which finishes at approximately 2.62 PPL despite being optimized solely for SlimPajama from initialization. The constant-LR baseline initially achieves the best downstream performance (lowest loss within the first ~100B tokens, reaching approximately 2.58 PPL at the 50B-token mark where the rewarmed models are still at ~2.62–2.67) but is eventually overtaken by all three rewarmed-and-decayed models as the decay phase allows finer convergence.
Upstream performance (Figure 3): The monotonic relationship inverts for Pile validation perplexity. MaxLr = 6 × 10⁻⁴ causes the most forgetting, finishing at approximately 2.90 Pile PPL. MaxLr = 3 × 10⁻⁴ finishes at approximately 2.65. MaxLr = 1.5 × 10⁻⁴ preserves the most upstream performance at approximately 2.58. All values represent substantial degradation from the pre-continual-training starting point (approximately 2.25 PPL). The constant-LR model shows the best upstream retention early in training (maintaining approximately 2.35 PPL through the first 50B tokens), but this advantage erodes over time — by 240B tokens, it has risen to approximately 2.72, performing worse than MaxLr = 3 × 10⁻⁴ on upstream retention. The from-scratch model, surprisingly, improves its Pile validation perplexity throughout SlimPajama training, dropping from approximately 2.90 to approximately 2.65, demonstrating that SlimPajama-trained models generalize backward to the Pile — strong evidence of positive cross-dataset transfer.
The Pareto tradeoff (Figure 4): The joint perplexity plot reveals the full trajectory of each condition in the upstream-downstream loss space. Key observations:
- All rewarmed models start from the same point (red dot, approximately (14.0 SlimPajama PPL, 9.0 Pile PPL in the plot's units). During warmup, they move up and right (both losses increase simultaneously), reaching their green points (end of warmup) at varying distances from the origin. Higher MaxLr pushes the model further from its starting point.
- After warmup, as the learning rate decays, all models curve down and left — downstream loss improves while upstream loss partially recovers, but neither returns to the starting point. The final resting points trace out a clear tradeoff frontier: MaxLr = 6 × 10⁻⁴ (furthest right, lowest downstream loss, highest upstream loss) to MaxLr = 1.5 × 10⁻⁴ (furthest up, best upstream retention, but worst downstream adaptation).
- The constant-LR model (black curve) follows a different trajectory: it initially stays close to the starting point on the upstream axis (good retention) while improving downstream, but then gradually drifts upward (more forgetting) without the guidance of decay to stabilize it.
- The from-scratch model (blue curve) starts from a completely different region of the space (high downstream loss, high upstream loss) and moves consistently downward and left as it learns both distributions simultaneously.
Convergence behavior detail: The paper notes that reaching the minimum learning rate at 240B tokens and holding constant thereafter produces a visible change in the loss curves. For the MaxLr conditions, the loss continues to improve slightly during this constant-LR plateau. The MaxLr = 6 × 10⁻⁴ condition benefits from a final learning rate of 6 × 10⁻⁵ (double the final rate of the MaxLr = 3 × 10⁻⁴ condition at 3 × 10⁻⁵), which may contribute to its better downstream convergence — the larger final learning rate may allow continued exploration of the loss landscape rather than premature convergence. This is not discussed explicitly by the paper but is a logical consequence of the schedule design: higher MaxLr produces higher final learning rate under the 10% decay ratio, meaning the "cooldown" phase differs across conditions in ways beyond the warmup peak.
4.3: Continual Pre-Training Outperforms Training from Scratch
Section 4.3 directly compares the rewarmed models against the from-scratch baseline on downstream performance (the from-scratch model is the blue curve in Figures 2 and 3, described above). The headline finding, visible in Figure 2, is that all three rewarmed models achieve lower SlimPajama validation perplexity than the from-scratch model by the end of training (~240B tokens). The gap is approximately 0.1 PPL between MaxLr = 6 × 10⁻⁴ (~2.52) and from-scratch (~2.62), with MaxLr = 3 × 10⁻⁴ (~2.55) and MaxLr = 1.5 × 10⁻⁴ (~2.60) also outperforming.
The paper emphasizes that this result holds despite the downstream dataset being approximately the same scale as the upstream dataset (297B vs. 300B tokens) and despite the from-scratch model having the advantage of being optimized solely for SlimPajama. This demonstrates positive transfer: the Pile provides knowledge that improves SlimPajama modeling even after 240B tokens of dedicated SlimPajama training.
Important qualification on the from-scratch vs. constant-LR comparison (Figure 2): After 200B tokens, the from-scratch model does overtake the constant-LR model on downstream perplexity. This means that while constant-LR training is competitive in the early stages (first ~100B tokens), it is eventually outperformed by both the rewarmed-and-decayed models and the from-scratch model. The paper interprets this as evidence that the decay phase is essential for convergence quality: the constant-LR model cannot settle into a deep minimum because the persistent learning rate prevents fine-grained optimization.
Cross-dataset generalization observation (Figure 3): The from-scratch model's improvement on Pile validation perplexity over the course of SlimPajama training — dropping from ~2.90 to ~2.65 — is presented as further evidence of positive transfer. A model trained exclusively on SlimPajama nevertheless learns representations that generalize backward to the Pile, indicating that the two datasets share a substantial common structure. This is not a result the experimental design set out to test, but the dual-validation tracking reveals it as an emergent property.
4.4: The Same-Data Ablation -- Rewarming Itself Causes Performance Degradation
Section 4.4 addresses a fundamental causal question: is the loss spike observed when switching to SlimPajama caused by the distribution shift (Pile → SlimPajama) or by the act of re-increasing the learning rate itself? The experiment re-warms the model (from the final Pile checkpoint, iteration 143,000) while continuing to train on the same Pile data for 50B tokens, using the same schedule parameterization as Section 4.2 (MaxLr values spanning 1.5 × 10⁻⁴ to 6 × 10⁻⁴, 1% warmup, cosine decay to 10%).
Headline result (Figure 5): Rewarming on the same data produces a loss spike qualitatively similar to the cross-dataset case. The Pile validation loss transiently increases after rewarming — at MaxLr = 6 × 10⁻⁴, it spikes from approximately 2.20 PPL to approximately 2.40 PPL during warmup — and then partially recovers as the learning rate decays, but does not return to its pre-rewarming level. All MaxLr conditions settle at higher final Pile PPL than before rewarming: MaxLr = 1.5 × 10⁻⁴ finishes at approximately 2.26 (closer to pre-rewarming 2.20), MaxLr = 3 × 10⁻⁴ at approximately 2.28, and MaxLr = 6 × 10⁻⁴ at approximately 2.30. The constant-LR condition (3 × 10⁻⁵) maintains stability with only a slight upward drift.
Diagnostic pattern — the joint plot difference (Figure 6 vs. Figure 4): The joint perplexity plot for the same-data ablation (Figure 6, tracking Pile PPL vs. SlimPajama PPL while training on Pile) reveals a linear relationship between the two losses across all conditions. All trajectories follow a common line: as Pile validation loss increases (due to rewarming), SlimPajama validation loss increases proportionally, and as Pile loss recovers (during decay), SlimPajama loss recovers along the same line. This is fundamentally different from the cross-dataset case (Figure 4), where the trajectories curve — downstream loss improves while upstream loss degrades, creating a convex tradeoff frontier. The paper interprets the linear pattern in Figure 6 as evidence that the model is "climbing out of a minimum during warmup and returning toward the same minimum as the learning rate is decayed," whereas in the cross-dataset case, the model settles at a different minimum that represents a compromise between the two data distributions.
Causal attribution: This ablation cleanly demonstrates that a substantial component of forgetting is caused by the rewarming operation itself, not by distribution shift. The learning rate increase mechanically pushes parameters away from the converged upstream minimum regardless of the data distribution being trained on. The distribution shift (present in Figures 2–4, absent in Figures 5–6) determines whether the model can recover downstream performance without fully recovering upstream performance (the tradeoff), or whether both losses move in lockstep (the same-data case). The paper does not attempt to quantify what fraction of total forgetting is attributable to optimization disruption versus distribution shift — the qualitative difference in trajectory shape is the main evidence.
4.5: Earlier Checkpoints Do Not Improve Downstream Adaptation
Section 4.5 tests the hypothesis that less-converged pre-training checkpoints might reside in loss landscape basins that are more amenable to further training (the "plasticity" hypothesis). Three checkpoints are selected from Pile pre-training: iteration 143,000 (fully converged, Pile PPL ~2.20–2.25), iteration 27,000 (chosen because its Pile PPL approximately matches the maximum Pile loss observed during rewarming experiments — ~2.50), and iteration 10,000 (roughly halfway to 27,000 in training progress). Each checkpoint is subjected to rewarming on SlimPajama with MaxLr = 3 × 10⁻⁴ for 50B tokens, with two warmup conditions tested: 0% (instant jump) and 1% (gradual).
Downstream result (Figure 7): The final checkpoint (iteration 143,000) achieves the best downstream SlimPajama validation loss, reaching approximately 2.70 PPL by 50B tokens for both warmup lengths. The iteration 27,000 checkpoint finishes at approximately 2.78 (0% warmup) and 2.80 (1% warmup). The iteration 10,000 checkpoint finishes at approximately 2.80 (0% warmup) and 2.82 (1% warmup). The ordering is monotonic: the more converged the checkpoint, the better the downstream performance, which is the opposite of what the plasticity hypothesis predicts. Warmup length has minimal effect within each checkpoint condition, consistent with Section 4.1 findings; the 0% warmup conditions show an initial spike (particularly visible for the iteration 27,000 and 10,000 checkpoints) but converge to similar final values.
Upstream result (Appendix A, Figure 8): The same monotonic ordering holds for Pile validation loss: iteration 143,000 achieves the best upstream retention (lowest Pile PPL at ~2.30 by 50B tokens), followed by iteration 27,000 (~2.40) and iteration 10,000 (~2.45). The 0% warmup conditions show the expected transient spike in upstream loss — reaching approximately 2.55–2.65 for the earlier checkpoints — before partially recovering. The fact that the earlier checkpoints both start and finish with worse upstream performance reinforces the conclusion that they are simply worse models, not more plastic ones.
Additional visualization (Appendix A, Figure 9): The paper provides a side-by-side comparison of Pile and SlimPajama validation losses for models trained from the various checkpoints versus a from-scratch baseline. The from-scratch model, as expected, starts with much higher loss on both datasets and converges more slowly. The iteration 143,000 checkpoint achieves lower Pile and SlimPajama validation loss faster than any other starting point.
Diagnosis: The paper concludes that "pre-training did not lead the model into a loss of plasticity that would make the model difficult to re-warm" (Section 4.5, results). The finding suggests that 300B tokens of upstream pre-training with a cosine-decayed schedule does not push the Pythia 410M model into a regime where continued training becomes hindered by sharp minima or exhausted learning capacity. This is a practically important negative result: practitioners should use their latest, most converged checkpoint, not hunt for an intermediate "optimal restart point."
Ablation Studies and Robustness Checks
Warmup length variation (Section 4.1, Figure 1): Varying warmup from 0% to 2% of the downstream dataset shows no significant effect on final perplexity for either upstream or downstream data. The only observable difference is a transient stability gap for the 0% warmup condition (loss spike in the first ~1B tokens), which fully resolves by 5B tokens. This ablation establishes that gradual warmup — considered essential during initial pre-training — is unnecessary during continual pre-training because the optimizer state is already well-initialized. The result is tested only at MaxLr = 3 × 10⁻⁴; whether it generalizes to the more aggressive MaxLr = 6 × 10⁻⁴ condition (where warmup might provide more stabilization benefits) is not explicitly tested, though the paper does test 0% and 1% warmup at MaxLr = 3 × 10⁻⁴ across different checkpoints in Section 4.5 and finds consistent behavior.
Maximum learning rate sweep (Section 4.2, Figures 2–4): Varying MaxLr across a factor-of-4 range (1.5 × 10⁻⁴ to 6 × 10⁻⁴) reveals a monotonic tradeoff: higher rates improve downstream adaptation at the cost of upstream forgetting. The relationship is smooth and continuous across the three tested values, suggesting that the Pareto frontier is well-behaved and does not exhibit threshold effects or regime changes within this range. Whether this monotonicity holds at more extreme values (e.g., 1 × 10⁻³, which would be ~3.3× the default peak) is not tested. The factor-of-two spacing provides sufficient resolution to demonstrate the trend but is too coarse to determine whether a specific intermediate value (e.g., 4.5 × 10⁻⁴) would offer a Pareto-optimal compromise that dominates both 3 × 10⁻⁴ and 6 × 10⁻⁴ — the convexity of the frontier in Figure 4 suggests such a point likely exists, but the paper does not attempt to locate it.
Checkpoint age (Section 4.5, Figures 7–8): Three checkpoints spanning 7% to 100% of Pile pre-training progress are tested under identical rewarming conditions (MaxLr = 3 × 10⁻⁴, both 0% and 1% warmup). The finding that earlier checkpoints perform worse on both upstream and downstream metrics contradicts the hypothesis that less-converged models retain more plasticity. This ablation is important because it rules out a plausible alternative strategy (restarting from an intermediate checkpoint) that practitioners might otherwise adopt based on intuition from the generalization and loss landscape literatures. However, the ablation has limited resolution: only three checkpoints are tested, and the spacing is uneven (iter 10,000, 27,000, 143,000). A finer-grained sweep — testing checkpoints at 50%, 75%, and 90% of training — would provide stronger evidence about whether the monotonic ordering holds throughout or whether there is a sweet spot near convergence. Additionally, the experiment is conducted only at MaxLr = 3 × 10⁻⁴; it is possible that earlier checkpoints would benefit from a different peak learning rate (e.g., a slower rewarming to avoid destabilizing their less-converged optimizer state), but this interaction is not explored.
Optimizer state inheritance (implicit ablation through the same-data experiment, Section 4.4, Figures 5–6): While not framed as an ablation, the same-data experiment effectively tests what happens when the optimizer state (Adam moments) accumulated over 300B tokens of Pile training is subjected to a rewarming schedule while training on the same data distribution. The fact that the model cannot recover its pre-rewarming loss level (Figure 5: all MaxLr conditions settle at higher Pile PPL than before rewarming) suggests that the optimizer state perturbation from rewarming causes permanent displacement from the upstream minimum, even when the data distribution does not change. This implies that resetting or partially resetting the optimizer state (e.g., zeroing the second-moment estimates while keeping the first moments) might mitigate the forgetting effect, but this ablation is not conducted. The paper also does not test whether a shorter rewarming period (e.g., 1B tokens instead of 50B) would allow recovery, or whether the non-recovery is caused by the number of steps taken at the elevated learning rate rather than the elevation itself.
Constant-LR baseline as an ablation of the decay phase (Section 4.2, Figures 2–4): The constant-LR condition (3 × 10⁻⁵ throughout) effectively ablates the cosine decay component of the schedule. The result — constant-LR performs well early but is overtaken by decayed models after 100B tokens — demonstrates that the decay phase is not merely a cooldown but an essential mechanism for achieving deeper convergence. This is a non-obvious finding: one might expect that a moderate constant learning rate would eventually reach the same minimum as a decayed schedule, just more slowly. The fact that it does not (the constant-LR model plateaus at a worse final loss) suggests that the shape of the schedule trajectory — specifically, the annealing from high to low learning rate — navigates the loss landscape in a way that constant-rate training cannot replicate. The paper does not explore whether restarting the decay from the constant-LR model's plateau point would close the gap, which would be a strong test of whether the benefit is from the exploration-exploitation sequence or from the absolute learning rate values visited.
Data mixing ratio (not ablated): The paper uses pure SlimPajama for downstream training with no replay of Pile data. An ablation that interleaves Pile and SlimPajama data during downstream training — at various mixing ratios — would test whether schedule engineering alone can match or exceed the forgetting mitigation provided by data mixing. The paper explicitly notes that replay is not used to maintain the focus on schedule-based interventions, but this means the relative contribution of data strategy vs. schedule strategy to the adaptation-forgetting tradeoff is never measured. A condition that mixes 10% Pile data into SlimPajama training at MaxLr = 6 × 10⁻⁴ would reveal whether the aggressive adaptation benefit of high MaxLr can be retained while reducing forgetting through data exposure rather than learning rate reduction.
Model scale and architecture (not ablated): All experiments use Pythia 410M with the GPT-NeoX architecture. There is no ablation testing whether the observed phenomena — warmup length insensitivity, MaxLr tradeoff monotonicity, checkpoint age ordering — hold at different model scales (e.g., Pythia 70M, 1.4B, 2.8B) or with different architectures (e.g., LLaMA-style with rotary embeddings and SwiGLU activations). The paper acknowledges this as a limitation (Section 5) and frames the current work as a "preliminary study." Whether the stability gap magnitude, the optimal MaxLr, or the recovery rate depend on model scale is entirely open. Larger models, which have more redundant capacity and potentially flatter minima, might exhibit different forgetting dynamics — or might forget less because their parameters are less specialized to the upstream distribution.
Schedule shape (not ablated): Only the linear-warmup + cosine-decay schedule is tested. Alternative schedule shapes — linear decay, inverse square root decay, exponential decay, constant-then-drop — are not explored. The paper inherits cosine decay from standard pre-training practice without testing whether it is optimal for the continual setting. Given that the function of the schedule is different in continual pre-training (managing a tradeoff rather than purely enabling convergence), other schedule families might offer better control over the adaptation-forgetting frontier. For example, a schedule with a shorter high-learning-rate phase and an earlier transition to decay might achieve the same downstream adaptation with less upstream displacement.
Critical Assessment
The paper makes four primary empirical claims, each with a designated "Takeaway," and a broader methodological contribution. Below, I assess each against the experimental evidence.
Takeaway 1: "The length of the warmup phase does not appear to have a significant effect on the Pile and SlimPajama validation losses."
The evidence from Figure 1 supports this claim for the specific conditions tested — Pythia 410M, MaxLr = 3 × 10⁻⁴, downstream training on SlimPajama for 50B tokens. The four warmup conditions (0%, 0.5%, 1%, 2%) produce visually indistinguishable curves after ~5B tokens. However, the claim is asserted more broadly than the evidence warrants. The experiment tests a single MaxLr value (3 × 10⁻⁴). At MaxLr = 6 × 10⁻⁴, the instantaneous jump to a much higher learning rate might produce a larger and more persistent stability gap, or might even cause training instability that gradual warmup prevents. The paper does test 0% and 1% warmup across different checkpoints at MaxLr = 3 × 10⁻⁴ in Section 4.5, and the pattern holds (0% and 1% curves converge), but still only at the default MaxLr. The generalization to "warmup length doesn't matter, period" is premature — the evidence supports the narrower claim that "warmup length doesn't matter at MaxLr = 3 × 10⁻⁴," which is the rate at which the model was originally pre-trained and for which the optimizer state is well-calibrated.
A stronger experiment would sweep warmup length at MaxLr = 6 × 10⁻⁴. If the stability gap at 0% warmup with the higher peak rate is larger but still transient, confidence in the claim's generality increases. If training diverges or the gap never closes, the claim requires qualification.
Takeaway 2: "Rewarming then decaying the learning rate appears necessary to learn well on the downstream task. Moreover, while keeping a constant learning rate is initially advantageous on Pile, this advantage vanishes when training long enough on SlimPajama."
The first sentence is supported by Figure 2: the rewarmed-and-decayed models outperform the constant-LR model on downstream SlimPajama perplexity at convergence. The constant-LR model plateaus at a worse final loss. However, the word "necessary" is too strong. The paper does not test alternative strategies that could achieve strong downstream performance without rewarming-and-decaying — for example, using a constant learning rate higher than 3 × 10⁻⁵ (the rate tested) or using a schedule that decays to zero rather than to a constant plateau. The constant-LR condition in the paper is only one specific rate (the final rate from upstream pre-training), and the fact that one constant rate underperforms does not prove that no constant rate can work.
The second sentence — about the advantage on Pile vanishing — is supported by Figure 3: the constant-LR model shows early upstream retention superiority (it maintains ~2.35 PPL at 50B tokens vs. ~2.45–2.60 for the rewarmed models), but by 240B tokens it has risen to ~2.72, worse than the MaxLr = 3 × 10⁻⁴ condition at ~2.65. The "advantage vanishes" claim is therefore accurate but conditional on training duration — at 50B tokens the constant-LR model is the better choice for a practitioner who cares primarily about minimizing Pile forgetting, while at 240B tokens it is not. The paper's framing of the 240B result as definitive understates the practical relevance of the early-stage advantage: if a deployment requires only modest downstream adaptation, stopping at 50B tokens with constant LR might be strictly preferable to rewarming at all.
An additional, unstated comparison: "A model that only learns on SlimPajama performs worse on SlimPajama than models pre-trained on Pile in spite of being optimized solely for the downstream task, highlighting positive transfer between the two datasets." This is supported by Figure 2: the from-scratch curve is above all three rewarmed curves at 240B tokens. However, the from-scratch model has not been trained for an equivalent total number of tokens — it has seen ~240B SlimPajama tokens, while the rewarmed models have seen ~300B Pile tokens + ~240B SlimPajama tokens. The comparison is not FLOPs-matched. The paper's implicit argument is that the Pile pre-training is a sunk cost and the relevant question is "given that I already have a Pile checkpoint, should I continue or restart?" but from an ab-initio planning perspective, the question "should I train on Pile then SlimPajama, or just train on SlimPajama for the same total compute?" is not answered by these experiments. A FLOPs-matched comparison — where the from-scratch model is trained for 540B tokens on SlimPajama (matching the combined Pile + SlimPajama budget) — would test whether the transfer benefit is truly about knowledge reuse or simply about total tokens seen.
Takeaway 3 (from Section 4.4): "Rewarming the learning rate appears to be a significant cause for the degradation of performance seen previously when starting to learn on the downstream task, as evidenced by rewarming then decaying the learning rate while training on the same dataset. The models do not appear to be able to recover from the performance hit due to rewarming the learning rate when training on the same dataset."
The first sentence is strongly supported by Figure 5: the same-data ablation produces a loss spike qualitatively matching the cross-dataset experiments. The attribution of forgetting to rewarming itself, at least in part, is rigorous because the distribution shift confound is eliminated.
The second sentence — about lack of recovery on the same dataset — is supported by the fact that all MaxLr conditions in Figure 5 settle at higher Pile PPL than before rewarming. This is a genuinely surprising result: one might expect that training on the same data at a higher learning rate temporarily displaces the model but that the decay phase would guide it back to the same minimum. The fact that it does not return suggests that the optimizer trajectory at the re-warmed learning rate visits regions of parameter space from which the decay path leads to a different (and worse) local minimum, even though the data distribution is unchanged. This is a finding about the non-reversibility of the optimizer trajectory — the path taken at high learning rate is not retraced during decay, and the destination depends on the path.
However, the ablation has a structural weakness: the model is rewarmed and then trained on the Pile for 50B tokens. The upstream pre-training ran for 300B tokens. It is possible that the model would recover to its original loss if trained for longer (e.g., 300B additional tokens on Pile), and the non-recovery observed at 50B tokens is an artifact of insufficient downstream budget. The paper does not test longer same-data runs, so the permanence of the displacement is not established.
Takeaway 4 (from Section 4.5): "Using an earlier checkpoint when pretraining on the Pile does not lead to learning faster on SlimPajama."
This is supported by Figures 7 and 8: the iteration 143,000 checkpoint achieves lower downstream loss and lower upstream loss than either iteration 27,000 or iteration 10,000. The ordering is monotonic and consistent across both warmup lengths tested. The claim is appropriately narrow — it says "does not lead to learning faster," not "leads to worse performance" (though the latter is also true).
The primary limitation is the small number of checkpoints tested (three) and the lack of MaxLr variation in this experiment. If earlier checkpoints have less well-developed optimizer statistics, the optimal MaxLr for rewarming them might be different from the default 3 × 10⁻⁴ used here — perhaps they need a lower peak rate to avoid destabilization, or perhaps a higher peak rate to compensate for their weaker starting position. The experiment only tests MaxLr = 3 × 10⁻⁴, so the conclusion that earlier checkpoints are worse is conditioned on using the same rewarming peak rate as the final checkpoint. This is not an unreasonable experimental choice — it tests the simplest strategy (identical rewarming) — but it leaves open the possibility that a checkpoint-specific MaxLr could reverse the ordering.
The broader methodological claim: that rewarming with appropriate MaxLr selection provides a controllable tradeoff between adaptation and forgetting. The evidence in Figures 2–4 shows a monotonic relationship between MaxLr and the upstream-downstream frontier, and the joint plot (Figure 4) visualizes the tradeoff clearly. This claim is well-supported for the specific model, dataset pair, and schedule family tested. The paper has demonstrated that the mechanism exists. What remains unestablished is the generality — whether the tradeoff curve shape, the monotonicity, and the optimal operating points transfer to different model scales, dataset pairs with varying degrees of overlap, or alternative schedule families.
Specific experimental weaknesses that limit confidence:
-
Single training run per condition. All curves in Figures 1–9 are single-run trajectories. Without error bars or replicate runs, it is impossible to distinguish systematic effects from run-to-run variance. A difference of 0.03 PPL between MaxLr = 3 × 10⁻⁴ and MaxLr = 6 × 10⁻⁴ at 240B tokens (Figure 2) could be within the noise floor of training stochasticity. The paper implicitly relies on the size of the validation sets (hundreds of millions of tokens) to produce low-variance perplexity estimates, but this addresses measurement noise, not trajectory noise — different random seeds would produce different training dynamics, and may converge to different final losses. The consistency of the monotonic ordering across experiments (MaxLr = 6 × 10⁻⁴ always outperforms MaxLr = 3 × 10⁻⁴ downstream) provides some informal evidence that the effects are real, but formal statistical support is absent.
-
No FLOPs-matched comparison between continual pre-training and from-scratch training. The from-scratch model is trained for 240B tokens on SlimPajama. The continual models are trained for 300B tokens on Pile + 240B tokens on SlimPajama. The from-scratch model sees less than half the total data. A fair comparison — training the from-scratch model for 540B tokens on SlimPajama (or a Pile+SlimPajama mixture) — would test whether the continual pre-training advantage is about total compute rather than transfer. The paper's framing (sunk cost of Pile pre-training is already paid) is practically reasonable — if you already have the checkpoint, you should continue rather than restart — but it overstates the positive transfer claim. It is possible that any model trained for 540B tokens on a sufficiently large corpus would outperform one trained for 240B, regardless of transfer. This experiment should be run to establish whether the transfer benefit is genuine or a data-volume artifact.
-
Only a single downstream dataset (SlimPajama) is tested. The paper motivates its work by the proliferation of new pre-training datasets, but only examines one downstream dataset, which is deliberately chosen for high similarity to the upstream data (the Pile). The paper acknowledges this in Section 5: results may not generalize to setups with larger distribution shifts, such as domain adaptation from general text to specialized domains (biomedical, legal, code). The absence of even a single dissimilar downstream corpus (e.g., a pure code dataset, a non-English language corpus) means the paper has characterized the easiest case for continual pre-training without establishing how performance degrades as the difficulty increases. The finding that rewarming causes performance degradation even on the same data (Section 4.4) predicts that more dramatic distribution shifts will amplify the forgetting problem, but the interaction between shift magnitude and optimal MaxLr is entirely unexplored.
-
No downstream benchmark evaluation beyond perplexity. Perplexity is a smooth, continuous signal that correlates with downstream task performance, but it is an imperfect proxy. A model could improve its perplexity on SlimPajama by learning superficial statistical patterns without acquiring the knowledge or reasoning capabilities that matter for applications. The paper acknowledges that evaluating on HELM or Harness is important future work, but the absence of such evaluation limits the practical interpretability of the results. A practitioner adopting the MaxLr = 6 × 10⁻⁴ recommendation because it produces the best SlimPajama perplexity might find that the model's factual accuracy, reasoning ability, or generation quality have not improved — or have degraded — relative to the MaxLr = 1.5 × 10⁻⁴ model with better upstream retention.
-
No multi-stage continual pre-training experiments. The paper studies a single upstream-to-downstream transition (Pile → SlimPajama). The introduction motivates continual pre-training as a response to the regular release of new datasets, which implies multiple sequential stages (Dataset A → Dataset B → Dataset C → ...). The paper does not test whether rewarming can be applied repeatedly, whether forgetting compounds across stages, or whether the optimal MaxLr for stage N depends on the schedule used in stages 1 through N-1. The conceptual framework (rewarming as a tradeoff dial) is compatible with multi-stage training, but the experimental evidence is limited to a single transition. This is a significant gap given that the paper's own argument for rewarming over progressive decrease is that progressive decrease "would cause [the learning rate] to eventually become too small if the number of training stages becomes high" — but no experiment tests whether rewarming on a 5-stage or 10-stage sequence does in fact outperform progressive decrease in that regime.
6. Limitations and Trade-offs
Single Model Scale and Architecture
The assumption or constraint. All experiments are conducted exclusively on Pythia 410M, a single model from a single architecture family (GPT-NeoX) at a single parameter count. The paper is explicit about this scope in Section 5:
"our investigation explores models of size 410M and fine-tuning dataset of size 297B tokens. While this is a preliminary study, in future work, we plan to verify whether our conclusions hold at different model scales (e.g., 3B and 7B) and different dataset scales (e.g., 100B and 600B)."
The paper frames the current work as a "preliminary study," which is an honest demarcation of scope, but it means the headline findings — warmup length does not matter, MaxLr creates a monotonic tradeoff, earlier checkpoints are worse — are conditioned on a specific regime of model capacity.
The consequence. There are at least three reasons to expect scale-dependence in these results, none of which the paper can rule out. First, the stability gap magnitude — the loss spike when rewarming — may scale with model size. Larger models typically operate at lower perplexities with sharper minima, so the same learning rate jump might cause proportionally larger displacement from the converged solution. A transient spike that resolves in <1B tokens for a 410M model might persist for tens of billions of tokens at 7B parameters, turning an ignorable artifact into a meaningful training cost. Second, loss of plasticity — which the paper concludes is absent at 410M scale ("pre-training did not lead the model into a loss of plasticity") — is a phenomenon that typically emerges with increasing model depth, training duration, and parameter count. The finding that the final checkpoint is best for continual pre-training at 410M parameters does not guarantee the same holds at 7B, 13B, or 70B. Third, the optimal MaxLr for the tradeoff may depend on model scale. Larger models have more capacity to absorb new data without overwriting old knowledge (less catastrophic forgetting due to greater representational redundancy), which could shift the Pareto frontier — the same MaxLr = 6 × 10⁻⁴ might cause less upstream forgetting in a 7B model than in a 410M model, changing the practitioner's calculus about which peak rate to choose.
What evidence exists in the paper. None — this is entirely unmeasured. There are no experiments at any scale other than 410M, and no theoretical argument for why the phenomena should be scale-invariant. The paper does not reference prior work that studies scale-dependence of learning rate dynamics in continual settings.
Mitigation status. The paper does not attempt to mitigate this limitation. It flags scale exploration as future work in Section 5 and treats the current results as establishing baselines. A practitioner deploying continual pre-training at production scale (billions of parameters) would need to replicate these experiments at their target scale before trusting the recommendations — the 410M results provide hypotheses to test (warmup length is irrelevant, MaxLr controls the tradeoff) but not verified predictions.
Perplexity as the Sole Evaluation Metric
The assumption or constraint. All conclusions are based entirely on validation perplexity — for both upstream (Pile) and downstream (SlimPajama) performance. The paper acknowledges this gap directly in Section 5:
"we plan to test our models throughout using benchmarks such as HELM or Harness instead of only loss or perplexity, as these benchmarks can provide important insight into the evolution of model capabilities."
Perplexity measures how well the model predicts held-out tokens from the same distribution. It is a smooth, continuous signal that correlates with broad language modeling quality, but it is not a direct measure of the capabilities that matter in deployment: factual accuracy, reasoning ability, generation coherence, resistance to hallucination, or performance on specific downstream tasks.
The consequence. The central practical claim of the paper — that rewarming with higher MaxLr "improves downstream performance" — is, strictly speaking, a claim about token-level prediction quality on SlimPajama validation data. A practitioner who chooses MaxLr = 6 × 10⁻⁴ because it minimizes SlimPajama perplexity might discover that the resulting model has degraded factual recall, makes more reasoning errors, or produces less coherent long-form text than a model trained with MaxLr = 1.5 × 10⁻⁴ (which preserves more upstream performance). This is not merely a theoretical possibility: catastrophic forgetting in the perplexity sense (increased Pile PPL from 2.25 to 2.90 at MaxLr = 6 × 10⁻⁴) implies that the model is systematically worse at predicting tokens from the Pile distribution, which includes domains like patent text, legal documents, and academic papers. If those domains contain factual knowledge that downstream tasks depend on — and they do, since the Pile includes Wikipedia, books, and scientific articles — then the perplexity-measured forgetting may correspond to genuine capability degradation that the perplexity-only evaluation framework cannot directly observe.
Furthermore, the positive transfer claim — that continual pre-training outperforms training from scratch — is also a perplexity claim. The from-scratch SlimPajama model is worse at predicting SlimPajama validation tokens than the rewarmed Pile-initialized model. But it is not obvious that a model with better SlimPajama perplexity is uniformly better at all downstream tasks one might care about. A model trained exclusively on SlimPajama might have learned different inductive biases that, while producing slightly worse token-level predictions on the SlimPajama validation distribution, generalize better to certain task formats or prompt styles. The paper provides no evidence either way.
What evidence exists in the paper. Only perplexity curves. No benchmark results, no few-shot evaluation, no task-specific accuracy measurements. The paper does include the joint-perplexity visualization (Figures 4, 6) which is informative about the optimization dynamics but still operates entirely within the perplexity framework.
Mitigation status. Not addressed. The paper acknowledges the gap explicitly and defers benchmark evaluation to future work. Until such evaluation is conducted, the practical recommendations (use higher MaxLr for better downstream performance, use lower MaxLr to preserve upstream knowledge) should be interpreted as hypotheses about capability transfers rather than verified prescriptions. A deployment decision based on this paper should include task-level evaluation as a mandatory validation step.
No FLOPs-Matched Comparison Between Continual Pre-Training and Training From Scratch
The assumption or constraint. The paper's headline finding — that continual pre-training with rewarming outperforms training from scratch on SlimPajama — compares a model that has seen ~300B Pile tokens + ~240B SlimPajama tokens (total: ~540B tokens) against a from-scratch model that has seen only ~240B SlimPajama tokens. The total compute is not matched: the continual model has consumed more than 2× the training FLOPs. The paper's framing (Section 4.3, and implicit throughout) is that the Pile pre-training is a sunk cost that has already been paid, and the relevant practical question is "given that I already have this Pile checkpoint, should I continue or restart?"
"This shows that finetuning instead of retraining might improve performance even when the downstream dataset is on the scale of the upstream dataset."
The consequence. There are two distinct questions a practitioner might ask, and the paper answers only one of them:
-
"I already have a pre-trained Pile checkpoint. Should I continue pre-training on SlimPajama, or retrain from scratch on SlimPajama?" The paper's results support "continue pre-training" as the answer — for the model scale, dataset sizes, and schedule tested, the continual model achieves better downstream perplexity than the from-scratch model at equivalent downstream compute.
-
"I have a fixed total compute budget. Should I spend it training on Pile then SlimPajama sequentially, or training on SlimPajama alone for the entire budget?" The paper does not answer this question, because the from-scratch baseline is not given a compute-matched budget. A from-scratch model trained on SlimPajama for 540B tokens (matching the combined Pile + SlimPajama budget of the continual model) might close or reverse the performance gap. The paper's observed ~0.1 PPL advantage for MaxLr = 6 × 10⁻⁴ over from-scratch at 240B downstream tokens could vanish if the from-scratch model were allowed an additional 300B tokens of training.
This matters for ab-initio planning: an organization deciding how to allocate a pre-training budget across multiple datasets needs to know whether sequential training genuinely produces transfer benefits, or whether the continual model's advantage is simply a data-volume artifact. The paper's evidence does not distinguish these interpretations.
What evidence exists in the paper. The from-scratch model is only trained for the ~240B tokens of the downstream phase. There is no experiment where the from-scratch model receives additional training tokens to match the total FLOPs of the continual model (either by training longer on SlimPajama or by training on a Pile+SlimPajama mixture for 540B tokens). The paper also does not report the final perplexity values numerically, making precise comparison difficult — all values must be read approximately from figure axes.
Mitigation status. The paper does not acknowledge this as a limitation, nor does it propose a FLOPs-matched comparison as future work. The implicit argument is that the sunk-cost framing is the practically relevant one, which is defensible for teams that already possess pre-trained checkpoints. But for researchers or organizations planning new training runs, the missing comparison is a significant gap in the evidence base for the positive transfer claim.
Data Similarity and Overlap Prevent Generalization to Genuinely Novel Domains
The assumption or constraint. The upstream (Pile) and downstream (SlimPajama) datasets are built from overlapping sources — both include Common Crawl, Wikipedia, GitHub, ArXiv, books, and StackExchange — and differ primarily in curation strategy, deduplication, and source weighting rather than in fundamental domain. The paper acknowledges this limitation in Section 5:
"Since in continual learning, different types of shifts can lead to variations in performance (Lesort et al., 2021), our results may not generalize to setups with different distribution shifts, such as language domain adaptation pre-training setups (Xu et al., 2019; Gururangan et al., 2020; Ke et al., 2023a)."
The paper also notes data overlap explicitly: "SlimPajama is built from similar sources as the Pile but with a higher quantity of data. Therefore, some upstream data may be repeated during downstream pre-training."
The consequence. The paper's findings characterize continual pre-training under essentially the most favorable possible conditions: high data similarity, substantial domain overlap, shared tokenizer, and positive transfer between corpora. The result that continual pre-training works well in this regime — that it outperforms training from scratch, that the MaxLr can be tuned to balance adaptation and forgetting — establishes an upper bound on what schedule engineering alone can achieve. It does not establish how performance degrades as the upstream-downstream shift becomes more severe.
Consider three progressively harder scenarios that the paper does not test:
-
Moderate domain shift (e.g., general web text → biomedical literature): the token distribution changes substantially, domain-specific terminology appears, and factual knowledge from the upstream corpus may be largely irrelevant. The optimal MaxLr might need to be higher to overcome the initial mismatch, but the resulting adaptation might cause more severe forgetting of general language capabilities.
-
Language shift (e.g., English → multilingual): the tokenizer, trained specifically on the Pile, may be poorly suited to the downstream language, limiting how effectively the model can adapt regardless of learning rate. Rewarming dynamics might interact with tokenizer adequacy in ways not present in the monolingual English setting.
-
Modality or task shift (e.g., text → code, or causal LM → instruction tuning): the loss function, architecture assumptions, and representation requirements may differ fundamentally, making the concept of "continuing pre-training" with a modified learning rate schedule insufficient as a strategy.
The same-data ablation (Section 4.4) provides suggestive evidence about what happens at the opposite extreme — no distribution shift at all — showing that rewarming itself causes permanent performance degradation even without domain change. This establishes a lower bound: some forgetting is inevitable even in the best case. What happens in between — at the moderate shift levels most common in practice — is entirely unmeasured.
What evidence exists in the paper. Only the Pile → SlimPajama transition, with the acknowledged high similarity and overlap. The paper does not include experiments with lower-similarity downstream datasets, controlled variation of distribution shift magnitude, or any measurement of how the optimal MaxLr or the tradeoff curve shape changes with shift severity.
Mitigation status. The paper explicitly acknowledges this limitation and notes that "our results may not generalize." It frames the current work as studying the fundamental optimization dynamics in a controlled setting before tackling harder domain shifts. This is a reasonable research strategy — establish baselines in the easy case first — but it means the paper's practical recommendations (rewarm with higher MaxLr for better adaptation) come with a large, unquantified caveat: they may fail entirely if the downstream data distribution is substantially different from the upstream one. A practitioner adapting a general English model to a specialized domain should treat the paper's findings as hypotheses to validate, not as verified transferable principles.
No Multi-Stage Continual Pre-Training Experiments
The assumption or constraint. The paper studies a single upstream-to-downstream transition: Pile → SlimPajama. However, the motivation presented in Section 1 is explicitly about an ongoing, multi-stage process:
"As the amount of data available for pre-training is ever-growing, new and improved datasets will continue to become available. Should practitioners always combine existing datasets and train from scratch to obtain the best performance? Doing so would quickly become prohibitively expensive."
This framing implies a sequence of dataset releases over time — Pile (2020), then RedPajama (2023), then SlimPajama (2023), then future datasets — each representing an opportunity to update the model. The paper's own critique of the progressive decrease strategy (Winata et al., 2023) is that it fails when the number of training stages becomes large:
"repeatedly decreasing the learning rate would cause it to eventually become too small if the number of training stages becomes high"
This critique is central to the paper's motivation for studying rewarming, yet the proposed alternative — rewarming at each stage — is never tested beyond a single stage.
The consequence. Several failure modes become possible in multi-stage continual pre-training that are invisible in the two-stage setting:
-
Compounding forgetting. After stage 2 (SlimPajama), the model has already forgotten some Pile knowledge. After stage 3 (a third dataset), does it forget both Pile and SlimPajama, or is forgetting concentrated on the most recently learned data? The paper provides no evidence about whether the forgetting curve is additive (each new stage degrades all prior stages equally), recency-weighted (newer learning crowds out older learning), or saturating (forgetting asymptotes after a few stages).
-
Optimizer state accumulation. The AdamW optimizer state (first and second moment estimates) accumulates statistics over the entire training history. After pre-training on Pile (300B tokens), the moments encode gradient information from the Pile distribution. After continual pre-training on SlimPajama (240B tokens), they encode a mixture of Pile and SlimPajama gradients. After a third stage, the moments become a mixture of three distributions. The effective learning rate — the actual parameter step size after Adam's adaptive scaling — may change as the second-moment estimates grow, making the same nominal MaxLr produce different effective step sizes at different stages. The paper's finding that MaxLr controls the tradeoff in stage 2 may not generalize to stage 3, stage 5, or stage 10 if the optimizer state dynamics shift.
-
Schedule reset questions. After completing stage 2 (SlimPajama with rewarming and cosine decay), the learning rate is again at its minimum (10% of MaxLr). For stage 3, should the MaxLr be the same as in stage 2? Higher (because the model is further from its original pre-training and needs more aggressive adaptation)? Lower (because the cumulative knowledge base is larger and more fragile)? The paper provides no guidance.
-
Cumulative cost of the stability gap. Each rewarming event produces a transient loss spike (the "chaotic phase" documented in Figure 1). In a single transition, this spike is negligible (the model recovers within ~5B tokens). In a 10-stage continual pre-training pipeline, the model would experience 10 such spikes, each requiring recovery tokens. The cumulative overhead might become non-trivial. The paper does not measure whether the spike magnitude changes across stages.
What evidence exists in the paper. Only the single Pile → SlimPajama transition. There are no experiments with three or more sequential datasets, no measurement of how optimizer state statistics evolve across stages, and no test of whether the optimal MaxLr changes when applied repeatedly.
Mitigation status. The paper does not acknowledge this as a limitation, despite the multi-stage motivation being central to the introduction. The critique of progressive decrease (which is fundamentally a multi-stage concern) is used to justify studying rewarming, but rewarming is never evaluated in the multi-stage setting that makes progressive decrease problematic. This is a structural gap between the paper's motivation and its experimental design: the problem is framed as requiring a solution that scales to many stages, but the solution is tested on exactly one stage. A practitioner genuinely facing a multi-dataset continual pre-training pipeline — which is exactly the scenario the introduction describes — would need to extrapolate from the two-stage results with no evidence about whether the extrapolation is valid.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper does not introduce a paradigm shift, a novel algorithm, or even a new training technique. What it does — and this is rarer and arguably more valuable for a workshop paper — is provide the first systematic diagnostic framework for what happens during the warmup phase of continual LLM pre-training, converting a routine optimizer detail into a strategic control point and mapping its behavioral landscape with sufficient clarity that practitioners can make informed, quantitative decisions.
The field's prior relationship with learning rate warmup was essentially ritualistic. Every major LLM pre-training run — GPT-3, Chinchilla, LLaMA, BLOOM, Pythia — used linear warmup followed by cosine decay because it was standard practice, not because anyone had carefully studied what warmup does in the continual pre-training context. The implicit model was: warmup prevents early-training instability from random initialization; therefore, when resuming training from a pre-trained checkpoint (where the optimizer state is already well-initialized), maybe warmup is optional or at least less important. The paper's contribution is to replace this implicit model with an explicit one, backed by controlled experiments that reveal what warmup actually controls in the continual setting.
The conceptual reframing is this: rewarming is not about stabilization — it is about deliberate, controlled destabilization for the purpose of enabling new learning. The evidence for this reframing comes from multiple converging findings. The same-data ablation (Section 4.4, Figures 5–6) demonstrates that rewarming itself causes a permanent loss degradation even with no distribution shift — the model climbs out of its converged minimum and cannot return, settling at a worse point on the same loss landscape. This is not a "stabilization" phenomenon; it is a "perturb and re-converge" phenomenon. The MaxLr sweep (Section 4.2, Figures 2–4) then shows that the magnitude of this perturbation — controlled by the peak learning rate — directly determines where on the adaptation-forgetting Pareto frontier the model ends up. Higher perturbation = more downstream adaptation + more upstream forgetting; lower perturbation = better upstream retention + less downstream improvement. The warmup phase, in this view, is the mechanism that administers the perturbation, and its length turns out to be irrelevant (Section 4.1) because the critical variable is the peak intensity, not the ramp-up speed.
This reframing has a specific practical consequence: it redirects optimization effort in continual pre-training from data-side interventions (replay, data mixing, domain weighting) toward schedule-side interventions. The paper's evidence that a substantial component of forgetting is mechanically induced by the learning rate increase itself, independent of what data the model sees, means that better schedule design — alternative warmup shapes, momentum buffer reset strategies, learning rate re-initialization rather than continuous accumulation — may be a more compute-efficient lever for managing forgetting than data mixing, which incurs storage and I/O overhead proportional to the number of upstream datasets. Prior work in continual learning for language (Ke et al., 2023a; Scialom et al., 2022; Winata et al., 2023) treated the schedule as a secondary concern to be set heuristically while focusing on data strategy. This paper suggests the opposite priority: nail the schedule first, because the schedule controls the fundamental optimization dynamics, and then use data strategy to fine-tune the forgetting-adaptation balance.
The paper also resolves a tension between two plausible intuitions that practitioners might hold. Intuition A: "I should use a lower learning rate when continuing pre-training to avoid catastrophic forgetting of the upstream data." Intuition B: "I should use a higher learning rate so the model can actually adapt to the new distribution." The paper's joint-tradeoff visualization (Figure 4) shows that both intuitions are partially correct — they correspond to different points on the same Pareto frontier — and that the optimal choice depends on the practitioner's relative preference for upstream retention vs. downstream adaptation. This is not a resolution in the sense of proving one side wrong; it is a resolution in the sense of providing the quantitative framework within which the tradeoff can be navigated deliberately rather than guessed at. A team that values upstream retention at 2:1 relative to downstream improvement can look at Figure 4 and pick MaxLr ≈ 1.5 × 10⁻⁴; a team with the opposite preference can pick MaxLr ≈ 6 × 10⁻⁴. Both are "correct" given their objectives — the paper provides the map, not the destination.
The negative result on checkpoint age (Section 4.5, Figures 7–8) also has a landscape-shifting effect, though a narrower one. The finding that fully converged checkpoints outperform earlier ones for continual pre-training contradicts a hypothesis — drawn from the loss landscape and plasticity literatures — that early-stopped models reside in wider basins that are more amenable to further training. By providing clean evidence that this hypothesis does not hold at the Pythia 410M scale with ~300B tokens of pre-training, the paper prevents researchers from wasting effort on checkpoint selection heuristics in this regime. It also reframes the research question: if plasticity loss is not the issue, what is the mechanism by which converged checkpoints achieve better downstream adaptation? The paper doesn't answer this, but it eliminates a false lead.
However, it is important to be precise about the magnitude of this contribution. This is an incremental diagnostic advance, not a paradigm shift. The paper maps the behavior of an existing technique (linear warmup + cosine decay) under a specific set of conditions (Pythia 410M, Pile → SlimPajama, AdamW, single transition). It does not propose a new schedule, demonstrate superiority over alternatives (beyond the constant-LR baseline), or show that its findings generalize across scales, architectures, or distribution shift severities. The field will not reorganize around this paper. But for the specific subcommunity working on efficient LLM lifecycle management — teams that need to update models on new data without re-training — this paper provides the first principled, empirically grounded guidance on how to set the learning rate schedule, replacing a practice that was previously guided by folklore and defaults.
Follow-Up Research This Work Enables
Multi-stage continual pre-training with schedule reset strategies. The paper's own motivation — that new datasets will continue to be released and progressive decrease fails when stage count is high — demands testing rewarming across more than two stages. A direct follow-up would train on Pile → SlimPajama → a third large-scale corpus (e.g., RedPajama-v2, FineWeb, or a temporally-shifted Common Crawl dump), applying the same rewarming procedure at each transition, and measure whether the MaxLr tradeoff curve shifts across stages. The specific hypothesis to test: does the optimal MaxLr for stage N depend on the cumulative training budget (suggesting a decay in the effective learning rate as optimizer moments accumulate), and if so, does resetting Adam's second-moment estimates (while preserving parameters) stabilize the tradeoff across stages? This experiment would directly address the gap between the paper's multi-stage motivation and its two-stage experimental design, and would test whether rewarming is genuinely scalable to the many-stage regime or whether it inherits the same "learning rate becomes too small" problem that the paper critiques in progressive decrease.
FLOPs-matched comparison between sequential and single-corpus training. The paper's headline finding that continual pre-training outperforms training from scratch on SlimPajama compares models with unequal total compute: the continual model sees ~540B tokens (300B Pile + 240B SlimPajama) while the from-scratch model sees ~240B tokens. A FLOPs-matched experiment would train a from-scratch model on SlimPajama for 540B tokens (or on a Pile+SlimPajama mixture for equivalent compute) and compare final perplexity against the continual model. This would distinguish whether the continual model's advantage is genuine positive transfer (knowledge from Pile that SlimPajama training alone cannot replicate) or simply a data-volume artifact (any model trained on 540B tokens outperforms one trained on 240B). The same-data ablation provides suggestive evidence that transfer is real — the model trained on Pile alone reaches lower SlimPajama PPL than random initialization, and the from-scratch SlimPajama model's Pile PPL improves during training — but the FLOPs-matched comparison would provide dispositive evidence and establish the magnitude of the transfer benefit in compute-equivalent terms. This is the single most important experiment the paper does not run, and it would directly inform ab-initio planning for organizations deciding between single-corpus and sequential training strategies.
Benchmark-level evaluation of the adaptation-forgetting tradeoff. The paper relies entirely on validation perplexity, which is a continuous signal but an imperfect proxy for the capabilities that matter in deployment. A natural follow-up would evaluate the models from the MaxLr sweep on downstream benchmarks — HELM scenarios, Harness tasks (MMLU, HellaSwag, ARC, GSM8K), and domain-specific evaluations for Pile sub-domains (patent classification for Freelaw, code generation for GitHub) — and map how the perplexity tradeoff translates into capability tradeoffs. The specific question: does the ~0.1 PPL improvement on SlimPajama from using MaxLr = 6 × 10⁻⁴ instead of MaxLr = 3 × 10⁻⁴ correspond to meaningful gains on downstream tasks, and does the ~0.25 PPL degradation on Pile correspond to measurable capability loss on tasks that depend on Pile-only knowledge? This experiment would determine whether the Pareto frontier in perplexity space maps monotonically to a Pareto frontier in capability space, or whether there are capability cliff edges — points on the perplexity tradeoff where a small additional degradation in Pile PPL corresponds to a catastrophic loss of specific factual or reasoning capabilities. The paper's practical recommendations cannot be fully actionable until this mapping is characterized.
Optimizer state reset and warmup-shape ablations. The same-data ablation (Section 4.4) reveals that rewarming causes permanent loss degradation even without distribution shift, which the paper attributes to the optimizer being pushed out of its converged minimum and following a non-reversible trajectory during decay. This suggests a specific causal mechanism: the Adam optimizer state (particularly the second-moment estimates, which normalize gradients and determine effective step sizes) encodes the geometry of the upstream loss landscape, and rewarming at a high learning rate interacts with this state to produce steps that the decay phase cannot retrace. A direct ablation would test whether resetting the optimizer state — zeroing the first and second moment buffers while keeping parameter values — before rewarming changes the recovery behavior. If the model can return to its original Pile loss after rewarming when the optimizer state is reset, the forgetting mechanism is optimizer-state-mediated rather than purely parameter-displacement-mediated, which would suggest different mitigation strategies (state reset vs. schedule smoothing). Additionally, testing alternative warmup shapes — exponential warmup (which concentrates the ramp near the peak), inverse-root warmup, or no warmup with a "start high, drop fast" schedule — against the linear warmup default would test whether the specific functional form of the warmup matters once the peak learning rate is controlled for, potentially revealing more efficient perturbation strategies that achieve the same downstream adaptation with less upstream displacement.
Scale-dependence of rewarming dynamics. The paper's findings are entirely at the 410M parameter scale, and the authors explicitly flag scale generalization as a key uncertainty. A systematic scale sweep — testing the same Pile → SlimPajama transition with Pythia models at 70M, 160M, 410M, 1.4B, and 2.8B parameters — would characterize whether the stability gap magnitude, the optimal MaxLr, the recovery rate, and the checkpoint age ordering are scale-invariant or scale-dependent. The paper's own hypothesis from the plasticity literature would predict that larger models, which are typically more overparameterized relative to their training data, might exhibit less forgetting (more redundant capacity to absorb new knowledge without overwriting old) and flatter MaxLr tradeoff curves (the Pareto frontier shifts). If true, this would mean the paper's 410M results are a lower bound on continual pre-training efficacy — larger models would be even better at the continual pre-training task, and the MaxLr choice would be less critical. If false — if larger models forget more aggressively or exhibit sharper tradeoffs — it would mean the 410M results overstate the ease of continual pre-training and that scale-specific schedule engineering is essential. Either outcome is scientifically informative and practically important.
Distribution shift magnitude as a controlled variable. The paper uses a single upstream-downstream dataset pair with acknowledged high similarity and overlap. A systematic study varying the severity of distribution shift — from near-identical (Pile → Pile, as in Section 4.4) through moderate (Pile → SlimPajama, Pile → C4 alone) to severe (Pile → pure code corpus like The Stack, Pile → non-English Wikipedia, Pile → domain-specific text like PubMed abstracts) — would characterize how the MaxLr tradeoff curve, the stability gap magnitude, and the recovery behavior change with shift severity. The paper's results already provide two points on this curve: zero shift (Section 4.4, linear recovery along a common trajectory, no tradeoff) and moderate shift (Section 4.2, curved tradeoff frontier). Filling in the intermediate and extreme cases would produce a phase diagram of continual pre-training that tells practitioners what to expect as a function of how different their new data is from their old data. The prediction from the paper's framework is that more severe shifts require higher MaxLr to achieve meaningful adaptation (because the initial parameter state is less well-suited to the new distribution) but also cause proportionally more forgetting (because the downstream gradients point in directions more orthogonal to the upstream minimum). Whether the tradeoff curve maintains its convex, monotonic shape or develops non-monotonicities (e.g., an intermediate MaxLr that performs worse on both axes than either extreme) at severe shifts is an open question with direct practical implications.
Practical Applications and Downstream Use Cases
Updating deployed LLMs when improved pre-training corpora are released. The most direct application of this work is for teams that maintain a deployed LLM and want to incorporate new, higher-quality pre-training data without retraining from scratch. The paper provides concrete, immediately actionable guidance: continue pre-training from the latest checkpoint using a rewarmed learning rate, do not waste tokens on gradual warmup (0% or near-0% warmup is sufficient), and choose your MaxLr based on how much upstream forgetting you can tolerate relative to downstream adaptation. For a team that prioritizes downstream performance and can accept some degradation on older data distributions (e.g., a model being updated for a new product release where the old data domains are less relevant), MaxLr = 6 × 10⁻⁴ (2× the original pre-training peak) provides the best SlimPajama adaptation, achieving ~2.52 downstream PPL versus ~2.62 for training from scratch — a meaningful improvement at zero additional pre-training cost beyond the compute already spent on the upstream model. For a team that must maintain performance on the original training distribution (e.g., a general-purpose API where users depend on consistent behavior across domains), MaxLr = 1.5 × 10⁻⁴ (0.5× the original peak) limits Pile forgetting to ~2.58 PPL (versus ~2.90 for the aggressive rewarming) while still providing some downstream improvement. The constant-LR baseline offers a third option: stop early (~50B tokens) at ~2.35 upstream PPL and ~2.58 downstream PPL if minimal upstream perturbation is the overriding goal and only modest downstream adaptation is needed. This is the paper's most directly deployable finding: the MaxLr dial maps cleanly to a tradeoff frontier that practitioners can navigate based on their specific deployment constraints.
Cost-efficient model lifecycle management for research labs and mid-scale organizations. Research groups and mid-scale companies that pre-train their own models but cannot afford to retrain from scratch every time a new dataset appears can use this paper's findings to implement a checkpoint update protocol. The protocol would be: (1) maintain the final (most converged) checkpoint from each pre-training run — the paper's Section 4.5 result shows earlier checkpoints are strictly worse for downstream adaptation; (2) when a new corpus becomes available, allocate a compute budget for continual pre-training (e.g., 240B tokens, matching the paper's decay period); (3) rewarm the learning rate to a peak value chosen based on the organization's upstream-downstream preference, with the warmup length set to near-zero (saving ~1% of the downstream compute budget, or ~3B tokens at the paper's scale, with no long-term penalty); (4) cosine-decay to 10% of the peak over the budget period. The paper demonstrates that this protocol produces a model that outperforms a from-scratch model trained on the new corpus alone, meaning the sunk cost of the original pre-training continues to pay dividends. For a lab that pre-trained on the Pile in 2021 and wants to update to SlimPajama in 2023, this protocol avoids a full retraining run (~300B tokens, costing tens of thousands of GPU-hours for a mid-scale model) while producing a better final model than if they had started over. The paper does not provide a dollar figure, but the compute savings are approximately the cost of the original pre-training run multiplied by the number of subsequent dataset releases — which, given the accelerating pace of corpus development, could be substantial over a model's operational lifetime.
Guiding data mixture decisions in multi-corpus pre-training. The paper's finding that a model trained sequentially on Pile then SlimPajama outperforms a model trained on SlimPajama alone (Figure 2, rewarmed curves vs. from-scratch), despite the sequential model having higher total Pile PPL (i.e., forgetting), has implications for how organizations design their data curricula. If sequential training on complementary corpora produces better downstream models than training on any single corpus — even when the single corpus is large enough to saturate the model's capacity from a data-volume perspective — then the optimal pre-training strategy may involve deliberately staging dataset exposure: broad-coverage corpus first (to build general linguistic and factual knowledge), then higher-quality or larger-scale corpus second (to refine and specialize), with rewarming at each transition. The paper's results suggest that this staged approach may be superior to the alternative of simply mixing all corpora into a single training run (though this comparison is not made explicitly — a head-to-head between sequential Pile→SlimPajama and a Pile+SlimPajama mixture trained for equivalent total tokens would be needed to confirm). The practical implication is that data procurement and cleaning efforts need not be synchronized: a team can train on the best available corpus today, and when a superior corpus arrives tomorrow, continue training rather than restarting, with confidence that the result will be better than training on either corpus alone.
Informing learning rate schedule defaults for continual fine-tuning pipelines. Beyond the specific Pile→SlimPajama setting, the paper's core finding — that warmup length is irrelevant in continual training because the optimizer state is already well-initialized — has implications for the much broader class of continual fine-tuning workflows. When adapting a pre-trained LLM to a new domain, task, or instruction format, practitioners routinely apply a learning rate schedule with warmup inherited from pre-training defaults (e.g., linear warmup over 1% of fine-tuning data). The paper's evidence suggests this warmup phase is unnecessary and can be eliminated, saving compute and simplifying schedule configuration. The caveat is that fine-tuning datasets are typically much smaller than SlimPajama (megabytes to gigabytes, not hundreds of billions of tokens), and the paper's warmup-length experiments were conducted at the 50B+ token scale — it is possible that at very small dataset sizes, warmup provides stabilization benefits that are invisible at scale. But for large-scale continual fine-tuning (domain adaptation with billions of tokens, instruction tuning with large prompt datasets), the paper's finding justifies dropping warmup as a default and allocating those tokens to decay-phase training instead.
When to Prefer This Method
The paper does not propose a named method or position rewarming against a specific set of alternative continual pre-training algorithms (beyond the constant-LR and from-scratch baselines). It studies a parameterization of the standard warmup-cosine schedule and characterizes its behavior, rather than introducing a novel technique that competes with alternatives. As such, a formal decision matrix ("prefer rewarming when X, prefer progressive decrease when Y, prefer replay when Z") would impose a structure the paper itself does not provide.
The closest the paper comes to a comparative recommendation is its implicit argument that rewarming with a chosen MaxLr is preferable to (a) not rewarming at all (which prevents downstream adaptation because the learning rate is too small), (b) using a constant learning rate (which underperforms at convergence on downstream loss compared to rewarming + decay), and (c) the progressive decrease strategy from prior work (which the paper critiques as eventually driving the learning rate to zero over many stages). These preferences are supported by the experimental evidence for the specific setting tested — single Pile→SlimPajama transition, 410M model, cosine decay schedule — but the paper does not test rewarming against progressive decrease, replay, or regularization-based methods, so explicit decision rules would be extrapolation beyond the evidence. The practical guidance is better framed as: within the warmup-cosine schedule family, the key variable is the peak learning rate, not the warmup length or the checkpoint age, and practitioners should select the peak rate to match their tolerance for upstream performance degradation.