Source-linked AI summary
Specializing Smaller Language Models towards Multi-Step Reasoning
Yao Fu, Hao Peng, Litu Ou, Ashish Sabharwal, Tushar Khot
TL;DR
Large language models’ chain-of-thought math reasoning was thought to require very large models, raising an accessibility problem for smaller systems. The paper specializes smaller T5 models by distilling teacher-generated reasoning and selecting for out-of-distribution math performance. It reports improved small-model math reasoning with a corresponding loss in generic abilities and a log-linear scaling curve after specialization.
Problem
Chain-of-thought math reasoning was believed to emerge mainly in models larger than 100B parameters, limiting access to this capability for smaller-model researchers and practitioners.
Method
The paper fine-tunes FlanT5 using correct GSM8K chain-of-thought solutions generated by GPT-3.5 code-davinci-002, aligns tokenizers dynamically, and selects models using three held-out math datasets.
Results
+10 average accuracy gain lifts CoT scaling for 250M, 760M, and 3B FlanT5 models across four math reasoning tasks, while specialized models lose generic BigBench Hard ability.
Takeaways & Limitations
Concentrating smaller models’ capacity on math reasoning can produce a log-linear scaling curve and move ability from generic tasks toward a specialized target.
Abstract
from arXiv · showhide
The surprising ability of Large Language Models (LLMs) to perform well on complex reasoning with only few-shot chain-of-thought prompts is believed to emerge only in very large-scale models (100+ billion parameters). We show that such abilities can, in fact, be distilled down from GPT-3.5 ($\ge$ 175B) to T5 variants ($\le$ 11B). We propose model specialization, to specialize the model's ability towards a target task. The hypothesis is that large models (commonly viewed as larger than 100B) have strong modeling power, but are spread on a large spectrum of tasks. Small models (commonly viewed as smaller than 10B) have limited model capacity, but if we concentrate their capacity on a specific target task, the model can achieve a decent improved performance. We use multi-step math reasoning as our testbed because it is a very typical emergent ability. We show two important aspects of model abilities: (1). there exists a very complex balance/ tradeoff between language models' multi-dimensional abilities; (2). by paying the price of decreased generic ability, we can clearly lift up the scaling curve of models smaller than 10B towards a specialized multi-step math reasoning ability. We further give comprehensive discussions about important design choices for better generalization, including the tuning data format, the start model checkpoint, and a new model selection method. We hope our practice and discoveries can serve as an important attempt towards specialized smaller models in the new research paradigm set by LLMs.
1. Introduction
The paper asks whether smaller models can acquire chain-of-thought math reasoning, an ability associated with very large models. It specializes smaller models toward this target and reports improved reasoning while trading away generic abilities.
- Smaller models’ access to multi-step math reasoning would broaden research and practical access to chain-of-thought abilities.
- Model specialization concentrates limited small-model capacity on a target task instead of spreading it across generic abilities.
- +10 average accuracy gain lifted CoT scaling for 250M, 760M, and 3B FlanT5 models across four math reasoning tasks.The tasks include one in-distribution and three out-of-distribution datasets.
- The approach distills GPT-3.5 chain-of-thought paths into FlanT5 and selects models using three held-out math reasoning datasets.The selection procedure is intended to support out-of-distribution generalization.
- The study frames specialization as a tradeoff in which gains in math reasoning can require decreased generic-task ability.It emphasizes both multi-dimensional ability tradeoffs and the design choices affecting generalization.
2. Background
The background presents language models as having multiple abilities that can be elicited and specialized. It positions this work among studies of specialized models, distillation, and smaller-model chain-of-thought reasoning.
- Large Language Models’ Abilities: Modern language-model training commonly proceeds from pretraining to instruction tuning, which elicits abilities from a base model.The paper identifies this sequence as the standard recipe for models such as GPT-3, PaLM, and OPT.
- Specialized Language Models: Specialized models focus on particular abilities, with code-davinci-002 and text-davinci-003 presented as examples of differing model emphases.
- Specialized Language Models: Small models may need to trade generic abilities for one special ability because their model capacity is limited.The paper contrasts this with large models, which can retain broad competence while specializing.
- Specialized Language Models: The central generalization challenge is that small specialized models may overfit tuning data and struggle when the data distribution shifts.
- Distillation and Data Augmentation: The paper uses teacher-generated data as a tool for specialization rather than presenting a direct contribution to distillation or data-augmentation research.It notes that specialization data could also come from human annotation.
- Most closely related works: Compared with related work, this study focuses on small models and explicitly examines the price paid for improving a target ability.Related work includes generic smaller-model CoT and large-model self-improvement.
3. Specializing Multi-Step Reasoning
The method specializes T5-family models for math chain-of-thought reasoning using teacher-generated GSM8K solutions and distribution matching. It addresses tokenizer mismatch with dynamic-programming alignment and evaluates generalization beyond the tuning data.
- The experiments use GSM8K as seed data, evaluate on MultiArith, ASDiv, and SVAMP, and measure generic reasoning with BigBench Hard.The additional math datasets test out-of-distribution generalization, while BigBench Hard measures generic reasoning ability.
- Distillation from Code-Davinci-002: Code-davinci-002 generates 40 chain-of-thought solutions per training question, and correct-answer solutions become specialization data.Each retained solution contains an answer and intermediate reasoning steps.
- Distillation from Code-Davinci-002: The study compares standard chain-of-thought outputs with in-context answer-only, in-context chain-of-thought, and zero-shot answer-only formats.In-context chain-of-thought adds reasoning to both demonstrations and the output.
- Distribution matching trains the student against the teacher’s per-step output distributions by minimizing their KL divergence.The paper chooses this objective because it reports faster convergence and better performance than sample matching.
- Aligning tokenizers by dynamic programming: Dynamic programming aligns GPT and T5 token sequences by minimizing the total edit cost between them.The alignment permits one-to-many and many-to-one token matches.
- Aligning tokenizers by dynamic programming: One-to-one token mappings reuse GPT distributions, whereas non-one-to-one mappings assign one-hot T5 distributions.
4. Experiments
The experiments specialize smaller T5-family models for multi-step math reasoning by distilling GPT-3.5 chain-of-thought data, revealing improved math performance alongside tradeoffs in generic abilities and generalization.
- Experimental setup: The study specializes raw T5 and instruction-tuned FlanT5 checkpoints across multiple model scales using augmented GSM8K chain-of-thought data.Evaluation covers in-distribution GSM8K, out-of-distribution MultiArith, ASDiv, and SVAMP, plus generic BigBench Hard tasks.
- Overall performance: +10 average accuracy gain follows GSM8K augmentation across all FlanT5 models.The specialized 3B model also outperforms current 11B and 6B distillation models on the GSM8K test set.
- Overall performance: 40.8 average accuracy on four math datasets is achieved by specialized FlanT5 11B versus 35.9 for LaMDA 137B.The specialized 11B model is on par with LaMDA 137B and slightly below PaLM 60B, while paying a clear BigBench Hard performance cost.
- Scaling behavior: Specialization changes smaller models’ scaling curves from flat to log-linear for math chain-of-thought performance.The result appears for directly CoT-trained raw T5 checkpoints and for specialized instruction-tuned FlanT5 models.
- Specialization dynamics: At early specialization checkpoints, BBH CoT accuracy falls to 0 while BBH answer-only accuracy drops from about 0.3 to about 0.1.Later tuning increases both GSM8K and combined MultiArith-ASDiv-SVAMP performance, but math performance can fluctuate afterward.
- Model selection: Selecting checkpoints on GSM8K does not necessarily maximize MultiArith-ASDiv-SVAMP performance, so validation should match the deployment goal.The authors recommend GSM8K validation for in-distribution generalization and an appropriate held-out set for out-of-distribution generalization.
- Distillation objective: Distribution matching converges faster with lower loss than sample matching, although their validation performances do not differ substantially.The faster convergence may matter when larger-model tuning is expensive.
- Data format: Training with in-context examples preserves both in-context and zero-shot ability, whereas zero-shot-only training increases zero-shot performance but gradually loses in-context learning.This demonstrates a second ability tradeoff beyond math versus generic reasoning.
5. Conclusion
The paper shows that specialization can redirect smaller models’ abilities toward multi-step math reasoning, producing log-linear scaling. This improvement involves tradeoffs in generic and generalization abilities.
- Specialization concentrates smaller models’ abilities from generic directions onto the target math reasoning task.
- After specialization, model performance increases smoothly with model scale, yielding a log-linear scaling curve.
- Instruction-tuned checkpoints generalize better than raw pretrained checkpoints as specialization base models.
- Specialization trades off math reasoning gains against losses in BBH performance and changes in generalization behavior.The reported balances include in-distribution versus out-of-distribution generalization and in-context versus zero-shot ability.