SlowGuess/NAACL_2022_OCR huggingface dataset:
Jiangang Bai $^{1,*}$ , Yujing Wang $^{2}$ , Hong Sun $^{2}$ , Ruonan Wu $^{2}$ , Tianmeng Yang $^{1}$ , Pengfei Tang $^{2}$ , Defu Cao $^{1,*}$ , Mingliang Zhang $^{1,*}$ , Yunhai Tong $^{1}$ , Yaming Yang $^{2}$ , Jing Bai $^{2}$ , Ruofei Zhang $^{2}$ , Hao Sun $^{2}$ , Wei Shen $^{2}$
<sup>1</sup>Peking University, Beijing, China
2Microsoft, Beijing, China
{pku_bjg,youngtimmy,cdf,zml24,yhtong}@pku.edu.cn
{yujwang,hosu,ruonan.wu,pengfeitang,yayaming,jbai,bzhang,hasun,sashen}@microsoft.com
Large-scale pre-trained language models have attracted extensive attention 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 pretrained 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.
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 (Jawahar et al., 2019; Hewitt and Manning, 2019) shows that pre-training enables a model to capture syntactic and semantic information in the self-attention mechanism. However, the attention maps, which
record the attention scores between tokens in a self-attention mechanism, are sometimes ineffective as they are learned implicitly without the guidance of explicit semantics (Jain and Wallace, 2019). If the knowledge can be leveraged in a reasonable way to guide the self-attention mechanism, we have a good chance to improve the quality of attention scores as well as the performance of downstream applications.
Recently, there have been multiple attempts for incorporating knowledge into transformer architectures. ERNIE (Zhang et al., 2019) and KE-PLER (Wang et al., 2019) 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 (Wang et al., 2020) 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
*The work was done when the author visited Microsoft.
Human Language Technologies, pages 107 - 115
Figure 1: Visualization of attention map from vanilla BERT for a case of query-ad matching
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.
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
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.
3.1 Self-Attention
The representation of a text sequence can be written as  $\mathbf{X} \in \mathbf{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 feedforward 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  $\mathbf{A}_{\mathrm{sa}}^{\mathrm{i}}$  of layer  $i$  is calculated by the dimension-normalized dot-product operation.
$$
\mathbf {A} _ {\mathrm {s a}} ^ {\mathrm {i}} = \text {S e l f - A t t e n t i o n} (\mathbf {X}) = \frac {\mathbf {Q K} ^ {\top}}{\sqrt {d}} \tag {1}
$$
where  $d$  is the dimension of representation vectors. In a vanilla transformer,  $\mathbf{A}_{\mathrm{sa}}^{\mathrm{i}}$  is then normalized by softmax and fed into position-wise feed-forward layers. In KAM-BERT, the self-attention map  $\mathbf{A}_{\mathrm{sa}}^{\mathrm{i}}$  is infused with semantic attention maps to calculate the final attention matrix, which will be described in the following sub-sections.
3.2 Semantic Attention Maps
Given a sequence of tokens  $\{x_0, x_1, \dots, 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} = \text{Relevance}(x_i, x_j)$ . Note that if  $x_i$  denotes a sub-word as in the BERT model, we define  $M_{i,j} = \text{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, \text{like}, \text{ta}, \# \# \cos, !\}$ , where "ta" and "##cos" are sub-words from "tacos", so both  $W(t_a)$  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.
3.3 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
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.
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  $\mathbf{A}_{\mathrm{sa}}^{\mathrm{i}}$  will be concatenated with each semantic attention map  $\mathbf{M}_{\mathrm{i}}$  separately along the channel dimension. Then, a multi-channel 2D-convolution is applied to generate an knowledge-infused attention map, denoted by  $\mathbf{A}_{\mathrm{sem}}^{\mathrm{i}}$ . The entire process can be formulated as below.
$$
\mathbf {A} _ {\mathrm {s e m}} ^ {\mathrm {i}} = \operatorname {C o n v} \left(C W C \left(\mathbf {A} _ {\mathrm {s a}} ^ {\mathrm {i}} \mid \mathbf {M} _ {\mathbf {1}. \mathbf {k}}\right)\right) \quad (2)
$$
where  $\mathbf{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  $\mathbf{A_{sa}}$ . If  $\mathbf{A_{sa}}$  has  $m$  attention heads, then  $\mathbf{A_{sem}}$  will also have  $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  $\mathbf{A}_{\mathrm{sa}}^{\mathrm{i}}$  and  $\mathbf{A}_{\mathrm{sem}}^{\mathrm{i}}$ .
$$
\mathbf {A} ^ {\mathrm {i}} = \operatorname {S o f t m a x} \left(\alpha \cdot \mathbf {A} _ {\text {s e m}} ^ {\mathrm {i}} + (1 - \alpha) \cdot \mathbf {A} _ {\text {s a}} ^ {\mathrm {i}}\right) \tag {3}
$$
After softmax activation, we get the final self-attention map  $\mathbf{A}^{\mathrm{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
$$
\mathbf {h} _ {\mathbf {j}} ^ {\mathbf {i}} = \mathbf {A} _ {\mathbf {j}} ^ {\mathbf {i}} \mathbf {V} ^ {\mathbf {i}}, \mathbf {H} ^ {\mathbf {i}} = \left(\bigoplus_ {j = 1} ^ {m} \mathbf {h} ^ {\mathbf {j}}\right) \mathbf {W} ^ {\mathbf {O}}, j \in m. \tag {4}
$$
In detail, we use the obtained attention map  $\mathbf{A}^{\mathrm{i}}$  to multiply the value matrix  $\mathbf{V}$  in the attention mechanism to get the representation  $\mathbf{h}^{\mathrm{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  $\mathbf{W}^{\mathrm{O}}$  to get the output representation of the  $i$ -th KAMBERT 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.
3.4 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
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,\dots,x_{n - 1}\}$  into a sequence of labels  $\{label_0, label_1, \ldots, 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
$$
\text {R e l e v a n c e} (W _ {i}, W _ {j}) = \left\{ \begin{array}{l l} 1 & W _ {i} \equiv_ {E} W _ {j} \\ 0 & \text {o t h e r w i s e} \end{array} \right. \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
$$
\text {R e l e v a n c e} (W _ {i}, W _ {j}) = \left\{ \begin{array}{l l} 1 & W _ {i} \equiv_ {T} W _ {j} \\ 0 & \text {o t h e r w i s e} \end{array} \right. \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
$$
P M I (x; y) = \log \frac {p (x , y)}{p (x) p (y)} \tag {7}
$$
$$
\operatorname {R e l e v a n c e} \left(W _ {i}, W _ {j}\right) = P M I \left(W _ {i}; W _ {j}\right) / 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:
$$
\operatorname {R e l e v a n c e} \left(W _ {i}, W _ {j}\right) = \left\{ \begin{array}{l l} 1 & E \left(W _ {i}\right) \equiv_ {K G} E \left(W _ {j}\right) \\ 0 & \text {o t h e r w i s e} \end{array} \right. \tag {8}
$$
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.
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.
4.1 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.
4.2 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 service<sup>1</sup> 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.
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.