Plug-and-Play Adaptation for Continuously-updated QA

Introduction

LM-as-KB is a new paradigm in which pre-trained language models (LMs) are used as implicit knowledge bases (KBs) Petroni et al. (2019). This is made possible by LMs’ impressive ability to memorize factual knowledge Heinzerling and Inui (2021); Brown et al. (2020). Recently, two tasks have been used to assess such ability: LAMA, a knowledge probing benchmark, challenges LMs to fill in masked words over relational knowledge Petroni et al. (2019); and closed-book QA (CBQA) examines whether LMs can correctly answer natural language questions Roberts et al. (2020).

For practical usage, LM-as-KB requires that LMs are updated periodically to stay current with the ever-evolving world. Thus, LMs’ ability to update knowledge should also be evaluated. To this end, we present Continuously-updated QA (CuQA), which tests the ability to continuously inject knowledge to update (or target knowledge), while retaining existing knowledge (or source knowledge). Specifically, we consider multiple large-scale knowledge updates (8k to 60k) covering two scenarios: injecting new knowledge (Scenario 1 in Figure 1) and updating existing knowledge (Scenario 2 in Figure 1) .

Our goal is to organize the implicit storage of knowledge, to add target knowledge (yellow box in Figure 1) and anchor to select target knowledge. A simple approach is to train updated LMs from scratch; however, this is far too expensive considering the parameter sizes of recent LMs, such as 175B for GPT-3 Brown et al. (2020) and about 11B for T5 Raffel et al. (2020). There has also been related work for the two scenarios. For Scenario 1, a method for continual learning can be adopted, constraining the distance between parameters before and after fine-tuning Chen et al. (2020). However, this approach still suffers from so-called catastrophic forgetting, where the LMs fail to retain large amounts of source knowledge. For Scenario 2, one may consider knowledge editing methods, where we see reasonable performances for a single knowledge edit while retaining the rest De Cao et al. (2021); Mitchell et al. (2021). However, this line of work does not perform well when multiple edits are accumulated, e.g., only 67% of 125 edits were updated, as reported in Mitchell et al. (2021).

We propose to efficiently extend LMs with plug-and-play modules that store target knowledge. More specifically, we adopt a parameter-expansion method in which the LM storing existing knowledge is extended with plug-in feed-forward modules storing updated knowledge. Depending on the input, the LM selectively uses either the original LM or a plug-in module. We stress that, by keeping the original LM intact, we retain (a) not only source knowledge, (b) but also those outdated from updates (red arrow in Figure 1). (a) is important to avoid catastrophic forgetting, while (b) is useful when updates need to be reverted due to ethical concerns—for example, there can be malicious attempts to override facts.

We evaluate our approach on zsRE Levy et al. (2017) and Natural Questions Kwiatkowski et al. (2019) to showcase successful updates of new knowledge and retention of existing knowledge. We measure the accuracies on both previous and updated knowledge and find that ours show x4 higher updates/forgets ratio, compared to fine-tuning. We will release our code and dataset.

We present CuQA, a novel task to assess LMs’ ability to continuously inject knowledge to update.

We propose a new methodology, plug-and-play adaptation, to continually learn new knowledge while better retaining existing knowledge.

Related Work

The relevant research can be categorized into three groups: Knowledge Editing, Continual Learning, and Adaptation. In Table 1, we compare these with our method.

