Enhancing Self-Attention with Knowledge-Assisted Attention Maps (2022)

Abstract: Large-scale pre-trained language models have attracted extensive attentions in the research community and shown promising results on various tasks of natural language processing. However, the attention maps, which record the attention scores between tokens in self-attention mechanism, are sometimes ineffective as they are learned implicitly without the guidance of explicit semantic knowledge. Thus, we aim to infuse explicit external knowledge into pre-trained language models to further boost their performance. Existing works of knowledge infusion largely depend on multi-task learning frameworks, which are inefficient and require large-scale re-training when new knowledge is considered. In this paper, we propose a novel and generic solution, KAM-BERT, which directly incorporates knowledge-generated attention maps into the self-attention mechanism. It requires only a few extra parameters and supports efficient fine-tuning once new knowledge is added. KAM-BERT achieves consistent improvements on various academic datasets for natural language understanding. It also outperforms other state-of-the-art methods which conduct knowledge infusion into transformer-based architectures. Moreover, we apply our model to an industry-scale ad relevance application and show its advantages in the real-world scenario.

## introduction language models pre-trained by a large text corpus have shown superior performances on a wide range of natural language processing tasks. many advanced models based on the transformer architectures achieve state-of-the-art results on various nlp benchmarks. existing literature @xcite shows that pre-training enables a model to capture syntactic and semantic information in the self-attention mechanism. however, the attention maps, which * the work was done when the author visited microsoft. recently, there have been multiple attempts for incorporating knowledge into transformer architectures. ernie @xcite and ke-pler @xcite utilize both large-scale textual corpora and knowledge graphs to train a representation model in a multi-task learning framework. they need to be retrained from scratch when injecting new knowledge, which is inefficient and can not benefit from existing pre-trained checkpoints. k-adapter @xcite integrates additional neural models to capture different kinds of knowledge. it enables adaptation based on pretrained language models. however, it does not instruct the self-attention mechanism directly and introduces a relatively large number of parameters to the original model. in this paper, we propose a novel and generic self-attention mechanism enhanced by explicit knowledge to address problems mentioned above. first, we show a failure case of query-ad matching, which motivates us to inject explicit knowledge into self-attention mechanism. in figure 1 , the attention map of a query-ad pair is visualized, and the goal is to judge if the query and ad text are semantically relevant. as shown in the figure, bert misclassifies this pair as irrelevant, probably because it does not understand the query word "glipizide", which rarely appears in the pre-training corpus. in fact, "glipizide" is a kind of medicine and highly related to the word "pharmacy" in ad text, so this case should be classified as relevant. in this case, if we know "glipizide" is semantically correlated to "pharmacy" as prior knowledge, we can enrich the attention maps accordingly. in addition, terms "fred" and "meyer" are from the same entity, so the attention scores between these two terms should be relatively high. based on this fact, we believe that simply using language models pre-trained by a general corpus is not enough to meet the satisfaction of a specific application. thus, our motivation is to inject explicit knowledge into the transformer architecture, which guides the pre-trained language model to perform better adaptation in an efficient fine-tuning procedure. to address the above motivation, we propose a novel architecture, namely kam-bert (knowledge-assisted attention maps for bert). first, it constructs semantic attention maps based on corresponding relevance functions defined by various kinds of semantic knowledge. specifically, we consider three kinds of semantic knowledge to guide the self-attention mechanism, i.e., entity, phrase segmentation, and term correlation. entity and phrase segmentation indicate the cohesion of continuous terms, while term correlation can help to enrich the semantic representation of a sentence. then, the knowledge infusion procedure is completed by concatenating these attention maps with vanilla self-attention and then performing 2dconvolution for integration. finally, the result attention maps are served as inputs for value projection, and the rest part is the same as a standard transformer. the kam-bert model can be fine-tuned on existing pre-trained checkpoints in a plug and play mode, which is highly efficient in practice. as illustrated in section 4, we compare kam-bert with bert and other knowledge-enhanced sotas on various natural language understanding tasks, where kam-bert shows consistent superiority. especially, we lift the average score of bert-base from 77.5 to 78.7 on the glue benchmark. we also demonstrate the advantage of kam-bert for knowledge injection on lama, a probing benchmark to analyze the factual and commonsense knowledge contained in a model. furthermore, kam-bert is successfully applied to the query-ad relevance scenario in a commercial search engine and shows significant lift in auc score. the major contributions of this paper are summarized as follows: ## related works after the nlp landmark models transformer @xcite and bert @xcite , many efforts have been devoted to pre-training representation models for natural language and utilizing them to benefit specific nlp tasks. most of them capture semantic information in the self-attention mechanism implicitly, while recent works demonstrated that injecting explicit knowledge significantly enhanced the performances of downstream tasks. ernie @xcite makes a preliminary attempt to utilize knowledge graph to improve the performance of knowledge-driven tasks. libert @xcite injects pairs of words with synonym and hyponym relations in wordnet. sensebert @xcite ) considered word-supersense knowledge by predicting the supersense of the masked word. @xcite incorporates knowledge bases into bert through knowledge attention and re-contextualization. wklm @xcite replaces entity mentions in the original document with names of other entities of the same type, and is trained to distinguish the correct entity mention from random ones. these models are supposed to be retrained when injecting new kinds of knowledge. k-adapter @xcite addresses this problem by plugging multiple ways of adapters for different kinds of knowledge. notably, most of methods above need to be pre-trained from scratch while our method do not. instead, we inject multiple kinds of knowledge directly into the self-attention maps and support efficient adaptation on a pre-trained language model. ## kam-bert as illustrated in figure 2 , kam-bert injects multiple kinds of knowledge into transformer-based pre-trained models in the form of multi-channel semantic attention maps. different kinds of knowledge can be extracted independently and infused together into one self-attention map in the transformer architecture. kam-bert can be fine-tuned directly from an existing checkpoint of bert, while additional parameters are initialized randomly and learned in the fine-tuning stage. thus, it is quite efficient and flexible to incorporate new kinds of knowledge into the model. below we first describe the standard selfattention mechanism. then, we will introduce the generic definition of semantic attention maps, as well as the methodology of multi-channel knowledge infusion which integrates semantic attention maps into transformer models. at last, the generation of different kinds of semantic attention maps will be presented. note that the time complexity of kam-bert is on-par with a vanilla bert. a detailed analysis can be found in the supplementary material. ## experiments we briefly introduced the extraction of semantic information in section 4.1. then we report experimental results on natural language understand and question answering tasks in section 4.2 and 4.3 respectively. in section 4.4, we show evaluation on lama, a benchmark especially designed to study how much semantic knowledge is contained in a language model. experiments on query-ad relevance is described in section 4.5. at last, we present ablation study in section 4.6 . ## conclusion in this paper, we proposed kam-bert, a flexible and efficient approach to inject knowledge into transformer-based pre-trained models. extensive experiments on glue and lama benchmark show that our approach outperforms all bert-style baselines and achieves new sota on qa tasks, suggesting that our models indeed integrate knowledge in an effective manner and have good generalization ability. in future work, we hope to investigate more types of knowledge which can be effectively integrated in our framework.