Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling

Pythia is a set of language models built to be studied rather than used. Stella Biderman, Hailey Schoelkopf and colleagues at EleutherAI published the paper in April 2023 (arXiv 2304.01373, ICML 2023), with compute donated by Stability AI. The suite is 16 decoder-only models at 8 sizes from 70M to 12B parameters. Every size was trained twice, once on the Pile and once on a deduplicated Pile, and every run saw its data in exactly the same order. Nathan Lambert's open-source-ai-reading-list puts it first among the "fully open" technical reports, ahead of the AI2 olmo line that took the same idea to competitive scale.

Why a suite, and why this one

The paper's starting point is that scaling laws describe how trained models behave as they grow, but almost nothing connects those laws to what happens during training, and the few studies that tried used private models or private checkpoints. The authors surveyed the public model families of early 2023 (GPT-2, GPT-3, GPT-Neo, OPT, T5, BLOOM) against three requirements and found none met all of them: models spanning several orders of magnitude of scale, all trained on the same data in the same order, and public data plus intermediate checkpoints.

The survey explains what each rival lacked. GPT-3 got half a mark for public access because its API costs money and OpenAI restricts the research allowed on it. GPT-Neo models are used as if they were a suite but differ in tokenizer, codebase and data order. OPT's training data is private and fewer than 10 of its checkpoints were public, for three sizes only. BLOOM handled training divergences by rewinding and skipping offending sequences, so its models differ in small undocumented ways. Counting checkpoints by the model in each suite that has the fewest, the best competitor was GPT-Neo with 30; Pythia offers 154 for every model.

Pythia deliberately trades performance for control. It uses parallel attention and feedforward layers at every size, even though that choice was believed to hurt models under 6B. The authors expected to pay for it and did not: the models match OPT at equal parameters and tokens.

What was trained

All models were trained for 299,892,736,000 tokens, about 300B, to token-match the original GPT-3 and OPT suites. The standard Pile is 334B tokens under the GPT-NeoX tokenizer, so the standard runs never see some of it. The deduplicated copy (MinHashLSH near-deduplication at a 0.87 threshold) is only 207B tokens, which means the deduped models run about 1.5 epochs. The paper turns that into an experimental feature: comparing checkpoints just before and just after the epoch boundary shows no evidence that the second epoch hurts benchmark scores.

The architecture follows GPT-3 with the changes that had become standard by 2023: fully dense attention instead of GPT-3's alternating sparse layers, Flash Attention, rotary embeddings (applied to 25% of dimensions per the config table), and the GPT-J parallel attention/MLP layout and initialization. Embedding and unembedding matrices are untied because earlier work suggested this makes interpretability research easier. The tokenizer is the BPE tokenizer trained on the Pile for GPT-NeoX-20B.

The training choice that goes most against received wisdom is batch size. Every model, including the 70M one, uses 1024 sequences of 2048 tokens, about 2M tokens per step. Small models were thought to need small batches to converge; the authors report no convergence problems at 4x to 8x the usual batch for sub-1B models, and wall-clock speedups of up to 10x on the smaller sizes. A checkpoint is saved every 1,000 steps (143 of them over 143,000 steps) plus step 0 and log-spaced steps 1, 2, 4 through 512, for 154 per model. Training ran on 40GB A100s with GPT-NeoX, DeepSpeed and ZeRO; the paper's whole compute bill, including the retrain described below, was 544,280 A100-hours.

The paper is candid that the suite was trained twice. A preliminary "v0" release had inconsistent batch sizes (4M tokens for three of the sizes) and a learning-rate schedule that decayed to 10% of peak for models up to 2.8B but to zero for 6.9B and 12B. The authors retrained everything with uniform settings, reran every analysis, and replaced the public models on 2023-03-31, keeping v0 available for replication. They also renamed the models in January 2023 so that parameter counts include embedding layers, matching GPT-2, BLOOM and OPT. They costed the alternative of training many random seeds per size and dropped it: 25 seeds would have run to about 10 million USD.

What was released

Everything is Apache 2.0: all weights and checkpoints on the HuggingFace Hub, the training code, the evaluation code with raw benchmark scores, and the data. The data release goes past a pointer to the Pile. EleutherAI ships the pre-tokenized files the dataloader actually read and a script that reconstructs the exact dataloader, so a researcher can read out the contents of any batch at any step. That is the property the case studies depend on, and the one that distinguishes Pythia from a set of open weights.

Evaluation uses EleutherAI's own LM Evaluation Harness, with every score rerun in-house instead of copied from other papers. On eight common benchmarks Pythia and Pythia-deduped perform very close to OPT and BLOOM. The per-benchmark result tables in the appendix did not survive extraction in the local copy of the source, and neither did Table 1 with the per-size hyperparameters, so exact scores and layer counts are not on this page.

Three things the evaluation contradicted

The authors flag three findings against the literature of the time. Deduplicating the Pile gave no clear benefit to language-modeling performance, which agrees with the GPT-NeoX-20B work but not with most papers; they suggest the Pile's deliberate upsampling of some subsets may not behave like ordinary duplication. Parallel attention did not cost small models anything. And BLOOM's "curse of multilinguality" showed up inconsistently: it underperformed on LAMBADA, PIQA and WSC but not on WinoGrande, ARC, SciQ or LogiQA.

Case studies

Each of the three case studies is chosen to need something only Pythia provides.

The gender-bias study retrains the tail of a run with the data changed. For the deduplicated 70M, 410M, 1.4B and 6.9B models, the authors take the checkpoint and optimizer state from 21B tokens (7%) before the end, and resume on exactly the same data with masculine pronouns swapped for feminine ones; the 1.4B model also gets a 63B-token (21%) version. Because order and content are otherwise identical, the pronoun frequency is the only variable. Stereotypical accuracy on WinoBias drops for every intervention and size, and the 6.9B model flips from pro-stereotypical to anti-stereotypical bias. CrowS-Pairs bias also falls, most for the larger models, which were also the most biased to begin with. LAMBADA accuracy barely moves.

The memorization study asks whether sequences seen late in training are memorized more. Using the Carlini et al. definition with 32-token prompts and 32-token continuations over every training sequence, the answer is no. The count of memorized sequences per batch fits a Poisson point process very well for the 12B model, so memorization is spread evenly across training. The practical advice follows: placing sensitive sequences at the start or end of the data will not reduce memorization, though putting them at the start at least lets you catch the problem in a partly trained checkpoint.

The term-frequency study counts, for every checkpoint, how often each arithmetic operand and each TriviaQA entity appeared in the data that checkpoint had actually seen. Models of 2.8B parameters and larger develop a correlation between task accuracy and term frequency after about 65,000 steps, 45% of the way through training; before that, and in smaller models, the correlation is mostly absent. Models under 1B rarely get the tasks right at all, even with 16 examples in the prompt. On multiplication, the gap between the most and least frequent operands widens as training goes on.

Where it sits

Pythia defines "fully open" as reproducibility down to the batch, which is a stricter standard than most later releases claimed. The OLMo paper cites Pythia and BLOOM as the most open models before it and compares its own checkpoints against Pythia-6.9B, which is the direct line from this report to olmo-accelerating-the-science-of-language-models. What Pythia did not try to be was competitive: 300B tokens was GPT-3's budget, and later open models trained on ten times that or more. Its lasting use is as an instrument for interpretability, memorization and data-attribution research, where fixed order and dense checkpoints matter more than benchmark rank.