In Table 1(a), knowledge editing methods De Cao et al. (2021); Mitchell et al. (2021); Dai et al. (2021) aim to efficiently edit model’s parameters on examples that have conflicts with old facts, while preserving the outputs of untargeted examples. Instead of directly updating gradients by fine-tuning, these methods transform the gradients for new edit parameters. As representative methods for knowledge editing, KnowledgeEditor (KE) De Cao et al. (2021) using LSTM produces gate vectors, then the gated sum of gradients is updated into the model, while MEND Mitchell et al. (2021) uses simple MLP layers and residual connections for the same purpose. Although these methods succeeded in updating the target examples less forgetting, their target scenario is a single edit, such that the cumulative effect of multiple edits does not reflect well, which disqualifies its use for our target task of update large-scale data (8K$\sim$60K). As reported in Mitchell et al. (2021), MEND successfully updates only 67% of edits when applying 125 edits, while our finding was consistent when none of the 125 edits was applied in our evaluation.In the case of KE, we reimplement the released code for testing: https://github.com/nicola-decao/KnowledgeEditor. In addition, for editing previous knowledge, KE and MEND simulate knowledge updates, by generating synthetic knowledge from LM. Such generations may not be realistic data and also give unfair advantages to LM-based methods, while we use actual up-to-date knowledge as new data, which were annotated on recent corpus Zhang and Choi (2021).

Continual Learning (CL) for NLP

For our task, we can adopt CL methods, learning a new task while preserving the accuracy on previous tasks. Kirkpatrick et al. (2017) proposed Elastic Weight Consolidation, alleviating catastrophic forgetting. This method regularizes learning on a new task, by constraining the parameters trained on the previous task. For NLP tasks, RecAdam Chen et al. (2020) uses the regularization and annealing technique, which is a CL baseline in our experiment. While CL approaches focusing on forgetting do not consider conflicts between old and new knowledge, our work deals with such a realistic scenario. Additionally, previous work Dhingra et al. (2021) proposed benchmarks for probing temporal language models, asking “Fill-in-the-Blank (FIB)" questions. Meanwhile, FIB questions are limited to evaluate masked language models, such as BERT and RoBERTa. We extend to evaluate arbitrary questions for a knowledge-intensive task; closed-book QA, which can evaluate generative LMs with broader applicability, to include T5 and GPT.

Task-aware Adaptation for Transformers

Recent works Hu et al. (2021); Wang et al. (2020); Lin et al. (2020) study LM adaptation to new labeled data in a new domain, which has a different data distribution from that at pretraining. These works show performance improvements on downstream tasks in the new domain, while fine-tuning a small number of parameters. However, these adaptation methods do not consider sequential training, and overwrite the new data into the parameters that store previous knowledge. In our experiment, it is observed that the adaptation methods are rapidly forgetting previously seen data, while performing well on new knowledge.

A Continuously-updated QA Task

In this section, we propose Continuously-updated QA (CuQA), a new continual learning task for knowledge updates in LMs based on closed-book QA (CBQA) Roberts et al. (2020). In CBQA, LMs answer factual questions with the implicit knowledge stored in the model, without any external context (i.e., in contrast to open-domain QA), so that LMs are required to adequately update their parameters to the target knowledge. In our CuQA, LMs learn source (original) knowledge first, then update them with target (new) knowledge without source knowledge access. For the above setting, source knowledge (to be retained) and target knowledge (to be added) in CuQA do not have any overlap of QA pairs (or paraphrases) for any given fact.

Specifically, we denote a factual pair of question and answer as $(q,a)$, source knowledge as $\mathcal{K}_{s}$, and target as $\mathcal{K}_{t}$. We first build an initial model $\theta^{old}$ pre-trained on source knowledge $\mathcal{K}_{s}$. Then, we inject target knowledge $\mathcal{K}_{t}$ into the pre-trained model and obtain the infused model $\theta^{new}$. Our goal is to memorize $\mathcal{K}_{t}$ on model $\theta^{new}$, with less forgetting $\mathcal{K}_{s}$. If knowledge in $\mathcal{K}_{t}$ conflicts one in $\mathcal{K}_{s}$, the model is required to adjust its parameters by reflecting the target knowledge. Note that multiple target knowledge can be sequentially updated to the model (see details in Section 4).

Research Questions

CuQA is designed to address the following research questions:

RQ1: Can the method learn target knowledge while retaining source knowledge?

RQ2: How does sequentially learning multiple target knowledge affect the performance?

RQ3: How does the size of each target knowledge affect the performance?

Metric

