Venue: NAACL 2022

# Enhancing Self-Attention with Knowledge-Assisted Attention Maps

 Jiangang Bai,1, Yujing Wang2, Hong Sun2, Ruonan Wu2, Tianmeng Yang1, Pengfei Tang2, Defu Cao1,*, Mingliang Zhang1,*, Yunhai Tong1, Yaming Yang2, Jing Bai2, Ruofei Zhang2, Hao Sun2, Wei Shen2

The work was done when the author visited Microsoft.



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 2D-convolution 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:

* First, we propose a novel self-attention mechanism enhanced by semantic attention maps, which incorporates knowledge from entity, phrase segmentation, and term correlation. Ablation study will demonstrate the effectiveness of these kinds of semantic attention maps.
* Second, KAM-BERT requires little extra memory and computation cost compared to vanilla BERT and other SOTAs. It can be finetuned efficiently on existing language models for a given application. We have successfully applied it to improve the performance of query-ad relevance in a commercial search engine.
* Last but not least, the proposed framework is generic and flexible for infusing various kinds of knowledge into the transformer architecture. Except for the three kinds of knowledge considered in this paper, we will also showcase how to incorporate other kinds of information, such as a knowledge graph. It opens up new opportunities for further exploration.

## 2 Related Works

After the NLP landmark models Transformer Vaswani et al. (2017) and BERT Devlin et al. (2019), 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

Figure 1: Visualization of attention map from vanilla BERT for a case of query-ad matching



performances of downstream tasks. ERNIE (Zhang et al., 2019) makes a preliminary attempt to utilize knowledge graph to improve the performance of knowledge-driven tasks. LIBERT (Lauscher et al., 2019) injects pairs of words with synonym and hyponym relations in WordNet. SenseBERT (Levine et al., 2020) considered word-supersense knowledge by predicting the supersense of the masked word. KnowBERT (Peters et al., 2019) incorporates knowledge bases into BERT through Knowledge attention and re-contextualization. WKLM (Xiong et al., 2019) 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 (Wang et al., 2020) 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.

## 3 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 self-attention 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.

### Self-Attention

The representation of a text sequence can be written as \({\bf X}\in{\bf R}^{N\times C}\), where \(N\) denotes the sequence length and \(C\) is the word embedding dimension size. A standard Transformer block is composed of a self-attention layer and a position-wise feed-forward layer, while each attention map is generated by a self-attention layer without any other prior knowledge introduced.

The self-attention mechanism plays an important role in the transformer-based model. In a vanilla Transformer, the self-attention map \({\bf A_{sa}^{i}}\) of layer \(i\) is calculated by the dimension-normalized dot-product operation.

\[{\bf A_{sa}^{l}=} \text{Self-Attention}({\bf X})=\frac{{\bf QK}^{\top}}{\sqrt{d}} \tag{1}\]

where \(d\) is the dimension of representation vectors. In a vanilla transformer, \({\bf A_{sa}^{l}}\) is then normalized by softmax and fed into position-wise feed-forward layers. In KAM-BERT, the self-attention map \({\bf A_{sa}^{i}}\) is infused with semantic attention maps to calculate the final attention matrix, which will be described in the following sub-sections.

### Semantic Attention Maps

