PaLM: Scaling Language Modeling with Pathways

In plain words

PaLM is Google's answer to a simple question from 2022: if you keep making a GPT-style language model bigger, does it keep getting better at tasks it was never trained for, given only a few examples in the prompt? The team built one with 540 billion parameters, plus smaller 8B and 62B siblings trained the same way, so they could see how abilities change with size.

The model itself is a plain "predict the next word" Transformer with a handful of tweaks. Most of them make training faster or more stable, or make text generation cheaper. Attention and the feed-forward part of each layer run side by side instead of one after the other. All attention heads share a single set of keys and values, which shrinks the memory needed when generating. Bias terms are removed, and an extra loss term keeps the output probabilities numerically tame. The engineering is the other half of the story. The model was trained on two TPU v4 "pods" of 3072 chips each, wired together by Google's Pathways system, and it kept those chips busy a larger share of the time than earlier models of similar size did. The authors also proposed a cleaner way to measure that utilization.

The results showed that scale still paid off. The largest model set new few-shot records on most standard English benchmarks. When prompted to write out its reasoning step by step (chain-of-thought prompting), it beat a fine-tuned system on grade-school math word problems. On some tasks the jump from 62B to 540B was far larger than the jump from 8B to 62B. The paper is also a candid record of what goes wrong at this size: the loss occasionally spiked during training, and the fix was to rewind and skip some data.

TL;DR

PaLM trains three dense decoder-only Transformers (8B, 62B, 540B) on one epoch of a 780B-token mixture of web, books, Wikipedia, news, code and conversation data §3. The architecture changes relative to a vanilla Transformer are SwiGLU, parallel attention/MLP blocks, multi-query attention, RoPE, tied embeddings, no biases and a 256k lossless SentencePiece vocabulary §2. The 540B model trains on 2×3072 TPU v4 chips with 12-way model and 256-way fully sharded data parallelism inside each pod and 2-way data parallelism across pods via Pathways, with no pipeline parallelism §4. The paper introduces model FLOPs utilization (MFU) and reports 46.2% for PaLM 540B versus 21.3–32.5% for earlier large models tab. 3. Author-reported results: few-shot SOTA on 28/29 English NLP tasks tab. 4, 69.3 on 5-shot MMLU tab. 6, 58% on GSM8K with 8-shot chain-of-thought plus a calculator tab. 10, and strong code results without a code-specific model tab. 12. The recipe is fully specified at the level of optimizer and schedule, but some details (RoPE settings, norm epsilon, exact vocab size, smaller-model batch schedules) are not stated, and no code or weights were released.

Problem

Few-shot prompting of large autoregressive LMs avoids task-specific fine-tuning, and earlier gains came from four levers: more parameters, more training tokens, cleaner and more diverse data, and sparse capacity §1. PaLM studies the first lever at a scale beyond earlier dense models. It asks whether few-shot quality keeps improving or saturates at 540B, and whether some abilities appear only at large scale §1. A second problem is systems-level. Earlier models at this scale either stayed on a single TPU system or used pipeline parallelism across pods or GPU clusters. Pipelining adds idle "bubble" time, extra weight reloads for each micro-batch and software complexity §4. PaLM aims to train pipeline-free across two pods at high efficiency §4.

Method

Section map of the compiled paper: §1 Introduction, §2 Model Architecture (§2.1 Model Scale Hyperparameters, §2.2 Model Card), §3 Training Dataset, §4 Training Infrastructure (§4.1 Training Efficiency), §5 Training Setup (§5.1 Training Instability), §6 Evaluation (§6.1 English NLP incl. §6.1.1 MMLU and §6.1.2 Finetuning, §6.2 BIG-bench, §6.3 Reasoning, §6.4 Code, §6.5 Translation, §6.6 Multilingual NLG, §6.7 Multilingual QA, §6.8 Analysis), §7 Memorization, §8 Dataset Contamination, §9 Exploring Explanations, §10 Representational Bias, §11 Ethical Considerations, §12 Related Work, §13 Open Questions in Scaling, §14 Conclusion. Appendices: A Contributions, B Compute Usage, C Dataset Analysis, D Datasheet, E Model Card, F Training for longer, G Sample outputs, H Additional results. The paper has no numbered equations, so the formulas below cite their section.

