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

Abstract: There is increasing interest in distilling taskspecific 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.

heading: 1 Introduction
text: 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 (Brown et al., 2020; Raffel et al., 2020; Chowdhery et al., 2022; Touvron et al., 2023). 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 (Touvron et al., 2023; OpenAI, 2023), in order to excel at downstream
tasks. Nevertheless, fine-tuning LLMs are challenging, primarily due to the demanding computational resources required (Wei et al., 2022; Chowdhery et al., 2022) and the limited access to LLMs, such as GPT-4 and Claude (Ouyang et al., 2022; OpenAI, 2023) 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 (Wang and Komatsuzaki, 2021; Touvron et al., 2023). KD has gained significant attention and led to models such as Alpaca (Taori et al., 2023) and Vicuna (Chiang et al., 2023). 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% (Ouyang et al., 2022), 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
265
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:
• We introduce a TA model and develop a threecomponent KD framework that leverages three signals to improve the sample efficiency given a annotation budget and imperfect teacher LLM.
• We further improve the three-component KD framework by introducing a two-stage training approach.
• We conduct extensive experiments on various datasets and public models and the proposed two-
stage framework increases performance by up to 20.79%, compared to fine-tuning without signals, suggesting the effectiveness of our framework.

heading: 2 Related Work

heading: 2.1 Language model prompting
text: Prompt-based learning refers to using prompts to induce the embedded knowledge in the language models to complete downstream tasks (Radford et al., 2019; Liu et al., 2023b; Raffel et al., 2020). Among various prompting methods, in-context learning (ICL) through adding a few demonstrations and instructions in the prompt to elicit the correct answer has been shown as an effective approach (Brown et al., 2020). Since the output from language models is sensitive to the instruction composition and demonstration selection, how to tune and design the prompts for ICL has attracted the attention of many researchers (Liu et al., 2023b). Chain-of-thought (CoT) prompting feeds the stepby-step examples and specifically designed instructions to ask the language models to generate the intermediate reasoning step of a complex task (Wei et al., 2022). By grounding on the reasoning steps, the language models perform better on a wide variety of complex tasks. Based on CoT methods, researchers also divide the reasoning process into multiple steps (Creswell and Shanahan, 2022; Creswell et al., 2022). Moreover, Yao et al. (2022) propose another prompting method to integrate the additional action step to obtain external knowledge, which has been shown to be more effective in reasoning and decision making tasks.

heading: 2.2 Signals and feedback in language models
text: Besides asking the LLM to generate more intermediate steps, some works utilize the signals or feedback generated from the language model output to continually refine the prompts and resolve the hallucination (Madaan et al., 2023; Shinn et al., 2023; Diao et al., 2023; Su et al., 2022; Peng et al., 2023; Welleck et al., 2022; Liu et al., 2023a; Li et al., 2024b). Diao et al. (2023) and Su et al. (2022) use the measured uncertainty or score to actively select the uncertain exemplars on the prompts to increase ICL performance. Madaan et al. (2023); Shinn et al. (2023) rely on the language model itself to compose the natural language feedback. These signal-based methods are usually utilized in language model prompting (ICL) without any model parameter updates, which are distinct from our work since we designed signals for collecting high-quality examples for model finetuning.

heading: 2.3 Knowledge distillation in language models
text: Knowledge Distillation (KD) can be framed as training small student models based on data generated from large teacher models to reduce model size and expense while retaining capabilities (Sanh et al., 2019; Hinton et al., 2015; Buciluǎ et al., 2006). For the KD pipeline with black-box LLMs, there are two lines of work. The first is to ask teacher models to generate the final answers and to do standard fine-tuning on the final answers (Yoo et al., 2021; Schick and Schütze, 2020, 2021; Zhou et al., 2023). Another line of work is to ask the student model to fine-tune the teacher-generated CoT or other prompting trajectories (rationales) (Yao et al., 2022; Ho et al., 2022; He et al., 2023; Shridhar et al., 2023; Hsieh et al., 2023; Wang et al., 2023; Feng et al., 2023). Through the developed prompting methods, researchers find that sequence-level distillation of reasoning steps generated by the teacher is more effective (Yao et al., 2022; Ho et al., 2022), enabling the reasoning or decision-making abilities in students. However, most sequence-level KD works utilize the known correct or incorrect trajectories to fine-tune the student (Li et al., 2023; Wang et al., 2023; Chen et al., 2023; Liu et al., 2023c; An et al., 2023). Moreover, they consider the scenario with sufficiently large fine-tuning set while the fixed teacher budget and the flawed teacher annotations in the unlabeled data are yet to be explored. Yu et al. (2023); Liang et al. (2023) boost the fine-tuning set by rephrasing noisy
rationales but do not actively select annotated samples to save teacher budget. Mirzadeh et al. (2020); Son et al. (2021) also leverage Teacher Assistant (TA) models for KD tasks but within the realm of computer vision, diverging from our approach that applies TA models in sequence-level KD tasks. Table 1 summarizes the similarity and uniqueness of our work with previous related ICL or KD approaches.