For evaluation, we measure the success of updates, retaining of source knowledge, and generality using exact match (EM) scores. Additionally, we measure the ratio of forgets to updates.

Accuracy on $\mathcal{K}_{t}$ : we evaluate how much model $\theta^{new}$ successfully updates examples in $\mathcal{K}_{t}$.

Accuracy on $\mathcal{K}_{s}$: how much model $\theta^{new}$ forgets examples in $\mathcal{K}_{s}$. This indicates performance degradation, when replacing $\theta^{old}$ with $\theta^{new}$.

Accuracy on $\mathcal{P}_{s}$, $\mathcal{P}_{t}$: how well model $\theta^{new}$ generalizes on semantically equivalent questions (or paraphrases).

F/U Ratio ($\#$ of forgets$/\#$ of updates): how many examples in $\mathcal{K}_{s}$ are forgotten per an update of one example in $\mathcal{K}_{t}$. ($\#$ of forgets) is equal to the difference of correct prediction cases in $\mathcal{K}_{s}$, between $\theta^{old}$ and $\theta^{new}$.

Method

In this section, we describe baseline approaches (Section 4.1), and introduce our proposed method for plug-and-play adaptation (Section 4.2).

We establish three baseline for (a), (b), and (c), in Table 1. Since we found that a knowledge editing approach is outperformed by fine-tuning, we exclude it as baselines, and add fine-tuning instead.

As a naive baseline, we start with the previous work Roberts et al. (2020) for CBQA, by fine-tuning T5 Raffel et al. (2020) with encoder-decoder structure. This baseline is to fine-tune the pre-trained model $\theta^{old}$ on facts in $\mathcal{K}_{t}$ to minimize the loss:

where $L$ refers to a seq2seq loss. This baseline is expected to optimize accuracy on target knowledge $\mathcal{K}_{t}$, thus increases the distance between the before- ($\theta^{old}$) and after-parameters ($\theta^{new}$) resulting in the risk of forgetting. For other baselines and our method, we adopt the same transformer: T5 as backbone network.

Regularized fine-tuning for CL

We adopt RecAdam Chen et al. (2020) aiming to reduce the forgetting risk by adding a constraint to minimize the distance between $\theta^{old}$ and $\theta^{new}$ as follow:

where $\|\cdot\|_{p}$ indicates $L_{p}$ norm. In addition, RecAdam uses an annealing technique, controlling the ratio between $\mathcal{R}$ and the fine-tuning loss (Eq. (1)) as follows:

where $k$ and $t_{0}$ are hyper-parameters.

Adapters for knowledge updates

2 Our Method

Motivated by the intuition of regularization to preserve source knowledge and that of adapters to inject target knowledge into new parameters, we show their strengths can be combined for our task. At the inference phase, our method selectively uses the plug-in modules to keep source knowledge intact, while tasks requiring target knowledge will be redirected to new plug-in modules.

Specifically, our distinction is augmenting function $f$ (in an original LM) with function $g$, representing source and target knowledge respectively. The function $f$ is a single layer in transformer trained on source knowledge $\mathcal{K}_{s}$, and $g$ is an augmented function with new parameters for $\mathcal{K}_{t}$. Existing work, such as LoRA, can be interpreted by adding the two functions:

However, the above add-aggregation has a limitation, as $g(x)$ can affect the model’s outputs, and increase the distance between hidden states in $\theta^{old}$ and $\theta^{new}$, which causes a forgetting problem.

Our key distinction is adding a selector, that is selectively activated for $q$ requiring the use of plug-in module $g$, as follows:

where $\sigma(q)$ is 1 or 0 depending on query $q$. While there can be various ways to train the selector in a sophisticated way, supervised either directly, or indirectly in an end-to-end manner, we show a simple unsupervised selector is already sufficient to show gains. Specifically, our selector is a key-value lookup where the key is $m_{i}$ and value is $g$. At inference time, when given query $q$ is based on facts in $\mathcal{K}_{t}$, we activate the augmented $g$ for generating its output. If $q$ is not from $\mathcal{K}_{t}$, we use only the original model $\theta^{old}$ for generation. To classify whether the input is from $\mathcal{K}_{t}$ or not, we build explicit memory with embeddings of $\mathcal{K}_{t}$ and leverage the distance with nearest neighbor (NN) in the memory.

where sim indicates cosine similarity. Based on $s_{q}$, if the score is greater than or equal to threshold $\delta$, we assume $q$ is from target knowledge $\mathcal{K}_{t}$. We build a indicator function as follows:

In other words, $s_{q}\geq\delta$ indicates that input $q$ is semantically similar with one fact in $\mathcal{K}_{t}$. At that time, our model is augmented with $g$ that stores new and updated knowledge.

Meanwhile, as shown in Figure 2, we apply the selective use of parameters to only a decoder in a transformer architecture, not a encoder. The switch $\sigma$ depends on query embedding $q$, and the embedding $q$ is extracted from T5 encoder. If we apply the switch $\sigma$ to hidden states in T5 encoder, this causes a recursion relation, or inefficient computations. By augmenting $g$ for the decoder, embedding $q$ is not changing during updating target knowledge, and depends on only pre-trained $\theta^{old}$.

Our new perspective has another benefit of naturally generalizing to sequential ($>$2) sources. Assume that there are multiple target knowledge to be sequentially updated, i.e., $\mathcal{K}_{t}^{1},\mathcal{K}_{t}^{2},...,\mathcal{K}_{t}^{M}$. We build multiple functions $g_{k}$ and memories $\mathcal{M}_{k}$ (where $k=1,...M$), according to each target knowledge. The new function considering the multiple knowledge is denoted as follows:

During training $j$-th target $\mathcal{K}_{t}^{j}$, the switch $\sigma_{k}(q)$ is activated where $1\leq k\leq j$. At inference time, our selector extracts top1-NN fact $m^{*}$, which is closest to a query $q$. If $m^{*}$ is in $\mathcal{M}_{k}$, the switch $\sigma_{j}(q)$ is activated where $1\leq j\leq k$, as follows:

If the NN fact $m^{*}$ is in $\mathcal{M}_{j}$, we estimate that its implicit knowledge is stored in the accumulated function $\sum_{k=1}^{j}g_{k}(x)$. That is, when $m^{*}$ is in $\mathcal{M}_{j}$, the activation is decided as follows:

An alternative adapter

We can replace LoRA with K-adapter Wang et al. (2020). In K-adapter, $f$ is a transformer layer (denoted as $\textrm{TRM}(x)$), and $g$ is multiple transformer layers with two projection layers (denoted as $\textrm{KIA}(x)$)). That is, $f(x)=\textrm{TRM}(x)$, consisting of one self-attention & two feed-forward layers. In the original paper Wang et al. (2020), $g(x)$ consists of multiple transformer layers and up&down projection layers. For K-adapter, we set a simple version with only a single transformer layer, as follows:

where the parameters in TRM are fixed and that in KIA is trainable on target knowledge.

Experiment

In this section, we demonstrate the effectiveness of our approach on CuQA.

We evaluate our method on the following closed-book QA datasets:

(1) Zero-shot Relation Extraction (zsRE): Levy et al. (2017) build relation-specific QA pairs, and De Cao et al. (2021) utilize this dataset for a closed-book QA task. This set provides question paraphrases based on the same fact and answer. We split this set into two groups ($\mathcal{K}_{s}$ and $\mathcal{K}_{t}$) that do not share the same facts. To validate generalization, we build held-out sets ($\mathcal{P}_{s}$ and $\mathcal{P}_{t}$) that are not used in training process. For this, we sample one QA pair among paraphrases based the same fact as $\mathcal{P}$.

(2) Natural Questions (NQ) + SituatedQA: Kwiatkowski et al. (2019) build NQ – a large-scale QA dataset based on user queries. We consider NQ as source knowledge $\mathcal{K}_{s}$ except outdated facts based on SituatedQA. Zhang and Choi (2021) proposed SituatedQA identifying temporal- and geographical-dependent questions on a subset of NQ. We use the temporal-dependent QA pairs as $\mathcal{K}_{t}$, which are annotated based on 2021 dump of Wikipedia. For $\mathcal{P}_{s}$ and $\mathcal{P}_{t}$, as both NQ and SituatedQA do not provide paraphrases, we follow De Cao et al. (2021) using back-translation for generating paraphrases.

