Scaling Instruction-Finetuned Language Models

In plain words

A pretrained language model learns to continue text, not to follow requests. One fix is to train it further on many examples where each task is written as an instruction. This paper asks how far that recipe goes when three things grow: the number of training tasks, the size of the model, and the share of training examples that show step-by-step reasoning.

The authors' tests say all three help. The resulting models do better on questions they never saw in training, whether they get no examples, a few examples, or are asked to reason out loud. Reasoning examples matter in particular: leave them out and the model gets worse at reasoning, even on tasks unlike its training data. The authors also release a family of smaller models trained the same way.

The models are not perfect. Some specialised systems still score higher on some of these tests, and the safety and bias checks cover only part of the risk.

TL;DR

Instruction finetuning ("Flan") trains a pretrained model on 1,836 tasks written as instructions, mixed from four sources §2.1. The paper scales the number of tasks, the model size (up to PaLM 540B), and chain-of-thought (CoT) data §1. On held-out benchmarks (MMLU, BBH, TyDiQA, MGSM) every size gains over the same model without finetuning §3 tab. 3. Flan-PaLM 540B reaches 75.2% five-shot MMLU with CoT and self-consistency, which the paper reports as state of the art tab. 4. Finetuning costs 0.2% of pretraining compute for Flan-PaLM 540B §2.2.

Problem

Pretrained language models can perform tasks described in natural language, but finetuning on instructions makes them respond better and reduces the need for few-shot exemplars §1. Earlier instruction-finetuning work scaled the number of tasks only on a 3B model, or the model size only on 62 tasks §7. The paper finds that finetuning on non-CoT instructions alone degrades CoT performance on held-out tasks §4.2. It notes that earlier studies evaluated only non-CoT held-out tasks, on models too small for CoT reasoning §4.2. A raw pretrained model also has usability problems: it may not stop generating, or may continue the user's input instead of answering it §6.

The evaluation assumes held-out tasks: none of the evaluation benchmarks appear in the finetuning mixture §2.1 §2.3. The benchmarks were chosen because current models still perform well below expert raters on them §2.3.

Method

Flan is the finetuning procedure, applied to a pretrained checkpoint. The same procedure is used for every model, with only learning rate, batch size, dropout and step count changed per model §2.2.

1. Data. Four mixtures: Muffin (80 tasks: 62 from Flan 2021 plus 26 new ones, including dialog and program synthesis data), T0-SF (193 tasks from T0 that do not overlap with Muffin), NIV2 (1,554 tasks after removing 44 MMLU-related ones), and CoT (nine datasets whose CoT annotations were written by human raters) §2.1. The total is 1,836 tasks from 473 datasets in 146 task categories fig. 2.
2. Formats. Each mixture uses the instruction templates given by its creators. CoT data gets about ten hand-written templates per dataset. Few-shot templates use exemplar delimiters such as "Q:" and "A:", applied at random per example. Four combinations exist: zero-shot or few-shot, each with or without CoT §2.1 fig. 3.
3. Mixing. Each task is weighted by its number of examples, subject to a maximum cap per mixture. The four mixtures are then balanced so that no underlying task is repeated more than once. The paper gives two proportion sets: A for the scaling and ablation experiments, B for the rest app. E.
4. Training. A constant learning rate, the Adafactor optimizer, and packing of several examples into one sequence, separated by end-of-sequence tokens. Attention is masked across packed boundaries §2.2. The framework is JAX-based T5X §2.2.
5. Checkpoint choice. One checkpoint per model is used for all evaluations. The step is chosen from periodic evaluations, every 2k to 10k steps depending on model size, on the held-out tasks §2.2.
6. Compute. Finetuning costs a small share of pretraining compute, from 0.2% to 1.6% per model. Flan-PaLM 540B used about 512 TPU v4 chips for 37 hours §2.2 tab. 2.
7. Evaluation. Few-shot exemplars: five for MMLU, three for BBH, one for TyDiQA, eight for MGSM. MMLU and BBH are scored both with direct answers and with CoT; TyDiQA only direct; MGSM only CoT §2.3. The normalized average is the mean of six normalized scores: direct and CoT for MMLU and BBH, direct for TyDiQA, and CoT for MGSM §2.3.

The paper has no algorithm box and no numbered equations. The procedure as implementers should read it:

# Build the training mixture  [§2.1] [app. E]
# Input:  mixtures: name -> list of tasks; caps: name -> max examples
# Output: list of (task, weight) pairs
def build_mixture(mixtures, caps):
    # Weight each task by its example count, capped per mixture  [app. E]
    weighted = []
    for name, tasks in mixtures.items():
        for task in tasks:
            count = min(task.num_examples, caps[name])
            weighted.append((task, count))
    # Balance the four mixtures to fixed proportions  [app. E]
    return balance_mixtures(weighted)