heading: 3 Methodology
text: 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 Ba: 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 Ba and return a high-quality fine-tuning set Bf : we want to filter out low-quality or wrong annotations generated by the teacher LLM.
We leverage the interactions among the student model Ms, the teacher model Me, and the TA model Mt to generate three types of signals and achieve the above goals. For a given example, we will use the student internal signal si as well as the TA-student signal from the TA model st to decide whether the example should be added to Ba. After collecting the fixed budget of samples for Ba, we ask the teacher model to annotate those examples and apply the TA-teacher signal se 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.

heading: 3.1 Signals for annotation
text: We propose two types of signals to decide whether an example meets the annotation criteria: student internal signal si and TA-student signal st.
Algorithm 1 Three-component Knowledge Distillation Framework
1: Input: Unlabeled dataset U , student model Ms, TA model Mt, teacher model Me, confidence prompt Pc, annotation criteria annotate(·), confidence set C, annotation bucket Ba, fine-tuning bucket Bf 2: Output: The desired fine-tuning bucket Bf 3: for each input x from dataset U do 4: ti(x) = Ms(x, Pi) for i ∈ {1, · · · , n} 5: si(x) = Uniq(t1(x) . . . tn(x)) 6: st(x) = Mt(x, t1, Pc) 7: if annotate(si, st) is True then 8: Add x into annotate bucket Ba 9: for each input x from dataset Ba do
10: te(x) = Me(x, Pi) 11: se(x) = Mt(x, te, Pc) 12: if se in C then 13: Add x into fine-tuning bucket Bf 14: return Bf

heading: 3.1.1 Student internal signal si
text: Motivated by self-consistency work (Wang et al., 2022), we utilize the disagreement number (selfconsistency number) as the measurement of the uncertainty of student models. For a given sample x and prompt Pi, we use stochastic temperature sampling with a fixed temperature and repeat the process for n times with answers t1, · · · tn:
ti(x) = Ms(x, Pi) for i ∈ {1, · · · , n}
then the internal signal can be calculated by counting the unique values among n answers
si(x) = Uniq(t1(x), · · · , tn(x))
where Uniq is counting the unique answers from t1 to tn, and si ∈ {1, 2, · · · , n}. In Table 14 in the Appendix, we present two examples with different numbers of student disagreement.

heading: 3.1.2 TA-Student signal st
text: We also utilize the TA model as an auxiliary signal to characterize the uncertainty of student’s generations, given the observation that the confidence estimated by a LLM itself is prone to be overconfident (Diao et al., 2023; Si et al., 2022). For a given example x and the student annotation t (rationales and answer), the TA model is provided with a crafted confidence prompt Pc to classify the
student annotation t
st(x) = Mt(x, t, Pc)
where st is a categorical variable belonging to the confidence set {very confident, confident, not confident, wrong answer}. Details and examples of the confidence prompt Pc and t are shown in Table 12 and 13 in Appendix.

heading: 3.1.3 Annotation criteria
text: Our annotation criteria are formulated by the student internal signal si and TA-student signal st. We compute the complexity score of an example x:
c(x) = α1si(x)∈C1 + β1st(x)∈C2 (1)
where α, β ∈ {0, 1} are the signal weights. When α = 1 and β = 1, it means both signal si and st are utilized, and conversely if either α or β is set to 0, the corresponding signal is not employed. 1 is the indicator function and C1, C2 are complexity sets to actively select the examples that bring the deepest learning curve for the student model. The example will be added to the annotation bucket Ba if c(x) ≥ α+ β until the predefined budget is met.
In our experiment setup, we use n = 5, C1 = {2, 3} and C2 = {confident, not confident} after the preliminary exploration about hyperparameter combinations (Appendix B.4), which ensures that we can choose questions that are not too easy or too hard to learn for the student model.

