[Paper Review] Beyond Answers: Transferring Reasoning Capabilities to Smaller LLMs Using Multi-Teacher Knowledge Distillation

Yijun Tian, Yikun Han+3 more|arXiv (Cornell University)|Feb 7, 2024

Natural Language Processing Techniques4 citations

TL;DR

TinyLLM proposes a multi-teacher knowledge distillation framework that transfers reasoning capabilities from multiple large language models (LLMs) to smaller student models by leveraging rationale-guided training. By using an in-context example generator and teacher-forcing Chain-of-Thought, the student learns high-quality, contextually grounded rationales, achieving up to 12.57% performance gain over full fine-tuning while using only 1.1% of the teacher model size.

ABSTRACT

Transferring the reasoning capability from stronger large language models (LLMs) to smaller ones has been quite appealing, as smaller LLMs are more flexible to deploy with less expense. Among the existing solutions, knowledge distillation stands out due to its outstanding efficiency and generalization. However, existing methods suffer from several drawbacks, including limited knowledge diversity and the lack of rich contextual information. To solve the problems and facilitate the learning of compact language models, we propose TinyLLM, a new knowledge distillation paradigm to learn a small student LLM from multiple large teacher LLMs. In particular, we encourage the student LLM to not only generate the correct answers but also understand the rationales behind these answers. Given that different LLMs possess diverse reasoning skills, we guide the student model to assimilate knowledge from various teacher LLMs. We further introduce an in-context example generator and a teacher-forcing Chain-of-Thought strategy to ensure that the rationales are accurate and grounded in contextually appropriate scenarios. Extensive experiments on six datasets across two reasoning tasks demonstrate the superiority of our method. Results show that TinyLLM can outperform large teacher LLMs significantly, despite a considerably smaller model size. The source code is available at: https://github.com/YikunHan42/TinyLLM.

Motivation & Objective

• To address the limited knowledge diversity and lack of contextual reasoning in single-teacher knowledge distillation for small LLMs.
• To improve the generalization and reasoning capability of compact student models by learning from multiple diverse teacher LLMs.
• To enhance distillation quality by incorporating rich, contextually grounded rationales instead of relying solely on ground-truth labels.
• To develop a scalable and efficient distillation pipeline that maintains high performance with minimal parameter overhead.

Proposed method

• Utilizes multiple large LLMs (e.g., FLAN-T5 and LLaMA 2) as co-teachers to provide diverse reasoning perspectives and knowledge.
• Employs an in-context example generator to provide task demonstrations, enabling teachers to generate contextually appropriate rationales.
• Applies a teacher-forcing Chain-of-Thought strategy to ensure that teacher-generated rationales are accurate and aligned with the input context.
• Trains the student model via multi-task instruction tuning that jointly optimizes for answer accuracy and rationale consistency.
• Introduces a trade-off weighting mechanism (α) to balance the contribution of rationales from different teachers during distillation.
• Uses a lightweight training setup with LoRA and full fine-tuning baselines for comparison.

Experimental results

Research questions

• RQ1 Can knowledge distillation from multiple diverse LLMs improve reasoning generalization in smaller student models compared to single-teacher approaches?
• RQ2 Does incorporating rationale signals from multiple teachers lead to better performance than relying only on ground-truth answers?
• RQ3 How effective is the in-context example generator and teacher-forcing Chain-of-Thought in producing contextually grounded rationales for distillation?
• RQ4 To what extent can a small student model outperform its larger teacher models in reasoning tasks?
• RQ5 How sensitive is the performance to hyperparameters like the trade-off weight α for rationale integration?

Key findings

• TinyLLM achieves an average performance improvement of +12.57% over full fine-tuning on 780M-parameter student models, despite using only 1.1% of the 7B teacher model size.
• A 780M student model outperforms both the 3B and 7B teacher models by +14.56% and +23.40%, respectively, across six datasets.
• A 250M-parameter student model outperforms the 3B teacher by +0.70% and the 7B teacher by +16.82%, using only 8.3% and 3.6% of the parameters.
• Ablation studies confirm that removing in-context examples or teacher rationales leads to significant performance drops, validating their critical role.
• Hyperparameter analysis shows that optimal α values vary by task: higher α is beneficial for concise commonsense reasoning, while lower α suffices for complex biomedical questions.
• The method consistently improves performance across both commonsense (OBQA, ARC, PIQA, RiddleSense) and biomedical (PubMedQA, BioASQ) reasoning tasks.