SlowGuess/NAACL_2022_OCR huggingface dataset:
where  $\mathcal{N}(e;\mathcal{G})$  is a set of neighboring entities of the entity  $e$ , AGG is the function that aggregates embeddings of neighboring entities of  $e$ , and UPDATE is the function that updates the representation of  $e$  with the aggregated messages from AGG.
However, simple aggregation (e.g., mean) cannot reflect the relative importance on neighboring nodes, thus we consider the attentive scheme (Velickovic et al., 2018; Brody et al., 2021) for neighborhood aggregation, to allocate weights to the target entity's neighbors by their importance. This scheme is helpful in filtering out less useful relations. Formally, we first define a scoring function  $\psi$  that calculates a score for every triplet  $(e_i, r_{ij}, e_j)$ , which is then used to weigh each node during aggregation:
$$
\begin{array}{l} \boldsymbol {e} _ {i} = \operatorname {E n t E m b e d} (e _ {i}), \boldsymbol {e} _ {j} = \operatorname {E n t E m b e d} (e _ {j}), \\ \boldsymbol {e} ^ {*} = \left[ \boldsymbol {e} _ {i} \parallel \boldsymbol {r} _ {i j} \parallel \boldsymbol {e} _ {j} \parallel \boldsymbol {h} _ {e _ {i}} \right], \\ \psi \left(\boldsymbol {e} _ {i}, \boldsymbol {r} _ {i j}, \boldsymbol {e} _ {j}, \boldsymbol {h} _ {e _ {i}}\right) = \boldsymbol {a} ^ {\top} \sigma (\boldsymbol {W} \cdot \boldsymbol {e} ^ {*}), \\ \end{array}
$$
where  $\sigma$  is a nonlinear activation,  $e^{*} \in \mathbb{R}^{4d}$  is concatenated vector where  $\|\cdot\|$  denotes the concatenation,  $\pmb{a} \in \mathbb{R}^{d}$  and  $\pmb{W} \in \mathbb{R}^{d \times 4d}$  are learnable parameters,  $\pmb{r}_{ij} \in \mathbb{R}^{d}$  is a embedding of the relation, and  $\pmb{h}_{e_i} \in \mathbb{R}^d$  is a context representation of the entity  $e_i$  obtained from the intermediate hidden states of the LM<sup>3</sup>.
The scores obtained from  $\psi$  are normalized across all neighbors  $e_j \in \mathcal{N}(e_i; \mathcal{G})$  with softmax:
$$
\begin{array}{l} \alpha_ {i j} = \operatorname {s o f t m a x} \left(\psi \left(\boldsymbol {e} _ {i}, \boldsymbol {r} _ {i j}, \boldsymbol {e} _ {j}\right)\right) \\ = \frac {\exp (\psi (\boldsymbol {e} _ {i} , \boldsymbol {r} _ {i j} , \boldsymbol {e} _ {j}))}{\sum_ {e _ {j ^ {\prime}} \in \mathcal {N} (e _ {i} ; \mathcal {G})} \exp (\psi (\boldsymbol {e} _ {i} , \boldsymbol {r} _ {i j ^ {\prime}} , \boldsymbol {e} _ {j ^ {\prime}}))}. \\ \end{array}
$$
Then, we update the entity embedding with a weighted average of the neighboring nodes with  $\alpha$  as an attention coefficient, denoted as follows:
$$
\boldsymbol {v} = \text {U P D A T E} \left(\sum_ {e _ {j ^ {\prime}} \in \mathcal {N} (e _ {i}; \mathcal {G})} \alpha_ {i j} \cdot \boldsymbol {e} _ {j ^ {\prime}}\right). \tag {3}
$$
By replacing the EntEmbed function in equation 2 with the above GNN in equation 3, we now represent each entity with its relational information in KG. This relational retrieval has several advantages over simple retrieval of a single entity from the entity memory. First, the relational retrieval with KG can consider richer interactions among entities, as described in Figure 3.
In addition, we can naturally represent an unseen entity – which is not seen during training but appears at test time – through neighboring aggregation, which is impossible only with the entity memory. In Figure 2, we provide an illustrative example of the unseen entity representation, where the unseen entity restenosis is represented with a weighted sum of representations of its neighboring entities myocardial_infarction, asthma, and pethidine, which is beneficial when the set of entities for training and test datasets have small overlaps.
We evaluate our model on two NLU tasks: Question Answering (QA) and Named Entity Recognition (NER). For QA, we use three domain-specific datasets: NewsQA (News, Trischler et al., 2017) and two subsets (Relation, Medication) of EMRQA (Clinical, Pampari et al., 2018). We use the Exact-Match (EM) and the F1 score as evaluation metrics. For NER, we use three datasets from different domains, namely CoNLL-2003 (News, Sang and Meulder, 2003), WNUT-17 (Social Media, Derczynski et al., 2017) and NCBI-Disease (Biomedical, Dogan et al., 2014). We use the F1 score as the evaluation metric. We report statistics and detailed descriptions of each dataset in Appendix B.2.
A direct baseline of our KALA is the adaptive pre-training, which is commonly used to adapt the PLM independent to the choice of a domain and task. Also, to compare ours against a more powerful baseline, we modify a recent method (Chen et al., 2020) that alleviates forgetting of PLM during fine-tuning. Details for each baseline we use are described as follows:
1. Vanilla Fine-Tuning (FT): A baseline that directly fine-tunes the LM on downstream tasks.
2. Fine-Tuning + more params: A baseline with one more transformer layer at the end of the
3The context representation of the entity is calculated with its mention as follows:  $h_e = \frac{1}{m^\omega - m^\alpha + 1}\sum_{i = m^\alpha}^{m^\omega}h_i^{l - 1}$
Method	NewsQA	Relation	Medication
Fine-Tuning	53.06 ± 0.63 | 67.20 ± 0.19	54.01 ± 1.14 | 61.43 ± 1.18	12.50 ± 0.28 | 43.31 ± 0.67
+ more params	53.59 ± 0.99 | 67.79 ± 0.67	54.06 ± 1.35 | 62.07 ± 1.44	12.46 ± 0.25 | 42.74 ± 0.91
TAPT	53.47 ± 1.69 | 67.59 ± 1.44	53.57 ± 2.05 | 60.87 ± 2.52	12.58 ± 0.42 | 43.82 ± 1.10
+ RecAdam	53.95 ± 1.02 | 67.89 ± 0.75	54.88 ± 1.94 | 62.54 ± 2.14	12.63 ± 0.30 | 43.86 ± 0.87
\( DAPT^† \)	53.68 ± 0.94 | 67.76 ± 0.61	55.29 ± 1.74 | 62.25 ± 1.80	12.67 ± 0.27 | 43.26 ± 0.88
KALA (point-wise)	53.41 ± 0.74 | 67.30 ± 0.45	56.13 ± 0.85 | 64.69 ± 0.92	12.01 ± 0.47 | 42.97 ± 0.70
KALA (relational)	54.25 ± 0.63 | 68.27 ± 0.63	55.96 ± 1.37 | 64.22 ± 1.15	12.75 ± 0.61 | 44.19 ± 0.46
Table 1: Experimental results of the extractive QA task on three different datasets with the BERT-base. The reported results are means and standard deviations of performances over five different runs with Exact Match / F1 score as a metric. The numbers in bold fonts denote the best score.  $\dagger$  indicates the method under an extremely high computational resource setting (See Figure 1).
Method	CoNLL-2003	WNUT-17	NCBI-Disease
Fine-Tuning	90.58 ± 0.19	45.70 ± 1.25	84.42 ± 0.58
+ more params	90.75 ± 0.23	46.42 ± 0.55	84.70 ± 0.49
TAPT	90.61 ± 0.73	45.39 ± 0.77	84.39 ± 0.73
+ RecAdam	90.69 ± 0.30	46.73 ± 0.94	84.99 ± 0.88
\( DAPT^† \)	90.30 ± 0.39	48.29 ± 1.08	84.68 ± 1.63
KALA (point-wise)	90.96 ± 0.21	47.33 ± 0.82	85.10 ± 0.73
KALA (relational)	91.02 ± 0.29	48.35 ± 0.92	85.77 ± 0.43
Table 2: Experimental results of the NER task on three different datasets with the BERT-base. The reported results are means and standard deviations over five different runs with an F1 score as a metric. The numbers in bold fonts denote the best score.  $\dagger$  indicates the baseline under an extremely high computational resource setting (See Figure 1).
LM. We use this baseline to show that the performance gain of our model does not come from the use of additional parameters.
3. Task-Adaptive Pre-training (TAPT): A baseline that further pre-trains the PLM on task-specific corpus as in Gururangan et al. (2020).
4. TAPT + RecAdam: A baseline that uses RecAdam (Chen et al., 2020) during further pre-training of PLMs (i.e., TAPT), to alleviate catastrophic forgetting of the learned general knowledge in PLMs from adaptive pre-training.
5. Domain-Adaptive Pre-training (DAPT): A strong baseline that uses a large-scale domain corpus outside the training set during further pretraining (Gururan et al., 2020), and requires extra data and large computational overhead.
6. KALA (pointwise): A variant of KALA that only uses the entity memory and does not use the knowledge graphs.
7. KALA (relational): Our full model that uses KGs to perform relational retrieval from the entity memory.
We use the uncased BERT-base (Devlin et al., 2019) as the base PLM for all our experiments on QA and NER tasks. For more details on training and implementation, please see the Appendix B.
Performance on QA and NER tasks On both extractive QA and NER tasks, our KALA outperforms all baselines, including TAPT and TAPT+RedcAdam (Gururangan et al., 2020; Chen et al., 2020), as shown in Table 1 and 2. These results show that our KALA is highly effective for the language model adaptation task. KALA also largely outperforms DAPT (Gururangan et al., 2020) which is trained with extra data and requires a significantly higher computational cost compared to KALA (See Figure 1 for the plot of efficiency, discussed in Section 5.3).
Effect of Using more Parameters One may suspect whether the performance of our KALA comes from the increment of parameters. However, the experimental results in Table 1 and 2 show that increasing the parameters for PLM during fine-tuning (+ more params) yields marginal performance improvements over naive fine-tuning. This result confirms that the performance improvement of KALA is not due to the increased number of parameters.
Importance of Relational Retrieval The performance gap between KALA (relational) and KALA (point-wise) shows the effectiveness of relational retrieval for language model adaptation, which allows us to incorporate relational knowledge into the PLM. The relational retrieval also helps address unseen entities, as discussed in Section 5.4.
We perform an ablation study to see how much each component contributes to the performance gain.
KFM Parameters We first analyze the effect of feature modulation parameters (i.e., gamma and beta) in transformers by ablating a subset of them in Table 3, in which we observe that using both
Method	NewsQA	Relation	WNUT-17	NCBI-Disease
Fine-Tuning	57.21 ± 0.56 | 71.91 ± 0.35	46.61 ± 2.75 | 53.89 ± 2.92	55.00 ± 1.66	86.91 ± 1.08
+ more params	58.07 ± 1.19 | 72.38 ± 1.04	45.12 ± 0.86 | 53.22 ± 1.27	56.62 ± 0.26	87.21 ± 0.26
TAPT	57.24 ± 0.53 | 71.77 ± 0.34	45.66 ± 2.20 | 53.23 ± 2.38	55.46 ± 1.90	86.24 ± 0.76
KALA (relational)	58.01 ± 0.57 | 72.70 ± 0.25	47.40 ± 1.67 | 55.13 ± 1.26	56.96 ± 0.27	87.72 ± 0.27
Table 5: Experimental results of the extractive QA and NER tasks on four different datasets - NewsQA, Relation, WNUT-17 and NCBI-Disease - with the RoBERTa-base. The reported results are means and standard deviations over five different runs. We use Exact Match and F1 score as a metric for QA, and F1 score for NER. The numbers in bold fonts denote the best score.
KFM (§3.2) Components	NewsQA
EM	F1
None (Fine-tuning)	53.06	67.20
+ Γ, Γ (gamma only)	54.10	67.98
+ B, B (beta only)	53.74	67.69
+ Γ, B (first only)	53.77	67.88
+ Γ, B (second only)	53.89	67.49
+ Γ, B, Γ, B (final)	54.25	68.27
Table 3: An ablation study of the KFM parameters  $\Gamma, B, \tilde{\Gamma}, \tilde{B}$ . We report the average results over five different runs.
Architecture
Variants (§5.2)	NewsQA
EM	F1
ERNIE	53.35	67.49
Adapter	53.32	67.38
KT-Net	53.15	67.01
EaE	53.00	67.40
ERICA	51.99	66.40
KALA (ours)	54.25	68.27
Table 4: Experimental results on knowledge integration architecture variants, averaged over five runs.
gamma and beta after both layer normalization on a transformer layer obtains the best performance.
Architectural Variants We now examine the effectiveness of the proposed knowledge conditioning scheme in our KALA framework. To this end, we use or adapt the knowledge integration methods from previous literature, to compare their effectiveness. Specifically, we couple the following five components with KALA: Entity-as-Experts (Févry et al., 2020), Adapter (Houlsby et al., 2019), KTNet (Yang et al., 2019), ERNIE (Zhang et al., 2019), and ERICA (Qin et al., 2021). Note that, most of them were proposed for improving pre-training from scratch, while we adapt them for fine-tuning under our KALA framework (The details are given in Appendix B.4). As shown in Table 4, our KFM used in KALA outperforms all variants, demonstrating the effectiveness of feature modulation in the middle of transformer layers for fine-tuning.
5.2 Robustness to Other PLMs
Although we believe our experimental results on Table 1, 2 with BERT (Devlin et al., 2019) are enough to show the effectiveness of KALA across different pre-trained language models (PLMs), one might be curious that KALA can work on even other PLMs such as RoBERTa (Liu et al., 2020). Thus, to address such concerns, we additionally conduct experiments on RoBERTa. As shown in Table 5, we observe that our KALA outperforms all baselines except for one case (Fine-Tuning +
Figure 4: Results on seen and unseen, where Seen denotes the context having less than 3 unseen entities, otherwise Unseen. Note that DAPT uses extra datasets in addition to the training dataset, thus the Unseen for other models could be considered as the Seen for DAPT.
more params on NewsQA). Thus, we believe that our KALA would be useful to any PLMs, not depending on specific PLMs.
Figure 1 illustrates the performance and training FLOPs of KALA against baselines on the NewsQA dataset. We observe that the performance of TAPT decreases with the increased number of iterations, which could be due to forgetting of the knowledge from the PLM. On the other hand, DAPT, while not suffering from performance loss, requires huge computational costs as it trains on 112 times larger data for further pre-training (See Appendix B.3 for detailed explanations on training data). On the other hand, our KALA outperforms DAPT without using external data, while requiring 17 times fewer computational costs, which shows that KALA is not only effective but also highly efficient.
To further compare the efficiency in various aspects, we report GPU memory, training wall time, and training FLOPs for baselines and ours in Table 6. Through this, we verify that our KALA is more efficient to train for language model adaptation settings than baselines. Note that the resource requirement of KALA could be further reduced by adjusting the size of the entity memory (e.g., removing less frequent entities). Therefore, to show the flexibility of our KALA on the typical resource constraint, we provide the experimental results on two different settings (i.e., tuning the number of entities in the entity memory) - KALA with memory size of 200 and 62.8k (full memory) in Appendix C.6.
Figure 5: A case study on one context of the NCBI-Disease dataset. A left table shows the context and its fact, and a right figure shows a visualization of token representations. Text in blue and red denote the seen and unseen entities, respectively.
T5-small	NewsQA
EM	F1
Fine-tuning	48.96	64.24
TAPT	48.66	64.30
+ RecAdam	48.37	63.41
KALA (ours)	51.78	66.88
Table 7: Experimental results on generative question answering with T5-small as a PLM and NewsQA as a dataset.
Method	GPU Mem.	Approx. Wall Time	FLOPs (1016)
Fine-Tuning	8 GB	3 hrs	9.5
+ more params	8.8 GB	3 hrs	10.1
TAPT	8 GB	3.8 hrs	10.1
DAPT	48 GB	40 hrs <	182.0
KALA (ours, 0.2k)	8.4 GB	3 hrs	9.97
KALA (ours, 62.8k)	9.2 GB	3 hrs	10.5
Table 6: Efficiency comparisons of GPU memory, Wall Time, and FLOPs on the NewsQA dataset. The number  ${0.2}\mathrm{k}$  and  ${62.8}\mathrm{k}$  indicate the size of entity memory used in our KALA.
5.4 Effectiveness on Unseen Entities
One remarkable advantage of our KALA is its ability to represent an unseen entity by aggregating features of its neighbors from a given KG. To analyze this, we first divide all contexts into one of Seen and Unseen, where Seen denotes the context with less than 3 unseen entities, and then measure the performance on the two subsets. As shown in Figure 4, we observe that the performance gain of KALA over the baselines is much larger on the Unseen subset, which demonstrates the effectiveness of KALA's relational retrieval scheme to represent unseen entities. DAPT also largely outperforms fine-tuning and TAPT as it is trained on an extremely large external corpus for adaptive pre-training. However, KALA even outperforms DAPT in most cases, verifying that our knowledge-augmentation method is more effective for tackling domain-specific tasks. The visualization of embeddings of seen and unseen entities in Figure 2 shows that KALA embeds the unseen entities more closely to the seen entities<sup>4</sup>, which explains KALA's good performance on the Unseen subset.
To better see how our KFM (§3.2) works, we show the context and its fact, and then visualize representations from the PLM modulated by the KFM.
As shown in Figure 5 right, the token '#on' is not aligned with their corresponding tokens, such as 'ex' (for exon) and 'cod' (for codon), in the baseline. However, with our feature modulation that transforms multiple tokens associated with the single entity equally, the two tokens (e.g., ('ex', '#on')), composing one entity, are closely embedded. Also, while the baseline cannot handle the unseen entity consisting of three tokens: 're', '#tina', and '#l', KALA embeds them closely by representing the unseen retinal from the representation of its neighborhood gene derived by the domain knowledge - (retinal, instance of, gene).
5.6 Extension to Generative Model
Our KALA framework is also applicable to encoder-decoder PLMs by applying the KFM to the encoder. Therefore, we further validate KALA's effectiveness on the encoder-decoder PLMs on the generative QA task (Lee et al., 2021) with T5-small (Raffel et al., 2020). Table 7 shows that KALA largely outperforms baselines even with such a generative PLM.
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.
<sup>4</sup>We quantitatively measure the mean of cosine distance of each unseen entity to its nearest seen entity, observing that KALA embeds unseen 1.5 times more closer to seen than DAPT (i.e., 0.07 for KALA vs 0.11 for DAPT for distance).
Ethical Statements
Enhancing the domain converge of pre-traind language models (PLMs) with external knowledge is increasingly important, since the PLMs cannot observe all the data during training and cannot memorize all the necessary knowledge for solving down-stream tasks. Our KALA contributes to this problem by augmenting domain knowledge graphs for PLMs. However, we have to still consider the accurateness of knowledge, i.e., the fact in the knowledge graph may not be correct, which affects the model to generate incorrect answers. Also, the model's prediction performance is still far from optimal. Thus, we should be aware of model's failure from errors in knowledge and prediction, especially on high-risk domains (e.g., biomedicine).
This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST) and No. 2021-0-02068, Artificial Intelligence Innovation Hub)), AITRICS, Samsung Electronics (IO201214-08145-01), and the Engineering Research Center Program through the National Research Foundation of Korea (NRF) funded by the Korean Government MSIT (NRF2018R1A5A1059921).
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. Layer normalization. arXiv preprint, arXiv:1607.06450.
Fan Bai, Alan Ritter, and Wei Xu. 2021. Pre-train or annotate? domain adaptation with a constrained budget. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021, pages 5002-5015.
Iz Beltagy, Kyle Lo, and Arman Cohan. 2019. Scibert: A pretrained language model for scientific text. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019, pages 3613-3618.
Shaked Brody, Uri Alon, and Eran Yahav. 2021. How attentive are graph attention networks? arXiv preprint, arXiv:2105.14491.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
Andrew Carlson, Justin Betteridge, Bryan Kisiel, Burr Settles, Estevam R. Hruschka Jr., and Tom M. Mitchell. 2010. Toward an architecture for never-ending language learning. In Proceedings of the Twenty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2010, Atlanta, Georgia, USA, July 11-15, 2010. AAAI Press.
Sanyuan Chen, Yutai Hou, Yiming Cui, Wanxiang Che, Ting Liu, and Xiangzhan Yu. 2020. Recall and learn: Fine-tuning deep pretrained language models with less forgetting. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pages 7870-7881.
Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning. 2020. ELECTRA: pretraining text encoders as discriminators rather than generators. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020.
Leon Derczynski, Eric Nichols, Marieke van Erp, and Nut Limsopatham. 2017. Results of the WNUT2017 shared task on novel and emerging entity recognition. In Proceedings of the 3rd Workshop on Noisy User-generated Text, NUT@EMNLP 2017, Copenhagen, Denmark, September 7, 2017, pages 140-147. Association for Computational Linguistics.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pages 4171-4186.
Rezarta Islamaj Dogan, Robert Leaman, and Zhiyong Lu. 2014. NCBI disease corpus: A resource for disease name recognition and concept normalization. J. Biomed. Informatics, 47:1-10.
Vincent Dumoulin, Ethan Perez, Nathan Schucher, Florian Strub, Harm de Vries, Aaron Courville, and Yoshua Bengio. 2018. Feature-wise transformations. Distill.