heading: 3.2 TA-Teacher signal se
text: With the desired annotation bucket, we utilize the teacher model Me to annotate the question x in the annotation bucket Ba with the few-shot ICL. To verify the correctness of teacher annotation te, we apply the similar TA-confidence as discussed in Section 3.1.2. The process of adding TA-confidence can be formulated as
te(x) = Me(x, Pi), se(x) = Mt(x, te, Pc)
where se is the TA-confidence of teacher annotations and serve as the TA-teacher signal. We calculate the confidence score for an example by
d(x) = 1se(x)∈C3 (2)
where C3 are the confidence set. Teacher annotations will be added to the fine-tuning set Bf only if d(x) ≥ 1. In our experiments, we set C3 = {very confident, confident}. Note that our
signal does not require an additional cost to call the teacher model, which makes our approach different from (Madaan et al., 2023; Shinn et al., 2023) where they propose filtering out low-quality samples by self-refine from the teacher model. Examples of questions, annotations, and the TA signal se are shown in Table 15 in Appendix.

heading: 3.3 Fine-tune the Student Model
text: After collecting the fine-tuning dataset Bf , we finetune the student model with questions and teachergenerated annotations from the fine-tuning bucket. We apply the autoregressive language modeling as well as the cross entropy loss, which is the same as the pretraining objective to fine-tune the student model (Wang and Komatsuzaki, 2021).

heading: 3.4 Extension: Two-stage Training
text: While the aim of using student internal signal si is to actively select the examples not too simple or challenging for the student to learn, the student answers t for calculating si could be quite inaccurate, for example, the accuracy of using GPT-J6B as the student is only 5.2% at HotpotQA task (Yang et al., 2018). With such a low accuracy, two challenges will emerge. First, the signal si cannot precisely reflect the difficulty and confidence of the student in the example. Second, the ratio of examples with si(x) ∈ C1 is small (most examples with si = 4/5), causing a low sampling efficiency.
To address these challenges, as shown in Figure 2, we develop a two-stage training. The first stage is the warm-up training, and we use 10% of the budget to fine-tune the student with our framework (using signals). For the second stage with a more
proficient student model, we use 90% of the budget to continue fine-tuning the student model with our framework. We expect that the student signals collected in the second stage are more reliable than the first stage, leading to a more performance boost.

heading: 4 Experiments
text: Through our extensive empirical analysis, we aim to address the following research questions:
• RQ1: How efficient is our three-component framework compared to a classical KD method?
• RQ2: How important is each type of signal to the framework?
• RQ3: How well does our framework work for student and TA models of different sizes?
• RQ4: Does extending the framework to twostage training boost more performance?
Dataset Since our work focuses on the scenario with sufficient unlabeled data but limited annotation budget, following the previous work (Wei et al., 2022; Yao et al., 2022), we evaluate our framework on four datasets with a large training set: HotpotQA (Yang et al., 2018), GSM8K (Cobbe et al., 2021), Aqua (Ling et al., 2017) and CommonSenseQA (Talmor et al., 2019). We choose HotpotQA, the closed-book question and answering (QA) dataset, for evaluation with ReAct prompting (Yao et al., 2022). For GSM8K and Aqua, the arithmetic reasoning task and CommonSenseQA, the commonse sense reasoning task, we evaluate performance with CoT prompting (Wei et al., 2022). In Table 2, we present the budget number, test number, prompting method, and evaluation metric of all four datasets. Backbone models For the teacher model, we use gpt-3.5-turbo based on InstructGPT 175B (Ouyang et al., 2022) to generate the CoT or ReAct trajectories. We use Vicuna-13B and Vicuna-65B (Chiang et al., 2023) as TA models to evaluate the answers
of the teachers and students. For the student models, we choose between the GPT-J-6B and Vicuna13B models (Wang and Komatsuzaki, 2021).
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 se as Eq (2). Throughout the experiments sections, we use “α = 1” and “fine-tune with si”, “β = 1” and “fine-tune with st” 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 HotpotQA dataset and CoT for the others (Wei et al., 2022; Yao et al., 2022). The other baseline model is Random-finetune, the classical sequence-level KD pipeline for LLMs (Ho et al., 2022; Yao et al., 2022; He et al., 2023), 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 (Li et al., 2023; Wang et al., 2023; Chen et al., 2023; Liu et al., 2023c; An et al., 2023; Zhao et al., 2023), not applicable in an unlabeled setting; those generate more augmented rationales (Liang et al., 2023; Shridhar et al., 2023; Feng et al., 2023) or more augmented questions (Yu et al., 2023), which are not achievable with a fixed budget; and those incorporate additional loss (Hsieh et al., 2023), which is orthogonal to our KD method. The details of the model configuration are included in Appendix A.

