PaLM 2 Technical Report

In plain words

PaLM 2 is a general-purpose language model from Google, the successor to PaLM. It predicts text, and it is used for many jobs: answering questions, working through word problems, writing code, translating, and writing in many languages. The report is less about one trick than about a set of choices in how the model was built. The team found that, for a fixed amount of training compute, the amount of training text and the model size should grow at about the same pace. Its training text covers far more languages than earlier models of this kind, and it mixes several different training goals instead of one. A small part of the text is labelled with how toxic it is, so the model can be nudged toward calmer output when it is used. The authors say a smaller, better-trained model can match or beat its much larger predecessor, which makes it cheaper to run. The sizes and internal design are not published, so the report is a source of results and training ideas, not a recipe you can rebuild exactly.

TL;DR

PaLM 2 is a Transformer language model in three sizes (S, M, L). Their sizes and architecture are withheld app. D. Its three main ingredients are compute-optimal scaling, where data and parameters grow at roughly equal rates §1 §2.1, a more multilingual and deduplicated pre-training mixture §3, and a tuned mixture of pre-training objectives that follows UL2 §1. PaLM 2-L is smaller than PaLM 540B but uses more training compute §1. Author-reported: BIG-Bench Hard 78.1 with chain-of-thought vs 65.2 for PaLM tab. 6, and GSM8K 91.0 with self-consistency vs 74.4 for PaLM tab. 7. Most training details (sizes, layers, objective weights, batch size, learning rate, token count) are not stated, so this card is a results-and-recipe reference, not an implementation spec.

Problem

Most earlier large language models used mostly English text and one training objective, either causal or masked language modelling §1. PaLM's non-code pre-training data was about 78% English §1. Earlier scaling studies had found that model size should grow faster than data §2. The report asks whether a smaller model trained on more and better data can match a larger one, which lowers inference cost and serving cost §1. It also asks how to measure memorization, toxicity and bias for a model that is evaluated as a base model rather than as a product §4 §5.

Method

The report describes four design decisions. Architecture details are not given.

1. Compute-optimal scaling §2.1. The team trained models at four compute budgets, 1e19, 1e20, 1e21 and 1e22 FLOPs. For each budget they chose token counts with the heuristic FLOPs ≈ 6ND §2.1. They fitted a quadratic to the final validation losses at each budget to find the loss-minimising model size, then fitted a power law across budgets to extrapolate app. A.1 fig. 4 fig. 5. Their conclusion is that D and N should grow in equal proportions, unlike the earlier trend in which model size grew about 3× faster than data §1 §2.1. The scaling models (400M to 15B parameters) are for the study only and are not the PaLM 2 models §2.1.

# Compute-optimal scaling fit  [§2.1] [app. A.1]
# Input:  runs: list of (flops, params, final_val_loss)
# Output: optimal params and tokens for a target compute budget

def fit_compute_optimal(runs, measured_budgets, target_flops):
    # 1. Group runs by compute budget (an isoFLOP band)  [§2.1]
    optimal_points = []
    for budget in measured_budgets:
        band = [run for run in runs if run.flops == budget]
        # 2. Fit a quadratic to the band's validation losses and take
        #    the minimum as the optimal size  [app. A.1]
        #    (smoothing method and size axis scale are not stated)
        best_params = quadratic_minimum(band)
        optimal_points.append((budget, best_params))
    # 3. Fit a power law of optimal size against FLOPs  [§2.1] [fig. 5]
    law = fit_power_law(optimal_points)
    params = law.predict_params(target_flops)
    # 4. Tokens follow from the compute heuristic C ≈ 6·N·D  [§2.1]
    tokens = target_flops / (6 * params)
    return params, tokens

2. Pre-training data mixture §3. The corpus includes web documents, books, code, mathematics and conversational data. It contains more non-English text than earlier large models and parallel data in hundreds of languages, with one side in English. Cleaning steps are de-duplication, removal of sensitive personal information and quality filtering. The report says no languages were explicitly kept or removed §3. The top-50 language shares of the multilingual web sub-corpus are tabulated in the appendix §3.

3. Objectives and architecture §1 app. D. PaLM 2 is Transformer-based and uses "a tuned mixture of different pre-training objectives", motivated by UL2 §1. The objectives, their weights and the schedule are not stated. The model card says the model uses stacked Transformer layers with size-dependent parameter counts, and that further details are withheld app. D.

4. Context length §3. The report says training increased the context length significantly over PaLM. No value is given.

5. Toxicity control tokens §3 §5.1. A small fraction of pre-training documents received a special control token marking their toxicity level. The labels came from a fixed version of the Perspective API. At inference the model can be conditioned on low, medium or high toxicity tokens. The report says this does not hurt unrelated tasks §3.

# Toxicity-conditioned generation  [§5.1] [§3]
# Input:  prompt (token ids), level ("low", "medium" or "high")
# Output: continuation token ids

def generate_toxicity_controlled(model, prompt, level):
    # 1. Look up the control token for the requested level  [§5.1]
    control = control_token_for(level)
    # 2. Insert it into the input; the position is not stated  [§5.1]
    conditioned = insert_control(prompt, control)
    # 3. Decode greedily, one continuation per prompt  [§5.1]
    return greedy_decode(model, conditioned)

