Large Language Models with Controllable Working Memory
Year: 2023

Abstract:
Large language models (LLMs) have led to a series of breakthroughs in natural language processing (NLP), partly owing to the massive amounts of world knowledge they memorize during pretraining. While many downstream applications provide the model with an informational context to aid its underlying task, how the model’s world knowledge interacts with the factual information presented in the context remains under explored. As a desirable behavior, an LLM should give precedence to the context whenever it contains taskrelevant information that conflicts with the model’s memorized knowledge. This enables model predictions to be grounded in the context, which then facilitates updating specific model predictions without frequently retraining the model. By contrast, when the context is irrelevant to the task, the model should ignore it and fall back on its internal knowledge. In this paper, we undertake a first joint study of the aforementioned two properties, namely controllability and robustness, in the context of LLMs. We demonstrate that state-of-theart T5 and PaLM models (both pretrained and finetuned) could exhibit low controllability and robustness that does not improve with increasing the model size. As a solution, we propose a simple yet effective method – knowledge aware finetuning (KAFT) – to strengthen both controllability and robustness by injecting counterfactual and irrelevant contexts to standard supervised datasets. Our comprehensive evaluation showcases the utility of KAFT across model architectures and sizes.

heading: 1 Introduction
text: Large language models (LLMs) pretrained on large scale datasets have shown promising results across natural language tasks (Vaswani et al., 2017; Devlin et al., 2019; Raffel et al., 2020a; Brown et al., 2020; Rae et al., 2021; Chowdhery et al., 2022; Smith et al., 2022). However, as models scale ever larger,
they become more expensive to train, making it unrealistic to frequently update model parameters. On the other hand, many real world applications often necessitate adjusting model behavior. This dilemma is especially sharp in the case of factual (world) knowledge that plays important role in realizing impressive performance of LLMs. It is well known that LLMs memorize large amounts of factual knowledge in their parameters (Petroni et al., 2019; Roberts et al., 2020; Geva et al., 2021), which could potentially be out-dated or incorrect. Even for moderate-size models, it is prohibitively expensive to retrain every time an update happens or a mistake is uncovered. Even if resources are ample, it is difficult to ensure that the modification of model parameters do not affect unrelated skills or knowledge.
In human cognition, working memory (George A. Miller, 1960) provides the biological brain with the ability to hold information temporarily to perform tasks such as conversation, reasoning, and mathematics in a way that is adaptive to the ever changing environment. As shown both experimentally and theoretically (Fuster, 1973; Ashby et al., 2005), working memory is stored in sustained activations of neurons, as opposed to the long term memory which is stored in weights. Working memory is also the immediate information buffer that is accessed while performing conscious tasks. In particular, it is where the fusion of perceptual inputs and long term memory happens (Fukuda and Woodman, 2017). This suggests that a potential method to solve LLMs’ pointwise knowledge update and correction problem is to control the working memory stored in activations, rather than editing the long term memory stored in the model weights.
As demonstrated by their powerful in-context few shot learning abilities (Brown et al., 2020), LLM could utilize different activation patterns resulting from different contexts during inference to solve a diverse set of tasks without any changes in
1774
the weights. It is natural to expect that the same would be true with factual knowledge. In particular, one could prepare a large list of natural language statements covering desired knowledge updates and corrections. At inference time, one can provide the relevant statements as context along with the input and hope that the model would perform the task based on the new knowledge presented in this context. Thus, if the model’s working memory is indeed controllable by context, then a single model with static long term memory can produce different results based on a varying set of facts available in different contexts. However, we demonstrate that this approach may fall short for existing LLMs as they have great tendencies to ignore the context and stick to their own parametric knowledge – the world knowledge stored in its model parameters. This raises a natural question:

