KALA: Knowledge-Augmented Language Model Adaptation (2022)

Abstract: Pre-trained language models (PLMs) have achieved remarkable success on various natural language understanding tasks. Simple fine-tuning of PLMs, on the other hand, might be suboptimal for domain-specific tasks because they cannot possibly cover knowledge from all domains. While adaptive pre-training of PLMs can help them obtain domain-specific knowledge, it requires a large training cost. Moreover, adaptive pre-training can harm the PLM’s performance on the downstream task by causing catastrophic forgetting of its general knowledge. To overcome such limitations of adaptive pre-training for PLM adaption, we propose a novel domain adaption framework for PLMs coined as Knowledge-Augmented Language model Adaptation (KALA), which modulates the intermediate hidden representations of PLMs with domain knowledge, consisting of entities and their relational facts. We validate the performance of our KALA on question answering and named entity recognition tasks on multiple datasets across various domains. The results show that, despite being computationally efficient, our KALA largely outperforms adaptive pre-training.

## introduction pre-trained language models (plms) @xcite have shown to be effective on various natural language understanding (nlu) tasks. although plms aim to address diverse downstream tasks from various data sources, there have been considerable efforts to adapt the plms to specific domains -distributions over the language characterizing a given topic or genre @xcite )-for which the acquisition of domain knowledge is required to accurately solve the downstream tasks (e.g., biomedical named entity recognition @xcite ). this problem, known as language model adaptation, can be viewed as a transfer learning problem @xcite under domain shift, where the model is pre-trained on the general domain and the labeled distribution is available for the target domain-specific task. the most prevalent approach to this problem is adaptive pre-training (figure 2a ) which further updates all parameters of the plm on a large domain-specific or curated task-specific corpus, with the same pretraining strategy (e.g., masked language modeling) before fine-tuning it on the downstream task @xcite @xcite . this continual pre-training of a plm on the target domain corpus allows it to learn the distribution of the target domain, resulting in improved performance on domain-specific tasks @xcite . while it has shown to be effective, adaptive pretraining has obvious drawbacks. first, it is computationally inefficient. although a plm becomes more powerful with the increasing amount of pretraining data @xcite , further pre-training on the additional data requires larger memory and computational cost as the dataset size grows @xcite . besides, it is difficult to adapt the plm to a new domain without forgetting the general knowledge it obtained from the initial pretraining step, since all pre-trained parameters are continually updated to fit the domain-specific corpus during adaptive pre-training @xcite . this catastrophic forgetting of the taskour kala framework embeds the unseen entities on the embedding space of seen entities by representing them with their relational knowledge over the graph, while the strong dapt baseline @xcite cannot appropriately handle unseen entities that are not given for task fine-tuning. the remaining step is how to integrate the knowledge into the plm during fine-tuning. to this end, we propose a novel layer named knowledgeconditioned feature modulation (kfm, §3.2), which scales and shifts the intermediate hidden representations of plms by conditioning them with retrieved knowledge representations. this knowledge integration scheme has several advantages. first, it does not modify the original plm architecture, and thus could be integrated into any plms regardless of their architectures. also, it only re-quires marginal computational and memory overhead, while eliminating the need of excessive further pre-training (figure 1 ). finally, it can effectively handle unseen entities with relational knowledge from the kg, which are suboptimally embedded by adaptive pre-training. for example, as shown in figure 2 , an entity restenosis does not appear in the training dataset for fine-tuning, thus adaptive pre-training only implicitly infers them within the context from the broad domain corpus. however, we can explicitly represent the unknown entity by aggregating the representations of known entities in the entity memory (i.e., in figure 2 , neighboring entities, such as asthma and pethidine, are used to represent the unseen entity restenosis). we combine all the previously described components into a novel language model adaptation framework, coined as knowledge-augmented language model adaptation (kala) (figure 3 ). we empirically verify that kala improves the performance of the plm over adaptive pre-training on various domains with two knowledge-intensive tasks: question answering (qa) and named entity recognition (ner). our contribution is threefold: ## related work language model adaptation nowadays, transfer learning @xcite is a dominant approach for solving natural language understanding (nlu) tasks. this strategy first pretrains a language model (lm) on a large and unlabeled corpus, then fine-tunes it on downstream tasks with labeled data @xcite . while this scheme alone achieves impressive performance on various nlu tasks, adaptive pre-training of the plm on a domain-specific corpus helps the plm achieve better performance on the domain-specific tasks. for example, @xcite demonstrated that a further pre-trained lm on biomedical documents outperforms the original lm on biomedical nlu tasks. also, @xcite showed that adaptive pre-training of the plm on the corpus of a target domain (domain-adaptive pre-training; dapt) or a target task (task-adaptive pre-training; tapt) improves its performance on domain-specific tasks. however, above approaches generally require a large amount of computational costs for pre-training. knowledge-aware lm accompanied with increasing sources of knowledge @xcite , some prior works have proposed to integrate external knowledge into plms, to enhance their performance on tasks that require structured knowledge. for instance, ernie @xcite and knowbert @xcite incorporate entities as additional inputs in the pretraining stage to obtain a knowledge-aware lm, wherein a pre-trained knowledge graph embedding from wikidata @xcite ) is used to represent entities. entity-as-experts (févry et al., 2020) and luke @xcite use the entity memory that is pre-trained along with the lms from scratch. erica @xcite further uses the fact consisting of entities and their relations in the pre-training stage of lms from scratch. previous works aim to integrate external knowledge into the lms during the pre-training step to obtain a universal knowledge-aware lm that requires additional parameters for millions of entities. in contrast to this, our framework aims to efficiently modify a general plm for the domain-specific task with a linear modulation layer scheme discussed in section 3.2, during fine-tuning. ## conclusion in this paper, we introduced kala, a novel framework for language model adaptation, which modulates the intermediate representations of a plm by conditioning it with the entity memory and the relational facts from kgs. we validated kala on various domains of qa and ner tasks, on which kala significantly outperforms relevant baselines while being computationally efficient. we demonstrate that the success of kala comes from both kfm and relational retrieval, allowing the plm to recognize entities but also handle unseen ones that might frequently appear in domain-specific tasks. there are many other avenues for future work, including the application of kala on pre-training of knowledge-augmented plms from scratch.