SlowGuess/NAACL_2022_OCR huggingface dataset:
Minki Kang $^{1,2*}$  Jinheon Baek $^{1*}$  Sung Ju Hwang $^{1,2}$  KAIST $^{1}$ , AITRICS $^{2}$
{zzxc1133, jinheon.baek, sjhwang82}@kaist.ac.kr
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. Code is available at: https://github.com/Nardien/KALA.
Pre-trained Language Models (PLMs) (Devlin et al., 2019; Brown et al., 2020) 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 (Gururanan et al., 2020)—for which the acquisition of domain knowledge is required to accurately solve the downstream tasks (e.g., Biomedical Named Entity Recognition (Dogan et al., 2014)).
Figure 1: F1 Score and Training FLOPs for different methods on Question Answering (NewsQA). Note that DAPT uses about 112 times larger data for adaptation. Details are in §5.3
This problem, known as Language Model Adaptation, can be viewed as a transfer learning problem (Yosinski et al., 2014; Ruder, 2019) 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 (Beltagy et al., 2019; Lee et al., 2020; Gururangan et al., 2020). 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 (Howard and Ruder, 2018; Han and Eisenstein, 2019).
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 (Gururangan et al., 2020), further pre-training on the additional data requires larger memory and computational cost as the dataset size grows (Bai et al., 2021). 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 (Chen et al., 2020). This catastrophic forgetting of the task
Human Language Technologies, pages 5144 - 5167
(Unseen data not in training dataset)
Context: non-ST-elevation myocardial infarction,  $100\%$  RCA, three stents,  $50\%$  mid LD. 2017-04-02 instant restenosis status post brachytherapy, (...) her shortness of breath and asthma flare, tetracycline, sulfa, Demerol. MEDICATIONS (...)
Question: Did the patient receive brachytherapies for instant restenosis?
Figure 2: Concepts (Left). (a) Adaptive Pre-training updates whole parameters of the PLM through further pre-training on the domain corpus. (b) Our method KALA integrates the external knowledge so that the PLM adapts to the target domain only with fine-tuning, which is realized by the affine transformation on the intermediate feature. Visualization of the contextualized representation from the PLM for seen and unseen entities (Right). Our 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 (Gururanan et al., 2020) cannot appropriately handle unseen entities that are not given for task fine-tuning.
KALA (Ours)
general knowledge may lead to the performance degradation on the downstream tasks. In Figure 1, we show that adaptive pre-training with more training steps could lead to performance degeneration.
Thus, it would be preferable if we could adapt the PLM to the domain-specific task without costly adaptive pre-training. To this end, we aim to integrate the domain-specific knowledge into the PLM directly during the task-specific fine-tuning step, as shown in Figure 2b, eliminating the adaptive pre-training stage. Specifically, we first note that entities and relations are core building blocks of the domain-specific knowledge that are required to solve for the domain-specific downstream tasks. Clinical domain experts, for example, are familiar with medical terminologies and their complex relations. Then, to represent the domain knowledge consisting of entities and relations, we introduce the Entity Memory, which is the source of entity embeddings but independent of the PLM parameters (See Entity Memory in Figure 2b). Then, we further exploit the relational structures of the entities by utilizing a Knowledge Graph (KG), which denotes the factual relationships between entities, as shown in Knowledge Graph of Figure 2b.
The remaining step is how to integrate the knowledge into the PLM during fine-tuning. To this end, we propose a novel layer named Knowledge-conditioned 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:
- We propose a novel LM adaptation framework, which augments PLMs with entities and their relations from the target domain, during fine-tuning without any further pre-training. To our knowledge, this is the first work that utilizes the structured knowledge for language model adaptation.
- To reflect structural knowledge into the PLM, we introduce a novel layer which scales and shifts the intermediate PLM representations with the entity representations contextualized by their related entities according to the KG.
- We show that our KALA significantly enhances the model's performance on domain-specific QA and NER tasks, while being significantly more efficient over existing LM adaptation methods.
Language Model Adaptation Nowadays, transfer learning (Howard and Ruder, 2018) 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 (Devlin et al., 2019). 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, Lee et al. (2020) demonstrated that a further pre-trained LM on biomedical documents outperforms the original LM on biomedical NLU tasks. Also, Gururangan et al. (2020) 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 (Vrandecic and Krötzsch, 2014), some prior works have proposed to integrate external knowledge into PLMs, to enhance their performance on tasks that require structured knowledge. For instance, ERNIE (Zhang et al., 2019) and KnowBERT (Peters et al., 2019) incorporate entities as additional inputs in the pretraining stage to obtain a knowledge-aware LM, wherein a pre-trained knowledge graph embedding from Wikidata (Vrandecic and Krötzsch, 2014) is used to represent entities. Entity-as-Experts (Févry et al., 2020) and LUKE (Yamada et al., 2020) use the entity memory that is pre-trained along with the LMs from scratch. ERICA (Qin et al., 2021) 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.
Our goal is to solve Natural Language Understanding (NLU) tasks for a specific domain, with a knowledge-augmented Language Model (LM). We first introduce the NLU tasks we target, followed by the descriptions of the proposed knowledge-augmented LM. After that, we formally define the ingredients for structured knowledge integration.
NLU tasks The goal of an NLU task is to predict the label  $\pmb{y}$  of the given input instance  $\pmb{x}$ , where the input  $\pmb{x}$  contains the sequence of tokens (Devlin et al., 2019):  $\pmb{x} = [w_{1}, w_{2}, \dots, w_{|\pmb{x}|}]$ . Then, given a training dataset  $\mathcal{D} = \{(\pmb{x}^{(i)}, \pmb{y}^{(i)})\}_{i=1}^{N}$ , the objective is to maximize the log-likelihood as follows:
$$
\max  _ {\theta} \mathcal {L} (\theta) := \max  _ {\theta} \sum_ {(\boldsymbol {x}, \boldsymbol {y}) \sim \mathcal {D}} \log p (\boldsymbol {y} | \boldsymbol {x}; \theta),
$$
$$
p (\boldsymbol {y} | \boldsymbol {x}; \theta) = g (\boldsymbol {H}; \theta_ {g}), \boldsymbol {H} = f (\boldsymbol {x}; \theta_ {f}),
$$
where  $f$  is an encoder of the PLM which outputs contextualized representation  $\pmb{H}$  from  $\pmb{x}$ , and  $g$  is a decoder which models the probability distribution  $p$  of the label  $\pmb{y}$ , with trainable parameters  $\theta = (\theta_{f},\theta_{g})$ . If the LM is composed of L-layers of transformer blocks (Devlin et al., 2019), the function  $f$  is decomposed to multiple functions  $f = [f^{0},\dots ,f^{L}]$ , where each block gets the output of the previous block as the input:  $\pmb{H}^{l} = f^{l}(\pmb{H}^{l - 1})$ .
Knowledge-Augmented Language Model The conventional learning objective defined above might be sufficient for understanding the texts if the tasks require only the general knowledge stored in PLMs. However, it is suboptimal for tackling domain-specific tasks since the general knowledge captured by the parameters  $\theta_{f}$  may not include the knowledge required for solving the domain-specific tasks. Thus, contextualizing the texts by the domain knowledge, captured by the domain-specific entities and their relations, is more appropriate for handling such domain-specific problems.
To this end, we propose a function  $h(\cdot ;\phi)$  which augments PLMs conditioned on the domain knowledge. Formally, the objective for a NLU task with
$f^0$  denotes a word embedding layer which gets  $\pmb{x}$  as an input, i.e.,  $H^0 = f^0 (\pmb {x})$ , for the sake of simplicity.
our knowledge-augmented LM is given as follows:
$$
\begin{array}{l} \max  _ {\theta , \phi} \mathcal {L} (\theta , \phi) := \max  _ {\theta , \phi} \sum_ {(\boldsymbol {x}, \boldsymbol {y}) \sim \mathcal {D}} \log p (\boldsymbol {y} | \boldsymbol {x}; \theta , \phi), \\ p (\boldsymbol {y} | \boldsymbol {x}; \theta , \phi) = g (\tilde {\boldsymbol {H}}; \theta_ {g}), \\ \tilde {\pmb {H}} ^ {l} = f ^ {l} (\pmb {H} ^ {l - 1}, h ^ {l} (\pmb {H} ^ {l - 1}, \mathcal {E}, \mathcal {M}, \mathcal {G}; \phi); \theta_ {f ^ {l}}), \\ \end{array}
$$
where  $\phi$  is parameters for the function  $h, \mathcal{E}$  is the set of entities,  $\mathcal{M}$  is the set of corresponding mentions, and  $\mathcal{G}$  is a knowledge graph. In the following, we will describe the definition of the knowledge-related inputs  $\mathcal{E}, \mathcal{M}, \mathcal{G}$ , and the details of  $h(\cdot, \phi)$ .
Definition 1 (Entity and Mention). Given a sequence of tokens  $\pmb{x} = [w_{1},\dots ,w_{|\pmb{x}|}]$ , let  $\mathcal{E}$  be a set of entities in  $\pmb{x}$ . Then an entity  $e\in \mathcal{E}$  is composed of one or multiple adjacent tokens within the input text:  $[w_{m^{\alpha}},\ldots ,w_{m^{\omega}}]\subseteq x^{2}$ . Here,  $m = (m^{\alpha},m^{\omega})$  is a mention that denotes the start and end locations for the entity within the input tokens  $\pmb{x}$ , which term is commonly used for defining entities (Févry et al., 2020). Consequently, for each given input  $\pmb{x}^{(i)}$ , there are a set of entities  $\mathcal{E}^{(i)} = \{e_1,\dots ,e_K\}$  and their corresponding mentions  $\mathcal{M}^{(i)} = \{m_1,\dots ,m_K\}$ . For example, given an input  $\pmb{x} = [\mathrm{New},\mathrm{York},\mathrm{is},\mathrm{a},\mathrm{city}]$ , we have two entities  $\mathcal{E} = \{\mathrm{New\_York},\mathrm{city}\}$  and their associated mentions  $\mathcal{M} = \{(1,2),(4,4)\}$ .
We further construct the entity vocabulary  $\mathcal{E}_{\mathrm{train}} = \bigcup_{i=1}^{N} \mathcal{E}^{(i)}$ , which consists of all entities appearing in the training dataset. However, at test time, we may encounter unseen entities that are not in  $\mathcal{E}_{\mathrm{train}}$ . To tackle this, we regard unknown entities as the null entity  $e_{\emptyset}$ , so that  $\forall e \in \mathcal{E}_{\mathrm{train}} \cup \{e_{\emptyset}\}$ .
Definition 2 (Entity Memory). Given a set of all entities  $\mathcal{E}_{\mathrm{train}} \cup \{e_{\emptyset}\}$ , we represent them in the continuous vector (feature) space to learn meaningful entity embeddings. In order to implement this, we define the entity memory  $E \in \mathbb{R}^{(|\mathcal{E}_{\mathrm{train}}| + 1) \times d}$  that comprises of an entity  $e \in \mathbb{R}$  as a key and its embedding  $e \in \mathbb{R}^d$  as its value. Also, to access the value in the entity memory, we define the point-wise memory access function EntEmbed which takes an entity as an input. For instance,  $e = \text{EntEmbed}(New_York)$  returns the embedding of the New_York entity, and  $e = \text{EntEmbed}(e_{\emptyset})$  returns the zero embedding. This entity memory  $E$  is the part of the parameter  $\phi$  used in function  $h$ .
Definition 3 (Knowledge Graph). Since the entity memory alone cannot represent relational information between entities, we further define a Knowledge Graph (KG)  $\mathcal{G}$  that consists of a set of factual triplets  $\{(h,r,t)\}$ , where the head and the tail entities,  $h$  and  $t$ , are the elements of  $\mathcal{E}$ , and a relation  $r$  is an element of a set of relations  $\mathcal{R}$ :  $h,t\in \mathcal{E}$  and  $r\in \mathcal{R}$ . We assume that a preconstructed KG  $\mathcal{G}^{(i)}$  is given for each input  $x^{(i)}$ , and provide the details of the KGs and how to construct them in Appendix A.
3.2 Knowledge-conditioned Feature Modulation on Transformer
The remaining problem is how to augment a PLM by conditioning it on the domain-specific knowledge, through the function  $h$ . An effective approach to do so without stacking additional layers on top of the LM is to interleave the knowledge from  $h$  with the pre-trained parameters of the language model (Devlin et al., 2019) consisting of transformer layers (Vaswani et al., 2017). Before describing our interleaving method in detail, we first describe the Transformer architecture.
Transformer Given  $|x|$  token representations  $\pmb{H}^{l - 1} = [h_1^{l - 1},\dots ,h_{|\pmb{x}|}^{l - 1}]\in \mathbb{R}^{| \pmb{x}| \times d}$  from the layer  $l - 1$  where  $d$  is the embedding size, each transformer block outputs the contextualized representations for all tokens. In detail, the  $l$ -th block consists of the multi-head self-attention (Attn) layer and the residual feed-forward (FF) layer as follows:
$$
\hat {\boldsymbol {H}} ^ {l} = L N (\boldsymbol {H} ^ {l - 1} + A t t n (\boldsymbol {H} ^ {l - 1}))
$$
$$
F F (\hat {\boldsymbol {H}} ^ {l}) = \sigma (\hat {\boldsymbol {H}} ^ {l} \cdot \boldsymbol {W} _ {1}) \cdot \boldsymbol {W} _ {2},
$$
$$
\pmb {H} ^ {l} = L N (\hat {\pmb {H}} ^ {l} + F F (\hat {\pmb {H}} ^ {l})),
$$
where  $LN$  is a layer normalization (Ba et al., 2016),  $\sigma$  is an activation function (Hendrycks and Gimpel, 2016),  $W_{2} \in \mathbb{R}^{d' \times d}$  and  $W_{1} \in \mathbb{R}^{d \times d'}$  are weight matrices, and  $d'$  is an intermediate hidden size. We omit the bias term for brevity.
Linear Modulation on Transformer An effective yet efficient way to fuse knowledge from different sources without modifying the original model architecture is to scale and shift the features of one source with respect to the data from another source (Dumoulin et al., 2018). This scheme of feature-wise affine transformation is effective on various tasks, such as language-conditioned image reasoning (Perez et al., 2018) or style-transfer in image generation (Huang and Belongie, 2017).
$^2 E\subseteq E'$  iff  $E = E^{\prime}$  , or  $E$  is included in  $E^{\prime}$  such that the order of elements in  $E$  and  $E^{\prime}$  is the same.
Figure 3: Framework Overview. (Left) The architecture of a knowledge-augmented LM with our method. Some of the input tokens are annotated as entities with their mentions. (Middle) Inside the transformer block, KFM (\$3.2) is applied after the layer normalization as in equation 1, to modulate the hidden representations of tokens in entity mentions. (Right) The retrieved embedding of an entity New_York is composed by the weighted aggregation of neighbors through the knowledge graph (\$3.3).
Motivated by them, we propose to linearly transform the intermediate features after the layer normalization of the transformer-based PLM, conditioned on the knowledge sources  $\mathcal{E},\mathcal{M},\mathcal{G}$ . We term this method as the Knowledge-conditioned Feature Modulation (KFM), described as follows:
$$
\boldsymbol {\Gamma}, \boldsymbol {B}, \tilde {\boldsymbol {\Gamma}}, \tilde {\boldsymbol {B}} = h ^ {l} (\boldsymbol {H} ^ {l - 1}, \mathcal {E}, \mathcal {M}, \mathcal {G}; \phi),
$$
$$
\hat {\boldsymbol {H}} ^ {l} = \boldsymbol {\Gamma} \circ L N (\boldsymbol {H} ^ {l - 1} + A t t n (\boldsymbol {H} ^ {l - 1})) + \boldsymbol {B},
$$
$$
\tilde {\boldsymbol {H}} ^ {l} = \tilde {\boldsymbol {\Gamma}} \circ L N (\hat {\boldsymbol {H}} ^ {l} + F F (\hat {\boldsymbol {H}} ^ {l})) + \tilde {\boldsymbol {B}}, \tag {1}
$$
where  $H^{l-1} \in \mathbb{R}^{|x| \times d}$  is the matrix of hidden representations from the previous layer,  $\circ$  denotes the hadamard (element-wise) product, and  $\Gamma = [\gamma_1, \dots, \gamma_{|\pmb{x}|}] \in \mathbb{R}^{|x| \times d}$ ,  $B = [\beta_1, \dots, \beta_{|\pmb{x}|}] \in \mathbb{R}^{|x| \times d}$ .  $\Gamma$  and  $B$  are learnable modulation parameters from the function  $h$ , which are conditioned by the entity representation. For instance, in Figure 3,  $\gamma$  and  $\beta$  for token 'New' are conditioned on the corresponding entity New_York. However, if tokens are not part of any entity (e.g., 'is'),  $\gamma$  and  $\beta$  for such tokens are fixed to 1 and 0, respectively.
One notable advantage of our KFM is that multiple tokens associated to the identical entity are affected by the same modulation (e.g., 'New' and 'York' in Figure 3), which allows the PLM to know which adjacent tokens are in the same entity. This is important for representing the tokens of the domain entity (e.g., 'cod' and 'on'), since the original PLM might regard them as separate, unrelated tokens (See analysis in §5.5 with Figure 5). However, with our KFM, the PLM can identify associated tokens and embed them to be close to each other.
Then, how can we design such functional operations in  $h$ ? The easiest way is to retrieve the entity embedding of  $e$ , associated to the typical to-
ken, from the entity memory  $\pmb{E}$ , and then use the retrieved entity embedding as the input to obtain  $\gamma$  and  $\beta$  for every entity (See Figure 3). Formally, for each entity  $e\in \mathcal{E}$  and its mention  $(m^{\alpha},m^{\omega})\in \mathcal{M}$
$$
\boldsymbol {v} = \operatorname {E n t E m b e d} (e) \tag {2}
$$
$$
\boldsymbol {\gamma} _ {j} = \mathbf {1} + h _ {1} (\boldsymbol {v}), \beta_ {j} = h _ {2} (\boldsymbol {v}),
$$
$$
\tilde {\gamma} _ {j} = \mathbf {1} + h _ {3} (\boldsymbol {v}), \tilde {\beta} _ {j} = h _ {4} (\boldsymbol {v}), \quad m ^ {\alpha} \leq j \leq m ^ {\omega},
$$
where  $\pmb{v}$  is the retrieved entity embedding from the entity memory,  $h_1, h_2, h_3$ , and  $h_4$  are mutually independent Multi-Layer Perceptrons (MLPs) which return a zero vector  $\mathbf{0}$  if  $e = e_{\emptyset}$ .
3.3 Relational Retrieval from Entity Memory
Although the simple access to the entity memory can retrieve the necessary entity embeddings for the modulation, this approach has obvious drawbacks as it not only fails to reflect the relations with other entities, but also regards unseen entities as the same null entity  $e_{\emptyset}$ . If so, all unseen entities are inevitably modulated by the same parameters even if they have essentially different meaning.
To tackle these limitations, we further consider the relational information between two entities that are linked with a particular relation. For example, the entity New_York alone will not give meaningful information. However, with two associated facts (New_York, instance of, city) and (New_York, country, USA), it is clear that New_York is a city in the USA. Motivated by this observation, we propose Relational Retrieval which leverages a KG  $\mathcal{G}$  to retrieve entity embeddings from the memory, according to the relations defined in the given KG (See Figure 3, right).
More specifically, our goal is to effectively utilize the relations among entities in  $\mathcal{G}$ , to improve
the EntEmbed function in equation 2. We tackle this objective by utilizing a Graph Neural Network (GNN) which learns feature representations of each node using a neighborhood aggregation scheme (Hamilton et al., 2017), as follows:
$$
\begin{array}{l} \boldsymbol {v} = \text {U P D A T E (E n t E m b e d} (e), \\ \operatorname {A G G} (\{\operatorname {E n t E m b e d} (\hat {e}): \forall \hat {e} \in \mathcal {N} (e; \mathcal {G}) \}), \\ \end{array}
$$