heading: 5 Results

heading: 5.1 RQ1: Comparison with Baseline Methods
text: Three-component KD framework outperforms Random-finetune and TA ICL We use GPT-J6B as the student model and Vicuna-13B as the TA model. We experiment with two variants of the proposed framework: TA-finetune (I) and TA-finetune (T). As shown in Table 3, our proposed framework shows around 16.7% improvement compared to the Random-finetune baseline. It is also interesting to see that the fine-tuned GPT-J-6B model with the
proposed framework can outperform the Vicuna13B ICL model, which is already fine-tuned with 70k ChatGPT conversations (Chiang et al., 2023).
Note that the teacher ICL results are only 57.15% on average, which means that the proportion of noise in the teacher annotations cannot be easily ignored. Another example in point is the multiple-choice QA Aqua dataset, where the finetuned student model without any signal (13.93%) actually performs worse than its ICL counterpart (21.30%) as well as the random guess baseline (20.00%), and is primarily caused by the noise from teacher annotations. An example of noisy teacher annotations in the Aqua dataset - the teacher model generates rationales with “no answer” as the final output when it cannot find a satisfactory result among the given choices. The student model imitates the teacher’s behavior and generates the similar “no answer” outputs, leading to a worse result than ICL. This case demonstrates the necessity of our annotation filtering signals to remove low-quality annotations. We also experiment with excluding the teacher annotations with a final “no answer” output and get 22.59%, and our methods with signals still achieve better performance.

heading: 5.2 RQ2: Ablation Study of Signals
text: After verifying the superiority of the overall framework on Random-finetune, we ask whether each component of signal presents its own functionality. In this section, we use the similar model and dataset setup in Section 5.1 and analyze the effects of each signal: TA-Teacher, TA-Student, and Student internal signal through ablation analysis.
Signal HotpotQA GSM8K Aqua CSQA Avg. si se N N 13.10 18.57 13.93 57.57 25.78 N Y 16.20 18.87 25.51 56.43 29.25 Y N 13.80 19.70 15.91 59.13 27.14 Y Y 16.80 19.33 23.86 60.36 30.09
(a) Effects of TA-Teacher signal. Results of fine-tuned students (GPT-J-6B) when considering our TA-Teacher signal. Signal st is not used for experiments in this table.
Signal HotpotQA GSM8K Aqua CSQA Avg. si st N N 13.10 18.57 13.93 57.57 25.78 Y N 13.80 19.70 15.91 59.13 27.14 N Y 14.20 20.30 12.44 59.71 26.66
(b) Effects of signals for annotation. Results of fine-tuned student models (GPT-J-6B) on four datasets when adding or not TA-Student signals or student internal signals. Signal se is not used in the experiment of this table.
Table 4: Ablation study of all three signals. “Y” and “N” represent “yes” and “no” to indicate the signal use.
TA-teacher signal brings salient performance boost First, we compare the fine-tuned performance when adding or not the TA-Teacher signal in the condition of with or without the student internal signals si and show the results in Table 4a.
From Table 4a, we can observe that on average, excluding the uncertain annotations effectively boosts the performance, leading to the growth of the average accuracy from 25.78% to 29.25% or from 27.14% to 30.09% for fine-tuned with student disagreement (si) or not, respectively. Furthermore, the effects are more significant on more challenging datasets, such as HotpotQA and Aqua. This could be attributed to a significant improvement in data quality for these challenging datasets. We find that the ratio of the annotations with correct answer increases from 27.54% to 49.80% and from 52.55% to 75.57% for HotpotQA and Aqua datasets.
Signals for annotation show their effectiveness We separately demonstrate the effects of two types of signals for annotation (student internal si and TA-student signal st) using only one of them to actively select the questions to annotate. We summarize the results in Table 4b. The average results in Table 4b show that the use of either of the signals for annotation can increase performance from 25.78% to 27.14% or 26.66%, respectively, demonstrating the effectiveness of selecting uncertain samples to fine-tune the student model.
From the ablation study, we find that all types of signal have their own functionalities and orthogonal to each other. The performance of using only
one type of signal can still exceed the results of the Random-finetune framework, and the combination of signals will lead to a more performance boost.

