Source-linked AI summary

The CoT Collection: Improving Zero-shot and Few-shot Learning of Language Models via Chain-of-Thought Fine-Tuning

Seungone Kim, Se June Joo, Doyoung Kim, Joel Jang, Seonghyeon Ye, Jamin Shin, Minjoon Seo

arXiv:2305.14045v2cs.CLcs.AIcs.LG

TL;DR

Smaller language models perform poorly on chain-of-thought reasoning for unseen tasks, while rationale data remains limited. The paper introduces the CoT Collection and fine-tunes Flan-T5 with it, producing CoT-T5. CoT-T5 improves zero-shot and few-shot performance over Flan-T5, including gains on BBH and domain-specific tasks.

  • Problem

    Smaller LMs do not reliably benefit from CoT prompting on unseen tasks, and existing rationale resources cover only a limited number of tasks.

  • Method

    The paper augments the Flan Collection with rationales across 1,060 tasks and fine-tunes Flan-T5 to create CoT-T5.

  • Results

    +4.34% and +2.60% average zero-shot accuracy improvements are reported for CoT-T5 3B and 11B over Flan-T5 across 27 BBH datasets.

  • Takeaways & Limitations

    CoT fine-tuning with the CoT Collection improves smaller LMs’ zero-shot generalization and few-shot adaptation capabilities.

  • Takeaways & Limitations

    CoT-T5 is not trained for dialogue with long-form responses and therefore does not necessarily support chat applications.

Abstract

from arXiv · show

Language models (LMs) with less than 100B parameters are known to perform poorly on chain-of-thought (CoT) reasoning in contrast to large LMs when solving unseen tasks. In this work, we aim to equip smaller LMs with the step-by-step reasoning capability by instruction tuning with CoT rationales. In order to achieve this goal, we first introduce a new instruction-tuning dataset called the CoT Collection, which augments the existing Flan Collection (including only 9 CoT tasks) with additional 1.84 million rationales across 1,060 tasks. We show that CoT fine-tuning Flan-T5 (3B & 11B) with CoT Collection enables smaller LMs to have better CoT capabilities on unseen tasks. On the BIG-Bench-Hard (BBH) benchmark, we report an average improvement of +4.34% (Flan-T5 3B) and +2.60% (Flan-T5 11B), in terms of zero-shot task accuracy. Furthermore, we show that instruction tuning with CoT Collection allows LMs to possess stronger few-shot learning capabilities on 4 domain-specific tasks, resulting in an improvement of +2.24% (Flan-T5 3B) and +2.37% (Flan-T5 11B), even outperforming ChatGPT utilizing demonstrations until the max length by a +13.98% margin. Our code, the CoT Collection data, and model checkpoints are publicly available.

1 Introduction

Smaller language models do not reliably gain the benefits of chain-of-thought prompting on unseen tasks, motivating broader rationale-based instruction tuning. The CoT Collection addresses this gap, and CoT-T5 improves zero-shot and few-shot performance over Flan-T5.

  • Smaller LMs often fail to benefit from CoT prompting, unlike models exceeding 100B parameters, creating computational and accessibility concerns.
  • The CoT Collection adds 1.84 million rationales across 1,060 tasks to support CoT instruction tuning.
  • +4.34% and +2.60% average accuracy improvements over Flan-T5 are reported for CoT-T5 3B and 11B in zero-shot BBH evaluation.The comparison covers 27 BBH datasets under CoT prompting.
  • +2.24% and +2.37% average improvements over Flan-T5 are reported for CoT-T5 3B and 11B across four domain-specific few-shot tasks.The four datasets are LEDGAR, Case Hold, MedNLI, and PubMedQA, with 64 randomly chosen instances each.
  • The authors examine CoT fine-tuning in lower-compute settings using fewer tasks or multilingual subsets of the collection.The ablations use 163 tasks for T0 and 80K instances translated into five languages.

2 Related Works

Prior work establishes CoT prompting and instruction tuning as tools for task generalization, but their benefits are less reliable for smaller models. This work responds with large-scale rationale-based instruction tuning across many task variants.

  • Chain-of-Thought Prompting: CoT prompting generates a rationale before the answer and improves reasoning for challenging tasks in large language models.
  • Chain-of-Thought Prompting: CoT prompting’s effectiveness does not necessarily transfer to smaller LMs, motivating instruction tuning on large amounts of rationales.
  • Instruction Tuning: Instruction tuning has been studied for generalization to unseen tasks, cross-lingual transfer, label generalization, and modular expert models.
  • Few-Shot Adaptation: A simple recipe combines LoRA with CoT fine-tuning on a target task to obtain strong few-shot performance.
  • Rationale Data: The collection expands rationale coverage from 9 publicly available CoT datasets to 1.84 million rationales and 1,060 task variants.

