PaperlayerView exact revision on arXiv

Source-linked AI summary

Exploring Knowledge Purification in Multi-Teacher Knowledge Distillation for LLMs

Ruihan Jin, Pengpeng Shao, Zhengqi Wen, Jinyang Wu, Mingkuan Feng, Shuo Yang, Chu Yuan Zhang, Jianhua Tao

arXiv:2602.01064v2cs.CL

TL;DR

Multi-teacher distillation is limited by conflicting rationales and high resource demands. The paper introduces Knowledge Purification, which consolidates teacher rationales through five methods, and reports improved distillation performance and conflict mitigation, with routing methods generalizing well out of domain.

  • Problem

    Multi-teacher distillation faces conflicting teacher rationales and high resource demands that hinder effective knowledge transfer.

  • Method

    Knowledge Purification consolidates rationales from multiple teacher LLMs into one rationale for distillation, using five methods including aggregation, routing, and RL-based selection.

  • Results

    Knowledge purification methods improve distilled-model performance and mitigate conflicts across commonsense and biomedical reasoning tasks; router-based methods show strong out-of-domain performance.

  • Takeaways & Limitations

    Knowledge Purification supports more efficient multi-teacher distillation, while LLM routing shows potential for broader generalization.

  • Takeaways & Limitations

    The study primarily focuses on NLP and LLMs, while broader machine-learning applications require further investigation.

Abstract

from arXiv · show

Knowledge distillation has emerged as a pivotal technique for transferring knowledge from stronger large language models (LLMs) to smaller, more efficient models. However, traditional distillation approaches face challenges related to knowledge conflicts and high resource demands, particularly when leveraging multiple teacher models. In this paper, we introduce the concept of \textbf{Knowledge Purification}, which consolidates the rationales from multiple teacher LLMs into a single rationale, thereby mitigating conflicts and enhancing efficiency. To investigate the effectiveness of knowledge purification, we further propose five purification methods from various perspectives. Our experiments demonstrate that these methods not only improve the performance of the distilled model but also effectively alleviate knowledge conflicts. Moreover, router-based methods exhibit robust generalization capabilities, underscoring the potential of innovative purification techniques in optimizing multi-teacher distillation and facilitating the practical deployment of powerful yet lightweight models.

1 INTRODUCTION

The paper motivates multi-teacher distillation by identifying knowledge conflicts and resource demands, then introduces Knowledge Purification to consolidate teacher rationales and improve distillation.

  • Large LLMs provide strong capabilities, but their computational cost motivates lightweight models that preserve those capabilities.
  • Multi-teacher distillation can suffer from conflicting rationales and escalating resource demands, hindering knowledge transfer.Conflicts may arise from hallucinations, divergent reasoning paths, or differing expertise domains.
  • As the teacher ensemble grows, TinyLLM performance declines rather than improves, indicating detrimental inter-teacher knowledge conflict.
  • Knowledge Purification consolidates multiple teacher rationales into one rationale used for distillation, mitigating conflicts and improving efficiency.
  • Five purification methods span aggregation, routing, and reinforcement-learning-based selection, with experiments showing improved distillation and conflict mitigation.Routing methods also show strong performance on out-of-domain datasets.

2 RELATED WORK

Related work positions multi-teacher distillation as a way to broaden student knowledge and routing as a way to activate suitable LLMs efficiently.

  • Multi-Teacher Knowledge Distillation: Multi-teacher distillation transfers broader knowledge and richer reasoning paths than single-teacher distillation, supporting student generalization.TinyLLM specifically trains small students from rationales generated by two teacher LLMs.
  • LLM Routing: LLM routing selects among candidate models for each question, aiming to activate ensembles efficiently while balancing cost and quality.Prior work explores hybrid and dynamic routing between strong and weak LLMs.

3 FORMULATION

The formulation defines rationale-based multiple-choice distillation, motivates purification from scaling problems, and expresses purification as consolidating teacher rationales before distillation.

  • 3.1 PRELIMINARIES: In multiple-choice QA, an LLM selects the correct option and generates rationales from a question and candidate-option set.Rationales are encouraged because they can enhance LLM performance.
  • 3.1 PRELIMINARIES: The answering formulation uses separate prompts for predicting options and generating rationales.
  • 3.1 PRELIMINARIES: Rationale-based distillation trains the student with ground-truth options and teacher rationales, balancing prediction and distillation losses with λ.
  • 3.1 PRELIMINARIES: Multi-teacher distillation expands the training data with rationales from n teachers and weights each teacher-specific distillation loss by λ_j.
  • 3.2 MOTIVATION ANALYSIS: The proposed purification methods include knowledge aggregation, LLM routing, and RL-based teacher selection.These categories comprise five methods defined for knowledge purification.
  • 3.2 MOTIVATION ANALYSIS: Increasing the teacher ensemble can reduce TinyLLM performance while increasing computational and hyperparameter-tuning challenges.
  • 3.3 KNOWLEDGE PURIFICATION: Knowledge Purification maps the set of teacher rationales R to a single consolidated rationale rP through a purification function f(·).The consolidated rationale is then used in distillation, with the resulting loss denoted LDL-KP.

