Exploring Knowledge Purification in Multi-Teacher Knowledge Distillation for LLMs

AI summary

Overview

• Research area: Natural Language Processing, specifically knowledge distillation of large language models (LLMs) and multi-teacher ensembles.
• Technical level: Intermediate. Readers need basic familiarity with knowledge distillation, multiple-choice reasoning benchmarks, and LLM routing, but the paper's central idea is intuitive.
• Scope: The paper proposes "knowledge purification" — merging the rationales produced by several teacher LLMs into one consolidated rationale — and evaluates five concrete purification methods against single-teacher and multi-teacher distillation baselines across commonsense and biomedical reasoning tasks.

What This Paper Is About

Distilling a small LLM from several strong teacher LLMs sounds better than distilling from one, but the teachers often disagree, hallucinate, or reason along incompatible paths, and the student inherits that conflict. The paper shows that with the TinyLLM framework, adding more teachers actually lowers the student's accuracy on OpenBookQA, ARC, RiddleSense and PubMedQA. The authors' goal is to fix this by condensing all teachers' rationales into a single purified rationale before distillation, and to find which purification strategy works best.

Key Contributions

1. Diagnosis of multi-teacher failure modes. The authors identify two obstacles in existing multi-teacher distillation frameworks: knowledge conflict among teachers (which grows with teacher count) and high resource demands (complex sampling and training pipelines).
2. The knowledge purification concept. They formalize replacing the set of teacher rationales with a single consolidated rationale, and rewrite the multi-teacher distillation objective around that consolidated rationale.
3. Five purification methods from three perspectives. Knowledge Aggregation (an LLM aggregator), three LLM routers (Plackett-Luce ranking, PLM classifier, similarity-based router, following RouterDC), and an RL-based Teacher Selection policy trained with policy gradients on a reward derived from student loss.
4. Empirical validation plus two new analyses. Beyond accuracy, they define Conflict Mitigation Value (CMV) to measure conflict reduction as teachers are added, and they test router-guided distillation on out-of-domain datasets (PIQA, BioASQ), where routing shows strong generalization.

Main Findings

• More teachers hurt the baseline, not help it. Extended experiments with TinyLLM showed distillation performance declining as the number of teacher LLMs increased, for the 77M, 248M and 783M students, contrary to the expectation that a bigger ensemble would improve the student.
• Routing-based purification leads on in-domain accuracy. For the FLAN-T5 small (77M) student, the similarity-based router reached the highest average accuracy of 45.66%, exceeding baselines by at least 4.9% as reported. For the FLAN-T5 base (248M) and FLAN-T5 large (783M) students, RL-based Teacher Selection performed best with average accuracies of 56.68% and 67.55%, surpassing the best baseline by 4.5% and 6.9% respectively as reported.
• Aggregation underperforms. Knowledge Aggregation showed "no significant improvement" and had the weakest overall results, despite using a strong aggregator (GPT-4); the authors state that the enhancing effect of a GPT-4-consolidated rationale on distillation "remains uncertain."
• Conflict Mitigation Value separates the methods. Across the three student sizes, Knowledge Aggregation produced negative CMV (−0.003, −0.007, −0.004). All routers and RL-based selection were positive: Plackett-Luce Ranking +0.001, +0.012, +0.010; PLM Classifier +0.018, +0.014, +0.021; Similarity-based Router +0.025, +0.020, +0.032; Teacher Selection +0.020, +0.019, +0.029. The similarity-based router had the highest CMV for all three students.
• The distilled student can beat most teachers. The distilled 783M student exceeded the average accuracy of three of the four teacher LLMs and ranked second only to Llama-3.1-8B-Instruct.
• Gains scale with student size. Knowledge purification helped larger students more, which the authors attribute to larger models being better able to learn from generated rationales, while smaller models mostly fit the final option.
• Routers generalize out of domain. Excluding aggregation and RL-based selection (limited transferability), the routers were tested on PIQA and BioASQ. The similarity-based router achieved the highest accuracy in most settings, and Plackett-Luce ranking outperformed the PLM classifier overall. Example: for the 77M student on PIQA, the similarity-based router scored 53.97 versus TinyLLM's 49.84; for the 783M student on BioASQ it scored 91.87 versus TinyLLM's 82.93.
• Routers are also the efficiency winners. Distilling FLAN-T5 large on ARC over a fixed 4000 epochs, total GPU hours were: Fine-tuning 0.7, Distilling-Step-by-Step 1.1, TinyLLM 2.6, Plackett-Luce Ranking 1.4, PLM Classifier 1.7, Similarity-based Router 1.8, Knowledge Aggregation 6.7, and RL-based Teacher Selection listed at 3.5. Training the router cost fewer resources than the distillation stage itself.
• Practical trade-offs differ sharply by method. Per the paper's practical comparison: Knowledge Aggregation needs the question and all rationales as prior input, more than 10B extra parameters, no training, is transferable, and has second-scale latency. Plackett-Luce Ranking needs only the question, about 278M parameters, no training, transfers to new datasets, and is second-scale. The PLM Classifier and Similarity-based Router need only the question, about 278M parameters, require training, transfer, and run at millisecond latency. RL-based Teacher Selection needs the question and rationales, about 278M parameters, requires training tied to distillation rewards, is not transferable without retraining, and has minute-scale latency.