heading: 5.3 RQ3: Generalization on Students and TAs
text: The ablation analysis presents the effectiveness of each signal, and we examine whether the effect of the framework can be generalized to student or TA models of different sizes. For student models, we additionally use Vicuna-13B as a student model and compare with the results in Table 3 with GPTJ-6B as the student model. For TA models, we will show the effects with various TA models in two aspects: TA-teacher signal and TA-student signal.

heading: 5.3.1 Generalizations on Students Models
text: Signals are effective for students of various sizes We try Vicuna-13B as the student, and use another Vicuna-13B as the TA and GPT-3.5 as the teacher. We fine-tune Vicuna-13B as three signal combinations: no signals, only student internal signal si, and student internal si + TA-teacher signals se. We show the results of GPT-J-6B and Vicuna-13B as the student model, respectively, in Table 5. Note that fine-tuning Vicuna-13B model on HotpotQA dataset will cause out-of-memory issues due to the large length of ReAct trajectories, so we only report the fine-tuned performance on other datasets.
We can observe that in Table 5, fine-tuning with the Vicuna-13B student model leads to much better performance than the GPT-J-6B model, indicating that a larger model size will cause better fine-tuning results, which is consistent with previous findings in Ho et al. (2022). The best results for Vicuna-13B student models achieved with both TA-Student and TA-Teacher signals also verifies the generalizability of our framework on larger student models.

heading: 5.3.2 Generalizations on TA Models
text: TA models are used to examine both student and teacher answers. An intuitive question could be:
How large could a model be an suitable choice of the TA? To explore this question, we use GPT-J-6B as the student, GPT-J-6B/Vicuna-13B/Vicuna-65B as the TA candidates, and GPT-3.5 as the teacher. To eliminate the other signals’ influence, we experiment with only one signal: TA-teacher signal se or TA-student signal st in this part. Choice of TA-Teacher signals Since the TA model in TA-teacher signal se filters incorrect teacher annotations to improve data quality in the fine-tuning bucket Bf , our hypothesis is that the data quality in Bf scales with the size of the TA, which influences the student fine-tuning result. To verify the hypothesis, we first collect the final teacher answer (the single answer) in different finetuning buckets. Then, we calculate the accuracy of the answer collection for each bucket, using the ground-truth answer from the original labeled dataset. We visualize the proportion of questions with correct teacher annotations with different TAs in Figure 3. Our observation is that with larger TAs, the fine-tuning data quality increases, but the improvement for model size from Vicuna 13B to 65B is smaller than that from GPT-J-6B to Vicuna-13B.
Then, we fine-tune the student model GPT-J6B with the bucket Bf evaluated by different TA models. The results can be found in Table 6. From Table 6, we obtain the highest fine-tuning accuracy with the best TA, but the marginal improvement in data quality from 13B to 65B revealed in Figure 3 leads to comparable performance. Choice of TA-student signals We vary the size of TA in the TA-student signals st to characterize the uncertainty of the student’s inference. We use Vicuna-13B or Vicuna-65B as the TA to evaluate the student’s inference from the GPT-J-6B model.
The fine-tuning results of the student model with only the TA-student signal are shown in Table 7.
From comparison with different TA models in the TA-student signal from Table 7, we can obtain comparable results when using a more advanced TA model to characterize the uncertainty of student’s generations. We conjecture that the difference in difficulty between the annotation buckets Ba of different TA models, especially from 13B to 65B, is not that large, which may not show much significant effect on the final fine-tuning results.
Results in Tables 6 and 7 indicate a trade-off between the TA model sizes and the fine-tuning performance. Using a more advanced TA can improve some performance. However, in some scenarios, the computation resource can not afford a large TA model, such as the 65B model, then using the 13B model, which can effectively improve the fine-tuning data quality, is also a reasonable choice.
Moreover, we use different budget numbers to fine-tune the student model and present the detailed results in Appendix B.1. The observation from model performance can verify the generalization of our framework with various budget numbers.