heading: Is it possible to design a mechanism to ensure that the context can reliably influence the model’s working memory?
text: Note that any such mechanism has to take into account the possibility of encountering noisy contexts. For example, any retrieval system that selects the task-relevant context from a large collection of contexts will be imperfect and occasionally provide irrelevant context. In such cases, it’s desirable that the model prediction does not get swayed by an irrelevant context. Interestingly, we show that the standard pretraining and finetuning methods do not ensure this behavior either. In fact, we demonstrate
that it’s the noise encountered during the training that often leads to the model ignoring the context.
In this work, we provide an affirmative answer to the aforementioned question and propose a novel approach – knowledge-aware finetuning (KAFT) – to make an LLM’s working memory controllable via relevant context while being robust against irrelevant context. Towards this, we aim to ensure that the model utilizes different types of information at its disposal in the following order:
relevant context
> model’s parametric knowledge (1)
> irrelevant context, (2)
where a > b indicates that a is prioritized over b. Thus, if the model decides that the context is relevant, it should ground its output in the context, ensuring the controllability of its working memory by the context. This is crucial when the context is in conflict with the model’s parametric knowledge. On the other hand, when the context is irrelevant, the model should instead stick to its parametric knowledge; thus ensuring robustness of its working memory against noise.
Our contributions. We develop first LLMs that utilize different knowledge sources with a predefined order of priorities. Along the way, we develop a systematic understanding of the working memories of LLMs and identify their shortcomings. Our key contributions are summarized below.
1. We undertake a systematic joint study of both controllability and robustness of the working memory of LLMs. Focusing on question answering (QA) tasks, we define the context-question relevance based on whether the context entails an answer to the question. We create a novel benchmark to measure the controllability by including contexts that imply an answer which contradicts the model’s pretrained knowledge.1 Similarly, we benchmark robustness against irrelevant contexts. We conduct an extensive evaluation of LLMs with different sizes across multiple architectures (encoderdecoder and decoder-only) from T5 (Raffel et al., 2020b) and PaLM (Chowdhery et al., 2022) family. We make the following key observations:
(a) LLMs could exhibit low controllability. Our experiments consistently show that both pretrained and QA finetuned LLMs tend to ignore a context when it contradicts with model’s world knowledge. We show that this problem persists and may intensify as the model becomes larger. We further show that the noise in the (QA) finetuning set plays an important role in emergence of this behavior (cf. Sec. 4.2). (b) LLMs may not be robust against context noise. We demonstrate that both pretrained and QA finetuned models are strongly interfered by irrelevant contexts, especially the ones that are on the same general topic as the underlying question (cf. Sec. 4.3).
2. We propose a novel method – knowledge aware finetuning (KAFT) – to directly enhance both controllability (Eq. 1) and robustness (Eq. 2) of an LLM. KAFT enhances the controllability by creating counterfactual data augmentations where the answer entity in the context is swapped to a different but plausible entity, in conflict with the ground truth (and potentially the model’s world knowledge). As for enhancing robustness, KAFT requires that the model should predict its pretrained closed-book answer rather than the ground truth answer whenever the context is irrelevant.
1We rely on in-context prompts in a closed book QA setup to measure the model’s parametric knowledge.
3. Through extensive empirical evaluation, we show that KAFT-based models successfully demonstrate the coexistence of controllability and robustness of model’s working memory (see Table 1 for an illustration).