# Finetune one pretrained model  [§2.2] [app. E]
# Input:  pretrained model, weighted mixture, hyperparameters
# Output: one checkpoint, used for every evaluation
def finetune(model, mixture, lr, batch_size, dropout, steps):
    # Pack several examples per sequence, split by end-of-sequence  [§2.2]
    sequences = pack_examples(mixture, separator=EOS)
    # Block attention across packed example boundaries  [§2.2]
    attention_mask = packing_mask(sequences)
    # Constant learning rate, Adafactor  [§2.2]
    optimizer = Adafactor(lr=lr)
    for step in range(steps):
        # batch_size is global, not per device  [app. E]
        batch = sample_batch(sequences, batch_size)
        loss = model.loss(batch, attention_mask, dropout=dropout)
        optimizer.step(loss)
        # Score the held-out tasks every 2k to 10k steps  [§2.2]
        if step % eval_every == 0:
            track_held_out_score(model, step)
    # Keep the step with the best held-out score  [§2.2]
    return best_checkpoint()

# Aggregate score on held-out benchmarks  [§2.3]
# Input:  raw scores for the six benchmark settings
# Output: one normalized average
def normalized_average(raw_scores):
    # Rescale each score against its random-guess baseline  [§2.3]
    # The exact formula is not stated; see Pitfalls.
    normalized = [normalize(score) for score in raw_scores]
    # Macro-average over the six settings  [§2.3]
    return mean(normalized)

Hyperparameters that the paper leaves to the implementer:

Per-model finetuning settings, copied from the paper:

Finetuning and pretraining compute, from the paper tab. 2:

Hyperparameters

See the Method section for the full list. The values the paper states:

• Mixture caps and proportions: app. E.
• Per-model batch size, dropout, LR and steps: app. E.
• Few-shot exemplars per benchmark: §2.3.
• Checkpoint interval: every 2k to 10k steps, depending on model size §2.2.
• Human-evaluation sampling: temperature 0.7, five responses §6.

Results

Author-reported. Setup: four held-out suites, where MMLU covers 57 tasks, BBH 23, TyDiQA eight languages and MGSM ten languages. Scores are few-shot prompted accuracy, averaged without weights over tasks §2.3 tab. 3. The paper's headline numbers:

Results by number of tasks and model size, with the full benchmark breakdown tab. 3:

Chain-of-thought results on held-out benchmarks tab. 4:

Across model families tab. 5. Norm. avg is the normalized average; the columns follow tab. 3:

Other findings, in words:

• Scaling. Finetuning helps at every size: the gain is 15.5 points at 8B, 10.4 at 62B and 9.4 at 540B. Most of the gain comes by 282 tasks; the paper reads the later gains as small §3 fig. 4.
• CoT data. Finetuning on non-CoT data alone lowers held-out CoT scores below the no-finetuning baseline at 540B. Adding the nine CoT datasets to the mixture recovers them, and at each size the combined model matches or beats the non-CoT-only model on non-CoT held-out scores §4.2 fig. 5.
• Limits on the claims. Flan-PaLM 540B does not beat code-davinci-002 on BBH algorithmic tasks, even with CoT and self-consistency, and is still below a finetuned ByT5 on TyDiQA §4.1.
• Zero-shot CoT. With CoT data in the mixture, Flan-PaLM 540B gains from the phrase "let's think step-by-step" on BBH; PaLM 540B without finetuning does not show this gain §4.3 fig. 6 fig. 7.
• Open-ended usability. On 190 prompts, raters preferred Flan-PaLM 540B over PaLM 540B in 79% of cases; the preference grows by about 10 points on CoT-trigger prompts. Each pair was rated by one rater §6 fig. 8 app. I.
• Responsible AI. In the authors' tests, instruction finetuning lowers the probability of toxic continuations on RealToxicityPrompts at every size, and raises zero-shot CivilComments toxicity-classification AUC for Flan-PaLM 540B to 86.5, against 71.4 for PaLM 540B app. C. The authors say disparities across identity groups remain app. C.

Paper vs code

The official repository is google-research/FLAN at commit e9e4ec6e2701182c7a91af176f705310da541277, licensed under Apache-2.0 README.md. Its README says the code was used to produce Flan-T5 and Flan-PaLM, and names Le Hou, an author of this paper, as a contact README.md. It generates the data mixtures only. It contains no model training code and no hyperparameters. Training ran in T5X §2.2; the hyperparameters are only in the paper app. E.