Given a sequence of tokens \(\{x_{0},x_{1},...,x_{n-1}\}\), the semantic attention map can be defined in a generic form \(M\in\mathbb{R}^{n\times n}\), where \(M_{i,j}\in[0,1]\) denotes the attention score from token \(i\) to token \(j\), and \(n\) is the number of tokens in the current sentence. Then, for a specific kind of knowledge, we need a corresponding relevance function to calculate the attention score, _i.e._, \(M_{i,j}=Relevance(x_{i},x_{j})\). Note that if \(x_{i}\) denotes a sub-word as in the BERT model, we define \(M_{i,j}=Relevance(W(x_{i}),W(x_{j}))\), where \(W(x_{i})\) denotes the entire word which contains the sub-word \(x_{i}\). For example, BERT will convert a sequence _"I like tacos !"_ into a sequence of sub-words, _i.e._, _[_I, like, ta, #cos, !_]_, where _"ta"_ and _"#cos"_ are sub-words from _"tacos"_, so both \(W(ta)\) and \(W(\#\# cos)\) denote the word _"tacos"_. In Section 3.4, we will introduce three kinds of semantic attention maps considered in this paper and the generation method for other knowledge-assisted attention maps.

### Multi-Channel Knowledge Infusion

In order to incorporate external knowledge into self-attention, we concatenate semantic attention maps with vanilla self-attention, and then infuse them into a single multi-head attention map using multi-channel 2D-convolutions. Applying 2D-convolution to a self-attention map is first pro 

posed by (Wang et al., 2021) and shows advantages in both NLP and CV tasks. Here we use 2D-convolution to infuse different kinds of knowledge.

First, we perform _Channel Wise Concatenation (CWC)_: the vanilla self-attention map \({\bf A^{i}_{sa}}\) will be concatenated with each semantic attention map \({\bf M_{i}}\) separately along the channel dimension. Then, a multi-channel 2D-convolution is applied to generate an knowledge-infused attention map, denoted by \({\bf A^{i}_{sem}}\). The entire process can be formulated as below.

\[{\bf A^{i}_{sem}}=Conv(CWC({\bf A^{i}_{sa}|M_{1..k})}) \tag{2}\]

where \({\bf M_{1..k}}\) is a set of semantic attention maps obtained by \(k\) different kinds of knowledge, including but not limited to the three ones considered in this paper; To infuse different types of knowledge, we apply a standard 2D convolution operation, the output dimension of which is the same as that of \({\bf A_{sa}}\). If \({\bf A_{sa}}\) has \(m\) attention heads, then \({\bf A_{sem}}\) will also has \(m\) attention heads. We adopt \(3\times 3\) kernel for the convolution empirically as it performs better than \(1\times 1\) and \(5\times 5\) kernels according to (Wang et al., 2021).

At last, we adopt a hyper-parameter \(\alpha\) to balance the importance of \({\bf A^{i}_{sa}}\) and \({\bf A^{i}_{sem}}\).

\[{\bf A^{i}}=Softmax(\alpha\cdot{\bf A^{i}_{sem}}+(1-\alpha)\cdot{\bf A^{i}_{ sa}}) \tag{3}\]

After softmax activation, we get the final self-attention map \({\bf A^{i}}\) with \(m\) heads for the \(i\)-th layer. Given the self-attention map, the rest components are identical to a vanilla Transformer, which can be calculated as

\[{\bf h^{i}_{j}}={\bf A^{i}_{j}V^{i}},{\bf H^{i}}=(\bigoplus_{j=1}^{m}{\bf h^{ j}}){\bf W^{O}},j\in m. \tag{4}\]

In detail, we use the obtained attention map \({\bf A^{i}}\) to multiply the value matrix \({\bf V}\) in the attention mechanism to get the representation \({\bf h^{j}}\) of the \(j\)-th attention head. Next, the outputs of all attention heads from each layer will be concatenated along the embedding dimension. Finally, we multiply this result with a linear transformation matrix \({\bf W^{O}}\) to get the output representation of the \(i\)-th KAM-BERT layer. Besides, we add a skip connection from the result attention map in the \(i\)-th layer to the self-attention map of the \(i+1\) layer to enhance the flow of information between layers.

### Generation of Semantic Attention Maps

The knowledge we inject into KAM-BERT includes entity information, phrase segmentation information, and term correlation information. We consider these three types of knowledge because they reflect language semantics from different perspectives. Entity and phrase represent coherence between adjacent words while term correlations build a semantic bridge for related words which may be far away or even unseen in the current sentence. The first two kinds of knowledge are presented as labeled sequences, and the last one is presented as relationship between tokens. As defined in Section 3.2, a specific kind of knowledge

Figure 2: Overview of KAM-BERT calculation flow. (a) is our KAM-BERT architecture, (b) (c) and (d) are specific knowledge-assisted attention maps calculated from entity, phrase segmentation and term correlation. We illustrate the first two transformer blocks here.

 

can be transferred to semantic attention maps once the corresponding \(Relevance\) function is defined. In the following paragraphs, we will demonstrate how to define \(Relevance\) functions for the three types of knowledge used in this paper. Also, we need to emphasize that the proposed framework is generic and is feasible to incorporate other semantic information like a knowledge graph. Thus, we will discuss how to generate other knowledge-assisted attention maps as our future work.

**Entity Attention Map** Named Entity Recognition (NER) Nadeau and Sekine (2007) is a standard task for natural language processing which has been studied for years. Mathematically, a entity extractor transforms the sequence of tokens \(\{x_{0},x_{1},...,x_{n-1}\}\) into a sequence of labels \(\{label_{0},label_{1},...,label_{n-1}\}\), where \(label_{i}\) falls into one of three classes, denoting non-entity words, starting words in entities and other words in entities. Based on the labeled sequence, the \(Relevance\) function for entity attention map can be defined as

\[Relevance(W_{i},W_{j})=\begin{cases}1&W_{i}\equiv_{E}W_{j}\\ 0&otherwise\end{cases} \tag{5}\]

where \(A\equiv_{E}B\) denotes \(A\) and \(B\) belong to the same entity.

**Phrase Segmentation Attention Map** Similar to the entity attention map, one can highlight the term correlations within the same phrase segmentation to emphasize the locality inductive bias. Syntax tree is a generic source to extract phrases in different semantic levels, which can be generated by a trained syntax parser. In a syntax tree, each internal node represents a phrase segment for a specific level. For example, we can select the parents of leaf nodes in the syntax tree as the root of each sub-tree which represents a phrase. We define the distance of an internal node \(i\) to the leaf node as \(level(i)\). Thus, the relevance function can be computed by

\[Relevance(W_{i},W_{j})=\begin{cases}1&W_{i}\equiv_{T}W_{j}\\ 0&otherwise\end{cases} \tag{6}\]

where \(A\equiv_{T}B\) denotes that \(A\) and \(B\) belong to the same sub-tree at \(level(i)\).

**Term Correlation Attention Map** In computational linguistics, Pointwise Mutual Information (PMI) has been widely used for finding associations between words Arora et al. (2016). In our work, we adopt PMI to measure the semantic correlations between terms. The PMI of a pair \((x,y)\) from discrete random variables \((X,Y)\) quantifies the discrepancy between the probability of their coincidence given joint distributions and individual distributions. We define the PMI-based relevance function as

\[PMI(x;y)= \log\frac{p(x,y)}{p(x)p(y)} \tag{7}\] \[Relevance(W_{i},W_{j})= PMI(W_{i};W_{j})/Z\]

where \(Z\) denotes the normalized factor of PMI matrix. In our experiments, PMI is calculated using a large web corpus. We calculate the probability \(p(x,y)\) of a word pair appearing jointly, and the probability of single word appearance is denoted as \(p(x)\) and \(p(y)\). Finally, we use \(\log p(x,y)-\log p(x)-\log p(y)\) to compute the PMI score.

To better incorporate semantic knowledge into attention maps, we further enrich each attention map by adding top \(k\) terms which do not appear in the current sentence but hold the highest average PMI scores with terms in the original sentence. Note that we should expand the selected \(k\) words to \(K\) subwords for BERT. Then the subwords will be appended to the original sentence. After augmentation, the input sentence has \(N+K\) words and the shape of an attention map becomes \((N+K)\times(N+K)\), where \(N\) and \(K\) denote the number of original terms and auxiliary terms respectively (see an example in Fig. 2(d)). In order to align the shapes of different attention maps (including the vanilla self-attention map), we add zero-padding for smaller ones. After passing one transformer layer, the output sequence length is still \(N+K\). Note that the auxiliary words are only utilized to enrich the semantics of original word representations, which is done within each transformer layer. Thus, we trim the output sequence length to \(N\) before taking it as input to the next transformer layer (while a new round of augmentation will be done in the next layer).

**Other Knowledge-assisted Attention Maps.** The KAM-BERT framework is generic and can be extended to other kinds of knowledge in future works. For each semantic type, we can define a specific \(Relevance\) function to transfer the corresponding information into semantic attention maps. For example, we can define the relevance function for a Knowledge Graph (KG) as:

\[Relevance(W_{i},W_{j})=\begin{cases}1&E(W_{i})\equiv_{KG}E(W_{ j})\\ 0&otherwise\end{cases}\] 

where \(E(W_{i})\) is the corresponding entity of word or sub-word \(W_{i}\) in a KG, and \(A\equiv_{KG}B\) represents that both \(A\) and \(B\) exist and are adjacent in a KG.

## 4 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 .

### Semantic Information Extraction

We use Stanza library Qi et al. (2020) to extract NER information and syntax information. Stanza NER takes one sentence as input and returns the start and end indices of the corresponding named entity in the sentence. While Stanza Parser can extract the corresponding syntax tree for each sentence. We use query-ad logs from a commercial search engine to calculate PMI matrix. These steps gain the knowledge required to generate the semantic attention maps mentioned in Section 3.4.

### GLUE Benchmark

The GLUE benchmark offers a collection of tools for evaluating the performance of models across a diverse set of NLP applications. It contains single-sentence classification tasks CoLA and SST-2), similarity and paraphrase tasks (MRPC, QQP and STS-B) and pairwise inference tasks (MNLI, RTE and QNLI). We use the default train/dev/test split for each dataset. The hyper-parameters are chosen based on the validation set (refer to appendix for details). After the model is trained, we make predictions on the test data and send the results to GLUE online evaluation service1 to get testing scores. Note that the original WNLI dataset in the GLUE benchmark is problematic, which causes the evaluation results to be 65.1. In order to make a fair comparison, most papers Devlin et al. (2019); Liu et al. (2019); Yang et al. (2019) choose to ignore the results of WNLI when calculating GLUE average score.