6. Memorization canaries §4, Memorization. Synthetic "canary" sequences were injected into the pre-training data to measure memorization. Two types were used: an interleave canary, which alternates blocks of N=50 tokens from two real documents in their original relative order, and a shuffle canary, which shuffles the tokens of one real document. Source documents had at least 500 tokens. Canary counts per language are given in a table §4, Memorization. The total share of canaries in the corpus is not stated.

7. Instruction tuning app. A.2. The report follows the general recipe of the Flan collection (more than 1,800 tasks, with instruction templates) app. A.2. Learning rate, batch size and number of steps are not stated.

8. Code model PaLM 2-S* §4.4. The S model was trained further on an extended, code-heavy and multilingual mixture. The token count of this extra training is not stated.

Hyperparameters

Results

Author-reported. Unless noted, the baseline is PaLM 540B. Results are for the base model under few-shot prompting, and the numbers are not those of any user-facing product §4 app. D.

Other reported results: PaLM 2 passes the C2-level language exams in every evaluated language, but these are simulated and not official grades fig. 1 app. C. In a 1-shot setting, the smallest PaLM 2 variant is said to be competitive with PaLM 540B §4.2.

Paper vs code

No official code.

The paper gives no code link, and the model is not released. The searches were: the method name and title words on the GitHub repository API, with and without the google-research organisation filter. Every hit was a third-party client, bot or wrapper, and none was used.

Pitfalls

• You cannot rebuild PaLM 2 from this report. Sizes, layers, width, heads, vocabulary, objective weights, token count, batch size and learning rate are withheld or not stated app. D §1. Use the report for recipe-level ideas and evaluation numbers, and treat any attempt to copy the architecture as guesswork.
• Scaling runs are not PaLM 2. The scaling models are a separate study and do not reflect the sizes or FLOPs of PaLM 2 §2.1 tab. 1.
• FLOP counts depend on which parameters you count. The study used total parameters for the 6ND estimate, but the non-embedding count fitted the optimal-size curve better. The embedding and output layers were shared in these runs app. A.1.
• Loss-optimal is not task-optimal. In the 1e22 study the 9.50B model has the lowest loss, but the 16.1B model has the higher downstream average (58.26 vs 57.68) §2.2 app. A.1.
• Extrapolation is thin. The power law uses four budgets, and the smoothing before the quadratic fit is not described app. A.1. Error bars come from a leave-one-out estimate with only three of four points app. A.1.
• Control tokens are a weak lever on their own. Dialog prompting alone reduced toxic responses more than control tokens. Adding control tokens to dialog prompting helped only on the standard dataset §5.1. The fraction of tagged data, the score thresholds for low, medium and high, and the token position are not stated §3 §5.1. The authors report no clear alignment cost on other evaluations but have no systematic ablation §5.1.
• Tags come from a specific API version. The labels use a fixed Perspective API version. The authors chose it to avoid drift as the API's signals change over time §3 §5.1.
• Some evaluations are contaminated or partly in training data. WinoGrande, ARC, DROP and COPA are in the Flan instruction mixture, so the instruction-tuned results may have seen training data §4.3 (footnote). BIG-Bench canary-marked data was removed from web pre-training §4.3 (footnote). Two language-exam sets were excluded for contamination, a German practice exam and a French writing exam app. C.
• Not every result improves. On WinoGrande the smallest model is below PaLM (77.9 vs 83.7) and PaLM 2-L is slightly below PaLM (83.0 vs 83.7) tab. 2. ARC-C is below GPT-4 tab. 5. MATH is below Minerva's self-consistency result tab. 7.
• Costly decoding. MATH uses 64 sampled paths, GSM8K 40, and the BIG-Bench Hard numbers use 250 examples per task with CoT. Compare only results with the same decoding setup §4.3.
• Checkpoint averaging. The large model's numbers average the five last checkpoints, so checkpoint variance is hidden §4 intro.
• Memorization is lower on average but not uniformly. PaLM 2 memorizes less than PaLM on average, and especially for text repeated fewer than about three times. For text repeated more often it memorizes more. The authors suspect this comes from de-duplication §4, Memorization. Tail-language canaries need fewer repetitions to be memorized §4, Memorization.
• Products are not the model. User-facing products add pre- and post-processing, and the underlying model may change. Expect product results to differ from the report §5 intro app. D.
• Language exam results are not official. Writing answers were graded by three contractors each, and the pass thresholds are the authors' reading of official guidelines app. C.

Lineage

• Builds on: PaLM, the predecessor whose data mix and benchmark tables PaLM 2 compares against §1.
• Builds on, no card: compute-optimal scaling (Hoffmann et al., 2022) §2; the UL2 objective mixture (Tay et al., 2023) §1; the Flan instruction-tuning recipe (Chung et al., 2022; Longpre et al., 2023) app. A.2; scaling-law work (Kaplan et al., 2020) §2.
• Compared against: PaLM 540B and PaLM-Coder-540B tab. 2 tab. 8; GPT-4 tab. 5 tab. 7; Minerva tab. 7; LaMDA on dialog toxicity §5.1; Google's production translation service §4.5.