heading: 2 Related Works
text: World knowledge in language models. Recent works established that LLMs memorize factual information present in the pretraining corpus. E.g., Petroni et al. (2019) utilize language model analysis (LAMA) probing to show that BERT models (Devlin et al., 2018) could act as knowledge bases. Roberts et al. (2020) reported similar findings for T5 models. It is therefore common practice to employ LLMs in tasks like closed book QA (Chowdhery et al., 2022).
Knowledge update in language models. Given that factual knowledge is ever-evolving, outdated memory of LLMs may lead to incorrect predictions (Lazaridou et al., 2021; Onoe et al., 2022). Furthermore, during deployment, one may unearth mistakes that need correction. Frequent retraining from scratch with an updated and corrected corpus would be prohibitively expensive. Ideas around finetuning (Zhu et al., 2020) and continued learning (Jang et al., 2022) train the model with less but still significant resources. Multiple recent efforts have studied how these models store the factual knowledge (Geva et al., 2021) and methods to update model parameters given new knowledge (De Cao et al., 2021; Dhingra et al., 2022; Mitchell et al., 2022; Meng et al., 2022a,b). These strategies change weights in response to single updates, risking inadvertently affecting unrelated skills or knowledge and creating burden to potentially store multiple versions of LLMs. We focus on updating the model behavior by providing a suitable context and ensuring that the model’s working memory is controllable by such contexts. Contextual and parametric knowledge. Guu et al. (2020); Joshi et al. (2020); Petroni et al. (2020) utilized retrieved context to assist language models in tasks such as QA. At the same time, LLMs memorize large amounts of knowledge in their parameters. Despite this dichotomy, only a few studies have previously addressed the relation between these two very different knowledge sources. Longpre et al. (2021) find that larger models have a greater tendency to ignore context in
favor of their own parametric knowledge, and that the noise in the context in the finetuning set plays a big role in causing this behavior. We incorporate the algorithms proposed by Longpre et al. (2021) for mitigating this problem as baselines in Sec. 4.4 (the Relevant Only Finetuning approaches), where we find such baselines lack robustness against irrelevant contexts (Fig. 1, 2). Kassner and Schütze (2020) showed that language models tend to be easily misled by certain types of irrelevant contexts. We observe similar phenomena in QA and show that our proposed KAFT leads to more robust models against irrelevant contexts. Finally, Pan et al. (2021) considers a scenario where some context sources may be less trustworthy than the model’s parametric knowledge. This scenario can be captured by an extension of our framework Eq.(1-2). For example, given three sources, one could enforce the following precedence order: source1 > source2 > model’s own knowledge > source3 > irrelevant contexts.
Notions of controllability and robustness. In control theory, controllability (Ogata, 1996) refers to the ability of using external inputs to manipulate system to reach all possible states. In the spirit of this definition, this paper measure the controllability of an LM’s working memory by external contexts. In the framework of controlled text generation (Zhang et al., 2022; Hu and Li, 2022), the notion of controllability explored here is a special type of fine-grained semantic control of the model’s behavior with the content of the context.
Notions of robustness. (Liang et al., 2022; Omar et al., 2022) survey many notions of robustness of language models around the notion of the invariance of model’s behaviors when the input is perturbed (for example, expressing similar semantic meanings in different ways). Our robustness benchmark is an extreme and input-dependent version under this framework. In our evaluations, the input contains two parts: the context and the question. In this work, a robust model’s response is invariant to large perturbations in the semantic content of the context, as long as these changes are not relevant to the question.
During the preparation of this manuscript, we were made aware of a parallel and independent investigation by Neeman et al. (2022) that shares some important aspects of our work.

heading: 3 Methods
text: For concreteness, consider a reading comprehension QA task where the model takes question q together with a context c as its input. The question has an answer label a. We also need a relevance label r denoting whether c entails a.
Starting with a pretrained LM M , we would like to build a model M ′ such that when the context c is relevant, its answer is always grounded in c, when c is irrelevant, it sticks to the pretrained model’s answer. In equations:
r = 1 : M ′(c+ q) = a (3) r = 0 : M ′(c+ q) = M(q) (4)
where + denotes string concatenation. This establishes the priority order of knowledge sources as in Eq. (1 & 2): if there is a conflict between a relevant context c and M ’s parametric knowledge, then the output should be consistent with c. In addition, irrelevant context should have no influence on the model’s output. Note that even though we are separating relevant vs irrelevant context here, the model does not know r a priori. It has to determine r based on the semantics of c and q.
In the KAFT data, r = 1 cases include relevant or counterfactual context, where a is the ground truth or counterfactual answer, respectively; r = 0 cases include empty or irrelevant contexts. Here the label is given by the pretrained model’s answer to the same question in a few-shot closed book setting, reflecting the model’s parametric knowledge. To provide more interpretability, we make the model output its classification of the context’s relevance along side the answer itself. See Table 3 for details.

heading: 3.1 Datasets
text: We construct KAFT based on several public datasets, including SQuAD 2.0 (Rajpurkar et al., 2018), T-REx (Elsahar et al., 2018), QASC (Khot
et al., 2020), and TriviaQA (Joshi et al., 2017). They cover several different QA formats, including multiple choice (QASC), Cloze (TReX), extractive (SQuAD), and open domain (TriviaQA). For each dataset, we may construct different types of context and corresponding labels as summarized in Table 4.

heading: 3.2 Models
text: We select two families of pretrained LLMs: T5 (Raffel et al., 2020b) representing the encoderdecoder architecture and PaLM (Chowdhery et al., 2022) representing the decoder only architecture. We include all three PaLM models (8B, 62B and 540B), while with T5 we restrict to the largest sizes (XL and XXL, with 3B and 11B parameters, respectively) because the smaller ones do not respond well to in-context few shot prompts, making it difficult to measure their parametric knowledge.

