Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes

The standard ways to get a small task model are fine-tuning on human labels or distillation: label unlabeled data with a big model and train on its labels. Both need lots of data to match the big model. This paper treats the big model as a source of reasoning, not just labels.

Step 1. Prompt the large model (540B PaLM) with about ten worked examples, each an input, a written rationale and a label, so that for every new input it writes a rationale and a label. For a word problem about carpeting a room, the rationale says area is length times width; the label is the expression.

Step 2. Train a small T5 model on two tasks over the same inputs. With the prefix [label] it learns to output the label; with [rationale] it learns to write the rationale. The loss is the label loss plus λ times the rationale loss. At inference only the label task runs, so the deployed model is as fast as any fine-tuned T5 and never needs the large model.

On four datasets (e-SNLI and ANLI for inference, CommonsenseQA, SVAMP word problems), the method beat standard fine-tuning and standard distillation at every data size and model size tested. It matched full-data fine-tuning with 12.5% of e-SNLI, 25% of ANLI, 75% of CQA and 80% of SVAMP. With human labels, T5 models from 220M to 11B beat PaLM's few-shot chain-of-thought on all four tasks; with only unlabeled data, on three of four. Putting rationale and label into one output sequence, instead of two tasks, sometimes did worse than no rationales.

A common deployment pattern: a customer prompts a large API model for a classification, extraction or routing task, it works, and then the bill and the latency arrive at production volume. The fix is a small fine-tuned model, but the customer has few labels. This paper is the recipe for that move.

Rationales stretch scarce labels. The large model is already being called; asking it to explain its answer costs extra output tokens per example, once, at training time. In return the small model needs a fraction of the labeled data, and at inference it emits only the label, so serving cost does not change.

The gain depends on the task. Rationales added 6.0 points on ANLI, a hard inference task, but about 1 point on e-SNLI and CQA. On e-SNLI, plain fine-tuning of a 220M model already beat PaLM's few-shot accuracy by about 15 points. Check whether the task needs reasoning before paying for rationales.

The student can beat the teacher. On e-SNLI and others, students trained only on PaLM's labels and rationales scored above PaLM's own few-shot accuracy. A small specialist trained on many teacher outputs can average out teacher mistakes; an FDE should measure on gold labels rather than assume the teacher is the ceiling.