Table of Contents
Fetching ...

Multistage Collaborative Knowledge Distillation from a Large Language Model for Semi-Supervised Sequence Generation

Jiachen Zhao, Wenlong Zhao, Andrew Drozdov, Benjamin Rozonoyer, Md Arafat Sultan, Jay-Yoon Lee, Mohit Iyyer, Andrew McCallum

TL;DR

This work tackles semi-supervised sequence generation under extreme data scarcity, showing that a student distilled from a few-shot prompted LLM can generalize better than the LLM itself. It introduces Multistage Collaborative Knowledge Distillation (MCKD), a cross-partition, multi-stage KD framework that iteratively refines pseudolabels by training paired students on disjoint unlabeled partitions. Across constituency parsing and semantic parsing tasks, MCKD achieves substantial gains, outperforming the LLM teacher and vanilla KD, and matching or exceeding supervised finetuning with many more labeled examples (e.g., 7.5% F1 gain on CRAFT with 50 labels, aligning with 500-label SFT). The method demonstrates strong label efficiency and scalable improvements through staged distillation, with practical implications for low-resource domains and specialized structured outputs.

Abstract

We study semi-supervised sequence generation tasks, where the few labeled examples are too scarce to finetune a model, and meanwhile, few-shot prompted large language models (LLMs) exhibit room for improvement. In this paper, we present the discovery that a student model distilled from a few-shot prompted LLM can commonly generalize better than its teacher to unseen examples on such tasks. We find that the student is able to learn a general pattern from the high-quality pseudolabels produced by the teacher during knowledge distillation (KD), and favorably not a general pattern from the low-quality pseudolables. Leveraging this discovery, we propose a new method, Multistage Collaborative Knowledge Distillation from an LLM (MCKD), for these tasks. MCKD first few-shot prompts an LLM to produce pseudolabels for unlabeled data. Then at each stage of an iterative KD process, a new pair of students is trained on disjoint partitions of the pseudolabeled data, and produces new and improved pseudolabels for their unseen partitions. We conduct extensive experiments on four syntactic and semantic parsing datasets and show the effectiveness of MCKD for low-resource semi-supervised sequence generation. On CRAFT biomedical parsing, for example, 3-stage MCKD with 50 labeled examples outperforms an LLM teacher and vanilla KD by 7.5% and 3.7% parsing F1, respectively, and matches the performance of supervised finetuning with 500 labeled examples.

Multistage Collaborative Knowledge Distillation from a Large Language Model for Semi-Supervised Sequence Generation

TL;DR

This work tackles semi-supervised sequence generation under extreme data scarcity, showing that a student distilled from a few-shot prompted LLM can generalize better than the LLM itself. It introduces Multistage Collaborative Knowledge Distillation (MCKD), a cross-partition, multi-stage KD framework that iteratively refines pseudolabels by training paired students on disjoint unlabeled partitions. Across constituency parsing and semantic parsing tasks, MCKD achieves substantial gains, outperforming the LLM teacher and vanilla KD, and matching or exceeding supervised finetuning with many more labeled examples (e.g., 7.5% F1 gain on CRAFT with 50 labels, aligning with 500-label SFT). The method demonstrates strong label efficiency and scalable improvements through staged distillation, with practical implications for low-resource domains and specialized structured outputs.

Abstract

We study semi-supervised sequence generation tasks, where the few labeled examples are too scarce to finetune a model, and meanwhile, few-shot prompted large language models (LLMs) exhibit room for improvement. In this paper, we present the discovery that a student model distilled from a few-shot prompted LLM can commonly generalize better than its teacher to unseen examples on such tasks. We find that the student is able to learn a general pattern from the high-quality pseudolabels produced by the teacher during knowledge distillation (KD), and favorably not a general pattern from the low-quality pseudolables. Leveraging this discovery, we propose a new method, Multistage Collaborative Knowledge Distillation from an LLM (MCKD), for these tasks. MCKD first few-shot prompts an LLM to produce pseudolabels for unlabeled data. Then at each stage of an iterative KD process, a new pair of students is trained on disjoint partitions of the pseudolabeled data, and produces new and improved pseudolabels for their unseen partitions. We conduct extensive experiments on four syntactic and semantic parsing datasets and show the effectiveness of MCKD for low-resource semi-supervised sequence generation. On CRAFT biomedical parsing, for example, 3-stage MCKD with 50 labeled examples outperforms an LLM teacher and vanilla KD by 7.5% and 3.7% parsing F1, respectively, and matches the performance of supervised finetuning with 500 labeled examples.
Paper Structure (37 sections, 8 figures, 4 tables)

This paper contains 37 sections, 8 figures, 4 tables.

Figures (8)

  • Figure 1: Overview of MCKD. (1) We use demonstrations from labeled data $\mathcal{D}^\text{labeled}$ to few-shot prompt an LLM teacher $t$ to produce pseudolabels for unlabeled data $\mathcal{D}^\text{unlabeled}$. We partition $\mathcal{D}^\text{unlabeled}$ into $\mathcal{D}^\text{unlabeled}_A$ and $\mathcal{D}^\text{unlabeled}_B$, and let $\mathcal{D}^{(x, \hat{y}_0)}_A$ and $\mathcal{D}^{(x, \hat{y}_0)}_B$ denote the same partitions but with teacher-generated pseudolabels. (2) At the $i$-th intermediate KD stage, students $s_A^i$ and $s_B^i$ are trained on previously pseudolabeled data $\mathcal{D}_A^{(x,\hat{y}_{i-1})}$ and $\mathcal{D}_B^{(x,\hat{y}_{i-1})}$, respectively, and leveraged to label the other partitions $\mathcal{D}_B^\text{unlabeled}$ and $\mathcal{D}_A^\text{unlabeled}$ and produce $\mathcal{D}_B^{(x,\hat{y}_i)}$ and $\mathcal{D}_A^{(x,\hat{y}_i)}$, which will be used to train the next-stage student(s). (3) In the final KD stage, a single final student $s^n$ is trained on both latest pseudolabeled partitions $\mathcal{D}^{(x, \hat{y}_{n-1})}_A$ and $\mathcal{D}^{(x, \hat{y}_{n-1})}_B$.
  • Figure 2: The test F1 scores of supervised finetuned (SFT) models increase with more gold annotated data. MCKD needs much less labeled data to match the performance of SFT.
  • Figure 3: The held-out performance of the stage-1 student (F1 against gold annotations) improves with its training performance (F1 against teacher pseudolabels). Its held-out performance also surpasses that of the few-shot prompted LLM teacher.
  • Figure 4: The student achieves a high F1 on held-out clean pseudolabels from an LLM teacher, but a low F1 on held-out noisy pseudolabels from the same teacher. The student tends to learn a general distribution over clean teacher pseudolabels while learning noisy ones via instance-wise memorization.
  • Figure 5: Performance (F1) on subsets of the held-out set $\mathcal{D}_{B}$ on which the teacher performs the best ($\mathcal{D}_{high}$) and the worst ($\mathcal{D}_{low}$). With sufficient amounts of pseudolabeled data, the student exhibits high agreement with the teacher on $\mathcal{D}_{high}$ while outperforming on $\mathcal{D}_{low}$.
  • ...and 3 more figures