• Caps match. The per-mixture caps for FLAN (30,000), T0 (20,000), CoT (100,000) and NIv2 (5,000) match the paper mixtures.py. The code adds a Dialog cap of 200,000 that the paper does not state mixtures.py.
• Mixing rates differ. Code defaults are FLAN 4, T0 3, CoT 1, NIv2 1, Dialog 1 (CoT-II and Dialog-II 0.3) mixtures.py. These do not reproduce the paper's Proportion A or B app. E. The weights in run_example.py for flan2022_submix (40 / 32 / 20 / 5 / 3 %) belong to the later Flan Collection, according to the README run_example.py README.md. Use the paper's numbers for the Flan-PaLM setting.
• Dialog and program synthesis. The paper puts dialog data (citing Taskmaster) and program-synthesis data into Muffin §2.1. The code's default Dialog mix excludes wiki_dialog, task_master and qrecc mixtures.py. The README says Task Master, DeepMind Code Contests and DrRepair could not be included for legal reasons, so there is no program-synthesis submix README.md.
• CoT data. The repo has nine CoT training files , matching the nine datasets in §2.1.
• Not checked. The total of 1,836 tasks, the ten CoT templates per dataset, and the per-model settings were not checked against code. The code is not a training recipe, so checkpoint selection and the model hyperparameters are paper-only.
• MMLU prompts. The paper links jasonwei20/flan-2 for its MMLU prompts. That repository's default branch has only a README and mmlu-cot.json, and no LICENSE file. It is not used here.
• Flan-T5 checkpoints. The paper says they are public; the weights and their licence were not checked.

Pitfalls

• Held-out is not untouched. The checkpoint step is chosen from periodic evaluations on the held-out tasks §2.2, which are the same benchmarks reported in the results. Treat held-out numbers as slightly optimistic.
• Two MMLU splits. Table 1 and Table 4 use the test split. Every other MMLU number uses the validation split app. D. Do not compare across them without checking.
• Two Flan-PaLM 540B normalized averages. Table 3 (scaling experiments, proportion A) gives 58.5; Table 5 (the rest of the experiments, proportion B) gives 58.4 for the same label app. E. The MMLU and BBH columns differ too. The 8B and 62B rows agree between the two tables.
• Text versus tables. §7 says Flan-T5 11B reaches 43.7 on BBH direct and PaLM 62B 37.5; Table 5 gives 45.3 and 37.4 §7 tab. 5. Use the table. The v5 versioning note says the PaLM 540B MMLU-CoT number was corrected app. K; this card uses v5 values.
• GSM8K is in training. The GSM8K training set is part of the CoT mixture. The paper's 83.9% GSM8K with CoT and self-consistency is therefore not a held-out result §4.1.
• Normalized average formula is not stated. The paper gives only the footnote example, where 55% against a 50% random baseline maps to 10 and 45% to -10 §2.3. This matches (accuracy - random) / (100 - random) × 100, but that formula is inferred from the example, not stated. Within each benchmark, tasks are averaged without weights §2.3.
• Compute is only partly stated. The 0.2% figure is for finetuning. The paper does not quantify the extra decoding cost of CoT or self-consistency. The number of sampled chains for self-consistency is not stated.
• Decoding settings for benchmarks are not stated. Only the human-evaluation sampling is specified (temperature 0.7) §6.
• Human evaluation is thin. 190 prompts, one rater per pair; raters were allowed to search the web to check correctness §6 app. I. The filter drops responses whose log probability is better than half the median. The paper's example (median -20, drop -3) fixes the direction, but for negative medians the exact cut-off is ambiguous. Implement it as the example reads and state that choice.
• Packing changes the effective batch. The listed batch sizes are global and the effective batch is larger because of packing app. E. The paper says attention is masked across packed boundaries §2.2; implement the mask.
• Learning rates differ by a factor of six across models. Flan-T5 uses 5e-4, Flan-PaLM 8B and 62B use 3e-3, and Flan-PaLM 540B uses 1e-3 app. E. Do not copy one value across models.
• Held-out is at the dataset level. MGSM is built from GSM8K problems §2.3, and the GSM8K training set is in the CoT mixture §4.1. The paper calls its benchmarks held-out tab. 3 §2.1, but it does not analyse this source-level overlap. The MMLU tasks were removed from NIV2 app. F.
• Responsible-AI results are narrow. Toxicity is scored with the Perspective API on English and a limited set of identity terms. The authors state that tail toxicity stays high and that the Winogender comparison uses a scoring change that ignores articles app. C.
• Model licences are not in this card. The paper's data sources are public, but the licences of the Flan-T5 and Flan-PaLM weights, and of each underlying dataset, were not checked here.

Lineage

• Builds on: instruction finetuning with task mixtures, from FLAN (Wei et al. 2021, arXiv: 2109.01652), T0 (Sanh et al. 2021, arXiv: 2110.08207) and NIV2 (Wang et al. 2022, arXiv: 2204.07705) §1 §2.1. Chain-of-thought prompting (Wei et al. 2022, arXiv: 2201.11903) §1. Self-consistency (Wang et al. 2022, arXiv: 2203.11171) §4. Models: T5, PaLM (Chowdhery et al. 2022, arXiv: 2204.02311), U-PaLM and the UL2 continued-pretraining method §2.2 §5.
• Compared against: PaLM and its finetuned versions tab. 1 tab. 3 tab. 5; Codex (code-davinci-002) and ByT5 as prior best tab. 4; text-davinci-003 in the appendix, added in v5 app. K. Prior work on instruction finetuning scaled smaller models: a 3B model and 1.6K tasks, or a 137B model on 62 tasks §7.
• Related: the Flan Collection, a later data paper whose generator code is the same repository (arXiv: 2301.13688).
• Superseded by: not stated in this paper.