Source-linked AI summary

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

Cheng-Yu Hsieh, Chun-Liang Li, Chih-Kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alexander Ratner, Ranjay Krishna, Chen-Yu Lee, Tomas Pfister

arXiv:2305.02301v2cs.CLcs.AIcs.LG

TL;DR

Deploying LLMs is difficult because of their memory and compute requirements, while finetuning and distillation of smaller models require substantial training data. Distilling step-by-step extracts LLM rationales as additional supervision in a multi-task framework. Across four NLP benchmarks, it achieves stronger results with fewer examples and smaller models, including a 770M T5 surpassing a 540B-parameter LLM using 80% of a labeled dataset.

  • Problem

    LLMs are costly to deploy, while finetuning and distillation require large amounts of data to train smaller models to comparable performance.

  • Method

    The method extracts LLM-generated rationales and uses them alongside labels to train smaller task-specific models in a multi-task framework.

  • Results

    Across four NLP benchmarks, the method outperforms finetuning and distillation with fewer examples and outperforms LLMs with much smaller models, including 770M T5 versus a 540B-parameter LLM using 80% of labeled data.

  • Takeaways & Limitations

    Distilling step-by-step reduces both the model size and training data required to achieve or surpass the original LLM’s performance.

  • Takeaways & Limitations

    The approach requires approximately 10-shot user-provided demonstrations for few-shot CoT prompting and adds slight training-time computation overhead.

Abstract

from arXiv · show

Deploying large language models (LLMs) is challenging because they are memory inefficient and compute-intensive for practical applications. In reaction, researchers train smaller task-specific models by either finetuning with human labels or distilling using LLM-generated labels. However, finetuning and distillation require large amounts of training data to achieve comparable performance to LLMs. We introduce Distilling step-by-step, a new mechanism that (a) trains smaller models that outperform LLMs, and (b) achieves so by leveraging less training data needed by finetuning or distillation. Our method extracts LLM rationales as additional supervision for training small models within a multi-task framework. We present three findings across 4 NLP benchmarks: First, compared to both finetuning and distillation, our mechanism achieves better performance with much fewer labeled/unlabeled training examples. Second, compared to few-shot prompted LLMs, we achieve better performance using substantially smaller model sizes. Third, we reduce both the model size and the amount of data required to outperform LLMs; our finetuned 770M T5 model outperforms the few-shot prompted 540B PaLM model using only 80% of available data on a benchmark, whereas standard finetuning the same T5 model struggles to match even by using 100% of the dataset. We release the code at: https://github.com/google-research/distilling-step-by-step .

1 Introduction

Distilling step-by-step addresses the deployment cost of LLMs and the data demands of training smaller models by using LLM-generated rationales as additional supervision. Across four NLP benchmarks, it improves performance with fewer examples and much smaller models, including a 770M T5 surpassing a 540B-parameter LLM with 80% of a labeled dataset.

  • 1 Introduction: 175 billion-parameter LLMs require at least 350GB of GPU memory, making low-latency deployment unaffordable for many product teams.State-of-the-art LLMs can exceed 500B parameters, increasing memory and compute requirements.
  • 1 Introduction: Traditional finetuning uses human-annotated data, while distillation uses LLM-generated labels, and both approaches can require substantial training data.Distillation may require large amounts of unlabeled data that are difficult to obtain.
  • 1 Introduction: Distilling step-by-step extracts LLM-generated rationales and trains smaller models on both label prediction and rationale prediction in a multi-task setup.Rationales provide task knowledge that may otherwise require many examples for small models to learn.
  • 1 Introduction: Over 50% fewer training examples on average, and up to over 85% fewer, produced better performance than traditional finetuning and distillation across datasets.The method also reduces the model size needed to outperform LLMs.
  • 1 Introduction: 770M T5 surpassed 540B-parameter LLM performance using only 80% of a labeled dataset.The comparison illustrates simultaneous reductions in model size and training data.

2 Related work