4 METHODOLOGY

The methodology defines knowledge purification through five methods that consolidate or select teacher rationales before distillation. These methods include aggregation, routing, and reinforcement-learning-based teacher selection.

  • Five knowledge-purification methods are defined to consolidate teacher rationales before distillation.The methods are introduced as alternatives for implementing knowledge purification.
  • Knowledge Aggregation: Knowledge aggregation uses a global LLM to combine all teacher-generated rationales into one consolidated rationale.Instruction-tuning prompts with in-context examples guide the aggregation process.
  • LLM Routing: LLM routing assigns each question to an appropriate candidate LLM and selects one rationale using router-predicted probabilities.Unlike aggregation, routing selects rather than combines rationales.
  • LLM Routing: Three routing methods are proposed: Plackett-Luce ranking, a PLM classifier, and a similarity-based router.The PL model ranks candidate teachers, while the classifier and similarity router use learned representations for routing.
  • Plackett-Luce ranking: The Plackett-Luce router models candidate-rationale selection probabilities with a softmax relationship and learns coefficients using cross-entropy loss.Question similarity weights can measure similarity between an input question and database questions.
  • RL-based Teacher Selection: RL-based teacher selection dynamically chooses teachers using question and rationale features, selecting the teacher with the highest prediction score.The selector is optimized with policy gradients, while the reward combines rationale-prediction and distillation losses; training alternates between distillation and RL.

5 EXPERIMENTS

Experiments evaluate knowledge purification across student models, teacher ensembles, practical criteria, conflict mitigation, and out-of-domain transfer. Routing and RL-based selection generally perform strongly, while aggregation shows weaker conflict mitigation and uncertain gains.

  • Experimental Setup: Experiments use four teacher LLMs, three FLAN-T5 student sizes, and four multiple-choice datasets spanning commonsense and biomedical reasoning.
  • Performance: 45.66% average accuracy is achieved by the similarity-based router when distilling FLAN-T5 small, exceeding baselines by at least 4.9%.
  • Practical Analysis: Routing methods require only the question as input, while RL-based selection requires retraining on new datasets and incurs substantially greater processing delay.The PLM classifier and similarity-based router operate at millisecond-level delay, unlike RL-based teacher selection.
  • Conflict Mitigation: Knowledge aggregation reports negative CMV for all three student models, whereas routing and RL-based selection report positive CMV.The similarity-based router achieves the highest CMV across all three student models.
  • Out-of-Domain Transfer: LLM routers generalize to PIQA and BioASQ, with the similarity-based router achieving the highest accuracy across most settings.PL ranking outperforms the PLM classifier overall and demonstrates robust generalization.
  • Efficiency: Knowledge aggregation has the highest GPU consumption, while RL-based teacher selection also incurs considerable computational consumption from iterative training.

6 LIMITATIONS

The study’s scope is constrained by limited teacher-model coverage and its focus on NLP and LLMs. Broader validation remains future work.

  • Four teacher LLMs were used because of limited computational resources, restricting assessment of knowledge purification with larger ensembles.A small-scale six-teacher evaluation was conducted, but further validation is needed.
  • The methods are tailored to LLM characteristics and were evaluated primarily in NLP, while broader machine-learning applications remain uninvestigated.LLM routing and teacher selection may generalize, but their implementation and evaluation require further investigation.

7 CONCLUSION

The paper addresses conflicts and resource demands in multi-teacher distillation through Knowledge Purification, which consolidates teacher rationales. Experiments report improved distilled-model performance, conflict mitigation, and strong router performance on out-of-domain data.

  • Knowledge Purification consolidates rationales from multiple teacher LLMs to reduce divergent reasoning paths and improve distillation efficiency.The approach targets knowledge conflicts and high resource demands in multi-teacher frameworks.
  • Five purification methods are proposed from distinct perspectives to address multi-teacher distillation challenges.
  • Experiments across commonsense and biomedical reasoning tasks show improved distilled-model performance and effective mitigation of knowledge conflicts.
  • LLM-router methods show exceptional performance on out-of-domain datasets, supporting their applicability within the studied setting.

B.1 KNOWLEDGE AGGREGATION

Knowledge aggregation uses a prompted aggregator to combine teacher rationales, while routing and selection methods choose or learn teacher contributions. The methods differ in training requirements, ranking objectives, and inference-time computation.

  • Knowledge Aggregation: GPT-4 aggregates multiple teacher rationales using an instruction-tuned generation prompt with one randomly selected in-context example from ten labeled samples.Llama-3.1-70b is also evaluated as a weaker aggregator using the same prompt format.
  • LLM Routing: Plackett-Luce routing learns coefficients for optimal rationale selection, while inference performs ranking without additional training.The optimal rationale is defined as the shortest description that still produces correct selections, favoring efficiency and cost reduction.
  • LLM Routing: Similarity-based routing uses question and LLM embeddings with sample-LLM and sample-sample contrastive losses to improve routing robustness.The question and LLM embeddings have dimension 768 in the implementation described.
  • Teacher Selection: The RL-based teacher selector alternates student distillation with selector training and jointly uses the question and rationale during optimization.

