Co-training and Co-distillation for Quality Improvement and Compression of Language Models (2023)

Abstract: Knowledge Distillation (KD) compresses computationally expensive pre-trained language models (PLMs) by transferring their knowledge to smaller models, allowing their use in resource-constrained or real-time settings. However, most smaller models fail to surpass the performance of the original larger model, resulting in sacrificing performance to improve inference speed. To address this issue, we propose Co-Training and Co-Distillation (CTCD), a novel framework that improves performance and inference speed together by co-training two models while mutually distilling knowledge. The CTCD framework successfully achieves this based on two significant findings: 1) Distilling knowledge from the smaller model to the larger model during co-training improves the performance of the larger model. 2) The enhanced performance of the larger model further boosts the performance of the smaller model. The CTCD framework shows promise as it can be combined with existing techniques like architecture design or data augmentation, replacing one-way KD methods, to achieve further performance improvement. Extensive ablation studies demonstrate the effectiveness of CTCD, and the small model distilled by CTCD outperforms the original larger model by a significant margin of 1.66 on the GLUE benchmark.

## introduction in recent years, the high computational cost of pretrained language models (plms) @xcite @xcite @xcite become a constraint in serving resource-limited or real-time applications. knowledge distillation (kd) @xcite is a popular model compression technique to tackle this issue, with a smaller model (student) learning (distilling) from a larger model (teacher). the ideal scenario would involve compressing the teacher to match the size of the small student without any performance drop. however, despite intense research in the area of kd @xcite @xcite @xcite @xcite , none of the existing methods have successfully avoided performance degradation during the distillation process. some approaches have attempted to mitigate the performance gap by incorporating external factors. for instance, @xcite incorporates data augmentation, while @xcite focuses on designing student architecture. however, these approaches have limitations as they only provide supplementary techniques rather than addressing the fundamental issue of performance loss in kd. this raises an interesting question: can we compress a model without scarifying the performance through kd? in this work, we propose a novel framework called co-training and co-distillation (ctcd) for improving the performance of a language model while compressing it through kd. ctcd involves jointly training the teacher and the student models, allowing them to transfer knowledge to each other bidirectionally, from the teacher to the student and vice versa. our work uncovers two key findings within the ctcd framework: firstly, we demonstrate that transferring knowledge from the smaller model to the larger model during co-training significantly improves the performance of the larger model. in other words, by employing knowledge distillation (kd), we enhance the performance of the teacher model compared to standalone training. this is significantly different from conventional one-way kd, where the teacher model cannot benefit from the distillation process since it is no longer trained or distilled. secondly, the improved performance of the larger model leads to further enhancements in the performance of the smaller model through kd from the teacher model to the student model. these novel findings enable the smaller model to surpass the independently trained larger model while maintaining its inference efficiency within the ctcd framework. moreover, the ctcd framework can be combined orthogonally with existing techniques for external factors such as student architecture design @xcite or data augmentation @xcite . we achieve this by replacing traditional one-way kd methods with the proposed ctcd framework, which holds promise for significant performance improvements in language modeling tasks. in addition to our ctcd framework, we introduce a community kd that utilizes ctcd framework with the pre-trained teacher model by combining ctcd with conventional kd. community kd distills knowledge from two students to each other, as well as from the pre-trained teacher to each student, during the co-training of the students. we show that distilling from the other student as well as the pre-trained teacher is better than only one-way distillation from the pre-trained teacher, as is done in conventional kd. note that the inference cost on a task is the same as in conventional kd, as we take one of the students distilled by community kd when deploying it to the task. we validate the effectiveness and efficiency of our ctcd on the glue benchmark @xcite , which contains different language understanding tasks. as shown in figure 1 , we focus on task-agnostic kd scenarios, where kd occurs during the pre-training stage. then, we fine-tune the distilled plms on each downstream task and evaluate their performance. this is more challenging compared to task-specific kd, as it requires substantial computational resources to train plms on a large text corpus, and the distilled knowledge should be transferable across diverse downstream tasks. our extensive ablation study revealed that the larger model benefits from the distillation of the smaller model, and the performance improvement of the larger model further enhances the performance of the smaller model. in our experiments, the student model compressed by ctcd framework obtained 1.66 higher gain than the original large model trained using a stand-alone method, demonstrating that our approach can improve model quality and inference efficiency concurrently. in summary, our contributions are as follows: one-way knowledge distillation knowledge distillation (kd) @xcite is a model compression technique in which a smaller student model distills knowledge from a larger, pre-trained teacher model. recently, many researchers have explored kd to address a high computational complexity of a language model resulting from its increasing size. this includes task-agnostic kd methods for the pre-training stage @xcite , task-specific kd method for we refer to the former as a "hard" loss and the latter as a "soft" loss. for example, the student soft loss captures the distance between the student's prediction and the teacher's prediction, and vice versa. this co-training and co-distillation approach improve the performance of the teacher model, which in turn benefits the performance of the student model. the fine-tuning stage @xcite , and both of pre-training and fine-tuning stages @xcite . additionally, @xcite have redesigned the architectures of student language models. however, the one-way kd can lead to a loss of knowledge, resulting in smaller models that generally have difficulty matching the performance of larger models, leading to performance degradation. in contrast, our ctcd can improve the quality of both teacher and student models, making it possible for the student to achieve the same quality as the original teacher model. ## co-training and co-distillation we first introduce the concepts of co-training and co-distillation briefly: co-training trains two (different-sized) models (e.g., a teacher and student) concurrently with the goal of achieving similar model quality. co-distillation transfers knowledge in both directions between two models (e.g., a teacher and student), during co-training. figure 2 illustrates how co-trained models learn together by comparing their prediction outputs against predictions from each other and to the hard labels (or "ground truth"). we refer to the former as a "soft" loss and the latter as a "hard" loss. for instance, the soft loss of the student model measures the accuracy of the student's prediction by considering the teacher's prediction as a soft label, and vice versa. task formulation suppose that we are given a classification task with k classes. for each training instance x and its ground truth label y, we denote that the ground truth distribution over the labels is q(k|x) (q(k) for simplicity) where for each label @xmath0...k}, q(y) = 1 and q(k) = 0 for all k ̸ = y. for each x, the teacher model t ϕ parameterized by ϕ and the student model s θ parameterized by θ predict the probability of each label k as p τ ϕ (k|x) and p τ θ (k|x), respectively as follows: the proposed objective l ct cd (θ, ϕ) consists of a normal kd objective l kd:@xmath1:@xmath2, during their co-training. @xmath3:@xmath4, ϕ) aims to train the student model by minimizing a weighted sum of the cross-entropy loss h(q, p θ ) between the ground truth q and student prediction p θ and kullback-leibler divergence (kl divergence) d(p τ ϕ , p τ θ ) between the predictions of the student and the teacher as follows: @xmath5, we introduce the reversed kd objective l rekd:@xmath6, ϕ) to train the teacher model t ϕ as follows: co-training with the equations ( 2 ) and (3), we get the final objective l ct cd (θ, ϕ) as follows: adapting to downstream task after model cotraining/-distillation, the trained smaller (student) model s θ * with trained parameter θ * can be deployed for multiple downstream tasks to improve inference efficiency. to fine-tune the model for a specific downstream task, we adapt the trained parameter θ * using the dataset for that task. we select one of the two students distilled by community kd (θ * = θ * 1 or θ * = θ * 2 ) and finetune the selected single student s θ * for downstream tasks, resulting that the inference cost does not increase compared with the conventional kd method. ## experiment we present a comprehensive analysis of the proposed ctcd method through empirical experiments. in section 5.1, we validate our ctcd method by comparing the performance of small models distilled by ctcd to the original large model on the glue benchmark @xcite . in section 5.2, we analyze the impact of co-distillation by adjusting loss weights for the soft losses of a student and a teacher. in section 5.3, we study the impact of training length on ctcd method, allowing us to determine the optimal training length for ctcd method. in section 5.4, we demonstrate the efficacy of community kd by comparing the performance of a model distilled by the community kd to a model distilled by the one-way kd. implementation details we use a learning rate of 5e-4, linear warm-up of 5%, adamw optimizer @xcite , and batch size of 128 with a100 gpus for pre-training. we train the teacher and student models from scratch for 20 epochs in section 5.1. to analyze the effectiveness of ctcd method, we train the teacher and student models for 10 epochs and 20 epochs in section 5.2 and section 5.3, respectively. in section 5.4, we train the models for 3 epochs after parameter remapping, which is the same as in conventional one-way kd method @xcite . training time for the one-way distillation, we need 1 gpu day to train the teacher for 10 epochs and 1.3 gpu days to distill knowledge from the pre-trained teacher to the student for another 10 epochs, which consumes a total of 2.3 days. for ctcd, it takes 3 gpu days to train both teacher and student models from scratch for 10 epochs. we use the automatic mixed precision (amp) of pytorch @xcite to accelerate training for all our models. dataset to validate our ctcd method, we use a reduced dataset (30m) created by uniformly sampling 1 out of every 4 sentences from the original pre-training dataset (bookcorpus @xcite + wikipedia (foundation)) used in the conventional one-way kd method @xcite . we evaluate our distilled models on the dev sets of the glue benchmark @xcite , which consists of nine sentence-level classification tasks. in section 5.4, we use the original pre-training dataset to train community kd method. ## limitations & future work limitations the proposed method co-train models from scratch and may require a longer pretraining time than the conventional kd method. however, as we described in section 5.3, when the student model is trained long enough with its teacher, it can outperform the models trained with the conventional kd on the downstream task. the proposed co-training method may increase the overall training cost compared with one-way distillation, and it may become a performance bottleneck depending on training resource constraints. however, note that ctcd can improve model quality while having the same inference cost as the oneway distillation on downstream tasks. future work 1) architecture sharing. models can share some of their architectures by reusing the output of such architectures and updating them together during back-propagation. this may help reduce the additional computing and memory overhead incurred by model co-training, while improving the model quality, especially for the student model. 2) integration of student architecture design and data augmentation. future research can focus on effectively combining the ctcd framework with student architecture design and data augmentation techniques. this integration provides a promising alternative to traditional one-way knowledge distillation methods, leading to significant improvements in language modeling tasks. ## conclusion the size and complexity of pre-trained language models (plms) can hinder their practicality for online downstream tasks. to address this, we introduced a novel framework called co-training and co-distillation (ctcd). by training models of different sizes together and extracting inter-model knowledge in both directions, the proposed ctcd framework improves both model efficiency and performance. the proposed framework overcomes the trade-off between efficiency and performance in traditional one-way knowledge distillation methods. notably, our compressed models achieved an impressive gain of 1.66 on the glue benchmark, outperforming large models trained using standalone methods.