SlowGuess/NAACL_2022_OCR huggingface dataset:
All experiments are constrained to be done with a single 12GB Geforce RTX 2080 Ti GPU for fairness in terms of memory and the availability on the academic budget, except for the DAPT and generative QA which use a single 48GB Quadro 8000 GPU. KALA training needs 3 hours in wall time with a single GPU. For all experiments, we select the best checkpoint on the validation set. For the summary of training setups, please see Table 10 and 12.
Fine-tuning Setup In the following three paragraphs, we explain the setting of fine-tuning for QA, NER, and generative QA tasks. For all experiments on extractive QA tasks, we fine-tune the Pre-trained Language Model (PLM) for 2 epochs with the weight decay of 0.01, learning rate of 3e-5, maximum sequence length of 384, batch size of 12, linear learning rate decay of 0.06 warmup rate, and half precision (Micikevicius et al., 2018).
For all experiments on NER tasks, we fine-tune the PLM for 20 epochs, where the learning rate is set to 5e-5, maximum sequence length is set to 128, and batch size is set to 32. We use AdamW (Loshchilov and Hutter, 2019) as an optimizer using BERT-base as the PLM.
For the generative QA task in Table 7, we fine-tune the T5-small (Raffel et al., 2020) for 4 epochs with the learning rate of 1e-4, maximum sequence length of 512, and batch size of 64. We also use the Adafactor (Shazeer and Stern, 2018) optimizer. Instead of training with the same optimizer as in BERT for QA and NER, we instead use the independent AdamW optimizer with the learning rate of 1e-4 and weight decay of 0.01 to train the KALA module with T5.
Adaptive Pre-training Setup In this paragraph, we describe the experimental settings of adaptive pre-training baselines, namely TAPT, TAPT (+ RecAdam), and DAPT. For QA tasks, we further pre-train the PLM for  $\{1,3,5,10\}$  epochs and then report the best performance among them. Specifically, reported TAPT result on NewsQA, Relation, and Medication are obtained by 1 epoch of further pre-training. We use the weight decay of 0.01, learning rate of  $5\mathrm{e} - 5$  maximum sequence length of 384, batch size of 12, and linear learning rate decay of 0.06 warmup rate, with a half-precision. Also, the masking ratio for the pre-training objective is set to 0.15, following the existing strategy introduced in the original BERT paper (Devlin et al., 2019).
For NER tasks, we further pre-train the PLM for 3 epochs across all datasets. In particular, the learning rate is set to 5e-5, batch size is set to 32, and the maximum sequence length is set to 128. We also use AdamW (Loshchilov and Hutter, 2019) as the optimizer for all experiments.
In the case of T5-small for generative QA in Table 7, we further pre-train the PLM for 4 epochs with the learning rate of 0.001, batch size of 64, maximum sequence length of 384, and Ada fac
tor (Shazeer and Stern, 2018) optimizer.
Regarding the setting of TAPT (+ RecAdam) on all tasks, we follow the best setting in the original paper (Chen et al., 2020) - sigmoid as an annealing function with annealing parameters:  $k = 0.5$ ,  $t_0 = 250$ , and the pretraining coefficient of 5000.
For training with DAPT, we need an external corpus having a large amount of data for adaptive pre-training. Thus, we first choose the datasets of two domains – News and Medical. Specifically, as the source of corpus for the News domain, we use the sampled set of 10 million News from the RealNews dataset used in Gururangan et al. (2021). As the source of corpus for the Medical domain, we use the set of approximately 100k passages from the Medical textbook provided in Jin et al. (2020). The size of pre-training data used in DAPT is much larger than TAPT. In other words, for experiments on NewsQA, TAPT only uses fine-tuning contexts containing 5.8 million words from the NewsQA training dataset, while DAPT uses more than a hundred times larger data – enormous contexts containing about 618 million words from the RealNews database. For both News and Medical domains, we further pre-train the BERT-base model for 50 epochs with the batch size of 64, to match the similar computational cost used in Gururangan et al. (2020). Other experimental details are the same as TAPT described above.
B.4 Architectural Variant Details
In this subsection, we describe the details of architectural variants reported in Section 5.1. For all variants, we use the same KGs used in KALA.
Entity-as-Experts (Févry et al. (2020); EaE) utilizes the entity memory similar to our work, but they use the parametric dense retrieval more like the memory neural network (Sukhbaatar et al., 2015). Similar to Févry et al. (2020); Verga et al. (2021), we change the formulation of query and memory retrieval by using the mention representation of the entity from the intermediate hidden states of PLMs, which is formally defined as follows:
$$
\boldsymbol {h} _ {e} = \frac {1}{m ^ {\omega} - m ^ {\alpha} + 1} \sum_ {i = m ^ {\alpha}} ^ {m ^ {\omega}} \boldsymbol {h} _ {i} ^ {l - 1}, \tag {4}
$$
$$
\boldsymbol {v} = \operatorname {s o f t m a x} \left(\boldsymbol {h} _ {e} \cdot \boldsymbol {E} ^ {\top}\right) \cdot \boldsymbol {E},
$$
where  $h_e$  represents the average of token representations of the entity mention  $m = (m^{\alpha}, m^{\omega})$ . We also give the supervised retrieval loss (ELLoss
in Févry et al. (2020)), when training the EaE model. With this retrieval, EaE also can represent the unseen entity  $e \notin \mathcal{E}_{train}$  if we know the mention boundary of the given entity on the context. We believe it is expected to work well, if the entity memory is pre-trained on the enormous text along with the pre-training of the language model from the scratch. However, it might underperform for the language model adaptation scenario, since it can fall into the problem of circular reasoning – the PLM does not properly represent the unseen entity, but it should predict which entity it is similar from the representation. Regarding the integration of the knowledge from the entity memory into the PLM, the retrieved entity representation  $v$  is simply added (Peters et al., 2019) to the hidden representations  $H$  after the transformer block as follows:
$$
\tilde {\boldsymbol {H}} ^ {l} = \boldsymbol {H} ^ {l} + h (\boldsymbol {v}) \tag {5}
$$
where  $h$  is Multi-Layer Perceptrons (MLPs).
Adapter (Houlsby et al., 2019) is introduced to fine-tune the PLM only with a few trainable parameters, instead of fine-tuning the whole parameters of the PLM. To adapt this original implementation into our KALA framework, we replace our Knowledge-conditioned Feature Modulation with it, where the Adapter is used as the knowledge integration module. We interleave the layer of Adapter after the feed-forward layer  $(FF)$  and before the residual connection of the transformer block. Also, instead of only providing the LM hidden states as an input, we concatenate the knowledge representation in Equation 3 to the LM hidden states. Note that we fine-tune the whole parameters following our KALA setting, unlike fine-tuning the parameters of only Adapter layers in Houlsby et al. (2019).
ERNIE (Zhang et al., 2019) is a notable PLM model that utilizes the external KB as an input for the language model. The key feature of ERNIE can be summarized into two folds. First, they use the multi-head self-attention scheme (Vaswani et al., 2017) to contextualize the input entities. Second, ERNIE fuses the entity representation at the end of the PLM by adding it to the corresponding language representation. We assume that those two features are important points of ERNIE. Therefore, instead of using a Graph Neural Network (GNN) layer, we use a multi-head self-attention layer to contextualize the entity embeddings. Then, we add it to a representation of the entity from the PLM, which is the same as the design in equation 5.
KT-Net (Yang et al., 2019) uses knowledge as an external input in the fine-tuning stage for extractive QA. Since they have a typical layer for integrating existing KB (Miller, 1995; Carlson et al., 2010) with the PLM, we only adopt the self-matching layer as the architecture variant of the KFM layer used in our KALA framework. The computation of the self-matching matrix in KT-Net is costly, i.e., it requires a large computational cost that is approximately 12 times larger than KALA.
ERICA (Qin et al., 2021) uses contrastive learning in LM pre-training to reflect the relational knowledge into the language model. We use the Entity Discrimination task from ERICA on the primary task of fine-tuning. We would like to note that, as reported in Section 5 of the original paper (Qin et al., 2021), the use of ERICA on fine-tuning has no effect, since the size and diversity of entities and relations in downstream training data are limited. Such limited information rather harms the performance, as it can hinder the generalization. In other words, contrastive learning cannot reflect the entity and relation in the test dataset.
B.5 FLOPs Computation
In this subsection, we give detailed descriptions of how the FLOPs in Figure 1 are measured. We majorly follow the script from the ELECTRA (Clark et al., 2020) repository to compute the approximated FLOPs for all models including ours. For FLOPs computation of our KALA, we additionally include the FLOPs of the entity embedding layer, linear layers for  $h_1, h_2, h_3, h_4$ , and GNN layer. Since the GNN layer is implemented based on the sparse implementation, we first calculate the FLOPs of the message propagation over one edge, and then multiply it to the average number of edges per node. Also, in terms of the computation on mentions, we consider the maximum sequence length of the context rather than the average number of mentions, to set the upper bound of FLOPs for our KALA. Note that, in NewsQA training data, the average number of nodes is 57, the average number of edges for each node is 0.64, and the average number of mentions in the context is 92.68.
In this section, we provide the analyses on the forgetting of TAPT, entity memory, number of entities and facts, location of the KLM layer, and values of Gamma and Beta.
Figure 7: Masked Language Model loss from Task-Adaptive Pre-Training on the domain-specific training dataset (Relation) and the general domain test dataset (Sampled wikipedia).
Figure 8: The performance (F1 score and Exact Match) and the GPU memory usage on NewsQA dataset with varying the size of elements in the entity memory.
C.1 Analysis on forgetting of TAPT
In Figure 1, we observe that the performance of TAPT decreases as the number of training steps increases. To get a concrete intuition on this particular phenomenon, we analysis what happens in the Pre-trained Language Model (PLM), when we further pre-train it on the task-specific corpus. Specifically, in Figure 7, we visualize the Masked Language Model (MLM) loss of TAPT on both domain-specific corpus from the Relation dataset and general corpus from the sampled Wikipedia documents during the adaptive pre-training. As Figure 7 shows, the test MLM loss increases while the training MLM loss persistently increases as the training step increases. This result indicates that TAPT on domain-specific corpus may yield the catastrophic forgetting of the general knowledge in the PLM.
C.2 Effects of the Size of Entity Memory
In this subsection, we analyze how the size of entity memory affects the performance of our KALA. In Figure 8, we plot the performance of KALA on the NewsQA dataset by varying the number of entity elements in the memory. Note that, we reduce the size of the entity memory by eliminating the entity appearing fewer times. Thus, the results are obtained by only considering the entities that appear more than [1000, 100, 10, 5, 0] times, e.g., 0 means the model with full entity memory. As shown in Figure 8, we observe that the size of the
Figure 9: Performance improvements of our KALA from simple fine-tuning, with varying the number of entities and facts in the context on Named Entity Recognition tasks.
Figure 10: The performance of our KALA with varying the location of the KFM layer inside the BERT-base model. y-axis denotes the F1 score on NewsQA and x-axis denotes the location of the KFM layer. For instance, 11 means the case where the KFM layer is appended in the 11th transformer layer of BERT-base.
entity memory is larger, the performance of our KALA is better in general. However, interestingly, we also observe that the smallest size of the entity memory shows decent performance, which might be due to the fact that some parameters in the entity memory are stale. For more discussions on it including visualization, please refer to Appendix D.2. Finally, we would like to note that, in Figure 1, we report the performance of our KALA in the case of [1000, 5, 0] (i.e., considering entities appearing more than [1000, 5, 0] times).
C.3 Effects of the Number of Entity and Fact
In this subsection, we aim to analyze which numbers of entities and facts per context are appropriate to achieve good performance in NER tasks. Specifically, we first collect the contexts having more than or equal to the  $k$  number of entities (or facts), and then calculate the performance difference from our KALA to the fine-tuning baseline. As shown in Figure 9, while there are no obvious patterns, performance improvements from the baseline are consistent across a varying number of entities and facts. This result suggests that our KALA is indeed beneficial when entities and facts are given to the model, whereas the appropriate number of entities and facts to obtain the best performance against the baseline is different across datasets.
Figure 11: Histogram of values of gamma and beta on the CoNLL-2003 dataset.
Figure 12: Visualization of contextual representations for seen and unseen entities on the NCBI-Disease dataset.
C.4 Effects of the Location of KFM
In the main paper and Appendix B.1, we describe that the location of the KFM layer inside the PLM architecture is the hyperparameter. However, someone might wonder which location of KFM yields the best performance, and what is the reason for this. Therefore, in this section, we analyze where we obtain the best performance in various locations of the KFM layer on the NewsQA dataset. Specifically, in Figure 10, we show the performance of our KALA with varying the location of the KFM layer inside the BERT-base model. The results demonstrate that the model with the KFM on the last layer of the BERT-base outperforms all the other choices. This might be because, as the final layer of the PLM is generally considered as the most task-specific layer, our KFM interleaved in the latest layer of BERT expressively injects the task-specific information from the entity memory and KGs, to such a task-specific layer.
C.5 Analysis on Values of Gamma and Beta
To see how much amount of value on gamma and beta is used to shift and scale the intermediate hidden representations in transformer layers, we visualize the modulation values, namely gamma and beta, in Figure 11. We first observe that, as shown in Figure 11, the distribution of values of gamma and beta approximately follow the Gaussian dis
tribution, with zero mean for beta and one mean for gamma. Also, we notice that the scale of values remain nearly around the mean point, which suggests that the small amount of shifting to intermediate hidden representations on transformer layers is enough to contribute to the performance gain, as we can see in the main results of Table 1, 2.
C.6 Detailed Efficiency Comparison
While we provide the efficiency on FLOPs in Figure 1, we further provide the efficiency on GPU memory, wall time, and FLOPs for training each method in Table 6. Specifically, we measure the computational cost on the NewsQA dataset with BERT-base, where we use the single Geforce RTX 2080 Ti GPU on the same machine. For our KALA, as we can flexibly manage the cost of GPU memory by reducing the number of entities in entity memory (See Figure 8 with Appendix C.2 for more analysis on the effects of the size of entity memory), we provide the experimental results on two settings - KALA with memory size  $0.2\mathrm{k}$  and  $62.8\mathrm{k}$  (full memory). As shown in Table 6, we confirm that the computational cost of our KALA with the full memory is similar to the cost of the more params baseline that uses one additional transformer layer on top of BERT-base. However, by reducing the number of entities in the memory, we can achieve better efficiency than more params in terms of GPU memory and FLOPs. Also, we observe that the training cost (i.e., Wall Time and FLOPs) of TAPT and DAPT is high, especially on DAPT, thus we verify that our KALA is more efficient to train for domain adaptation settings.
D Additional Visualization Results
Here we provide the frequency distribution of entities, additional case studies, and more illustrations of textual examples and embedding spaces.
D.1 Additional Representation Visualization
While we already show the contextualized representations of seen and unseen entities in the latent
Figure 13: Distribution of frequency of entities on QA datasets: NewsQA, Relation, and Medication, where almost all entities appear less than 10 times, while an extremely few numbers of entities appear very frequently.
space in Figure 2 right, we further visualize them including the missing baselines of Figure 2, such as Fine-tuning or TAPT, in Figure 12 on the NCBI-Disease dataset. Similar to Figure 2, we observe that all baselines fail to closely embed the unseen entities in the representation space of seen entities. While this visualization result does not give a strong evidence of why our KALA outperforms other baselines, we clearly observe that KALA is beneficial to represent unseen entities in the feature space of seen entities, which suggests that such an advantage of our KALA helps the PLM to generalize over the test dataset, where the context contains unseen entities.
D.2 Entity Frequency Distribution
We visualize the frequency of entities in Figure 13 and 14. The entity frequency denotes the number of mentions of their associated entities within the entire text corpus of the training dataset. As shown in Figure 13 and 14 of QA and NER datasets, the entity frequency follows the long-tail distribution, where most entities appear a few times. For instance, in the NewsQA dataset, more than  $20\mathrm{k}$  entities among entire  $60\mathrm{k}$  entities appear only once in the training dataset, whereas one entity  $(CNN^{10})$  appears approximately  $20\mathrm{k}$  times. This observation suggests that most of the elements in the entity memory are not utilized frequently. In other words, only few entities are accurately trained with many training instances, whereas there exists the stale embeddings which are rarely updated. This observation raises an interesting research question on the efficient usage of the entity memory, as we can see in Figure 8 that the small size of entity memory could result in the better performance (See Appendix C.2). We leave the more in-depth analysis on the entity memory as the future work.
Figure 14: Distribution of frequency of entities on NER datasets: CoNLL-2003, WNUT-17, and NCBI-Disease, where almost all entities appear less than 10 times, while an extremely few numbers of entities appear very frequently.
D.3 Additional Case Study
In addition to the case study in Figure 5, we further show the case on the question answering task in Figure 15, like in Section 5.5. With this example, we explain how the factual knowledge in KGs could be utilized to solve the task via our KALA. The question in the example is "who was kidnapped because of her neighbor". We observe that DAPT answers this question as Araceli Valencia. This prediction may come from matching the word 'her' in the question to the feminine name 'Araceli Valencia' in the context. In contrast, our KALA predicts the Jaime Andrade as an answer, which is the ground truth. We suspect that this might be because of the fact "(Jaime Andrade, spouse, Valencia)" in the knowledge graph, which relates the 'Valencia' to the 'Jaime Andrade'. Although it is not clear how it directly affects the model's performance, we can reason that KALA can successfully answer the question by utilizing the existing facts.
D.4 Additional Data Visualization
In Figure 16 and 17, we visualize the examples of the context with its seen and unseen entities and its relational facts. We first confirm that the quality of facts is moderate to use. For instance, in the first example of Figure 16, the fact in the context that Omar_bin_Laden is son of Osama_bin_Laden, is also appeared in the knowledge graph. In addition, we observe that there are facts that link unseen entities to the seen entities in both Figure 16 and 17. Thus, while some of the facts in the knowledge graph are not accurate, we can represent the unseen entities with their relation to the seen entities. We expect that there is a still room to improve in terms of the quality of KGs, allowing our KALA to modulate the entity representation more accurately. We leave the study on this as the future work.
10 Almost every context in NewsQA includes the text 'CNN' since they are originated from the CNN News.
DAPT
Context PHOENIX, Arizona (CNN) – Jamie Andrade had just gotten out of the shower when the men came to snatch him. Jamie Andrade was kept in this closet for three days without food or water, police say. His wife, Araceli Valencia, was mopping the kitchen in ... (ellipse)...
Facts (Sampled)
Valencia, spouse, Jaime Andrade Jr.)
(Jamie Andrade Jr., spouse, Valencia)
Question who was kidnapped because of her neighbour?
Answer Jaime Andrade
DAPT prediction Araceli Valencia,
KALA prediction Jaime Andrade
Figure 15: A textual example from NewsQA with predictions from each method (DAPT and KALA), and also the T-SNE plot of contextualized representations from the last layer of BERT obtained by each method. Grey dots indicate tokens without any mentions, and dots in other colors indicate tokens with mentions to the entity. We also represent sampled facts in Knowledge Graph we used. Blue text indicates the mention of seen entities and red text indicates the mention of unseen entities. The fact is represented as the format of (head, relation, tail). Text with yellow background indicates the ground truth answer span.
Context MADRID, Spain (CNN) – One of Osama bin Laden’s sons has been denied asylum in Spain, an Interior Ministry spokeswoman told CNN on Wednesday. Omar bin Laden pictured earlier this year during television interview in Rome, Italy. Omar bin Laden, who is in his late 20s, stepped off a plane at Madrid’s Barajas International Airport during a stopover late Monday and informed authorities that he planned to request political asylum, the spokeswoman said. Bin Laden has publicly called on his father to abandon terrorism. He prepared his formal asylum request Tuesday at the airport with the help of a translator, filing it around 1 p.m., the spokeswoman said. The Interior Ministry, which had 72 hours to reply to the request, was required to seek the opinion of the U.N. High Commissioner for Refugees on the matter. The UNHCR recommended … (ellipse) …
(Bin Laden, significant event, Flight)
(International Airport, country, Spain)
(International Airport, [UNK], Madrid)
(Omar bin Laden, father, Osama Bin Laden)
(Spain, diplomatic relation, Italy)
Osama Bin Laden, child, Omar Bin Laden)
(Italy, diplomatic relation, Spain)
Question 1 Where was Omar previously denied?
Answer 1 asylum in Britain.
Question 2 Did Spain give a reason for turning down the asylum?
Answer 2 was given
Question 3 Who was denied asylum in Britain?
Answer 3 Omar bin Laden
Question 4 What family member of Omar bin Laden was associated with terrorism?
Answer 4 his father