heading: 5.4 RQ4: Effects of Two-stage Training
text: As we described in Section 3.4, we extend our one-stage training to two-stage training to better utilize the student internal signals. For the model and dataset setup, we use the same setting as in Section 5.1, which means that for the two-stage training, we use 200 examples for the first warmup stage and the remaining 1,800 budgets for the second stage. For signal setup, we choose the student internal signal si and experiment under two conditions: with or without TA-Teacher signals se. We summarize the results in Table 8.
Two-stage extension obtain better results in most cases From the results in Table 8, we observe that extending the one-stage to the two-stage framework obtains a relative improvement of 1.86% and 3.48% for fine-tuning with or without TATeacher signals, respectively, and compared with the Random-finetune baseline, the two-stage framework brings the highest relative improvement of 20.79%. Our extension obtains better results in 6/8 comparison cases, which verifies the effectiveness of the extension. The improvement can be attributed to more reliable signals si extracted by the student model in the second stage. We show the intermediate results in Appendix B.2 to interpret the improvement of the two-stage training. Moreover, we extend our fine-tuning to more stages to a curriculum learning pipeline, and the results of four-stage learning are presented in Appendix B.3.

heading: 6 Conclusions
text: 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.

heading: 7 Limitations
text: In our framework, we apply several types of signals: student intenal signal, TA-student and TA-teacher signal to refine the knowledge distillation of LLM.
In addition to these internal and external signals, we can try other signals developed in existing
works about active learning (Zhang et al., 2022; Settles, 2009; Li et al., 2024a). For example, we can aggregate the model confidence on each generated token in the rationales as the student or teacher’s confidence on their annotations. Moreover, in addition to querying the TA model of confidence on the final answers, we can also ask the TA to evaluate the intermediate reasoning steps of the teacher or student models on each question, which could be another type of external signal.
We have verified the effectiveness of our framework on multiple models, prompting methods, and datasets. In related work, we have disccused multiple sequence-level KD methods, using more advanced fine-tuning loss (Hsieh et al., 2023) and more complex generated rationales (Shridhar et al., 2023) to improve KD performance. These methods are complementary to our three-component framework, aiming to increase the results with sufficiently large teacher-annotated data. Our method could further enhance their efficiency in fixedbudget and unlabeled data settings. We leave the exploration of the combination with existing works to the future.
The student models (GPT-J-6B and Vicuna-13B) in the experiments are fine-tuned on 8 NVIDIA A100 GPUs, which may not be accessible to everyone and have a negative environmental impact. For further experiments, we can extend our experiments to smaller language models with fewer than 1 billion parameters as the student model. Furthermore, we can also use more advanced teacher models, such as GPT-4 (OpenAI, 2023), to observe whether a better teacher model leads to a greater improvement in performance.

heading: Acknowledgments
text: We would like to thank the anonymous reviewers as well as Paiheng Xu for reviewing the article and for providing helpful comments and suggestions.

