Source-linked AI summary
Teaching Small Language Models to Reason
Lucie Charlotte Magister, Jonathan Mallinson, Jakub Adamek, Eric Malmi, Aliaksei Severyn
TL;DR
The paper addresses whether reasoning capabilities associated with large language models can be transferred to smaller models. It generates chain-of-thought data with large teachers and finetunes smaller students on those outputs. The method improves task performance across arithmetic, commonsense, and symbolic reasoning datasets, including a GSM8K increase for T5 XXL from 8.11% to 21.99% with PaLM 540B-generated CoT.
Problem
CoT prompting improves large models, but reasoning capabilities do not reliably emerge in models below 100 billion parameters.
Method
The method generates CoT annotations for supervised datasets with large teacher models and finetunes smaller student models on the resulting data.
Results
The method improves task performance across arithmetic, commonsense, and symbolic reasoning datasets; T5 XXL GSM8K accuracy rises from 8.11% to 21.99% with PaLM 540B-generated CoT.
Takeaways & Limitations
CoT knowledge distillation can improve smaller models’ task accuracy, although the magnitude of improvement is task dependent.
Takeaways & Limitations
Experiments are limited to English, one task at a time, non-public model access, large-scale inference, and the original CoT approach without self-consistency.
Abstract
from arXiv · showhide
Chain of thought prompting successfully improves the reasoning capabilities of large language models, achieving state of the art results on a range of datasets. However, these reasoning capabilities only appear to emerge in models with a size of over 100 billion parameters. In this paper, we explore the transfer of such reasoning capabilities to models with less than 100 billion parameters via knowledge distillation. Specifically, we finetune a student model on the chain of thought outputs generated by a larger teacher model. Our experiments show that the proposed method improves task performance across arithmetic, commonsense and symbolic reasoning datasets. For example, the accuracy of T5 XXL on GSM8K improves from 8.11% to 21.99% when finetuned on PaLM-540B generated chains of thought.
1 Introduction
Chain-of-thought prompting improves reasoning in large language models but does not reliably transfer to smaller models. This work asks whether knowledge distillation can transfer those capabilities by finetuning smaller students on larger teachers’ generated reasoning.
- CoT prompting breaks reasoning tasks into intermediate steps and improves large language models across commonsense, symbolic, and mathematical datasets.
- Smaller language models generally produce illogical CoT, and prompting can reduce accuracy below 10 billion parameters.
- The central research question is whether LLM reasoning capabilities can be transferred to smaller LMs through finetuning.
- 21.99% GSM8K accuracy was achieved by T5 XXL after finetuning on PaLM 540B-generated CoT, up from 8.11%.Using GPT-3 175B-generated CoT produced 18.42% accuracy.
2 Related Work
Prior work established CoT prompting and later extensions, while concurrent studies explored related teacher–student approaches. This paper distinguishes itself through broader comparisons of teachers, datasets, and model and dataset sizes.
- CoT prompting uses exemplars to encourage intermediate reasoning, and self-consistency improves accuracy by selecting the majority-voted reasoning path.
- Concurrent studies explored CoT distillation through diverse sampling, rationalization prompting, separate decomposition and solving models, and markup-and-mask explanations.
- This work compares more teacher models and studies how student model size and dataset size affect accuracy across multiple datasets.
3 Method
The proposed pipeline generates chain-of-thought annotations for existing supervised datasets and finetunes smaller language models on them. Teacher prompts include the target answer to improve the quality of generated reasoning.
- The method has two steps: annotate an existing supervised dataset with teacher-generated CoT, then finetune a smaller student model on that data.
- PaLM 540B and GPT-3 175B serve as teachers because CoT reasoning improves with model scale.
- Few-shot prompts use eight exemplars and provide the target after the question but before example CoT.The target guidance is intended to allow correction of small reasoning mistakes.
4 Experimental Setup
The experiments evaluate distilled reasoning on arithmetic, commonsense, and symbolic tasks using large teachers and T5 students of different sizes. Accuracy is measured from final answers, with calculator-assisted evaluation additionally accounting for arithmetic errors.
- The benchmark covers arithmetic, commonsense, and symbolic reasoning tasks.
- Arithmetic evaluation uses GSM8K, MAWPS, and ASDiv with official or cross-validation splits and checks whether the final CoT answer matches the target.
- Calculator-assisted accuracy is also computed to account for arithmetic mistakes despite correct chains of thought.
- StrategyQA is evaluated using an 80% training, 10% validation, and 10% testing split because no testing split is available.
- Symbolic evaluation includes last-letter concatenation and coinflip state tracking, with emphasis on out-of-distribution generalizability.
- PaLM 540B and GPT-3 175B are teacher models, while differently sized T5 models are students trained on their generated CoT data.
5 Results
The proposed CoT distillation method improves T5 performance across arithmetic, commonsense, and symbolic reasoning tasks, while results vary with teacher, student size, and data amount. Ablations show benefits from target-conditioned CoT generation and substantial data and model-size effects.
- Arithmetic reasoning: The proposed method improves arithmetic task accuracy across all evaluated datasets, with calculator-assisted results approaching 8-shot PaLM 540B.The gains may remain limited by the mathematical abilities of small models.
- CoT generation ablation: 79.37% GSM8K accuracy with target-conditioned CoT generation exceeds 59.98% without the target.The analysis attributes most benefits to correcting CoT with a missing or incorrect step, rather than merely copying the answer.
- Commonsense reasoning: 71.98% StrategyQA accuracy with CoT finetuning improves over the 68.12% baseline using only 1319 of 1648 examples.The improvement is smaller than on arithmetic datasets, which the authors relate to missing factual knowledge in smaller models.
- Symbolic reasoning: CoT finetuning significantly improves Coinflip accuracy for three-coinflip OOD generalization but is slightly weaker than baseline for four coinflips.Both traditional finetuning and the proposed method fail to generalize Last Letter Concatenation to longer sequences.
- Teacher models: The method also improves results when GPT-3 175B supplies the teacher CoT, while the strongest teacher source differs by task.Golden CoT performs best on StrategyQA, whereas PaLM-generated CoT performs best on GSM8K.
- Model and dataset size: T5 base, with 44 times fewer parameters than T5 XXL, matches baseline T5 XXL on GSM8K after CoT training, while T5 small outperforms it with a calculator.With only 20% of the examples, the method reaches 11.22% versus the 8.11% baseline, but a quality gap remains.
6 Discussion
The discussion finds that CoT distillation can substantially improve task accuracy, sometimes with few generated examples, but gains depend on task requirements and support trade-offs between model size, dataset size, and accuracy.
- CoT data from models over 100 billion parameters can significantly improve smaller-model task accuracy.The discussion states that even a small number of CoT examples may suffice.
- Improvements are task dependent, with limited gains on StrategyQA attributed to smaller models lacking required factual knowledge.The authors also suggest that some residual improvement may reflect learning how to approach the task.
- The CoT distillation pipeline allows model size and dataset size to be traded off against accuracy.
7 Conclusion
The paper proposes transferring reasoning from very large language models to smaller models through CoT knowledge distillation and reports improved accuracy across multiple benchmark datasets.
- The two-step pipeline generates CoT for existing datasets with large language models, then finetunes smaller models on that CoT.
- Finetuning smaller models on generated CoT improves task accuracy across arithmetic, commonsense, and symbolic reasoning benchmarks.
8 Limitations
The experiments are limited to English and single-task settings, require access to non-public models and substantial inference, and evaluate only the original CoT prompting approach.
- Experiments cover only English and one task at a time, leaving other languages and multitask settings unexplored.
- Replicating the results requires access to non-public models and inference over large amounts of data.
- The work evaluates the original CoT prompting approach but not later improvements such as self-consistency.
9 Ethical Considerations
The paper identifies ethical concerns arising from text generation by both teacher and student models, including potentially non-factual or offensive outputs.
- Teacher and student models may generate non-factual or offensive text during the text-generation process.The authors relate this risk to the input data, which consists of standard peer-reviewed NLP benchmark tasks.
A Dataset Usage and Licenses
This section documents dataset licensing, intended usage, symbolic-dataset generation, and the hardware and APIs used for experiments.
- Dataset licenses: GSM8K uses the MIT license, MAWPS uses CC BY 4.0, and ASDiv uses CC BY-NC 4.0.
- Dataset usage: The authors state that they follow the intended usage of the listed datasets.
- Dataset licenses: StrategyQA uses the MIT license, while its open-domain BIG-bench version uses the Apache License 2.0.
- Dataset construction: Symbolic reasoning datasets are generated as described in Wei et al. (2022).
- Experimental implementation: Inference and finetuning use different T5 sizes on TPUs, while PaLM 540B inference also runs on TPUs.
- Experimental implementation: Results can be replicated through public PaLM and GPT-3 APIs.