MiCoTA: Bridging the Learnability Gap with Intermediate CoT and Teacher Assistants

TL;DR

MiCoTA introduces intermediate-sized models and reasoning sequences to enhance small language models' ability to learn long-chain reasoning, significantly improving their performance on complex reasoning benchmarks.

Contribution

The paper proposes MiCoTA, a novel distillation framework using intermediate models and reasoning sequences to bridge the learnability gap for small language models.

Findings

SLMs distilled with MiCoTA outperform baseline models on reasoning benchmarks.

MiCoTA produces data more aligned with small model distributions.

Significant improvements in reasoning scores on multiple benchmarks.

Abstract

Large language models (LLMs) excel at reasoning tasks requiring long thought sequences for planning, reflection, and refinement. However, their substantial model size and high computational demands are impractical for widespread deployment. Yet, small language models (SLMs) often struggle to learn long-form CoT reasoning due to their limited capacity, a phenomenon we refer to as the "SLMs Learnability Gap". To address this, we introduce \textbf{Mi}d-\textbf{Co}T \textbf{T}eacher \textbf{A}ssistant Distillation (MiCoTAl), a framework for improving long CoT distillation for SLMs. MiCoTA employs intermediate-sized models as teacher assistants and utilizes intermediate-length CoT sequences to bridge both the capacity and reasoning length gaps. Our experiments on downstream tasks demonstrate that although SLMs distilled from large teachers can perform poorly, by applying MiCoTA, they achieve significant improvements in reasoning performance. Specifically, Qwen2.5-7B-Instruct and Qwen2.5-3B-Instruct achieve an improvement of 3.47 and 3.93 respectively on average score on AIME2024, AMC, Olympiad, MATH-500 and GSM8K benchmarks. To better understand the mechanism behind MiCoTA, we perform a quantitative experiment demonstrating that our method produces data more closely aligned with base SLM distributions. Our insights pave the way for future research into long-CoT data distillation for SLMs.