Prior work uses distillation, human rationales, and LLM-generated reasoning to transfer knowledge or improve predictions, but Distilling step-by-step uses rationales as supervision for smaller deployable models. The paper studies data size, downstream model size, and fully unlabeled datasets.

  • 2 Related work: Knowledge distillation transfers knowledge from larger teacher models into smaller models and can support learning from limited labeled data through pseudo-label generation.Prior approaches commonly train students on noisy labels generated by larger models.
  • 2 Related work: Human-generated rationales have been used to regularize behavior, guide predictions, improve performance, and support interpretability, but they are expensive.Their uses span both supervision and explanation generation.
  • 2 Related work: LLM-generated reasoning steps have been used to improve few-shot or zero-shot prompting and to provide additional finetuning data for self-improvement.Their utility for test-time deployment remains limited by LLM size.
  • 2 Related work: Distilling step-by-step instead uses generated rationales as informative supervision for smaller task-specific models that avoid large computation or memory costs at deployment.The approach distills rationales rather than only labels to reduce the need for large unlabeled datasets.
  • 2 Related work: The paper extends concurrent rationale-supervision work by varying training dataset size and downstream model size and evaluating fully unlabeled datasets.These studies examine both method effectiveness and deployment-oriented scaling choices.

3 Distilling step-by-step

Distilling step-by-step uses LLM-generated rationales as additional supervision to train smaller task-specific models through multi-task learning. The framework elicits rationales and labels from an LLM, then trains a smaller model to predict labels and generate rationales without requiring an LLM at deployment.

  • 3.1 Extracting rationales from LLMs: Distilling step-by-step prompts an LLM to generate labels and natural-language rationales for examples, then uses both outputs to train smaller downstream models.The method uses few-shot Chain-of-Thought prompting with demonstrations containing inputs, labels, and user-provided rationales.
  • 3.2 Training smaller models with rationales: Standard finetuning uses human-annotated labels, whereas task-specific distillation uses LLM-generated pseudo-labels to train smaller models with label-prediction loss.Both paradigms train the smaller model by minimizing cross-entropy between predicted and target tokens.
  • 3.2 Training smaller models with rationales: Feeding generated rationales as model inputs requires an LLM to produce a rationale before each prediction, limiting deployability.Distilling step-by-step instead learns rationale generation during training so the deployment model does not depend on an LLM.
  • 3.2 Training smaller models with rationales: The smaller model is trained as a multi-task system that predicts task labels and generates corresponding rationales from the input.Task prefixes can specify whether the model should produce a label or rationale.
  • 3.2 Training smaller models with rationales: Rationale-generation loss teaches intermediate reasoning steps that could guide the model toward better label predictions, while rationales are unnecessary at test time.This design removes the need for an LLM during deployment.

4 Experiments

Across four NLP benchmarks, Distilling step-by-step improves data efficiency and enables smaller task-specific models to outperform standard baselines and few-shot LLMs. It uses LLM-generated rationales in multi-task training and benefits from higher-quality rationales and additional unlabeled data.

  • Reducing training data: Distilling step-by-step achieves better performance than standard finetuning with substantially fewer labeled examples across datasets.It reduces the required training examples by 12.5% on e-SNLI, 75% on ANLI, 25% on CQA, and 20% on SVAMP when outperforming standard finetuning.
  • Reducing training data: Distilling step-by-step outperforms standard task distillation across all four datasets while using much less unlabeled data.On e-SNLI, it surpasses standard task distillation trained with 100% of the data using only 12.5% of the unlabeled dataset.
  • Reducing model size: Distilling step-by-step outperforms Few-shot CoT and PINTO Tuning with much smaller T5 models when human-labeled data is available.It exceeds 540B PaLM Few-shot CoT using 220M T5 on e-SNLI, 770M T5 on ANLI and SVAMP, and 11B T5 on CQA.
  • Reducing model size: With only unlabeled data, Distilling step-by-step surpasses the teacher LLM on three of four datasets and can exploit augmentation more effectively than standard distillation.On SVAMP, ASDiv augmentation closes the gap to Few-shot CoT for Distilling step-by-step, while standard task distillation still underperforms.
  • Minimum model size and data: Distilling step-by-step can reduce both model size and data requirements for surpassing LLMs.On e-SNLI, it exceeds Few-shot CoT with a 220M T5 model using 0.1% of the full dataset; on ANLI, it does so with a 45× smaller model and 50% of the unlabeled set, while standard methods often need more data or larger models.
  • Method and ablations: LLM-generated rationales are used as additional supervision in a multi-task setup, rather than treating rationale and label prediction as one task.The ablation finds that multi-task training consistently performs better, while single-task rationale training can underperform standard finetuning on ANLI and CQA.