1. Architecture §2

A decoder-only Transformer with causal attention, where each position sees itself and earlier positions §2, changed as follows:

• SwiGLU MLP. The hidden activation is \textrm{Swish}(xW) \cdot xV, which needs three matmuls in the MLP instead of two §2. d_\textrm{ff} = 4 \times d_\textrm{model} tab. 1.
• Parallel block. Instead of the serial form y = x + \textrm{MLP}(\textrm{LayerNorm}(x + \textrm{Attention}(\textrm{LayerNorm}(x)))) each layer computes y = x + \textrm{MLP}(\textrm{LayerNorm}(x)) + \textrm{Attention}(\textrm{LayerNorm}(x)) §2 (the paper's serial formula as typeset is missing one closing parenthesis; balanced here). The input matmuls of MLP and attention can then be fused, giving about 15% faster training at large scale. Ablations showed a small quality loss at 8B and none at 62B. Quality neutrality at 540B is extrapolated, not measured §2.
• Multi-query attention. Queries are projected to [k, h] (k heads, head size h), keys and values to [1, h], shared by all heads §2. Head size is 256 at every scale tab. 1. The authors report neutral quality and training speed but large savings in autoregressive decoding cost §2.
• RoPE position embeddings instead of absolute or relative embeddings §2.
• Tied input/output embeddings §2.
• No biases in any dense kernel or layer norm, reported to improve stability for large models §2.
• Vocabulary. SentencePiece, 256k tokens, trained on the training data. It is lossless and reversible: whitespace is kept (important for code), out-of-vocabulary Unicode falls back to one token per UTF-8 byte, and numbers are always split into single digits ("123.5" becomes "1 2 3 . 5") §2.

2. Data §3

780B tokens, one epoch for all three models, shuffled identically, with mixture weights chosen so that no subcorpus repeats §3. Sources and proportions are in the table below tab. 2. Webpages are sampled in proportion to a quality-classifier score (from GLaM), so low-quality pages are down-weighted but not removed §3. Code comes from GitHub repositories with non-copyleft licenses, filtered to 24 languages by file extension (196GB), and deduplicated by Levenshtein distance between files §3. About 22% of the tokens are non-English §1.

tab. 2

3. Training setup §5

• Init. Kernels use fan-in variance scaling, written as W \sim \mathcal{N}(0, 1/\sqrt{n_{in}}). Input embeddings use E \sim \mathcal{N}(0, 1) because no layer norm is applied to the embeddings. Because embeddings are tied, pre-softmax logits are multiplied by 1/\sqrt{n}, with n the embedding size §5.
• Optimizer. Adafactor without factorization, which the authors describe as Adam with "parameter scaling": the learning rate is multiplied by the RMS of the parameter matrix. Embeddings and layer-norm scales therefore get a different effective step size than kernels §5.
• Schedule. LR 10^{-2} for the first 10,000 steps, then decayed proportionally to 1/\sqrt{k} (k = step) §5. \beta_1 = 0.9. \beta_2 = 1 - k^{-0.8}, which the authors find more stable than a fixed 0.99 because rare embedding tokens otherwise get poor second-moment estimates §5. Global-norm gradient clipping at 1.0. Weight decay is dynamic, equal to lr^{2.0} §5.
• Loss. Mean next-token log-likelihood, no label smoothing, plus an auxiliary z-loss that pushes the softmax normalizer toward \log Z = 0 for stability §5: \mathcal{L}_z = 10^{-4} \cdot \log^2 Z §5
• Sequences. Length 2048 for all models. Examples are concatenated and cut into exact 2048-token chunks (no padding; examples may be split) and separated by an [eod] token §5.
• Batch-size ramp. 540B: 512 sequences (1M tokens) to step 50k, 1024 (2M) to step 115k, 2048 (4M) to step 255k, the end of training. Smaller models follow "similar" schedules §5 tab. 1. Two reasons are given: small batches are more sample-efficient early on, and large batches give better TPU utilization §5.
• No dropout in pretraining. Dropout 0.1 is used for most fine-tuning §5.
• Bitwise determinism. Restarting from any checkpoint reproduces the run exactly. This relies on JAX+XLA+T5X determinism and on a data pipeline in which each batch is a function of the step number only §5.

4. Loss-spike mitigation §5.1

The 540B run showed about 20 loss spikes at irregular times despite clipping. The smaller models showed none §5.1.

Procedure: spike recovery  [§5.1]
1. Detect a loss spike.
2. Reload the checkpoint from ~100 steps before the spike began.
3. Skip ~200–500 data batches (those seen just before and during the spike).
4. Resume training.

Replaying the same batches from a different, earlier checkpoint did not spike. The authors therefore attribute spikes to a batch combined with a particular parameter state, not to bad data alone §5.1. Step 2 relies on the deterministic step→batch mapping of §5.

5. Parallelism and Pathways §4

• Within each pod (3072 TPU v4 chips, 768 hosts), every weight tensor is sharded with 12-way model parallelism × 256-way fully sharded data parallelism ("2D finalized" from GSPMD) §4.
• Forward: weights are all-gathered over the data axis, and one fully sharded activation per layer is saved. Backward: the remaining activations are rematerialized, chosen because it gives higher throughput at larger batch sizes §4.
• Across pods: 2-way data parallelism. A single Python client sends half the batch to each pod. Each pod computes gradients, exchanges them with the other pod (1:1 between corresponding hosts, about 1.3 GB per host pair per step, an 81 Tbps aggregate burst), sums local and remote gradients, and applies an identical update §4 fig. 2.
• Gradient transfers are split into smaller chunks routed over many DCN links to reduce congestion. Two-pod throughput is about 1.95× a single pod (97% of perfect weak scaling). The remaining gap comes from not overlapping the backward pass with the cross-pod reduction §4.

6. Model FLOPs utilization §4.1 app. B

MFU is observed tokens/s divided by the theoretical maximum tokens/s at peak matmul FLOPs. It counts only the FLOPs the model requires, not rematerialization §4.1. For an N-parameter decoder with L layers, H heads, head dim Q and sequence length T, and peak throughput P FLOP/s app. B:

R = \frac{P}{6N + 12LHQT} \qquad \textrm{MFU} = \frac{\textrm{observed tokens/s}}{R} app. B

Here 6N is 2N forward plus 4N backward matmul FLOPs, and 12LHQT = 6LH(2QT) is the attention matmuls app. B. Worked example for PaLM 540B: 238.3K tokens/s at batch 2048 gives (238.3 \times 6 \times 540)/(275 \times 6144) = 45.7\% without attention, 46.2% with it. Hardware FLOPs utilization (including rematerialization) is 57.8% app. B.

7. Code fine-tuning (PaLM-Coder) §6.4

Two stages from each pretrained model: (a) 6.5B tokens of 60% Python from a new 5.8B-token "ExtraPythonData" GitHub set not used in pretraining, 30% multi-language code (same source as pretraining code but held out) and 10% natural language; (b) 1.9B more tokens of ExtraPythonData Python. In total 7.75B tokens, 5.9B of them Python §6.4.

Hyperparameters

Compute per model app. B:

app. B. The 540B run used 6144 chips for 1200 h and 3072 chips for 336 h, including downtime and repeated steps, with 271.43 tCO2e estimated app. B.

Results

All numbers are author-reported. "Prior SOTA" is the best published few-shot result from a single pretrained checkpoint, taken from GLaM, GPT-3, Megatron-Turing NLG, Gopher, LaMDA or Chinchilla tab. 4.

Training efficiency tab. 3

English NLP, PaLM 540B vs prior SOTA tab. 4 (few-shot shot counts in parentheses)

PaLM 540B beats prior SOTA on 24/29 tasks 1-shot and 28/29 few-shot §6.1. Averages across the 29 tasks, 1-shot tab. 5:

MMLU, 5-shot tab. 6

SuperGLUE fine-tuned. Dev average 92.6 (T5-11B 89.9, ST-MoE-32B 93.2) tab. 7. Test leaderboard 90.4 vs ST-MoE-32B 91.2 and the best decoder-only LM 71.8 tab. 9. The gap between few-shot and fine-tuned remains large on dev, e.g. WiC 64.6 vs 78.8 and RTE 81.2 vs 95.7 tab. 8.

GSM8K, 8-shot tab. 10

The 8-shot CoT exemplars are exactly those of Wei et al. §6.3. Across 7 reasoning datasets, PaLM 540B + CoT reports SOTA on GSM8K, MAWPS, SVAMP and StrategyQA, and close to SOTA on ASDiv, AQuA and CommonsenseQA §6.3 fig. 10. The external calculator never added more than 5% on any dataset §6.3. Of 150 random GSM8K problems, PaLM 62B got 45 wrong, and scaling to 540B fixed a substantial share in every error category fig. 9.

BIG-bench §6.2. On the 58 tasks shared with GPT-3, Gopher and Chinchilla, 5-shot PaLM 540B beats the prior SOTA on 44 and scores above the average human rater in aggregate §6.2 fig. 3. Over 150 tasks, 25% show a "discontinuity" above +10% and 15% above +20%. The discontinuity is the 540B score minus a log-linear extrapolation from 8B→62B, e.g. logical_sequence 13%/25%/87% gives +50% §6.2. The average human still beats PaLM 540B on 35% of individual tasks §6.2.

Code tab. 12 (shots in parentheses; Davinci Codex numbers measured by the authors via the API)

PaLM's pretraining saw about 2.7B Python tokens, against 100B for Codex 12B §6.4 tab. 11. DeepFix compile rate with a revised prompt rose from 82.1% to 83.4% §6.4.

Translation (BLEU, PaLM 540B) tab. 14

The 0-shot prompt names the languages, while 1-shot and few-shot prompts do not tab. 14.

TyDiQA-GoldP (EM, validation): PaLM 540B fine-tuned averages 80.0, against mT5 XXL 79.1 and ByT5 XXL 81.4. Few-shot it averages 60.5 tab. 17.

Chinchilla comparison tab. 21: on 9 shared English tasks the average is PaLM 8B 51.2, 62B 64.8, 540B 70.3, Chinchilla 65.2, Gopher 59.5, at 37.4 / 295.7 / 2527.2 / 588.0 / 504.0 zettaFLOPs. Training PaLM 62B on to 1325B tokens on a refreshed mixture raises that average to 66.5, MMLU to 62.8 and GSM8K (no calculator) from 33.0 to 48.7 app. F. On 890B tokens, repeated data underperformed refreshed data on most of 8 tasks app. F.

Memorization §7: with a 50-token prompt and a 50-token greedy continuation, exact-match rates on training spans are 1.6% (8B) and 2.4% (540B) §7 fig. 18. Spans seen once are memorized 0.75% of the time by 540B, and spans seen more than 500 times over 40% §7.

Contamination §8: 10 of the 29 English tasks were flagged as contaminated. On their clean subsets (an example counts as dirty if ≥70% of its 8-grams appear in training), accuracy deltas are split evenly between positive and negative tab. 18.

Paper vs code

No official code. The model weights and training code were not released with the paper.

Pitfalls

• Initialization notation is ambiguous. "Fan-in variance scaling" implies variance 1/n_{in} (std 1/\sqrt{n_{in}}), but the paper writes \mathcal{N}(0, 1/\sqrt{n_{in}}), which under the usual (mean, variance) convention means std n_{in}^{-1/4} §5. The phrase points to std = 1/\sqrt{n_{in}}. Treat that as an interpretation, not a stated fact.
• Adafactor's "parameter scaling" changes the meaning of LR. The 1e-2 LR is multiplied by each matrix's RMS, so plugging 1e-2 into plain Adam would diverge §5. Factorization is off. Adafactor's update clipping and ε are not stated.
• LR schedule form. "1e-2 for 10k steps, then decayed at 1/√k" §5 is consistent with lr = 1/\sqrt{\max(k, 10^4)}, since 1/\sqrt{10^4} = 10^{-2}. That continuous form is our reading. The paper does not write a formula.
• Weight decay of lr² is tiny (1e-4 at peak, by our arithmetic) and decays with the LR. The paper does not say whether it is decoupled or which parameters receive it §5.
• β2 depends on the step. 1 - k^{-0.8} is 0 at k = 1, so the very first steps effectively normalize by the current squared gradient. Implementations must count k from 1, not 0, to avoid 0^{-0.8} §5. The exact indexing is not stated.
• Parallel layers were not ablated at 540B. The quality cost seen at 8B (small, none at 62B) is assumed to vanish at 540B by extrapolation §2. Small reproductions should expect a slight loss.
• The 540B attention width differs from d_model. With head size 256, 48 heads give 12,288 query dims against d_model 18,432, so the Q and output projections are not square. At 8B (16 × 256 = 4,096) and 62B (32 × 256 = 8,192) they match d_model tab. 1. This is arithmetic on tab. 1; the paper does not discuss it.
• Unstated architecture details: RoPE base and the fraction of head dims rotated, norm epsilon, whether "LayerNorm" means scale-only LayerNorm or RMSNorm (no biases are used), exact vocabulary size, and how the z-loss interacts with the 1/\sqrt{n} logit scaling.
• The 62B model is not trained "identically". §2.1 says all three models used the same data and training apart from batch size, but the reported 62B checkpoint is at 795B tokens, not 780B, "due to an oversight in checkpoint selection" §6.8 app. B. The 8B and 62B batch-switch steps and total steps are not given.
• Loss spikes are handled by hand. About 20 manual rewinds of ~100 steps, skipping 200–500 batches each §5.1. Reproducing this needs bitwise-deterministic, step-indexed data loading. The paper offers no principled fix. The compute figures include these repeated steps app. B.
• Numbers disagree across tables. The Chinchilla comparison lists PaLM 540B "0-shot" HellaSwag 83.6 and BoolQ 88.7 tab. 21, which match the 1-shot entries of tab. 4 (0-shot there: 83.4 and 88.0). GSM8K with CoT+calculator is 58% in tab. 10, while §6.4 quotes 58.1 from fig. 10. Quote the table you mean.
• The checkpoint matters for some tasks. WebQuestions 1-shot for 540B varies a lot across checkpoints 2B tokens apart and peaked at 770B tokens. All reported results use the final checkpoint §6.8.
• pass@k is not the unbiased estimator. The paper reports the raw fraction of problems solved by k samples, not the Chen et al. estimator, so numbers are not directly comparable to papers that use it §6.4. Davinci Codex numbers come from an API model with unknown size, data and post-processing §6.4.
• The translation 0-shot vs few-shot comparison is confounded. 0-shot prompts include language names and few-shot prompts do not, which may explain why 0-shot sometimes wins (fr→de: 25.2 0-shot vs 9.5 1-shot) tab. 14 tab. 15.
• A broken reference in the paper. §6.4 points to a table of code data by programming language, but the appendix file defining it is not included in the compiled paper, so the reference does not resolve. §6.4 also cites tab. 12 where tab. 11 (code token counts) is meant.
• The cost hidden by the abstract. 2.56×10²⁴ training FLOPs (29,600 PF-days) for 540B, and about 1.95× rather than 2× scaling across pods because the cross-pod gradient transfer does not overlap the backward pass app. B §4.

Lineage

• Builds on: the Transformer decoder as used by GPT-3 (Brown et al., 2020), which supplies the few-shot evaluation protocol and the 29-task suite §6.1. SwiGLU (arXiv:2002.05202), multi-query attention (arXiv:1911.02150), RoPE (arXiv:2104.09864), GPT-J's parallel block (Wang & Komatsuzaki, 2021), Adafactor (arXiv:1804.04235) §2 §5. Data pipeline from GLaM (arXiv:2112.06905) and LaMDA (arXiv:2201.08239) §3. Pathways (Isard et al., 2022) and GSPMD sharding §4. Chain-of-thought prompting (arXiv:2201.11903) §6.3.
• Compared against: GPT-3, GLaM, Gopher (arXiv:2112.11446), Chinchilla (arXiv:2203.15556), Megatron-Turing NLG 530B (arXiv:2201.11990), LaMDA tab. 4. Codex (arXiv:2107.03374) and Davinci Codex tab. 12. T5-11B and ST-MoE-32B on SuperGLUE tab. 7. mT5 and ByT5 on TyDiQA tab. 17.
• Memorization methodology follows Carlini et al. (arXiv:2202.07646) §7.