Implementation

For T5 model, we use a large version with total 770M parameters. In our experiment, we assume that the old model $\theta^{old}$ storing source knowledge is available. For NQ, we used the open-source pre-trained modelhttps://huggingface.co/google/t5-large-ssm-nq as the model $\theta^{old}$. For zsRE, we load and train T5 modelhttps://huggingface.co/google/t5-large-ssm on source knowledge. For training, we set batch size 64 on 4 RTX3090 GPUs, and used Adam Kingma and Ba (2015) optimizer with learning rate 4e-4. For development set, we sample each 1K from $\mathcal{K}_{s}$, $\mathcal{K}_{t}$, and select the maximum harmonic mean of their accuracies as a best model. As a hyper-parameter, we search $\delta$ in a range of with 0.05 step size, and found the best value ($\delta$=0.9) based on development set. As embedding memory $\mathcal{M}$, we used additional parameters: 60M for zsRE and 8.5M for NQ. The size of the memories can be reduced by several techniques, such as random projection Luan et al. (2020) and binary encoding Yamada et al. (2021), which is left out of our focus.

Comparison with baselines

We compare our method with baselines, as mentioned in Section 3.2; Fine-tuning (B-I), RecAdam (B-II), LoRA (B-III), and K-adapter (B-IV). When re-implementing K-adapter, we do not freeze the parameters of decoder, unlike in the original paper Wang et al. (2020), because the performance is not changing when freezing. We train each model until 80 epochs and select a best model by the harmonic mean of source/target knowledge in development set.

1 R1: Comparing Ours with Baselines

Table 3 shows our main experimental results on two CBQA datasets. First, the model $\theta^{old}$ memorizes the source knowledge $\mathcal{K}_{s}$ well and generalizes on the paraphrase set $\mathcal{P}_{s}$ as well, showing high accuracy on both datasets. After training on $\mathcal{K}_{t}$, all models perform well on $\mathcal{K}_{t}$ and $\mathcal{P}_{t}$. These results indicate that these models are at least appropriate for memorizing training data in the current task.

Meanwhile, while acquiring $\mathcal{K}_{t}$, the models show variant results on $\mathcal{K}_{s}$ and $\mathcal{P}_{s}$, which have the different ability of retaining previous knowledge against forgetting. In Fine-tuning (B-I), its performances on source knowledge $\mathcal{K}_{s}$ and $\mathcal{P}_{s}$ decrease as training epochs (see Figure 3). RecAdam (B-II) alleviates the forgetting problem of fine-tuning, but the performance gains are marginal on two datasets. K-adapter (B-III) shows the strong performance on $\mathcal{K}_{s}$ with less forgetting, however, does not perform well on $\mathcal{P}_{s}$ and $\mathcal{P}_{t}$ showing low generalization. Because LoRA (B-IV) has the fewest trainable parameters, its forgetting is more aggravated, showing the worst performance on $\mathcal{K}_{s}$ and $\mathcal{P}_{s}$ in both zsRE and NQ. Ours with either K-adapter or LoRA shows the best performance on $\mathcal{K}_{s}$ and $\mathcal{K}_{t}$. In terms of the F$/$U ratio, our method also shows the lowest loss when updating one new example. Figure 3 shows how the performance of each model changes over training epochs, on the development set.