heading: 3.3 Relevant context
text: We define the relevance of a context by whether it logically entails an answer to the question, which is a strong requirement - even if a piece of context is on the same topic of the question or contain the answer label, it might still be irrelevant. In practice, this happens often among retrieved results. In Sec 4.4, we show that if the model is still required to fit on to the ground truth label when given an irrelevant context, then the model becomes more likely to ignore relevant contexts. It is therefore crucial to strive towards precise logical entailment when building relevant context. We apply several techniques to improve the semantic connection between the context and the QA pair as shown in Table. 4. More details can be found in Appendix A.1.

heading: 3.4 Irrelevant Context
text: An irrelevant context is any context that does not entail the answer. An easy irrelevant context is completely off topic. We obtain them with random sampling for all datasets. A hard irrelevant context is on the same topic, sometimes discussing the same entities involved in the QA pair but does not logically entail the answer. SQuAD 2.0 already contains human labels on whether the answer can be derived from the context, thus providing hard irrelevant contexts. TriviaQA provides somewhat extensive paraphrases for each answer. We filter the retrieved contexts to find ones that do not contain any answer paraphrase, and use them as hard irrelevant context.

heading: 3.5 Probing pretrained knowledge
text: We first use the pretrained model to generate M(q) in Eq. 4, which are then used to assemble the KAFT finetuning dataset according to Eq. 4. We use handengineered few-shot knowledge probing prompts that condition the model to answer a question according to its world knowledge acquired during pretraining. In Appendix A.3, we provide more details on the construction of these prompts.

heading: 3.6 Counterfactuals
text: To train the model to be controllable by the context, we explicitly engineer plausible training data where the context is in conflict with the model’s pretrained world knowledge. Given a triple of question, answer, and relevant context, we use a pretrained T5 XXL model to generate a triple of question, counterfactual answer, and counterfactual context with prompt engineering. We apply several filtering and postprocessing techniques to ensure the quality. Details are given in Appendix A.4.

heading: 3.7 Metrics
text: In this section, we define metrics that measures controllability and robustness. All results are from single runs.
Controllability. To measure controllability, we supply the model with a counterfactual context and examine whether it can output the corresponding counterfactual answer. For a fair comparison, we select questions which all five pretrained models can answer correctly in a closed book few-shot setting, which are referred to as head questions. Since they are likely well represented in the pretraining set, such questions are challenging as we swap the answer to counterfactuals. Since we don’t have any paraphrases of the counterfactual answer, we choose to use thresholded unigram recall to measure the performance. In particular, a model output is rated positive if the output of the model contains > 80% of the answer unigrams, with stop-words removed.
Robustness. To measure robustness, we use the human labeled "impossible" slice of SQuAD 2.0, since SQuAD 2.0 contains many examples where the context is on the same general topic of the question but does not contain the answer. We measure the rate when the model successfully avoids extracting answers from such irrelevant contexts. The avoidance is considered successful if the con-
text contains less than 50% of the unigrams in the model’s prediction, removing stop words.

heading: 3.8 Baselines
text: Pretrained. We evaluate the pretrained model’s controllability and robustness in a zero shot reading comprehension QA setup. The context is concatenated with the question in input sequence.
Noisy finetuning. In this approach, the label is the ground truth answer whether the context is relevant or not. This is a standard method implicitly used in most QA datasets.2 In this work, we construct this baseline for KAFT by first removing all counterfactual augmentations and then replace all labels with the ground truth label.
Relevant only finetuning. The approach where only relevant context and the corresponding ground truth label are used during finetuning, which is shown to improve controllability in (Longpre et al., 2021). As a baseline for KAFT we remove all counterfactual and irrelevant augmentations and only keep the relevant slice of our finetuning data.
UQA V2. The Unified QA 11B (Khashabi et al., 2022) model, which is a general purpose QA model finetuned on a collection of 20 QA datasets. We take the largest model (11B) in the UQA V2 family as a baseline and compare with KAFT T5 XXL which is of similar size in Fig. 2. Since UQA V2 contains SQuAD 2.0 in its training set, where the label for irrelevant context is an empty string, it
2As a notable exception, SQuAD 2.0 has empty strings as labels for its irrelevant context.
does not completely follow the noisy finetuning prescription introduced earlier.
KAFT noCF. The KAFT method with no counterfactual augmentations.
KAFT noCF and noTQA. The KAFT method with no counterfactual augmentations and no TriviaQA slice.
We include more details on the hyper parameters of model finetuning, prompts, post processing, data filtering, and metric computations in the Appendix A.2.

