Multitask Prompted Training Enables Zero-Shot Task Generalization

In plain words

Large language models often handle tasks they were never trained on, if the task is described in the prompt. One explanation is that they pick up many tasks by accident while predicting text. The authors ask whether the same ability can be built on purpose, by training one model on many tasks at once, each phrased as a plain-language request.

They take a pretrained text-to-text model and fine-tune it on a large set of existing datasets. Each dataset is rewritten with several prompt wordings, so the model sees the same task in many forms. They then test it on kinds of tasks that were left out of training, such as deciding whether one sentence follows from another. The model does well on these held-out tasks even though it is much smaller than the language models it is compared with.

TL;DR

Sanh, Webson, Raffel and co-authors build a system that turns NLP datasets into prompted text pairs, and fine-tune an encoder-decoder T5 model on 62 datasets and 12 task types. Four task types are held out: natural language inference, coreference, sentence completion and word sense disambiguation. T0 (11B) matches or exceeds GPT-3 (up to 175B) on 9 of 11 held-out datasets, and the authors say it is about 16 times smaller (author-reported). Read this if you want to build multitask instruction data or check zero-shot claims. The Pitfalls section lists where the paper and its code disagree.

Problem

• Large language models show zero-shot generalization to new tasks, usually attributed to implicit multitask learning during pretraining §1.
• The question is whether explicit multitask training on prompted datasets induces zero-shot generalization directly §1.
• Earlier multitask work generalizes to held-out datasets, not held-out tasks, and single-prompt formats do not extend to new tasks §1 §2.
• Zero-shot ability depends on model size and is sensitive to prompt wording §1. Choosing a prompt by its validation score leaks evaluation information, so the authors do not select prompts this way §5.
• Tasks are a fuzzy partition of NLP datasets. The authors group by task format rather than by skill, and they hold out whole tasks §3.

Method

Partition. The authors collect 62 English datasets in 12 task types with publicly contributed prompts. The 11 held-out evaluation datasets are WSC, Winogrande, ANLI R1 to R3, CB, RTE, COPA, HellaSwag, StoryCloze and WiC §3 app. D.4, tab. 5. Held-out task types are NLI, coreference, sentence completion and word sense disambiguation §3. Datasets in the four held-out types are excluded from T0 and T0+. T0++ adds SuperGLUE WSC and WiC to its training data, which is why those two are greyed out in its results §5 tab. 5 tab. 7. LAMBADA is dropped from evaluation because GPT-3 used a different variant app. D.3.

Prompts. A prompt is an input template and a target template, plus metadata such as answer choices and the evaluation metric. Templates map a data example to text §4. Each dataset is written as several prompts, and the collection (P3) has 2073 prompts for 177 datasets, 11.7 per dataset on average §4. Prompt templates are written in Jinja in the PromptSource tool; an empty rendered prompt is skipped app. C. Prompts that are not the dataset's original task, for example generating a document from its summary, are used in training but not reported in evaluation §4.