C.1 DETAILS OF MODEL SELECTION

Experiments use FLAN-T5 students, a four-model teacher ensemble, six multiple-choice datasets, and comparisons against established distillation baselines. Conflict mitigation is evaluated through performance-based CMV rather than rationale-level information metrics.

  • Experimental Setup: FLAN-T5 small, base, and large serve as students, while the teacher ensemble contains FLAN-T5 xlarge, Llama 2-chat, BioMistral-7B, and Llama-3.1-8B-Instruct.The ensemble combines homogeneous, heterogeneous, biomedical-specialized, and updated-model roles.
  • Experimental Setup: The evaluation covers four in-domain datasets—OBQA, ARC, RiddleSense, and PubMedQA—and two out-of-domain datasets, PIQA and BioASQ.
  • Baselines: Baselines include direct inference, fine-tuning, Distilling-Step-by-Step, and TinyLLM, with additional evaluation against the single-teacher ABKD method.Standard knowledge distillation is excluded because the authors consider it functionally similar to fine-tuning for LLM distillation.
  • Baselines: The study emphasizes comparisons with multi-teacher methods because prior work reported stronger cross-task adaptability for multi-teacher distillation.TwT is discussed but excluded as both a baseline and a knowledge-purification method for stated methodological reasons.
  • Conflict Evaluation: Conflict Mitigation Value measures average accuracy improvement from purification over TinyLLM across incremental teacher-ensemble sizes.The metric is performance-based because rationales are intermediate states and rationale-level metrics do not apply uniformly to routing.

D.1 EXTENDED EXPERIMENTS OF THE TINYLLM FRAMEWORK

Extended experiments show that TinyLLM performance declines as more teacher LLMs are added, contrary to the expectation that larger ensembles improve students. The best overall performance occurs with two teachers, motivating knowledge purification experiments.

  • TinyLLM distillation incorporates rationales from multiple teacher LLMs through a loss function weighted by each teacher's importance.The framework is extended beyond its original two-teacher setting to study incremental teacher additions.
  • As the number of teacher LLMs increases, distilled student-model performance declines rather than improving.The experiments add teachers incrementally from an initial FLAN-T5 xlarge teacher.
  • Across all three student models, the best overall performance is achieved with two teacher LLMs.With four teachers, FLAN-T5 base performs worse than with a single teacher.
  • The extended experiments evaluate knowledge distillation with increasing teacher counts using the same setup described in Appendix D.1.Complete results are reported in Tables 10–14.

D.3 GENERALITY TOWARD BROADER TASK DOMAINS AND MORE TEACHER LLMS

The paper evaluates knowledge purification across broader task domains and larger teacher ensembles using MMLU and supplementary experiments. Router-based purification generalizes strongly, while the study remains bounded by computational limits and planned future evaluations.

  • D.3 GENERALITY TOWARD BROADER TASK DOMAINS AND MORE TEACHER LLMS: The broader-domain evaluation uses MMLU's 57 tasks and expands the teacher ensemble with two additional LLMs.The existing assessment focuses mainly on commonsense and biomedical reasoning with four teachers.
  • D.3 GENERALITY TOWARD BROADER TASK DOMAINS AND MORE TEACHER LLMS: 65.19% average accuracy is achieved by the similarity-based router, surpassing baselines by at least 7.3%.On MMLU, a routing-based method reaches 55.26% accuracy.
  • D.3 GENERALITY TOWARD BROADER TASK DOMAINS AND MORE TEACHER LLMS: Routing-based methods attain the strongest reported results in the broader evaluation, supporting their generalization across task domains and larger teacher ensembles.The reported evaluation includes commonsense reasoning, biomedical reasoning, and multitask language understanding.
  • D.3 GENERALITY TOWARD BROADER TASK DOMAINS AND MORE TEACHER LLMS: The study cannot yet test knowledge purification with eight, ten, or more teachers because of computational resource limitations.The authors identify broader NLP applications and more advanced distillation sampling as future evaluation targets.
  • D.3 GENERALITY TOWARD BROADER TASK DOMAINS AND MORE TEACHER LLMS: The paper includes a detailed OBQA case study and visualization of the knowledge purification process.Figure 5 presents an example involving five proposed purification methods.

E THE USE OF LARGE LANGUAGE MODELS

The paper states that LLMs were not used for conceptualization or drafting, while presenting an OBQA example containing teacher-generated rationales. The section documents post-drafting language editing as the only reported LLM use.

  • The authors state that no LLMs were used during the paper's original conceptualization or drafting.They attribute the core content and findings to their own research and critical evaluation.
  • LLM use was limited to post-drafting linguistic refinement, including grammar, typo, and sentence correction.
  • The OBQA example presents a multiple-choice question asking which animal only eats plants.The listed options are rat, moth, chimpanzee, and pig.
  • The example displays rationales generated by FLAN-T5 xlarge, Llama 2-chat, BioMistral-7B, and Llama-3.1-8B.The figure is described as an example using five proposed knowledge-purification methods.
Loading 2602.01064v2…