heading: 4 Results
text: In this section we measure the controllability and robustness of KAFT with the metrics defined in Sec. 3.7 and compare with baselines in Sec. 3.8.

heading: 4.1 Larger models may ignore more contexts
text: Most benchmarks improve as a function of model size, including TriviaQA exact match (EM) accuracy, as shown in the first row of Fig. 1. However, we found that larger models may ignore the context more. This may happen for the pretrained model, but the behavior is especially severe for models finetuned on QA tasks using baseline approaches. We demonstrate this effect in the second row of Fig. 1. This highlights a need for designing new methods to improve the controllability of LLMs.

heading: 4.2 KAFT and controllability
text: One of the most striking phenomenon observable from Fig. 1 is that KAFT achieve immense improvements in controllability while maintaining
performance on standard QA. For example, the KAFT PaLM 540B model achieves 24X better controllability compared to the noisy finetuning when the context is in conflict with the model’s pretrained factual knowledge, while performing similarly on regular contexts. In addition, KAFT is the only finetuning approach that consistently achieves better controllability than the pretrained models. Most of this gain originates from the counterfactual augmentation where the model explicitly learns the priority order in Eq. 1 when a conflict does appear. However both relevant only finetuning and KAFT without counterfactual augmentations also exhibit stronger controllability compared to noisy finetuning, even when there is no explicit counterfactual augmentations in both cases. The reason is that both approaches avoid irrelevant contexts that does not imply an answer. Thus the model is less prone to ignore the context compared to noisy finetuning.

heading: 4.3 KAFT and robustness
text: For the pretrained model, the robustness decreased slightly from T5 XL to XXL and from PaLM 8B to 62B (see third row in Fig. 1). But the difference is small. Relevant only finetuning suffers the most loss because it does not have irrelevant contexts
during training. Noisy finetuning only alleviates this loss slightly, still vastly underperforming the pretrained model.
KAFT, on the other hand, significantly boosts robustness. For example, the KAFT PaLM 540B model achieves 6X better robustness compared to noisy finetuning and 1.6X better robustness compared to the pretrained model. Adding the counterfactual augmentation slightly reduces robustness, but the difference is comparably small.

heading: 4.4 Analysis and ablation studies
text: We perform ablation studies to understand the effect of different augmentations in KAFT, as well as the general effect of added context noise.
Effect of KAFT data augmentations. In Fig. 2, we systematically reduce the sampling rate of different data augmentation slices when training KAFT-T5 XXL models. We observe that reducing or removing the counterfactual and irrelevant data augmentations severely reduces controllability and robustness, respectively. In addition, KAFT models significantly out-perform the very strong baselines of Unified QA V2 on both controllability and robustness, showing that KAFT cannot
be replaced by simply adding more supervised data.
KAFT models memorize few counterfactual. One potential risk of adding counterfactual context-answer pairs in the training set is unwanted memorization. We check whether KAFT models memorizes the counterfactual answers in the
training set using the same prompts we used to probe the pretrained model’s closed book answers. We find very little memorization: e.g., the KAFT-PALM 540B model only memorized 0.1% more counterfactuals compared to the pretrained PALM model after KAFT finetuning. Results for other models are similar (cf. Table. 6). The model learns the desirable correlation between the context and the output, rather than memorizing the counterfactual answers.
Context noise reduces controllability. By context noise we refer to the subset of training data where the model is required to produce an answer that is not implied by the provided context, or required to ignore the context while it actually imply the answer. On the flip side, we find that it is possible to achieve good controllability without explicit counterfactual augmentations if we can reduce context noise in the training data.
Table. 5 shows how different amounts of context noise impact the model’s controllability. In particular, because TriviaQA contexts are produced by a retrieval system, it is not guaranteed that a context logically implies the answer. This is even true when the context contains exact matches of the answer. On the other hand, TReX, SQuAD and QASC contains much less context noise given our KAFT construction methods Sec. A.1. Due to this intrinsic noise, including TriviaQA in KAFT caused a negative impact on controllability, especially when there are no explicit counterfactual augmentations. The first row shows noisy finetuning, which contains the most noise. The last row shows that KAFT with TriviaQA data removed. Even though this model is not finetuned on TriviaQA, it has the best controllability. The second row uses a simpler and more noisy filter than KAFT by considering a context to be relevant if it contains the answer.

