Multi-Stage Balanced Distillation: Addressing Long-Tail Challenges in Sequence-Level Knowledge Distillation (2024)

Abstract: Large language models (LLMs) have significantly advanced various natural language processing tasks, but deploying them remains computationally expensive. Knowledge distillation (KD) is a promising solution, enabling the transfer of capabilities from larger teacher LLMs to more compact student models. Particularly, sequence-level KD, which distills rationale-based reasoning processes instead of merely final outcomes, shows great potential in enhancing students’ reasoning capabilities. However, current methods struggle with sequence-level KD under long-tailed data distributions, adversely affecting generalization on sparsely represented domains. We introduce the Multi-Stage Balanced Distillation (BalDistill) framework, which iteratively balances training data within a fixed computational budget. By dynamically selecting representative head domain examples and synthesizing tail domain examples, BalDistill achieves state-of-the-art performance across diverse long-tailed datasets, enhancing both the efficiency and efficacy of the distilled models.

## introduction large language models (llms) like gpt-4 and llama have revolutionized tasks ranging from text generation to language translation through their deep understanding and generation of human-like text @xcite @xcite . despite their success, the deployment of these models is hindered by their substantial size and computational demands, especially in environments with limited resources. knowledge distillation (kd) offers a viable solution by transferring knowledge from expensive teacher models to smaller, efficient student models. specifically, sequence-level kd focuses on distilling rationale-based reasoning processes rather than final outcomes. it leverages the teacher's reasoning processes, encapsulated in chain-of-thought (cot) rationales, to enhance the student models' generative capabilities @xcite @xcite . however, there are a few challenges to fully leverage the power of sequence-level kd, as follows. (c1) sequence-level kd encounters significant challenges when training with long-tailed data distributions, which are prevalent in real-world scenarios -data often follows a power-law distribution with a few dominant classes (head) and many rare classes (tail) @xcite . such distributions feature a few dominant classes and many underrepresented ones, leading to models that generalize poorly on sparsely represented domains. (c2) traditional kd methods in the text area to solve long-tail challenges, often reliant on direct access to model weights or loss adjustment primarily suited for straightforward classification tasks @xcite @xcite @xcite , falter under the complexities of sequence-level kd, especially when the teacher model is a black box and the task is generative, which is our target. (c3) addressing this imbalance is critical, yet resourceintensive, as it typically requires generating a large volume of synthetic data to balance the dataset @xcite . moreover, naively up-sampling the long-tailed dataset may dramatically increase the number of calls to the teacher models. budget constraints play a crucial role in kd for black-box llms, as querying the teacher for rationales can be costly and time-consuming @xcite . our proposed solution, the multi-stage balanced distillation (baldistill), tackles all the challenges above by strategically generating balanced training sets within budget constraints and itera-figure 1 : overview of the proposed iterative baldistill framework. the framework is composed of multiple stages. for each stage, we apply the balancing policy to decide the data distribution in the training batch. for head domains with sufficient data, we actively extract the examples by ifd metrics using the student model. for the tail domains, we call the teacher model to generate the synthetic examples and the corresponding rationales. the teacher model finally annotates the balanced training batch and fine-tunes the student model. our contributions are summarized as follows: long-tail learning focuses on long-tail distributed data and has been an emerging topic of interest in the nlp community @xcite @xcite @xcite @xcite . approaches to solving the long-tail problem include rebalancing, information augmentation, and module improvement @xcite @xcite @xcite . despite the importance of long-tail learning, studies have shown that llms struggle to learn long-tail knowledge @xcite . in this work, we propose to improve llms' ability to learn longtail knowledge via multi-stage distillation over balanced datasets. active learning aims to reduce labeling effort by selecting only the most useful examples. traditional active learning can be categorized into uncertainty-based methods @xcite @xcite and diversity-based methods @xcite . in the llm era, active learning has been used to reduce human annotation costs by (1) strategically selecting the most informative examples for human feedback or annotation @xcite @xcite and ( 2 ) integrating language models as annotators within an active learning framework without human supervision @xcite @xcite @xcite . in this work, we propose to solve the long-tail problem in the student llm by actively distilling knowledge from a black-box teacher llm to meet the budget requirement. ## experiment through our extensive empirical analysis, we aim to address the following research questions: dataset to verify the effectiveness of our framework on various reasoning tasks, we evaluate our method on five long-tailed datasets, following previous work @xcite @xcite . these include text classification: r52 and reuters (hayes and weinstein, 1990), question answering: abstractiveqa and multiple-choiceqa @xcite and arithmetic: math @xcite . for text classification datasets, we treat the label of inputs as the domain; for other datasets, the domain information of inputs is annotated as metadata from the data provider. the detailed construction process and domain information for these datasets can be found in appendix a. we also show two example distributions of the datasets in figure 5 in appendix a. for each dataset, we prepare two budget settings for the experiment. in evaluation metrics since we are dealing with long-tailed imbalanced data, for each dataset, we choose to use both the micro-and macro-averages to evaluate the method robustness @xcite . for the classification datasets (r52 and reuters), we report micro-f1 and macro-f1, where micro-f1 is a global average f1 score and macro-f1 is computed by taking the unweighted mean of all the per-class f1 scores @xcite . for other datasets, we also report the micro-/macro-f1 for abstractiveqa datasets and micro-/macro-accuracy for math and multi-choiceqa datasets. note that the f1 score for the abstractqa is the word-level f1 score between the token list of ground truth answer and the generated answer, different from the f1 for the classification task. model setup for the teacher model, we use gpt-4 (openai, 2023) to generate the cot rationales for each dataset. we choose between llama2-7b and llama3-8b as our student models @xcite . we include the detailed configurations and implementations of the model in appendix b. ## conclusions in this paper, we propose a novel framework baldistill to enhance performance on long-tail datasets in the current teacher-student knowledge distillation process. our framework is a multi-stage pipeline, and at each stage, we call the student models to actively select the representative examples from head domains while prompting the teacher to generate synthetic examples for tail domains. with a fixed budget restriction for calling the teacher, our extensive empirical evaluations show that our framework can significantly increase fine-tuning results across multiple datasets. furthermore, we demonstrate the effectiveness of all framework components through ablation studies. furthermore, our experiments only focus on the decoder-only student models: llama3 and llama2. incorporating more encoder-decoder models such as @xcite would benefit future studies. another future direction for our paper is to explore the application of knowledge distillation in large vision-language models (lvlms) @xcite @xcite @xcite @xcite . in this paper, we have focused on experiments related to knowledge distillation in large language models (llms). in future work, we aim to use knowledge distillation to further address the issue of hallucination @xcite @xcite in small lvlms such as llava-7b @xcite and vila-7b @xcite .