Meta-Learning Adaptive Knowledge Distillation for Efficient Biomedical Natural Language Processing (2022)

Abstract: There has been an increase in the number of large and high-performing models made available for various biomedical natural language processing tasks. While these models have demonstrated impressive performance on various biomedical tasks, their training and run-time costs can be computationally prohibitive. This work investigates the use of knowledge distillation, a common model compression method, to reduce the size of large models for biomedical natural language processing. We further improve the performance of knowledge distillation methods for biomedical natural language by proposing a meta-learning approach which adaptively learns parameters that enable the optimal rate of knowledge exchange between the teacher and student models from the distillation data during knowledge distillation. Experiments on two biomedical natural language processing tasks demonstrate that our proposed adaptive meta-learning approach to knowledge distillation delivers improved predictive performance over previous and recent state-of-the-art knowledge distillation methods.

## introduction while there has been an increase in the number of large, pre-trained language models with impressive performance on various biomedical tasks @xcite @xcite @xcite , the training and deployment of these models can be computationally prohibitive and time-consuming, especially in resource-constrained settings. the inference latencies and storage costs of these models make their deployment for real-word biomedical applications a challenge. knowledge distillation @xcite @xcite , a model compression technique which aims to transfer the performance of a large and computationally inefficient teacher model to a smaller and more efficient student model, has been proposed as a way to reduce the size of large models while retaining their predictive performance. while a variety of knowledge distillation approaches have been proposed in the literature @xcite @xcite , their effectiveness have largely not been evaluated on biomedical natural language processing tasks. in this work, we evaluate the effectiveness of the proposed approaches for knowledge distillation on biomedical nlp tasks. to further enhance performance, we propose an adaptive meta-learning method for distilling large and inefficient biomedical models into more efficient and smaller ones. in experiments conducted on two biomedical natural language processing tasks, we find that our proposed meta-learning approach to knowledge distillation delivers improved predictive performance over previous and recent state-of-the-art knowledge distillation methods. ## knowledge distillation knowledge distillation is a model compression method which aims to transfer knowledge from large and accurate but computationally inefficient models to smaller and more efficient models without significant loss in task performance. this is usually achieved by training a smaller and computationally efficient student model to imitate the outputs of a larger and inefficient teacher model with a knowledge distillation objective. for instance, the knowledge distillation objective proposed in @xcite uses the final output logits produced by the teacher model to transfer its hidden knowledge to the student model. concretely, given a teacher model t parametrized by θ t , a student model s parametrized by θ s and a dataset d containing n instances @xmath0, y i )} n @xmath1, the knowledge transfer between teacher and student can be achieved by training the student with a knowledge distillation objective l kd of the form: ## meta-learning meta-learning, also known as learning to learn @xcite @xcite aims to develop algorithms and models that are able to learn more efficiently with experience, by generalizing from the knowledge of related tasks. these models are able to learn how to learn, by improving their own learning process over time. various approaches to meta-learning have been proposed and applied in various areas. these approaches include specific-architectures for learning to learn @xcite , learning to update model parameters from background knowledge @xcite , and gradient-based model-agnostic meta-learning methods @xcite @xcite . example natural language processing tasks to which meta-learning has been applied include machine translation @xcite and quality estimation @xcite . gradient-based model-agnostic meta-learning algorithms such as maml @xcite often involve a bi-level optimization objective where feedback from the performance of an inner-learner (student model) is used to optimize a meta-learner (teacher model) with the aid of a meta-objective. in other words, in contrast with the teacher model in common knowledge distillation approaches which does not take into account feedback from the student model, the teacher model in meta-learning is able to receive and utilize feedback from the student model in order to improve itself. additionally, in knowledge distillation the teacher and student models are usually trained one after the other, with the teacher model trained first and then fixed during the student training. on the other hand, the student and teacher models in metalearning are trained jointly together in order for them to improve each other. ## knowledge distillation with meta-learning some works have investigated the use of the bilevel optimization framework in meta-learning to improve knowledge distillation, that is, to employ meta-learning to explicitly optimize the teacher for better knowledge transfer during the knowledge distillation process. for instance, @xcite trained a teacher network that can be adapted across several domains with meta-learning, and then perform standard knowledge distillation to distil the knowledge present in the teacher network into a student network. however, @xcite utilize meta-learning only to train a teacher model, and not throughout knowledge distillation training, thus limiting the generalizability of their approach. in order to enable the teacher model to better transfer knowledge to the student, @xcite proposed the use of a meta-learning pilot update mechanism which improves the alignment between the student and the teacher in knowledge distillation. in their approach, @xcite update both the teacher and student throughout the knowledge distillation training process, resulting in improved knowledge distillation performance. ## meta-learning adaptive knowledge distillation an important limitation in all aforementioned knowledge distillation methods, including those that make use of meta-learning, is that they treat the rate of knowledge exchange between teacher and student (α and β in equation 1 ) as fixed during training. this is not ideal, as the optimal rate and level of knowledge exchange between teacher and student should be updated during training to account for their current state. a relevant and analogous human analogy is that school teachers teach and students learn different curricula depending on the student's educational level (e.g. nursery, primary, secondary, or university student). in most circumstances, it would not be appropriate for a human teacher to be teaching university-level knowledge to primary school students, and vice-versa. therefore, α and β in knowledge distillation also need to be adaptive and learnable. as a solution to the aforementioned issue, in this work we propose to treat α and β as learnable parameters which are updated during training. our work builds on that of @xcite and further enhances it with learnable α and β. this would allow the values of α and β to change to reflect the needs of the student throughout training. as we demonstrate in the experiments, this change results in improved knowledge distillation performance. we refer to our adapted approach as meta-learning adaptive knowledge distillation (metaadaptivekd), and our overall training algorithm is illustrated in algorithm 1. ## results and discussion the results obtained by our approach and the other knowledge distillation methods on the two biomedical datasets are as shown in in terms of f1 performance of the student models, we find that just finetuning the student model (finetune) without any knowledge distillation objective underperforms all other distillation methods on the gad dataset and also underperforms all other methods except patientkd on the chemprot dataset, which demonstrates the effectiveness of knowledge distillation in general. patientkd outperformed kd on the gad dataset but not on the chemprot dataset, while metadistil outperforms kd and patientkd on both datasets. finally, we find that our approach metaadap-tivekd, which adaptively learns α and β with metalearning, outperforms all previous distillation methods on both datasets. the fact that our approach outperforms metadistil (a meta-learning method which uses fixed α and β) demonstrates the importance of not keeping α and β fixed during knowledge distillation, but instead learning their optimal values from the distillation data during training, as done in our approach. ## conclusion in this work, we proposed a new meta-learning approach to knowledge distillation. in contrast to previous methods which manually set the rate of knowledge exchange between student and teacher and keep them fixed throughout training, our approach learns their optimal values adaptively from the distillation data during training. in experiments conducted on two biomedical datasets, we demonstrated that our approach outperforms previous knowledge distillation methods.