TL;DR
We propose PromptKD to distill student-friendly knowledge through prompt tuning and demonstrate for the first time that such knowledge is effective even in generative language models.
Figure. Comparison of instruction-following performance of KD methods using the GPT-2 model family. Owing to the student-friendly knowledge, our PromptKD outperforms others with only an additional 11K parameters. Dashed reference line represents the performance of the teacher model.
Abstract
Recent advancements in large language models (LLMs) have raised concerns about inference costs, increasing the need for research into model compression.
While knowledge distillation (KD) is a prominent method for this, research on KD for generative language models like LLMs is relatively sparse, and the approach of distilling student-friendly knowledge, which has shown promising performance in KD for classification models, remains unexplored in generative language models.
To explore this approach, we propose PromptKD, a simple yet effective method that utilizes prompt tuning - for the first time in KD - to enable generative language models to transfer student-friendly knowledge.
Unlike previous works in classification that require fine-tuning the entire teacher model for extracting student-friendly knowledge, PromptKD achieves similar effects by adding a small number of prompt tokens and tuning only the prompt with student guidance.
Extensive experiments on instruction-following datasets show that PromptKD achieves state-of-the-art performance while adding only 0.0007% of the teacher's parameters as prompts.
Further analysis suggests that distilling student-friendly knowledge alleviates exposure bias effectively throughout the entire training process, leading to performance enhancements.
Method Overview
Our PromptKD is a knowledge distillation method designed for instruction-following tasks in generative language models.
It introduces soft prompts that are prepended to the teacher's input, guiding the teacher to extract knowledge at a level similar to the student's for adaptive teaching.
Each iteration in the training phase consists of three steps.
First, the student generates responses that are treated as training data (pseudo-targets) to address exposure bias by incorporating the model's own outputs into training.
Second, the prompt is updated by minimizing the KL divergence loss, encouraging the teacher to generate responses at a similar level to the student.
Since the teacher's output distribution changes due to the emergence of prompts can lead to instability in the early stages of prompt training, apply a regularization loss here.
Finally, the updated prompt facilitates distillation by minimizing the discrepancy between teacher and student outputs, ensuring the student learns effectively from the teacher's student-friendly knowledge.
Main Results
Table. Evaluation results on 5 instruction-following datasets. †Results surpass those of the teacher.
PromptKD demonstrates state-of-the-art performance in instruction-following tasks across five datasets, outperforming other knowledge distillation baselines and showcasing strong generalization ability on four additional datasets not used in training.
It consistently surpasses baseline methods across different model sizes and families, including outperforming MiniLLM, which additionally uses language modeling loss computed by the corpus used for pre-training (PromptKD does not use this loss).
Notably, PromptKD is the only method to outperform the teacher on all datasets, highlighting the effectiveness of modifying the teacher to extract student-friendly knowledge for generation tasks.
Analysis
Figure. The measurement of exposure bias. Excess accumulated error (ExAccErr) is measured with respect to generation steps and training progress, where values closer to 0 indicate alleviation of exposure bias.
The above figure illustrates ExAccErr over generation steps and its variation up to 50 steps during training.
In part (a), ExAccErr increases for most methods as generation length grows, indicating accumulated exposure bias.
GKD reduces this error by using student-generated responses but still shows a gap for the teacher's oracle response.
PromptKD, however, maintains near-zero ExAccErr, demonstrating its superior ability to alleviate exposure bias.
In part (b), ExAccErr is measured across early training stages, where PromptKD, MiniLLM, and GKD consistently show lower ExAccErr than other methods, with PromptKD maintaining the most stable and near-zero values, effectively reducing exposure bias throughout training.
Table. Qualitative results of generated response from the Dolly validation set with and without using prompts for the Llama-13B teacher. A teacher with a prompt generates a response more similar to that of the student.
To understand how the prompt affects the teacher model, we compared the responses generated by the teacher with and without the prompt, alongside the student model's output.
The results show that the original teacher generates more complex responses, while the prompt-prepended teacher produces simpler, student-friendly responses that closely align with the student's output.
Despite being simpler, these responses remain accurate, demonstrating how the prompt effectively tailors the teacher's output to be more comprehensible for the student.
Additionally, quantitative results also confirm that the teacher's output becomes more similar to the student's while maintaining sentence quality. For more details, please refer to the paper.
BibTeX
If you find our work useful, please cite our paper:
@inproceedings{kim-etal-2024-promptkd,
title = "PromptKD: Distilling Student-Friendly Knowledge for Generative Language Models via Prompt Tuning",
author = "Kim, Gyeongman and Jang, Doohyuk and Yang, Eunho",
booktitle = "Findings of the Association for Computational Linguistics: EMNLP 2024",
year = "2024"
}