Plug-and-Play Knowledge Injection for Pre-trained Language Models (2023)

Abstract: Injecting external knowledge can improve the performance of pre-trained language models (PLMs) on various downstream NLP tasks. However, massive retraining is required to deploy new knowledge injection methods or knowledge bases for downstream tasks. In this work, we are the first to study how to improve the flexibility and efficiency of knowledge injection by reusing existing downstream models. To this end, we explore a new paradigm plug-and-play knowledge injection , where knowledge bases are injected into frozen existing downstream models by a knowledge plugin . Correspondingly, we propose a plug-and-play injection method map-tuning , which trains a mapping of knowledge embeddings to enrich model inputs with mapped embeddings while keeping model parameters frozen. Experimental results on three knowledge-driven NLP tasks show that existing injection methods are not suitable for the new paradigm, while map-tuning effectively improves the performance of downstream models. Moreover, we show that a frozen downstream model can be well adapted to different domains with different mapping networks of domain knowledge. Our code and models are available at https://github.com/THUNLP/Knowledge-Plugin .

## introduction recent years have witnessed rapid development in enhancing pre-trained language models (plms) with various external knowledge bases, i.e., knowledge injection for plms (levine et al., 2020; zhou et al., 2020; zhang et al., 2019; peters et al., 2019; * equal contribution † corresponding authors figure 1 : illustration of plug-and-play knowledge injection, where knowledge bases and models are decoupled. @xcite . knowledge injection improves the performance of plms on a wide range of tasks such as information extraction @xcite , question answering @xcite , and text generation @xcite . existing injection methods commonly inject knowledge by knowledge-aware pre-training or fine-tuning @xcite @xcite . however, rarely studied is how to inject knowledge into a downstream model that is already adapted to a specific task. if we want to apply a new knowledge injection method to enhance models on a specific task, we have to discard task-specific downstream models and retrain them. in addition, one downstream model working with multiple knowledge bases requires retraining itself to inject each knowledge base. retraining models is time-consuming and resource-intensive, leading to the need for a flexible and efficient injection paradigm. toward flexible and efficient injection, we explore a novel paradigm plug-and-play knowledge injection, where knowledge bases are injected into frozen existing downstream models by knowledge modules. the knowledge module bridges the knowledge base and the model, and we call it a plugin vividly. under this paradigm, a downstream model would have multiple plugins, each corresponding to a combination of an injection method and a knowledge base, which ensures flexibility. moreover, knowledge plugins should be small enough to ensure efficiency. intuitively, as shown in figure 1 , we treat models and knowledge bases as computers and flash disks, respectively. in this work, we study two settings for the plugand-play knowledge injection paradigm. the first is general plug-and-play knowledge injection, aiming to inject knowledge into all downstream models (trained from a particular plm) by a general plugin without any task-specific training. in this setting, all downstream models share exactly one plugin for one combination of an injection method and a knowledge base. the second is task-specific plugand-play knowledge injection, where knowledge plugins are trained to better adapt to downstream tasks while keeping downstream models frozen. by our pilot study, we find that existing methods @xcite @xcite that can be used directly can not be well applied to the plug-andplay injection paradigm. to this end, we propose map-tuning, a preliminary exploration of learning knowledge plugins. specifically, we train a lightweight mapping network that augments model inputs with mapped knowledge representations, e.g., transe @xcite . to meet the general and task-specific injection requirements, we design general map-tuning and task-specific map-tuning, respectively. general map-tuning adopts language modeling as its objective to learn knowledge plugins and seeks better generalizability. task-specific map-tuning adopts task targets for plugin learning and seeks better task adaptation. we use three typical knowledge-driven nlp tasks to evaluate our plug-and-play knowledge injection, including relation classification @xcite , entity typing @xcite , and question answering @xcite . the experimental results show that: (1) after adapting plms to downstream tasks through full-parameter finetuning or parameter-efficient tuning, also known as delta tuning @xcite , injecting knowledge into these downstream models by general map-tuning leads to performance improvements in almost all cases; (2) using taskspecific map-tuning to inject domain knowledge further enables a frozen downstream model to work well in different domains. we hope our contribution can draw more attention to the plug-and-play knowledge injection paradigm and inspire more future research. two injection settings. as shown in figure 2(a), plug-and-play knowledge injection decouples knowledge injection from model training, which is different from existing paradigms. for general plug-and-play knowledge injection, m is obtained based on only p and b, and then it is directly plugged into all downstream models, d 1 , d 2 , . . ., without any additional training. for task-specific plug-and-play knowledge injection, it is allowed to train m 1 , m 2 , . . . for d 1 , d 2 , . . . respectively while keeping d 1 , d 2 , . . . frozen. challenges. the general plug-and-play knowledge injection poses serious challenges to methods designed for it. m is expected to improve the performance of d, yet d has never seen m or been seen by m during training. the only prior condition is that p and b are visible during training m. therefore, the designed methods for general injection need to endow m with enough generalizability such that m can adapt to unknown d 1 , d 2 , . . . well. even though the knowledge base b may have rich knowledge, without a good adaptation of m, useful information brought to d will be less than disruptive noise. the task-specific plug-and-play knowledge injection relaxes the restrictions, where m i is allowed to be trained with frozen d i . compared to injection during fine-tuning, the training of m i should be fast and the parameter number of m i should be small compared to that of d i . otherwise, the methods would be meaningless. hence, it requires simple and efficient architecture designs and informative training objectives for m i . ## map-tuning in this section, we first present the overall framework of map-tuning, which is designed for plugand-play knowledge injection. then, we show how to use it for general injection and task-specific injection, where the methods are called general maptuning and task-specific map-tuning respectively. ## related work to enhance plms with external knowledge, there are two mainstream paradigms: injection during pre-training and injection during fine-tuning @xcite . for injection during pre-training, researchers usually construct new knowledge-aware objectives, such as entity prediction @xcite , entity discrimination @xcite , entity and relation discrimination @xcite , and link prediction @xcite . in this way, knowledge will be implicitly stored in the parameters of plms. injection knowledge during pretraining can simultaneously improve performance on a range of downstream knowledge-driven tasks. however, the training cost of this paradigm is expensive. taking the typical knowledge-enhanced plms luke @xcite and ke-pler @xcite as an example, it takes more than 3,000 gpu hours to train them. injection knowledge during fine-tuning is a relatively lightweight paradigm, where external knowledge is often used to augment model inputs for specific tasks @xcite @xcite @xcite . when injecting unstructured textual knowledge, some methods retrieve task-related information from external corpora to augment the original input text @xcite . when using structured knowledge, such as knowledge graphs, existing methods usually apply knowledge representation learning methods @xcite to encode structured knowledge into embeddings, and then fuse these knowledge embeddings with input token embeddings using knowledge injection methods @xcite @xcite . in general, existing knowledge injection methods mainly target plms and adopt paradigms where knowledge and models are highly coupled. toward flexible and efficient injection, we study a new paradigm, plug-and-play knowledge injection, where we decouple models and knowledge sources, and then inject knowledge into downstream models without retraining the models. this work is also related to parameter-efficient tuning @xcite and plugins for large language models @xcite @xcite @xcite @xcite while we are the first to study knowledge injection in a parameter-efficient and plug-and-play way. in this work, we propose a new paradigm of injection toward flexible and efficient knowledge injection. in this paradigm, downstream models can be enhanced with little computational cost, which benefits large amounts of models. we first systematically evaluate existing knowledge injection methods and find that they are not suitable for plugand-play injection. then, we propose map-tuning for this paradigm, which effectively injects knowledge into downstream models to enhance them. there are four promising directions for future investigation into plug-and-play knowledge injection. (1) how can we reduce the performance gap between methods for this novel paradigm and those for the previous injection paradigms, while maintaining superior flexibility and efficiency? (2) besides factual knowledge, how can we effectively plug diverse knowledge bases, such as text corpora, voice, images, and even other plms? (3) after injecting the knowledge in a plug-and-play way, how can the plms do various types of complex reasoning based on the injected knowledge @xcite ? (4) can the plug-and-play knowledge injection methods for these sources be unified, so we can plug a combination of multiple sources? we hope this work can attract attention to and inspire research on these problems.