Methodology in Plain English

The setup is multiple-choice question answering. Each teacher LLM is asked to produce a rationale for a question, and the student is trained both on the correct answer (prediction loss) and on the teacher's rationale (distillation loss), combined by a weight lambda.

Instead of feeding the student every teacher's rationale, the authors insert a purification step that outputs a single rationale. Three families of purifiers are tried. The first is aggregation: a powerful LLM (GPT-4) is prompted, with an in-context example, to write one rationale from all of them. The second is routing: train a small model to pick which teacher's rationale to use for each question — either by ranking teachers with a Plackett-Luce model learned via cross-entropy weighted by question similarity, by classifying the question embedding from a CLS token through a two-layer perceptron, or by computing cosine similarity between the question embedding and trainable per-teacher embeddings as in RouterDC. The third is reinforcement learning: a policy scores each teacher's state (question embedding plus its rationale embedding, gated by whether that teacher answered correctly) and selects the teacher with the highest predicted score; the reward is the negative of the student's losses, so the selector and the student are trained alternately.

Data is split so that 80% of each training set is used for distillation and the remaining 20% forms a public set; the public sets of all four datasets are combined into a joint dataset used to train the routers. Training used AdamW with learning rate 5×10⁻⁵, batch size 8, maximum input length 512, and lambda = 4, on four NVIDIA A100 80GB GPUs. Teachers were FLAN-T5 xlarge (2.85B), Llama 2-chat (7B), BioMistral-7B and Llama-3.1-8B-Instruct; students were FLAN-T5 small (77M), base (248M) and large (783M). mDeBERTaV3-base served as the language encoder for the PLM classifier, the similarity-based router and the RL teacher selector.

Why This Matters

Impact on research. The paper reframes multi-teacher distillation from "collect everything" to "condense first," and it provides a measurable quantity (CMV) for knowledge conflict rather than treating it as a qualitative complaint. It also connects the LLM routing literature to distillation, showing that a router trained only on questions can replace expensive pre-sampling of teacher outputs.

Real-world applications.

• Deploying small, cheap models for commonsense question answering on edge or on-premise hardware, using a distilled 77M–783M student instead of a 7–8B teacher.
• Domain-specialized assistants, such as biomedical question answering, where one teacher (BioMistral-7B) supplies domain knowledge and others supply general reasoning.
• Cost-controlled pipelines that route queries to the most suitable expert model rather than always invoking the largest one.
• Rapid adaptation of a distillation pipeline to new data, since routers require only the question as input and can guide rationale sampling without first querying every teacher.

Industry relevance. The efficiency table is the practical argument: routers cut total GPU hours relative to TinyLLM (1.4–1.8 versus 2.6 in the reported ARC setting), while aggregation is far more expensive (6.7 total, with the purification-stage figure footnoted as measured when using Llama-3.1-70b as the aggregator). For teams constrained by compute, that difference decides whether multi-teacher distillation is viable at all. The work is supported by the National Natural Science Foundation of China under Grant U2436210.

Future Directions

• Scale the teacher ensemble. The authors note that limited computational resources restricted them to four teachers, and that they could only run a small-scale evaluation with six teachers (Appendix D.3); larger ensembles are needed to fully test purification, especially for domain specialization.
• Extend beyond NLP multiple-choice tasks. The methods are tailored to LLM characteristics; the authors state that routing and teacher selection may generalize to broader machine learning tasks but that implementation and evaluation remain for future work.
• Make RL-based selection transferable. Its CMV is strong, but it is not transferable without retraining and carries minute-scale latency; closing that gap is an open engineering question.
• Understand why aggregation fails. A strong GPT-4 aggregator did not reliably help and produced negative CMV, leaving open how to prompt or train an aggregator that genuinely consolidates rather than dilutes teacher knowledge.

Target Audience

Researchers and graduate students working on model compression, distillation and LLM ensembles; practitioners who need to build smaller task models from several strong teachers under compute constraints; and readers interested in LLM routing or reinforcement-learning-based model selection, who will find the practical comparison table and out-of-domain results the most directly reusable part of the paper.

Authors’ abstract

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.