heading: 5 Conclusion
text: In this work, we analyzed the interaction between LLMs’ parametric knowledge (stored in its model parameters) and knowledge contained in informational contexts provided as a part of the input sequence. We find that models are prone to ignoring the context, especially when the context is in conflict with the parametric knowledge. In addition, the model’s output can be swayed by irrelevant context even when there is no logical link between such context and the model’s task at hand. We quantitatively characterize these behaviours as controllability and robustness of LLMs when one attempts to control their working memory with noisy context. We proposed a new finetuning method, KAFT, that utilizes data augmentations to substantially boost the controllability and robustness of an LLM without significantly affecting its performance on standard QA tasks. With KAFT, we can build LLMs with a clear order of priority when utilizing information from difference sources, including its own parametric knowledge.

heading: 6 Limitations

heading: 6.1 Multiple sources
text: In this work, we trained a model that can utilize two sources of information with predefined priority order, with one of them being the model’s own parametric knowledge. While this is the first step towards LLM’s information utilization with clear, predefined priorities, we acknowledge that real world applications could be more nuanced. For example, KAFT may need to be expanded to treat multiple sources of information with different trustworthiness which may translate to the following desired priority order:
relevant context 1 > relevant context 2 (5)
> model’s parametric knowledge (6)
> relevant context 3 (7)
> all irrelevant context (8)
This orders of priority determines the handling of conflicts. In addition, any irrelevant context should have no influence on the model’s output.

heading: 6.2 Multitask / in-context learning
text: KAFT currently only explores QA tasks. We acknowledge that the applications of LLMs go far beyond a single style of tasks. We have not yet
achieved controlled utilization of information in a task agnostic way. Ideally, the model should learn to prioritize retrieved relevant information in any task that LLMs are capable of, including in-context few-shot or zero-shot scenarios.

heading: 6.3 Dynamically enforce "learning to ignore"
text: In this work, it was necessary to build a different KAFT dataset for each model. Because in Eq. 4, whenever the context is irrelevant, the model fits on to the pretrained model’s answers which depends on the model. This presents additional workload when applying KAFT to new models. In future, it’s worthwhile to explore a dynamic methods that generates closed booked answers during training. At each training step involving irrelevant context, we could run the forward pass twice, one with the provided context and another without. Then we can compute a new loss:
r = 1 : Loss = CE(M ′(c+ q), label) (9) r = 0 : Loss = CE(M ′(c+ q),
stop_gradient(M ′(q))) (10)
where + denotes string concatenation. This is different from Eq. 4 as it fits on to the closed book answers of the current version of the finetuned model, rather than that of the pretrained model. It’s not yet clear whether this would achieve better robustness. It’s also more expensive because two forward passes are necessary for each training example. However it might be justified by the improved simplicity in directly applying KAFT with minimal prepossessing.
This approach is somewhat similar to classifier free guidance (Ho and Salimans, 2022), which has been successfully applied to image generation models. One added benefit of classifier free guidance is the ability to tune the strength of context conditioning after the model is trained, which is another interesting direction to explore here.

heading: 7 Ethics statement: Broader impacts and potential risks
text: In this work, we study approaches to finetune LLMs to make them more grounded and faithful to provided contexts. If our method is applied broadly, it has the potential to correct the unwanted or biased behavior of LLMs with a carefully curated set of natural language instructions without expensive retraining. This provides one feasible avenue
towards improving language models to correct a potential bias that is embedded in the pretraining corpus. At the same time, we acknowledge that our method does not completely address such issues on its own, because 1) instances where the model’s working memory is not controllable by the context even after KAFT is applied may remain; 2) the finetuning dataset used in KAFT may inadvertently introduce or strengthen certain biases. For example, we acknowledge that all KAFT datasets used in this study are English datasets, and so it is a valuable future work direction to extend KAFT to be more representative of all languages.
In addition, we acknowledge that the use of LLMs can be expensive in terms of energy usage. We utilize existing pretrained LLMs such as T5 and PaLM. KAFT’s energy usage is small compared to the pretraining process, but it still leaves a significant energy footprint. In particular, the most expensive training, KAFT-PaLM 540B, takes 12190 TPU v4 hours. It is our hope that methods such as KAFT will provide a way for reducing the need for frequently retraining LLMs, and thus could lead to a more environmentally friendly experimentation.