3 The COT COLLECTION

The CoT Collection creates rationale-augmented instruction data from diverse tasks, using large-model in-context generation and manually designed task-family demonstrations. The pipeline applies selection and filtering criteria to make the data suitable for training.

  • Dataset Scope: The collection addresses scarce rationale data by augmenting 1.84 million rationales across 1,060 tasks.
  • Rationale Augmentation: Rationales are generated for existing instances with instructions, inputs, and answers rather than by generating new task instances.
  • Source Dataset Selection: The source Flan Collection contains 1,836 diverse NLP tasks, from which 1,060 tasks are selected.
  • Source Dataset Selection: The selection excludes long-output, non-public, and input-output-misaligned datasets.Long outputs are excluded because appending rationale and answer can exceed the 512-token training limit.
  • Filtering: Tasks producing very short, uninformative rationales are excluded to prevent negative transfer during multitask learning.
  • Creating Demonstrations for ICL: Task-family prompts share 6–8 manually created demonstrations, with three annotators producing and comparing rationales from sampled instances.Demonstrations are organized across 26 task groups sharing formats such as multiple-choice QA and closed-book QA.
  • Rationale Augmentation: OpenAI Codex performs rationale augmentation through in-context learning, with generated rationales conditioned on task demonstrations and answers.For multiple-choice QA, the prompt conditions rationale generation on the ground-truth label.

4 Experiments

The experiments evaluate CoT-T5, created by CoT fine-tuning Flan-T5 on the CoT Collection, across zero-shot generalization, multilingual adaptation, and few-shot domain-specific tasks. Results show improvements over Flan-T5 and several larger or instruction-tuned baselines, with limitations around cross-lingual transfer.

  • Experimental setup: CoT-T5 is obtained by CoT fine-tuning Flan-T5 on the CoT Collection and evaluated in zero-shot and few-shot settings.CoT fine-tuning trains models to generate rationales before answers using the trigger phrase “Let’s think step by step.”
  • Zero-shot generalization: CoT-T5 improves average BBH CoT accuracy over Flan-T5 by +4.34% for 3B and +2.60% for 11B models.The BBH evaluation covers 27 unseen datasets in the zero-shot setting.
  • Zero-shot generalization: CoT fine-tuning enables smaller models to outperform approximately four-times-larger T0-11B and Tk-Instruct-11B models in Direct and CoT Evaluation.The authors attribute the broader gains to training on diverse tasks and combining Flan and CoT Collection data.
  • Multilingual adaptation: CoT fine-tuning produces gains across five languages, including mathematical reasoning in Korean, Japanese, and Chinese using 60k–80k language-specific instances.These experiments train and evaluate within a single target language rather than testing cross-lingual transfer.
  • Few-shot generalization: CoT-T5 yields stronger few-shot adaptation than Flan-T5 and outperforms ChatGPT and Claude baselines using maximum-length demonstrations.The comparison covers four domain-specific datasets, while the authors note that long inputs make in-context learning computationally expensive.

5 Analysis of of CoT Fine-tuning

The analysis tests whether CoT fine-tuning benefits more from broader task coverage or additional instances, and whether it preserves in-domain performance. Results favor diverse task coverage, while in-domain gains are observed under a setting that limits the catastrophic-forgetting test.

  • The analysis compares expanding task diversity against increasing instances and examines whether CoT fine-tuning preserves in-domain accuracy.The research questions address task-versus-instance scaling and possible catastrophic forgetting.
  • Task and instance scaling: 10K instances across 1,060 tasks outperform 180K instances across 9 tasks on BBH, indicating broader task coverage is more crucial than instance count.The comparison uses a baseline trained on the existing 9 CoT tasks and a CoT-T5 configuration sampled from the CoT Collection.
  • In-domain task accuracy: CoT-T5 consistently improves in-domain accuracy on the five evaluated learned tasks.The evaluation uses ANLI-R1, ANLI-R2, ANLI-R3, RTE, and Winogrande.
  • In-domain task accuracy: The in-domain result may reflect using the same tasks for Flan-T5 and CoT-T5 fine-tuning, leaving broader catastrophic-forgetting effects unresolved.The authors state that adding tasks not used to train either model could produce different results and defer further exploration.

6 Conclusion

The paper constructs a large CoT rationale collection and fine-tunes Flan-T5 to improve smaller language models’ reasoning. CoT-T5 shows better zero-shot generalization and serves as a stronger base model when trained with few instances.

  • COT COLLECTION contains 1.84M CoT rationales extracted across 1,060 NLP tasks.
  • CoT fine-tuning Flan-T5 produces CoT-T5, which improves zero-shot generalization and few-instance training performance.