heading: B Supplementary Experimental Results
text: In this section, we will show the supplementary experimental results, including our framework with different budgets and the performance of extending our framework to curriculum learning.
B.1 Generalizations on Different Budgets
In Section 5.1, we have verified the effectiveness of our framework with 2,000 as the budget number, and to show the generalization of our framework on different budget numbers, we will repeat the experiments with the same model and dataset setup with the budget numbers 1,000 and 3,000, respectively. We present the results in Table 9.
From Tables 9 and 3, we find that the performance of all methods improves a lot when the budget number increases from 1,000 to 2,000 and will get a comparable performance when the budget increases from 2,000 to 3,000. Our proposed threecomponent KD framework can outperform the random-finetune method in all three cases, which suggests that our framework’s effectiveness can be generalized to different numbers of budgets.
B.2 Details of Two-stage Training
For the interpretation of the improvement in twostage training in Section 5.4, we conjecture that it could be attributed to two reasons. First, after the warm-up stage, the better student model is more expert in the questions, so the student’s internal signal si should be more reliable. Second, samples extracted by student internal signals si (the uncertain questions for the student) should be more challenging in the second stage and continually fine-tuning the student with increasingly difficult
data, leading to a deeper learning curve and a higher performance boost. To verify these two hypotheses, we present the accuracy of the student model and the difficulty level of the samples in the annotation bucket Ba before and after the warm-up stage.
We first present the student ICL accuracy (ICL Acc) in the first stage and the student accuracy in the second stage (Warm-up Acc) in Figure 4. We find that after fine-tuning on a small dataset to warm up, the student model becomes more promising and expert in answering complex questions in
most cases. Then, we also quantify the difficulty of samples in the annotation bucket Ba of in the first stage and second stage, respectively, by measuring the teacher model accuracy on questions in Ba. We visualize the accuracy of the teacher model of the samples in the annotation bucket in two stages in Figure 5. We find that in most datasets, the teacher model’s accuracy on the questions in the annotation bucket for the second stage is lower than that for the first stage and suggests more difficult and challenging questions in Ba. This observation indicates that continually fine-tuning the student model on the slightly more difficult data causes a deeper learning curve, leading to a better performance.
B.3 Explorations on Curriculum Learning
Inspired by the results of the two-stage training, our next question is whether the performance improvement will continue to increase when using more stages. We can frame multiple-stage training into the curriculum learning scope, which means that at each stage, we will have better student models and the samples extracted by student signals should be more difficult and challenging (Xu et al., 2020). For our experiment setup, we divide our training into four stages and use 200, 400, 600, and 800 budgets (a total of 2,000) in each stage. For the final step, we will integrate all the annotations in every stage and continue fine-tuning the student model for a few epochs. We present the results of one-stage training and curriculum learning in Table 10.
From Table 10, we do not observe improvement as in the results of the two-stage training, and we conjecture that there should exist multiple reasons to influence the results of curriculum learning, such as the learning rate for each stage and the number of budgets in each stage. Here, we keep the learning rate for each stage as the same value, which may cause the fine-tuning to a sub-optimal point. We
leave the further exploration of curriculum learning of our framework as a future direction.
B.4 Preliminary Explorations of Annotation Criteria
In Section 3.1.3, we choose C1 = {2, 3} as the annotation criteria of the internal signal si, and this hyperparameter choice is decided by our preliminary experiment. We compare the fine-tuning result with only si (without signal st and se) by choosing C1 = {2, 3} and C ′1 = {4, 5} and without si (randomly sampled) by fixing budget number 300 in the HotpotQA dataset. The model configuration is the same as in Section 5.1. The result is shown in Table 11.
From Table 11, we find that the student model fine-tuned with C1 = {2, 3} achieves the best performance, indicating that fine-tuning with C1 = {2, 3} can bring the deepest learning curve, so we choose this hyperparameter choice in the following experiments. Furthermore, we observe that the model fine-tuned with the annotation criteria C ′1 cannot outperform random selection without signal si, suggesting that it is hard for the student to learn useful information from hard questions.

heading: Future Work
text: We
leave the further exploration of curriculum learning of our framework as a future direction. B. 4 Preliminary Explorations of Annotation Criteria
In Section 3. 1. 3, we choose C1 = {2, 3} as the annotation criteria of the internal signal si, and this hyperparameter choice is decided by our preliminary experiment. We compare the fine-tuning result with only si (without signal st and se) by choosing C1 = {2, 3} and C ′1 = {4, 5} and without si (randomly sampled) by fixing budget number 300 in the HotpotQA dataset. The model configuration is the same as in Section 5. 1. The result is shown in Table 11. From Table 11, we find that the student model fine-tuned with C1 = {2, 3} achieves the best performance, indicating that fine-tuning with C1 = {2, 3} can bring the deepest learning curve, so we choose this hyperparameter choice in the following experiments. Furthermore, we observe that the model fine-tuned with the annotation criteria C ′1 cannot outperform random selection without signal si, suggesting that it is hard for the student to learn useful information from hard questions