Model. The model is a T5-style encoder-decoder. Input text goes to the encoder, and the decoder is trained by maximum likelihood to produce the target text. The model is never trained to generate its input §5. The starting point is the LM-adapted T5 of Lester et al., trained further on a language-model objective §5. Variants: T0 (11B), T0+ (also trains on GPT-3's evaluation datasets), T0++ (also trains on SuperGLUE except RTE and CB), and T0 (3B) §5.

# Build one prompted example from a raw example  [§4] [app. C]
# Input:  example (dict of fields), template (input and target templates)
# Output: (input_text, target_text), or None if the prompt is empty

def materialize(example, template):
    # Fill the input template with the example's fields  [§4]
    input_text = render(template.input_template, example)
    # Fill the target template; it is the text the model must produce  [§4]
    target_text = render(template.target_template, example)
    # Skip examples where the template does not apply  [app. C]
    if input_text == "" or target_text == "":
        return None
    return input_text, target_text

# Sampling cap per template in the training mixture  [§5]
# Input:  train_size (examples in the dataset), num_templates
# Output: number of examples drawn for each template

def template_cap(train_size, num_templates):
    # Large datasets share a budget of 500,000 across their templates  [§5]
    if train_size > 500000:
        return 500000 // num_templates
    # Smaller datasets keep their full size for every template  [code:
    # t0/seqio_tasks/tasks.py#L205-L214]
    return train_size

# Training loss for one example  [§5]
# Input:  input_ids (encoder tokens), target_ids (decoder tokens)
# Output: scalar negative log-likelihood of the target

def training_loss(input_ids, target_ids):
    # Score each target token given the input; only targets are scored  [§5]
    token_log_probs = decoder_log_probs(input_ids, target_ids)
    # Sum vs mean over tokens is not stated in the paper  [§5]
    return -sum(token_log_probs)

Training. All training examples are combined and shuffled, and the training mixture is sampled with the cap above §5 [code: t0/seqio_tasks/tasks.py#L205-L232]. Sequences are packed to the maximum length §5. Hyperparameters are in the next section.

Checkpoint. The checkpoint with the highest score on the validation splits of the training datasets is kept. No held-out data is used for this choice §5.

Evaluation. For multiple-choice tasks, rank classification: score every answer option by its log-likelihood under the model, pick the highest, and do not normalize by length §5. The metric is accuracy §5. For each dataset, the reported number is the median over all its prompts, with the interquartile range (Q3 minus Q1) as the robustness measure §5. Validation splits are used §5; ANLI uses its dev rounds tasks.py.

# Rank classification for a multiple-choice task  [§5]
# Input:  input_ids (encoder tokens), choices (list of token lists)
# Output: index of the chosen answer

def predict(input_ids, choices):
    scores = []
    for choice in choices:
        # Log-probability of every token of this option  [§5]
        # [code: t0/model.py#L61-L64]
        token_log_probs = score_tokens(input_ids, choice)
        # Plain sum, no length normalization  [§5]
        scores.append(sum(token_log_probs))
    return argmax(scores)

Prompt robustness experiments. The main T0 model averages 8.03 prompts per dataset (the non-original prompts included). Ablations use 1 prompt per dataset, 5.7 original prompts per dataset, and no prompted training §6.2 fig. 6. The dataset-count ablation uses 39, 49 and 55 datasets for T0, T0+ and T0++ §6.2 fig. 7.

Hyperparameters

Hardware and compute: about 27 TPU v3-512 hours for one T0 run and about 270 hours for all experiments in the paper app. B.1 tab. 1.

Results

Author-reported. Zero-shot evaluation on validation splits; per dataset, median over all its prompts §5. GPT-3 reports one prompt per dataset, while the T0 numbers cover all prompts §6.1. GPT-3 numbers appear only in fig. 4 and are not tabulated here.

• Held-out datasets: T0 matches or exceeds all GPT-3 models on 9 of 11 held-out datasets. Winogrande and HellaSwag are the two exceptions §6.1 §7.
• NLI: T0 outperforms GPT-3 on all NLI datasets, though the T5+LM baseline does not §6.1.
• Against the T5+LM baseline with the same architecture and prompts, T0 improves on every held-out dataset §6.1 fig. 4.
• BIG-bench subset: at least one T0 variant outperforms all baselines on all tasks except StrategyQA. The baselines are decoder-only models from the BIG-bench maintainers, each evaluated with one prompt per dataset §6.1 fig. 5.
• Prompt robustness, p=5.7 versus p=1: median improves on 8 of 11 datasets, and the IQR shrinks on 7 of 11 §6.2 fig. 6.
• Prompt robustness, all prompts versus 5.7: median rises on 9 of 11 datasets and the IQR shrinks on 8 of 11 §6.2.
• Datasets ablation: median rises for all 5 held-out datasets from 39 to 49 datasets, but the IQR shrinks on only 1 of 5 §6.2 fig. 7.
• GPT-3 robustness on RTE: 10 prompts, median 52.96%, IQR 1.28 §6.2. One of them reproduces GPT-3's reported prompt at 58.8%, below the 63.5% GPT-3 reports §6.2.
• Without instructions, the HellaSwag median rises from 33.65% to 57.93%. For Winogrande, FLAN's prompt without instructions gives 62.15%, with no large change §7.

Table 7 gives per-dataset medians for T5+LM and the T0 variants. Values are medians across prompts, as reported in tab. 7 app. F. An asterisk marks a cell the paper greys out as not zero-shot for that model, because the model trained on that dataset.

BIG-bench subset (accuracy, one BIG-bench prompt per dataset): tab. 8 app. F. Accuracies for T5+LM, T0, T0+ and T0++; baseline models appear only in fig. 5.

Bias and contamination (not in the headline numbers). On CrowS-Pairs and WinoGender, T0 mean accuracy is 59.2 and 84.2 respectively tab. 3 app. B.3. WinoBias pro-minus-anti gaps are reported in tab. 4. Contamination checks found 9.12% of HellaSwag examples and 5.2% of RTE hypotheses in C4 app. E tab. 6.

Paper vs code

Checked against the official repository at code_commit, which contains the training and evaluation scripts and the mixture definitions. Discrepancies:

• Checkpoint selection. The paper selects the checkpoint by validation score on training datasets §5. The training README evaluates a fixed step, 1,112,200 README.md, and the repo contains no selection script.
• Sampling weight. The paper says the mixture samples each dataset in proportion to its number of examples, with the 500,000 cap for large datasets §5. The code sets a cap per template, so each template of a dataset under 500,000 examples contributes its full size [code: t0/seqio_tasks/tasks.py#L205-L232]. A dataset with several templates therefore gets more weight than its example count would give.
• Target length in evaluation. The paper truncates targets to 256 tokens §5. The evaluation script defines --target_max_length (default 256) run_eval.py, but the answer-choice tokenization uses max_length (default 1024) instead [code: evaluation/run_eval.py#L268-L276].
• Special tokens in evaluation. Inputs are tokenized with add_special_tokens=False run_eval.py. Answer options are not, so the end-of-sequence token of the tokenizer's default setting is scored as part of each option. This is a reading of the code, not a check of the output.
• Training steps and optimizer details. The paper gives neither the number of steps nor the dropout's location. The repo supplies 1,112,200 steps README.md and a learning-rate gin file, but dropout, Adafactor and packing settings live in external T5 gin files that are not in this repo.
• Training code. The released training path is Mesh TensorFlow on TPUs. The README says the PyTorch replication is still a work in progress README.md.
• Hotpot QA. The paper recasts HotpotQA as closed-book QA fig. 2. The code drops the generative HotpotQA prompts because they exceed 1024 tokens tasks.py. The paper does not list the kept prompts.
• Mixture membership. The T0+ and T0++ additions match the paper's description tasks.py.

Pitfalls

• Zero-shot is only as clean as the held-out split. Tasks are a fuzzy partition §3, and the authors argue that NLI and paraphrase are distinct while holding out sentence completion as a precaution app. D.2. A reader who disagrees should treat the held-out numbers as one partition, not a fact.
• Contamination is partial. 9.12% of HellaSwag and 5.2% of RTE hypotheses appear in C4 app. E tab. 6. The search used exact matches on 16-token windows, and flagged matches were checked by hand app. E.
• Median over prompts is not GPT-3's protocol. GPT-3 reports one prompt per dataset, which the authors believe is a manually tuned best prompt §6.1 fn. 1. The T0 medians cover all prompts, so the comparison is not like for like.
• Validation splits, not test. Every number is on validation or dev splits §5. Compare with the same split.
• LAMBADA was dropped. The unprompted form gave T5+LM 6.2% and T0 18.7%; the cloze form gave T0 27.8%, and the GPT-3 variant gave 40.5% app. D.3.
• Answer scoring has no length normalization. Long options are penalized when the option is scored by its summed log-probability §5 [code: t0/model.py#L61-L67].
• fp16 is not safe at inference. The code's data collator warns that fp16 predictions can differ greatly from fp32 data_collator.py.
• Prompts are external. The 2073 prompts are in the P3 collection and the PromptSource repository, not in the paper text §4 app. G. Reproducing the mixture needs both.
• Bias. The authors report that T0++ produces biased or false answers, for example "yes" to "Is the earth flat?" app. B.3 tab. 2. The training data excluded forum and social-media sources, but the models are not bias-free.
• Compute cost. Training used about 270 TPU v3-512 hours across all experiments app. B.1 tab. 1. A reproduction needs that hardware class, and the repo's training command is for TPUs.

Lineage

• Builds on: the T5 encoder-decoder §5 and the LM-adapted T5 of Lester et al. §5. Prompting as the input format builds on GPT-3 §2.
• Compared against: GPT-3 §6.1; the BIG-bench baselines §6.1.
• Concurrent work: FLAN ("Finetuned Language Models Are Zero-Shot Learners"), which trains decoder-only models, one held-out task per model §7. The paper names no arXiv id for it.
• Superseded by: not stated in the paper.