VE - KD : Vocabulary-Expansion Knowledge-Distillation for Training Smaller Domain-Specific Language Models (2024)

Abstract: We propose VE-KD, a novel method that balances knowledge distillation and vocabulary expansion with the aim of training efficient domain-specific language models. Compared with traditional pre-training approaches, VE-KD exhibits competitive performance in downstream tasks while reducing model size and using fewer computational resources. Additionally, VE-KD refrains from overfitting in domain adaptation. Our experiments with different biomedical domain tasks demonstrate that VE-KD performs well compared with models such as BioBERT (+1% at HoC) and PubMedBERT (+1% at PubMedQA), with about 96% less training time. Furthermore, it outperforms DistilBERT and Adapt-and-Distill, showing a significant improvement in document-level tasks. Investigation of vocabulary size and tolerance, which are hyperparameters of our method, provides insights for further model optimization. The fact that VE-KD consistently maintains its advantages, even when the corpus size is small, suggests that it is a practical approach for domain-specific language tasks and is transferrable to different domains for broader applications.

## introduction language models such as bert @xcite and roberta @xcite have provided significant performance improvements in solving natural language processing (nlp) tasks, enabling many practical applications that have increased productivity, understanding, and accessibility in diverse industries. these traditional models still hold value in terms of cost-effectiveness and ease of deployment, even though large language models (llms) demonstrate remarkable few-shot capabilities in nlp tasks. one reason is that training or fine-tuning llms such as gpt-3 requires an immense amount of data and computational resources. another reason is the growing demand for artificial intelligence (ai) applications that run on local machines because some applications require independence from network connectivity or have concerns about information security and confidentiality when using llm-based application programming interface (api) services such as gpt-4. many industrial and academic fields use specialized terminology and concepts that general language models might not fully understand. these potential gaps in understanding may result in less effective or even erroneous solutions, making it essential to adapt language models to specific domains. however, llms such as gpt-3 and gpt-4 are difficult to use because it is expensive and challenging to obtain high-quality labeled data for additional pre-training and because domain knowledge must be added through the api. in contrast, general bert models have the advantage of easy of finetuning and specialization in different domains. for example, bert performs better in named-entity recognition tasks compared with gpt-family models such as biogpt @xcite . in industrial applications, operational efficiency is often the primary concern. for example, high latency can be detrimental for applications that require real-time response or that process large amounts of input data, such as monitoring systems or predictive analytics. larger models need more powerful and thus more expensive hardware setups but typically have capacity constraints imposed to manage costs. this also limits the model size that can feasibly be realized. therefore, reducing resource consumption by compressing a model improves its deployment adaptability. although the need for domain adaptation and model compression is particularly prominent in industrial applications within a specific domain, given the complexities inherent in these processes, a simplistic sequential approach may not yield the best results. first, both tasks require high-quality data, which can be difficult to obtain. second, using general methods such as domain-adaptation followed by distillation or distilling an already domain-adapted model requires two or more steps or training and hyperparameter tuning @xcite , which makes the learning process difficult to optimize. during the domain-adaptation phase (secondarystage unsupervised pre-training), there is a significant risk of losing general knowledge due to overfitting when a small corpus is used. moreover, 2-step training requires more computational resources and time, possibly requiring further iterations to achieve the most effective outcomes. therefore, a method that can proficiently perform domain adaptation and model compression simultaneously is needed to overcome these issues. in this paper, we propose ve-kd, a novel method that can simultaneously perform domain adaptation and model compression from a teacher model such as bert. we also show that our method significantly outperforms the teacher model on related tasks with a corpus, is robust and easy to optimize, and has lower requirements in terms of computational resources and time. ## related work large pre-trained models such as bert @xcite and roberta @xcite have become ubiquitous in the nlp field @xcite . in terms of domain shifts, secondarystage unsupervised pre-training on a new domain has proven to be advantageous. contextualized tokenizations are adapted to text from the target domain through masked language modeling, as introduced by @xcite and @xcite . meanwhile, @xcite performed continual pre-training to adapt the bert model to the biomedical domain, utilizing both pubmed abstracts and pmc full-text resources. the use of contrastive learning also increases the representation ability for specific domains. @xcite investigated the use of contrastive learning to develop discriminative entity representations in the field of cross-domain named entity recognition. however, many specialized domains contain unique terms that are not included in the vocabulary of pre-trained language models. @xcite proposed a biomedical pre-trained model called pubmedbert in which the vocabulary was constructed from scratch and the model was pretrained from scratch. furthermore, in many specialized domains, sufficiently large corpora may not be available to support pre-training from scratch. general domain vocabulary can be extended with in-domain vocabulary @xcite in order to solve this out-of-vocabulary issue. knowledge distillation (kd) @xcite aims to transfer the knowledge from a large teacher model to a small student model. distil-bert @xcite uses soft labels and embedding outputs to supervise the student model. tinybert @xcite and mobilebert @xcite includes self-attention distributions and hidden-states for training the student model. minilmv2 @xcite avoids restrictions on the number of student layers and supervises the student model by using the self-attention distributions and value relation of the teacher's last transformer layer. the ad-kd approach @xcite explores the token-level rationale behind the teacher model based on integrated gradients and transfers attribution knowledge to the student model. several frameworks for general knowledge distillation using llms have also been proposed. gkd @xcite is a general-knowledge distillation framework that supports distillation on larger-scale plms using various distillation methods and f -distill @xcite formulates sequence level knowledge distillation through minimization of a generalized f-divergence function. @xcite and @xcite have proved that distillation using extracted rationales from black box llms is effective. in this paper, we focus on task-agnostic knowledge distillation approaches, where a smaller distilled pre-trained model can be directly fine-tuned on downstream tasks. ## methods in this study, we propose ve-kd, a knowledge distillation method with vocabulary expansion, as shown in figure 1 . unlike adapt-and-distill @xcite , which requires 4-step training, our approach simultaneously lightens the model and resolves the adaptability issues of special domains, which have been a problem in general-purpose models pre-trained on large corpora. by continually distilling knowledge from the teacher model, ve-kd effectively avoids overfitting, a common issue that arises during the domain adaptation phase, especially when working with small corpora. in the knowledge distillation aspect of ve-kd, a larger bert model serves as the teacher model, instructing a smaller student model. through the distillation process, the student model learns to mimic the behavior of the larger teacher model in general terms. simultaneously, the vocabulary expansion aspect broadens the model's vocabulary to capture domain-specific terms, thereby enhancing the method's ability to adapt to domain-specific tasks. ## experiment details and results in this section, we conduct our experiments in the biomedical domain. ## analysis in this section, we analyze the impact of training time and various settings on performance. ## conclusion in this paper, we proposed ve-kd, a novel method that merges vocabulary expansion and knowledge distillation. we also showed that our method achieves competitive performance on various downstream tasks. our experimental results demonstrate that ve-kd is effective; its performance is competitive with well-known models such as biobert and pubmedbert, and its training efficiency is noteworthy. it outperforms distilbert and adaptand-distill method, especially in document-level tasks. furthermore, ve-kd is more robust compared to general domain adaptation. ve-kd using distillation mechanism which can avoid overfitting, especially work with small corpora. we thoroughly investigated the effects of vocabulary size and tolerance and obtained insights that can help us configure more efficient models. because of its efficiency across various domainspecific nlp tasks, ve-kd lays the groundwork for further research in task-specific model optimization and application across diverse domains. one limitation of our study is that we did not evaluate the model's generalization abilities on outof-domain tasks. another limitation is that we have not yet fully explored the applicability of the ve-kd method to other model structures. for instance, we are looking for more efficient and fitting methods of knowledge segmentation and aggregation for generative language models.