5 Discussion

Distilling step-by-step extracts rationales from LLMs as informative supervision for smaller task-specific models. The approach reduces both the data and model size required to achieve or surpass the original LLM’s performance.

  • 5 Discussion: Distilling step-by-step extracts rationales from LLMs as informative supervision for training smaller task-specific models.The method uses rationales as additional supervision rather than relying only on predicted labels.
  • 5 Discussion: The method reduces the training dataset required to curate task-specific smaller models.This establishes a resource-efficient alternative to existing training-to-deployment methods.
  • 5 Discussion: The method reduces the model size required to achieve, and even surpass, the original LLM’s performance.The paper frames this reduction in model size as part of a resource-efficient training-to-deployment paradigm.

Limitations

The approach has limitations involving demonstration requirements, training computation, and the reasoning capabilities of teacher LLMs on complex tasks.

  • Limitations: The method requires approximately 10-shot user-provided demonstrations to use few-shot chain-of-thought prompting.The authors note that recent methods may elicit rationales without user-annotated demonstrations.
  • Limitations: Training task-specific models with rationales incurs slight training-time computation overhead.At test time, the multi-task design avoids this overhead by predicting labels without generating rationales.
  • Limitations: LLMs may have limited reasoning capability on more complex reasoning and planning tasks.The authors identify rationale quality as an area for future study.

Ethics statement

Downstream smaller models can inherit biases from the larger teacher LLM. The authors envision applying progress in reducing antisocial LLM behaviors to smaller language models.

  • Ethics statement: Downstream smaller models are subject to biases inherited from the larger teacher LLM.The authors connect this issue to efforts to reduce antisocial behaviors in LLMs.
  • Ethics statement: Progress in reducing antisocial behaviors in LLMs could also be applied to improve smaller language models.

A.1 Implementation

The implementation trains T5 models with specified learning rates, batch sizes, input lengths, and maximum training steps. Results are reported over four random runs with standard errors in the presented plots.

  • A.1 Implementation: T5-Base (220M) and T5-Large (770M) use learning rate = 5 × 10−5, batch size = 64, max input length = 1024, and maximum 10000 steps.
  • A.1 Implementation: T5-XXL (11B) uses learning rate = 5 × 10−5, batch size = 32, max input length = 1024, and maximum 4000 steps.
  • A.1 Implementation: All results are reported over 4 random runs.
  • A.1 Implementation: The presented plots include the standard error.

A.2 Datasets

The experiments use five datasets—e-SNLI, ANLI, CQA, SVAMP, and ASDiv—with documented sources and dataset statistics. Validation and test-set handling varies by dataset availability, including a 10% training subsample for validation when needed.

  • Data provenance: The datasets are described with their experiment sources and original author-released sources, with readers referred to those sources for licensing terms.
  • Data characteristics: To the authors’ knowledge, the datasets contain neither information naming or uniquely identifying individuals nor offensive content.
  • Datasets: The experiments use e-SNLI, ANLI, CQA, SVAMP, and ASDiv, obtained from publicly available dataset repositories.ANLI uses the R1 split, while CQA includes human-labeled explanations from prior work.
  • Dataset splits: 10% of each original training set is randomly subsampled for validation when no validation set is provided.The dataset statistics are reported in Table 3.
  • Dataset splits: CQA uses its original validation set as the test set because ground-truth labels are unavailable for its original test set.
Loading 2305.02301v2…