Teaching-Assistant-in-the-Loop: Improving Knowledge Distillation from Imperfect Teacher Models in Low-Budget Scenarios (2024)

Abstract: There is increasing interest in distilling task-specific knowledge from large language models (LLM) to smaller student models.Nonetheless, LLM distillation presents a dual challenge: 1) there is a high cost associated with querying the teacher LLM, such as GPT-4, for gathering an ample number of demonstrations; 2) the teacher LLM might provide imperfect outputs with a negative impact on the student’s learning process. To enhance sample efficiency within resource-constrained, imperfect teacher scenarios, we propose a three-component framework leveraging three signal types. The first signal is the student’s self-consistency (consistency of student multiple outputs), which is a proxy of the student’s confidence. Specifically, we introduce a ”teaching assistant” (TA) model to assess the uncertainty of both the student’s and the teacher’s outputs via confidence scoring, which serves as another two signals for student training. Furthermore, we propose a two-stage training schema to first warm up the student with a small proportion of data to better utilize student’s signal. Experiments have shown the superiority of our proposed framework for four complex reasoning tasks. On average, our proposed two-stage framework brings a relative improvement of up to 20.79% compared to fine-tuning without any signals across datasets.

## introduction large language models (llms) have demonstrated state-of-the-art (sota) performance across a wide spectrum of tasks, and their efficacy is primarily attributed to their substantial model size, enabling them to possess capabilities that smaller models lack @xcite @xcite . in most cases, a pre-trained llm undergoes several specialized fine-tuning stages, such as instruction tuning, self-supervised tuning, and reinforcement learning with human feedback @xcite , in order to excel at downstream tasks. nevertheless, fine-tuning llms are challenging, primarily due to the demanding computational resources required @xcite and the limited access to llms, such as gpt-4 and claude @xcite prevents further fine-tuning. to resolve the above difficulties, knowledge distillation (kd) uses the outputs of a larger llm (teacher) to train a smaller model (student), such as gpt-j-6b or llama-7b @xcite . kd has gained significant attention and led to models such as alpaca @xcite and vicuna @xcite . the kd pipeline first utilizes in-context learning (icl) on the teacher model to generate outputs, forming distillation sets and then use the teacher generations to fine-tune the student model. however, current work of distilling llms presents two major difficulties: first, it can be prohibitively expensive to collect sufficiently large distillation sets, especially when querying proprietary llms like gpt-4, so a budget of asking the teacher is required in most real-life scenarios; second, as llms may not have seen task-specific data, the quality of their demonstrations might be low. for instance, the zero-shot accuracy of instructgpt on the aqua task is only 34.25% @xcite , and these suboptimal demonstrations negatively affect the student's performance, as elaborated in section 4. effectively excluding flawed annotations, especially for unlabeled sets, is essential. to address these kd challenges, we introduce a novel three-component kd framework for llms to learn efficiently from an imperfect teacher within budget constraints. in addition to the student and the teacher (as in standard kd pipeline), our framework introduces a teaching assistant (ta) model which mimics an ta in the real-life class and acts as an intermediary to communicate with students and teachers. the ta can have a larger size and different architecture than the student, but smaller size than the teacher, leading to significantly reduced inference cost compared to the teacher. with the help of the ta model, our framework utilizes three types of new signals to refine the distillation. first, the student internal signal (selfconsistency score) gauges the student's confidence and determines whether it should be "forwarded" to the teacher for annotation. second, the ta model independently generates its ta-student signal for a question, aiding in the decision to annotate. lastly, the ta model assesses annotations from the teacher (ta-teacher signal), deciding whether they merit inclusion in the student model's training dataset. moreover, the signal from the student model on a complex task is of lower quality for active selection. therefore, we propose an extension of the two-stage training. initially, we allocate 10% of the annotation budget to fine-tune the student. after this "warm-up" stage, we utilize the remaining 90% of the annotation budget. the student in both stages are fine-tuned within the proposed framework. our experiments demonstrate that this approach further improves performance in various tasks. our work introduces three main contributions: ## methodology figure 1 shows the overview of our framework. our three-component framework consists of two pipelines: 1): collect a fixed budget of examples that meet annotation criteria and add them to the bucket b a : the intuition here is that we want to choose the challenging examples which the student model has a chance to learn correctly instead of choosing examples randomly that will include oversimple or complicated examples beyond the student's learning ability 2): annotate the examples from b a and return a high-quality fine-tuning set b f : we want to filter out low-quality or wrong annotations generated by the teacher llm. we leverage the interactions among the student model m s , the teacher model m e , and the ta model m t to generate three types of signals and achieve the above goals. for a given example, we will use the student internal signal s i as well as the ta-student signal from the ta model s t to decide whether the example should be added to b a . after collecting the fixed budget of samples for b a , we ask the teacher model to annotate those examples and apply the ta-teacher signal s e from the ta model to filter out low-quality demonstrations. we introduce these three signals in detail below, and the algorithms outlining this procedure are also presented in algorithm 1. ## experiments through our extensive empirical analysis, we aim to address the following research questions: dataset since our work focuses on the scenario with sufficient unlabeled data but limited annotation budget, following the previous work @xcite , we evaluate our framework on four datasets with a large training set: hotpotqa @xcite , gsm8k @xcite , aqua @xcite and commonsenseqa @xcite . we choose hotpotqa, the closed-book question and answering (qa) dataset, for evaluation with react prompting @xcite . for gsm8k and aqua, the arithmetic reasoning task and commonsenseqa, the commonse sense reasoning task, we evaluate performance with cot prompting @xcite . in proposed framework we present two variants of the three-component framework with different signal weights as discussed in section 3.1.3: in our first framework, we use α = 1 and β = 0 (ta-finetune (i)); in our second framework (tafinetune (t), we use α = β = 1 to compute annotation criteria. both frameworks use the ta-teacher signal s e as eq ( 2) . throughout the experiments sections, we use "α = 1" and "fine-tune with s i ", "β = 1" and "fine-tune with s t " interchangeably. baseline models we design several baseline experiments to show the superiority of our proposed methods. the first group of baseline methods is the few-shot icl methods: student-icl, ta-icl, and teacher-icl with the corresponding prompting methods, that is, react for the hot-potqa dataset and cot for the others @xcite . the other baseline model is random-finetune, the classical sequence-level kd pipeline for llms @xcite @xcite , where the examples are sampled randomly until the budget is met, and all annotations from the teacher model are used for student fine-tuning. although there are other sequencelevel kd methods, some depend on ground-truth labels @xcite @xcite @xcite , not applicable in an unlabeled setting; those generate more augmented rationales @xcite @xcite or more augmented questions @xcite , which are not achievable with a fixed budget; and those incorporate additional loss @xcite , which is orthogonal to our kd method. the details of the model configuration are included in appendix a. ## conclusions in this paper, we propose a novel unified framework to resolve two challenges in the current teacherstudent knowledge distillation process. our framework utilizes a third ta model and signals from the student and teacher side to actively select samples and improve data quality in teacher annotations. additionally, to better utilize student signals, we also extend our framework to two-stage training. with a limited budget for teacher annotations on an unlabeled dataset, our extensive empirical evaluations show that the proposed framework can significantly increase kd results with signals and all proposed signals show its own effectiveness.