heading: Acknowledgements
text: We would like to thank Slav Petrov for his insightful comments on an early draft of the paper that significantly helped improve the presentation of our work. We would also like to thank the reviewers and meta-reviewer for their thoughtful feedback on our submission.

heading: A Appendix
text: A.1 Details on relevant context construction SQuAD 2.0 has human labels for this particular aspect. But most datasets do not. For TReX, the question is cloze style where we mask a certain entity within the triples statement. We build a relevant context by concatenating the original statements with a number of sampled irrelevant statements, after randomly shuffling their order. This ensures the relevance of the context while keeping it challenging. The training set of QASC provides 2 gold statements that implies the answer via a two hop reasoning. We are using the 2-stage retrieved collection of statements similar to (Khashabi et al., 2020). We find that the gold statements, or semantically equivalent ones, often exist in the retrieved results. To improve relevance we will randomly add one of the two golden statements and mix it in
the retrieved context to build a relevant context for the KAFT training set. We manually checked on a random small subset that this ensures a relevance ratio around 90%.
TriviaQA is especially challenging because there is no human labeled gold context, while all existing contexts are obtained by a retrieval system. We filter the context by whether they contain the answer. This turned out to be insufficient and leaves a large fraction of irrelevant contexts that do not logically entail the answer. We apply additional filters based on the unigram overlaps of the context with the question, as well as based on the output of a logically entailment model.
A.2 Training Details
We use a learning rate of 0.0002 on all models. The batch size is 32 for all PaLM models and 16 for T5 models. For T5 XL we pick the checkpoint at 100000 finetune steps and for T5 XXL models we pick the checkpoint at 90000 steps. For PaLM 8B and 62B, we pick the checkpoint at 40000 finetuning steps. For PaLM 540B we pick the checkpoint at 15000 steps. These steps are generally determined by avoiding overfitting. However for larger models we are also constrained by compute resources.
A.3 Knowledge Probing Prompts
In this section we provide details on how the knowledge probing prompts in Table 7-9 are constructed. In particular, our goal is to make the model only answer questions where it knows the answer. To do this, we construct prompts that contains two types of QA pairs:
1. Regular QA pairs if the model can answer the specific question correctly in multiple few-shot in-context settings.
2. QA pairs where the answer is "I don’t know" for T5 models or "?" for PaLM models, if the model cannot answer the question correctly in most few-shot in-context settings.
With such specially designed prompts, we encourage the model to abstain if it does not know the answer. The counterfactual context used in the controllability benchmark is constructed using the same method. However we ensure no entities overlaps exist between the prompts that generates the training data vs the test data.
A.4 Counterfactual Generation
To train the model to be controllable by the context, we explicitly engineer plausible training data where the context is in conflict with the model’s pretrained world knowledge. This is done in 3 steps:
1. We apply a diverse set of few-shot prompts similar to Table 10 to condition a pretrained T5 XXL model to generate plausible counterfactual answers.
2. We remove examples if the generation is unsuccessful, when it’s either too long or have a large overlap with the original answer.
3. We replace all occurrences of the original answer with the counterfactual answer in the original context to build the counterfactual context.
With this approach, we build a new QA data set where the answer implied by the context is likely to be in conflict with the model’s existing knowledge.
A.5 Evaluations for Counterfactual memorization and relevance classification
One potential danger of adding counterfactual context-answer pairs in the training set is unwanted memorization. We check whether KAFT models memorizes the counterfactual answers in the training set using the same prompts we used to probe the pretrained model’s closed book answers. The results in Table 6 show that KAFT has little unwanted memorization of counterfactual answers. Instead the model learns the desirable correlation between the context and the output, as demonstrated in Figure 1.
As illustrated in Table 1 and described in Table 3, we require the model to generate its judgements on whether the provided context is relevant. As a sanity check, we evaluated this part of the output on 1000 class-balanced SQuAD2 validation questions, the relevance prediction from KAFT-T5-XXL has 84% precision and 98% recall.
A.6 Postprocessing
After we obtain the output from the pretrained model to the question, which is concatenated after the knowledge probing prompt, we need to postprocess it and removed unwanted components. We do two types of post-processing on the pretrained predictions:
1. Truncation: We truncate the model’s output on special tokens such as < extra_id_1 >, punctuation, line change symbols and question/context initialization symbols such as "Q:", "Question:", "CONTEXT:". These symbols are a frequent in the pretrained model’s responds to our QA style knowledge probe prompts and indicate that the model is ready to move on to the next question that is unrelated to the answer of the current question.
2. Abstain: We normalize all abstain symbols. Whenever the model indicate abstaining using either "I don’t know", "unsure" or "?" in the output as responses to our prompt, we record "unsure" as its answer when constructing the label in the irrelevant slices of KAFT.
A.7 Dataset and task details KAFT mixes together a number of datasets, each with multiple augmentation slices. During training, data from these difference sources are sampled in a round-robin style according to predefined mixture weights. We list these weights as well as the corresponding dataset stats as in Table 11. The sampling ratio from each slice is computed using a product of the normalized dataset level rate and the normalized slice level rate as follows:
R(d, s) = rd∑ d′ rd′ rds∑ s′ rds′
(11)
where d, d′ denote different datasets and s, s′ denote difference slices within each dataset. For ex-
ample, the sampling ratio from the QASC relevant slice is given by:
R(QASC, relevant)
= 0.3
1.3 + 0.3 + 0.1 + 0.2
0.5
0.5 + 0.25 + 0.02
= 0.0831 (12)
The KAFT-TriviaQA training set contains 45593 relevant examples and 72697 irrelevant examples. The KAFT-QASC training set contains 8134 relevant examples and the same number of irrelevant examples. The KAFT-SQuAD2 dataset contains 78125 relevant examples and 117287 irrelevant examples. The KAFT-TReX training set contains 75365 relevant examples and 47503 irrelevant examples.
A.8 Licensing and scientific artifacts
In this work, we used the following scientific artifacts: TriviaQA is licensed under Apache License 2.0. The SQuAD 2.0 dataset is licensed under CC BY-SA 4.0. T-REx is under a Creative Commons Attribution-ShareAlike 4.0 International License. QASC is under CC BY license. T5 models are under Apache License 2.0. Unified QA models are under Apache License 2.0. The PaLM models are proprietary. All these artifacts are properly cited when we mention them the first time. Our use for these artifacts are consistent with their licenses.
We create the following scientific artifacts and we will partly release them after this paper is pub-
lished: The KAFT finetuning method will be released under Apache License 2.0. The KAFT-T5 models will be released under Apache License 2.0. The KAFT-PaLM models will be proprietary.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
section 6
3 A2. Did you discuss any potential risks of your work? section 7
3 A3. Do the abstract and introduction summarize the paper’s main claims? abstract and section 1
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts? Section 1-5
3 B1. Did you cite the creators of artifacts you used? section 1-5
3 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? Appendix A.8
3 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Appendix A.8
7 B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? We use well established public datasets that are constructed based on publicly available information.
3 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? Sec 1-5, Appendix A.8
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Appendix A
C 3 Did you run computational experiments? Sec 3-4
3 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? Sec 1-4, Sec 7, Appendix A
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found hyperparameter values? Appendix A
3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Section 4, Appendix A
3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Jax
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: For example, we acknowledge that all KAFT datasets used in this study are English datasets, and so it is a valuable future work direction to extend KAFT to be more representative of all languages. In addition, we acknowledge that the use of LLMs can be expensive in terms of energy usage. We utilize existing pretrained LLMs such as T5 and PaLM. KAFT’s energy usage is small compared to the pretraining process, but it still leaves a significant energy footprint. In particular, the most expensive training, KAFT-PaLM 540B, takes 12190 TPU v4 hours. It is our hope that methods such as KAFT will provide a way for reducing the need for frequently retraining LLMs, and thus could lead to a more environmentally friendly experimentation