Minimal Distillation Schedule for Extreme Language Model Compression (2024)

Abstract: Recent studies have revealed that language model distillation can become less effective when there is a significant capacity gap between the teacher and the student models. In order to bridge the gap, teacher assistant-based distillation has been introduced, in which the selection of the teacher assistant plays a crucial role in transferring knowledge from the teacher to the student. However, existing approaches for teacher assistant-based distillation require numerous trials to find the optimal teacher assistant.In this paper, we propose a novel approach called Minimal Distillation Schedule (MiniDisc), which enables the scheduling of an optimal teacher assistant in just one trial for extreme model compression (e.g, to 5% scale). In particular, we empirically show that the performance of the student is positively correlated with the scale-performance tradeoff of the teacher assistant. We then introduce a new \lambda -tradeoff metric that quantifies the optimality of the teacher assistant without the need for trial distillation to the student. By employing a sandwich framework, MiniDisc can select the optimal teacher assistant with the best \lambda -tradeoff.We extensively evaluate MiniDisc through a series of experiments on the GLUE benchmark. The results demonstrate that our approach achieved an improved efficiency compared to various state-of-the-art baselines. Furthermore, we showcase the scalability of MiniDisc by applying it to a language model with billions of parameters.

## introduction but are inapplicable to those requiring limited computational resources @xcite . to address this issue, lms can be compressed using a range of strategies such as model quantization @xcite , pruning @xcite , etc., among which knowledge distillation @xcite has gained significant attention. it operates within the teacher-student framework, where a large model acts as the teacher, transferring its knowledge to a smaller student model. recent advances @xcite ) have shown a significant performance decline in conventional distillation methods when dealing with a substantial capacity gap between the teacher and the student models. to alleviate this, teacher assistantbased distillation @xcite has been proposed. this approach involves distilling the teacher model into an intermediate-scale teacher assistant, which then serves as an intermediary to transfer knowledge to the student model. while teacher assistant-based distillation generally lifts the performance of the student @xcite , the performance of the student is largely impacted by the choice of the teacher assistant as illustrated in figure 1 . in fact, we observe there is potentially a turning point of the student performance, indicating a scale-performance (i.e., x-v.s. y-axis) tradeoff in scheduling the teacher assistant. however, existing studies schedule the teacher assistant in an enumeration manner, resulting in an inferior solution that requires maximally many trials to meet the optimal teacher assistant (maximal distillation schedule, in short maxidisc). to this demand, we propose a minimal distillation schedule (minidisc) that enables the identification of the optimal teacher assistant in just a single trial. we define a λ-tradeoff metric to empirically measure the tradeoff between scale and performance for a given teacher assistant, as depicted in figure 1 . this allows us to determine the optimality of the teacher assistant without requiring multiple trial distillations to the student model. to efficiently obtain the optimal teacher assistant based on the λ-tradeoff metric, we introduce minidisc within a sandwich framework, consisting of three stages. in the specification stage, we utilize gridding and pruning techniques to generate a series of teacher assistant candidates with varying scales. in the optimization stage, we demonstrate that the generated candidates adhere to the incremental property and the sandwich rule. furthermore, we present two approximations that enable the computation of the λ-tradeoff for each teacher assistant candidate at a lower computational cost. in the selection stage, we choose the optimal teacher assistant by selecting the candidate with the highest λ-tradeoff value. it is worth noting that minidisc can be directly extended to scenarios involving multiple sequential teacher assistants by recursively applying the mini-disc procedure. however, this work focuses on a single teacher assistant as it is sufficiently effective. to verify the effectiveness of minidisc, we conduct experiments on glue @xcite . experimental results exhibit the competitive performance of minidisc compared to several state-ofthe-art baselines, with improved efficiency (10×) of minidisc compared to maxidisc. further, minidisc is applied to large lms enct5 xl @xcite and llama2 7b @xcite to show its scalability. ## related work model pruning model pruning @xcite spans from unstructured pruning @xcite @xcite to structured pruning @xcite @xcite @xcite . unstructured pruning prunes parameters at neuron level referring to parameter magnitude @xcite or learning dynamics @xcite , while structured pruning @xcite prunes parameters at module level relying on parameter sensitivity. although unstructured pruning enjoys a finer-grained pruning, it can only fit specialized devices. in contrast, structured pruning generally fits modern acceleration devices. in our work, we adopt structured pruning for deriving the structures of candidates for its benefits for distillation. pruning also offers an opportunity to optimize the efficiency and effectiveness of our method due to its merits @xcite @xcite @xcite @xcite . knowledge distillation knowledge distillation @xcite can be divided into two categories: task-specific @xcite @xcite and task-agnostic @xcite @xcite @xcite distillation. task-specific methods distill finetuned models with task-specific data, while task-agnostic methods distill pretrained models directly with task-agnostic data. learning bjective is central to distillation, and distilling logits @xcite is the most common way. recently, hidden states @xcite , attention distributions @xcite @xcite , and high-order relations @xcite are taken into consideration for better abstraction. teacher assistant-based distillation @xcite @xcite is showcased to trade in teacher scale for student performance by inserting an intermediate teacher assistant. however, setting an optimal teacher assistant for the student is nontrivial. in this work, we aim to achieve this goal. minidisc uses only one trial while maxidisc uses many trials to schedule the optimal teacher assistant. ## conclusions in this paper, we propose minidisc to identify an optimal teacher assistant for teacher assistantbased distillation in minimally one trial in contrast to maxidisc. having observed that the scaleperformance tradeoff of the teacher assistant is of great importance to the performance of the student, we introduce a λ-tradeoff measure that quantifies the scale-performance tradeoff of the teacher assistant, and show that it is positively correlated with the student performance. to efficiently compute the measures for teacher assistant candidates and select the optimal one, we design a sandwich optimization for these candidates. comprehensive results demonstrate the improved efficiency of minidisc.