Limitations

The limitations concern chat applicability, multilingual step-by-step reasoning, and reproducibility and rationale quality tied to the augmentation models. These boundaries leave broader dialogue, language, and data-construction questions for future work.

  • COT-T5 is not necessarily suitable for chat applications because it was not trained on long-form dialogue responses.
  • Because COT-T5 uses Flan-T5 as its base, it lacks step-by-step reasoning in diverse languages.The authors report nearly zero accuracy on math problems in different languages for smaller models, with only slight improvement from CoT fine-tuning.
  • Reproducibility is constrained by the shutdown of Codex, the proprietary LLM used for rationale augmentation.The paper also identifies more powerful augmentation models and improved prompting as opportunities to improve dataset quality.

A Analysis of COT COLLECTION

The analysis examines rationale quality, diversity, and reproducibility in COT COLLECTION. It reports diverse formats, strong quality under ROSCOE metrics, and high-quality rationales from multiple LLMs.

  • COT COLLECTION rationales are analyzed for quality, diversity, and reproducibility against human-authored or existing CoT rationales.
  • Diversity of Rationales: COT COLLECTION includes more diverse textual rationale formats than the 9 existing CoT datasets.The comparison parses rationales and contrasts their structural formats.
  • Quality of Rationales: COT COLLECTION rationales are reported as faithful, less repetitive, informative, and logical, including in comparison with human-authored rationales.The evaluation uses 13 label-free ROSCOE metrics.
  • Reproducibility: Different LLMs produce high-quality rationales in terms of semantic alignment and language coherence.The reproducibility comparison includes Bard, Claude, and OpenAI Codex.

B Filtering COT COLLECTION

COT COLLECTION filters augmented rationales to remove inconsistent, overlong, duplicate, repetitive, and code-contaminated outputs while retaining high-quality reasoning.

  • The filtering stage is applied after generating multiple rationales to ensure high-quality instances.
  • Filtering removes rationales that omit the ground-truth answer, exceed the combined rationale-and-answer limit of 512 tokens, duplicate earlier outputs, or contain repetitive sentences.
  • Additional trigger-token filtering excludes outputs where Codex begins generating code after the rationale.
  • Nucleus sampling with p =0.8 and no_repeat_n_gram =3 was effective for generating good-quality rationales.

C Training and Evaluation Details of CoT-T5

CoT-T5 is trained on COT COLLECTION using task-mixture sampling and one-epoch fine-tuning, with separate optimizer choices for the 3B and 11B models. The section also documents training settings, rationale examples, and analyses of rationale diversity and augmentation quality.

  • Training: CoT-T5-3B uses AdamW and CoT-T5-11B uses Adafactor; both models train for 1 epoch, requiring 1 day and 7 days respectively on 8 A100 GPUs.
  • Training: Training samples combine FLAN, P3, existing CoT, SNI, and other dialogue-and-code datasets using proportions of 23.94%, 30.85%, 7.89%, 25.47%, and 11.85%, respectively.
  • Analysis: Figure 5 compares the top 20 common root verbs and their top 4 noun objects between 9 existing CoT tasks and COT COLLECTION’s 1,060 tasks.
  • Analysis: Figure 6 evaluates rationale augmentation quality across Semantic Alignment, Semantic Similarity, Logical Inference, and Coherence using reference-free ROSCOE scores.
  • Examples: Table 8 provides rationale examples spanning Multiple choice QA, Extractive QA, Closed-book QA, Formal Logic, NLI, and Arithmetic.
  • Filtering analysis: Table 9 presents filtered rationale examples where missing ground-truth tokens can indicate inconsistent generation and degeneration, especially in math datasets.

D Prompts used for Rationale Augmentation

Rationale augmentation prompts use task-family demonstrations and varied instruction formats, including multiple-choice, process, dialogue, and open-ended questions. The examples show rationales connecting contextual evidence to selected answers.

  • Prompt demonstrations: Each task family uses shared demonstrations, with 26 demonstrations used for each family during rationale augmentation and 4 examples included in the described prompt set.
  • Multiple-choice prompts: Multiple-choice prompts pair an instruction and question with answer options, then elicit a rationale and selected answer.
  • Process prompts: Process prompts present ordered steps and ask how an altered condition, such as no available water, affects the outcome.
  • Dialogue prompts: Dialogue prompts ask models to infer relationships or conclusions from conversational context and available choices.
  • Other prompt formats: Other prompts ask direct questions about settings, sports injuries, generators, and fermentation, with rationales explaining the selected answer.
  • Prompt demonstrations: Prompt examples also include demonstrations organized explicitly by task family, such as MCQA, with separate instruction, options, answer, and rationale fields.
Loading 2305.14045v2…