Footnote 1: [https://gluebenchmark.com](https://gluebenchmark.com)

The scores on all datasets in GLUE benchmark are listed in Table 1. We report test scores on BERT-Base, BERT-Large, RoBERTa related models and their corresponding enhanced models. The performances of BERT-Base, BERT-Large and RoBERTa-Large are reproduced using the official checkpoints provided by corresponding authors.

As shown in the table, our models outperform all corresponding baselines. KAM-BERT-Base achieves an average GLUE score of 78.7, lifting 1.2 scores from standard BERT-Base model with only a few extra parameters introduced to the baseline model. Particularly, the improvements on CoLA datasets are fairly large, showing that our knowledge integration method has good generalization performance for natural language inference and understanding. ERNIE have also added external information such as entity and knowledge graph, but it needs much more time for a joint re-training. As for BERT-Large and its counterpart KAM-BERT-Large, the average improvement on GLUE benchmark is 0.9. We can see that the improvement becomes smaller when the model grows larger, because larger models often capture more semantic knowledge in the pre-training phrase. But incorporating explicit knowledge is still indispensable for achieving a superior performance.

### Question Answering

We conduct experiments on two kinds of question answering tasks, i.e., commonsense QA and open-domain QA. Commonsense QA aims to answer questions with commonsense. We adopt CosmosQA for evaluation, which requires commonsense-based reading comprehension formulated as multiple answer selection. Open-domain QA aims to answer questions using external resources such as collections of documents and webpages. We consider two public datasets for this task, i.e., Quasar-T and SearchQA.

The results of CosmosQA are shown in Table 2. Compared with BERT-Large, KAM-BERT-Large achieves 10.6% improvement in accuracy. KAM-RoBERta-Large further improves the accuracy of RoBERTa-Large by 5.4%, which indicates that our models has better knowledge inference ability. For open-domain QA, our model also achieves better results compared to corresponding baselines. This because that KAM-based models can make full use of the infused knowledge. At the same time, one can notice that KAM-based models have fewer parameters than K-Adaptor, demonstrating its effectiveness for knowledge infusion. WKLM Xiong et al. (2019) forces the pre-trained language model to incorporate knowledge from a knowledge graph. This makes WKLM to achieve a better score on QA tasks, but KAM-BERT performs even better than WKLM.

 



### LAMA Benchmark

To further verify whether KAM-BERT better integrate internal knowledge into pre-trained language models, we conduct experiments on LAMA, a widely used benchmark for knowledge probing. LAMA examines models' abilities on recalling relational facts by cloze-style questions. The first place micro-averaged accuracy is used as evaluation metrics. The evaluation results are shown in Table 3. KAM-BERT consistently outperforms corresponding baselines on all tasks. It indicates that KAM-BERT can generate better attention maps with semantic guidance.

### Query-Ad Relevance

Query-ad relevance measures how relevant a search ad matches with a user's search query. Very often queries and ads have words with special meanings, which are not easily understood well by traditional NLP techniques but can benefit from the knowledge-assisted mechanism proposed in this work. Besides, user queries and ads text often contain noises, so evaluation on query-ad relevance task would test our model's robustness and resistance of noise. We compare BERT and KAM-BERT on a large-scale internal dataset of a commercial search engine. As shown in Table 5, our model outperforms corresponding baselines by a large margin, which is statistically significant under 95% confidence interval. One thing to call out is that, although NER and syntax parsing results are nosier comparing to the ones in academic datasets, we still have good improvements on this dataset. This indicates the way we combine those knowledge together makes our model more robust to noisy inputs.

### Model Analysis

In this section, we explore the sensitivity of hyper-parameter \(\alpha\), and then conduct ablation experiments on three types of added knowledge.

**Hyper-parameter Analysis** The optimal \(\alpha\) value after grid search is 0.2, which means that the original attention maps still dominate the token relationships. We chose three tasks from different fields to do ablation study for \(\alpha\). Our model is KAM-BERT-Base, and its performance is shown in Table 4. In three different tasks, setting \(\alpha\) to 0.2 achieves the best results. An intuitive understanding is that when \(\alpha\) is small, external knowledge plays an unimportant role and cannot participate in

\begin{table}
\begin{tabular}{l c c c c c c c c c c} \hline \hline
* \\ WKLM & 348M & 58.763/3 & 43.749/9 &
- \\ WKLM+Ranking & 348M & 61.766/76.458/52.2 & – \\
* \\ RoBERTa multitask & 355M & 59.966/74.4 & 465.12 & 81.2 &
- \\ K-Adapter & 384M & 62.067/3.4 & 46.353/0.3 & 81.8 &
**KAM-RoBERa-Large (Ours)** & 361M & **64.468** & **46.653** & **81.9** \\ \hline \hline \end{tabular}
\end{table}
Table 4: Results for Ad Relevance

\begin{table}
\begin{tabular}{l l c c c c c c c c c} \hline \hline
**Model** & **\#Params** & **CoLA** & **SST-2** & **MRPC** & **STS-B** & **QQP** & **MNLI-m/mm** & **QNLI** & **RTE** & **WNLI** \\ \hline BERT-Base & 109.5M & 77.5 & 50.1 & 92.6 & 88.784/3.8 & 85.784/6.7 & 71.0/88.9 & 83.6/83.2 & 89.4 & 67.9 & 65.1 \\ ERNIE & 114M & 77.5 & 52.3 & **93.5** & 88.2/– & 83.2/– & 71.2/– & 84.0/83.2 & **91.3** & **68.8** & 65.1 \\ SenseBERT & 133M & 77.9 & **54.6** & 92.2 & **89.2/82.8** & 83.5/82.3 & 70.3/88.8 & 83.6/– & 90.6 & 67.5 & 65.1 \\ CorefBERT & 110M & 77.6 & 51.5 & 93.7 & 89.1/– & 85.8/– & 71.3/– & 84.2/**83.5** & 90.5 & 67.2 & 65.1 \\
**KAM-BERT-Base** & 112M & **78.7** & 53.7 & 93.2 & **89.2**/85.0 & **86.1/84.8** & **71.2/89.4** & **84.3/83.5** & 90.6 & 67.9 & 65.1 \\ \hline BERT-Large & 345M & 80.5 & 60.5 & 94.9 & 89.3/85.4 & 87.6/86.5 & 72.1/89.3 & 86.8/85.9 & 92.7 & 70.1 & 65.1 \\
**KAM-BERT-Large** & 350.4M & **81.4** & **61.5** & **95.3** & **91.4/88.1** & **88.5/87.3** & **72.5/89.5** & **87.1/86.6** & **93.0** & **72.0** & 65.1 \\ \hline RoBERTa-Large & 355M & 83.9 & 63.8 & 96.3 & 91.0/89.4 & 89.4/87.9 & 72.7/90.1 & 89.5/89.7 & 94.2 & 84.2 & 65.1 \\ KEPLER & 360M & 84.4 & 63.6 & 94.5 & 89.3/— & 91.2/– & 91.7/– & 87.2/86.5 & 92.4 & 85.2 & 65.1 \\
**KAM-RoBERa-Large** & 361M & **84.6** & **65.1** & **96.6** & **92.4/90.7** & **90.2/88.4** & **73.2/90.5** & **89.8/89.7** & **94.6** & **85.4** & 65.1 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Comparison of different model backbones on GLUE benchmark

\begin{table}
\begin{tabular}{c c c c} \hline \hline
**Model** & **SQuAD** & **Google-RE** & **T-REx** & **ConceptNet** \\ \hline BERT-Large & 14.1 & 9.8 & 31.1 & 15.6 \\
**KAM-BERT-Large** & **14.5** & **10.1** & **32.2** & **16.0** \\ \hline RoBERTa-Large & 15.9 & 11.3 & 33.7 & 17.1 \\
**KAM-RoBERa-Large** & **16.3** & **11.9** & **34.5** & **17.3** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Results for knowledge probing benchmark LAMA

\begin{table}
\begin{tabular}{l c c c} \hline \hline
**Model** & **SoAuD** & **Google-RE** & **T-REx** & **ConceptNet** \\ \hline BERT-Large & 14.1 & 9.8 & 31.1 & 15.6 \\
**KAM-BERT-Large** & **14.5** & **10.1** & **32.2** & **16.0** \\ \hline RoBERTa-Large & 15.9 & 11.3 & 33.7 & 17.1 \\
**KAM-RoBERa-Large** & **16.3** & **11.9** & **34.5** & **17.3** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Results for Question Answering

the entire training process well. With the gradual increase of \(\alpha\), the intervention of external knowledge on the attention map will increase, and the attention relationship in the original sequence will be gradually lost, resulting in the decline of model performance.

**Ablation Study** For a comprehensive understanding of our model design, we conduct ablation study with the following settings in Table 6. (1) _KAM-BERT w/o PMI_: the PMI-based attention maps are removed; (2) _KAM-BERT w/o Phrase_: the phrase-based attention maps are removed; (3) _KAM-BERT w/o Entity_: the entity-based attention maps are removed. (4) _KAM-BERT w/o Convolution_: the convolution layers for knowledge-assisted attention map integration are removed. Instead, we merge all the attention maps through average aggregation.

The average scores of all ablation experiments are better than BERT, but are relatively worse than KAM-BERT, demonstrating all the components are beneficial for the final performance. At the same time, we observe that after deleting the entity attention map, the score of KAM-BERT drops drastically from 78.7 to 77.7. This shows that the gain brought by entity information is the greatest. In addition, the convolution layer is indispensable for achieving a superior performance.

### Case Study

In Figure 3, we visualize an example of query-ad relevance, where the query is _"buy glipizide"_ and ad text is _"Fred Meyer Pharmacy Near Me"_. The darker color in the figure represents a higher attention score. Figure 3(a) is the attention map of vanilla BERT without adding explicit knowledge. When encountering rare words like _"glipizide"_, the self-attention mechanism cannot do a good job to decide which terms should _"glipizide"_ attend to. But in Figure 3(b), the attention map of KAM-BERT uses term correlations to learn that _"slipizzide"_ is a medicine, so it focuses on the medicine-related tokens like _"Pharmacy"_ and _"antidiabetic"_. Figure 3(c) and 3(d) are the attention maps for the sentence _"The Academy of Fine Arts is located in Northern Maidan."_ We only visualize the key parts due to space limitation. Figure 3(c) shows the attention map for vanilla BERT, and 3(d) is the attention map for KAM-BERT. It can be observed that the terms in the same entity are highly correlated with each other in KAM-BERT, which is more reasonable than vanilla BERT.

## 5 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.

Figure 3: Visualization of attention maps in BERT and KAM-BERT.

 



## References

* Arora et al. (2016) Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. 2016. A latent variable model approach to pmi-based word embeddings. _TACL_, 4:385-399.
* Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. Bert: Pre-training of deep bidirectional transformers for language understanding. In _NAACL-HLT (1)_.
* Hewitt and Manning (2019) John Hewitt and Christopher D Manning. 2019. A structural probe for finding syntax in word representations. In _Proceedings of the 2019 NAACL-HLT(1)_, pages 4129-4138.
* Jain and Wallace (2019) Sarthak Jain and Byron C Wallace. 2019. Attention is not explanation. In _Proceedings of the 2019 NAACL-HLT(1)_, pages 3543-3556.
* Jawahar et al. (2019) Ganesh Jawahar, Benoit Sagot, and Djame Seddah. 2019. What does bert learn about the structure of language? In _Proceedings of the 57th Annual Meeting of the ACL_, pages 3651-3657.
* Lauscher et al. (2019) Anne Lauscher, Ivan Vulic, Edoardo Maria Ponti, Anna Korhonen, and Goran Glavas. 2019. Informing unsupervised pretraining with external linguistic knowledge. _arXiv preprint arXiv:1909.02339_.
* Levine et al. (2020) Yoav Levine, Barak Lenz, Or Dagan, Ori Ram, Dan Padnos, Or Sharir, Shai Shalev-Shwartz, Amnon Shashua, and Yoav Shoham. 2020. SenseBERT: Driving some sense into BERT. In _Proceedings of the 58th Annual Meeting of the ACL_.
* Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. Roberta: A robustly optimized BERT pretraining approach. _CoRR_, abs/1907.11692.
* Nadeau and Sekine (2007) David Nadeau and Satoshi Sekine. 2007. A survey of named entity recognition and classification. _Lingvisticeae Investigatione_, 30(1):3-26.
* Peters et al. (2019) Matthew E Peters, Mark Neumann, Robert Logan, Roy Schwartz, Vidur Joshi, Sameer Singh, and Noah A Smith. 2019. Knowledge enhanced contextual word representations. In _Proceedings of EMNLP-IJCNLP'19_, pages 43-54.
* Qi et al. (2020) Peng Qi, Yuhao Zhang, Yuhui Zhang, Jason Bolton, and Christopher D. Manning. 2020. Stanza: A Python natural language processing toolkit for many human languages. In _Proceedings of the 58th Annual Meeting of the ACL: System Demonstrations_.
* Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. _Advances in neural information processing systems_, 30:5998-6008.
* Wang et al. (2020) Ruize Wang, Duyu Tang, Nan Duan, Zhongyu Wei, Xuanjing Huang, Cuihong Cao, Davin Jiang, Ming Zhou, et al. 2020. K-adapter: Infusing knowledge into pre-trained models with adapters. _arXiv preprint arXiv:2002.01808_.
* Wang et al. (2019) Xiaozhi Wang, Tianyu Gao, Zhaocheng Zhu, Zhiyuan Liu, Juanzi Li, and Jian Tang. 2019. Kepler: A unified model for knowledge embedding and pretrained language representation. _arXiv preprint arXiv:1911.06136_.
* Wang et al. (2021) Yujing Wang, Yaming Yang, Jiangang Bai, Mingliang Zhang, Jing Bai, Jing Yu, Ce Zhang, Gao Huang, and Yunhai Tong. 2021. Evolving attention with residual convolutions. _arXiv preprint arXiv:2102.12895_.
* Xiong et al. (2019) Wenhan Xiong, Jingfei Du, William Yang Wang, and Veselin Stoyanov. 2019. Pretrained encyclopedia: Weakly supervised knowledge-pretrained language model. In _ICLR_.
* Yang et al. (2019) Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. 2019. XLnet: Generalized autoregressive pretraining for language understanding. In _Advances in Neural Information Processing Systems_, volume 32. Curran Associates, Inc.
* Zhang et al. (2019) Zhengyan Zhang, Xu Han, Zhiyuan Liu, Xin Jiang, Maosong Sun, and Qun Liu. 2019. Ernie: Enhanced language representation with informative entities. In _Proceedings of the 57th Annual Meeting of the ACL_, pages 1441-1451.