In an ablation study, we test which component has the higher impact on memorizing implicit knowledge, on paraphrase set $\mathcal{P}_{s}$ and $\mathcal{P}_{t}$. In our method with LoRA, the function $f$ in Eq. (8) can be applied to any projection layer in transformers. While the original work Hu et al. (2021) applies to query- and value-matrices ($W_{Q},W_{V}$) in self-attention, we consider feed-forward layers ($W_{FF}$), as well as self-attention. In addition, we observe how does the performance vary when the number of parameters increases by controlling rank $r$. In Table 4, we empirically found applying feed-forward layers is more effective than query and value projection, especially on target knowledge $\mathcal{P}_{t}$. These results indicate that memorizing factual knowledge is more relevant with a feed-forward module, which is consistent with the views in Sukhbaatar et al. (2019); Geva et al. (2020).

To evaluate the scalability of our method on multiple $\mathcal{K}_{t}$ ($>$2), we assume multiple updates (five-phase) with smaller amount of examples, by splitting target knowledge $\mathcal{K}_{t}$ in zsRE (Large, 60K), into four sets, from $\mathcal{K}_{t}^{1}$ to $\mathcal{K}_{t}^{4}$ (each 15K). In this experiment, we train models during 40 epochs/phase. To generalize for LoRA baseline, we aggregate multiple $g_{k}$ by addition, by activating all the switches at inference, i.e., $\sigma_{1:M}(x)=1$ in Eq. (13). This setting assumes that this baseline cannot leverage our selector to organize the storage of implicit knowledge. Figure 4 shows the performances of Fine-tuning, LoRA, and Ours, over training epochs. In fine-tuning, the accuracy on source knowledge keeps dropping during the whole training process. In LoRA, multiple updating deteriorates memorizing target knowledge stored in adapters, faster than source knowledge stored in the original parameters. This indicates that the fewer parameters, the faster the forgetting. In contrast, our method consistently outperforms the baselines, by retaining five knowledge, with forgetting less. To summarize these results, sequential updates aggravate forgetting of the fine-tuning method, which can be overcome through the selective use of adapters.

As the size of target knowledge increases, it makes LMs suffer from more forgetting, increasing the distance between before- and after-parameters. In this section, we observe how does the performance of each model vary as different sizes of $\mathcal{K}_{t}$. Figure 5 shows the the accuracies of zsRE datasets (Large-60K, Medium-30K, Small-15K), over training epochs. On source knowledge $\mathcal{K}_{s}$, the performance of fine-tuning and LoRA keeps dropping, and the accuracy drops are proportional to the size of target knowledge. Meanwhile, our method with LoRA consistently maintains high performance, which is not sensitive to training epochs. On target knowledge $\mathcal{K}_{t}$, the performances of three models reach high accuracy. However, our method on Large zsRE shows unstable performance at the end of training, which may need to use early stopping.

4 Analysis of Selector

In Table 5, we show the distribution of selector’s predictions and the ground-truths, in our experiment on zsRE (Large). Nearest Neighbor-based selector successfully classifies 88.9% of examples, while 11.1% failed. In our method, if the selector classifies an input as target knowledge, the plug-in $g$ is activated. Instead of the use of $g$, we can retrieve answers aligned with questions in $\mathcal{M}$, not generate them. We compare our generation with the retrieval in each case of Table 5. Table 6 shows the accuracy of predicting the answers, where the numbers in each cell indicate EM of our generation (retrieval: in parentheses). If an example in source knowledge is incorrectly classified as target, there is no relevant fact in $\mathcal{M}$, thus the accuracy in this case is zero. In contrast to Retrieval, our generative method is robust in this case, achieving 70.8% EM, because ours with $g$ learned the source knowledge.

Conclusion

This paper studies how to accumulate new knowledge to LMs that stores existing knowledge. We propose a simple yet effective method to update target knowledge into new parameters, preventing from forgetting source knowledge. On two datasets: zsRE and NQ, our empirical results show that our proposed method can improve existing approaches for continual learning or task adaptation.

Acknowledgement

This research was supported by SNU-NAVER Hyperscale AI Center, and IITP grants funded by the Korea government (MSIT) [2021-0-02068 SNU AIHub, IITP-2022-2020-0-01789].

References