PaperlayerView exact revision on arXiv

Source-linked AI summary

Large Language Models with Controllable Working Memory

Daliang Li, Ankit Singh Rawat, Manzil Zaheer, Xin Wang, Michal Lukasik, Andreas Veit, Felix Yu, Sanjiv Kumar

arXiv:2211.05110v1cs.CLcs.AIcs.LG

TL;DR

The paper asks how LLMs should reconcile pretrained world knowledge with relevant or irrelevant input context. It jointly studies controllability and robustness, proposes KAFT with counterfactual and irrelevant-context training, and finds that KAFT improves both properties across evaluated models. A remaining limitation is that the method requires a separate dataset for each model and may require additional forward passes in a dynamic alternative.

  • Problem

    LLMs may ignore relevant contexts that conflict with pretrained knowledge or be swayed by irrelevant contexts, complicating context-based knowledge updates.

  • Method

    KAFT augments supervised training with counterfactual contexts and trains on pretrained closed-book answers when contexts are irrelevant.

  • Results

    KAFT substantially boosts controllability and robustness without significantly affecting performance on regular metrics.

  • Takeaways & Limitations

    KAFT can produce LLMs with a clear priority order for using contextual information, including their pretrained world knowledge.

  • Takeaways & Limitations

    The method requires a different KAFT dataset for each model, while a dynamic alternative would require two forward passes per training example and its robustness benefit is uncertain.

Abstract

from arXiv · show

Large language models (LLMs) have led to a series of breakthroughs in natural language processing (NLP), owing to their excellent understanding and generation abilities. Remarkably, what further sets these models apart is the massive amounts of world knowledge they internalize during pretraining. While many downstream applications provide the model with an informational context to aid its performance on the 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 task-relevant information that conflicts with the model's memorized knowledge. This enables model predictions to be grounded in the context, which can then be used to update or correct specific model predictions without frequent retraining. 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-the-art T5 and PaLM (both pretrained and finetuned) could exhibit poor controllability and robustness, which do not scale with increasing model size. As a solution, we propose a novel method - Knowledge Aware FineTuning (KAFT) - to strengthen both controllability and robustness by incorporating counterfactual and irrelevant contexts to standard supervised datasets. Our comprehensive evaluation showcases the utility of KAFT across model architectures and sizes.

1 Introduction

The paper studies whether LLMs can prioritize relevant contextual information over conflicting pretrained knowledge while ignoring irrelevant context. It introduces KAFT, a finetuning approach designed to improve both controllability and robustness.

  • Motivation: LLM working memory can be controlled through context without changing model weights, supporting pointwise knowledge updates and corrections.The paper frames this as controlling activations rather than editing long-term knowledge stored in weights.
  • Working-memory behavior: The desired priority order is relevant context over pretrained knowledge, while irrelevant context should be ignored in favor of pretrained knowledge.This distinction defines controllability for conflicting relevant contexts and robustness against noisy contexts.
  • Evaluation: The benchmark jointly evaluates controllability and robustness using contradictory contexts and irrelevant contexts across model architectures and sizes.For QA, relevance is defined by whether the context entails an answer to the question.
  • Findings: Pretrained and QA-finetuned LLMs often ignore contexts that conflict with world knowledge, with the problem becoming more severe as models grow larger.The paper also reports that noise in QA finetuning data contributes to this behavior.
  • KAFT: KAFT strengthens controllability with counterfactual context augmentations and robustness by training on pretrained closed-book answers for irrelevant contexts.The method swaps context answer entities with plausible alternatives for counterfactual examples and uses model answers when contexts are irrelevant.
  • Results: KAFT models demonstrate coexistence of controllability and robustness against conflicting and irrelevant contexts.Table 1 illustrates improved behavior relative to baseline methods.

2 Related Works

Related work examines how language models store factual knowledge, how that knowledge can be updated, and how models use contextual versus parametric information. Prior studies also document susceptibility to irrelevant or untrustworthy contexts.

  • World knowledge in language models: Prior probing studies show that BERT and T5 models memorize factual world knowledge in their parameters.These findings motivate treating pretrained models as repositories of factual information.
  • Knowledge update in language models: Because factual knowledge changes and can be outdated, retraining language models from scratch is prohibitively expensive.Related work therefore investigates how factual knowledge is stored and how specific parameterized knowledge can be updated efficiently.
  • Contextual and parametric knowledge: Previous work uses retrieved context to improve downstream tasks while LLMs also retain substantial knowledge in their parameters.The paper studies the interaction between these contextual and parametric knowledge sources.
  • Contextual and parametric knowledge: Earlier studies found that language models can be misled by certain irrelevant contexts, while this work evaluates the phenomenon in QA and applies KAFT for robustness.Related work also considers settings where context may be untrustworthy and should be ignored.
  • Contextual and parametric knowledge: An extension is proposed in which multiple information sources receive an ordered priority relative to model knowledge and irrelevant contexts.The cited formulation places source1 and source2 above model knowledge, followed by source3 and irrelevant contexts.
  • Related investigations: The authors note an independent investigation sharing important aspects of their work.This establishes overlap with contemporaneous research.

3 Methods

The method defines how a finetuned model should prioritize relevant context over pretrained knowledge while ignoring irrelevant context, then constructs KAFT data and evaluation settings to enforce these behaviors.

  • Method objective: KAFT targets outputs grounded in relevant context and unchanged by irrelevant context, with the model determining relevance from context-question semantics.Relevant-context labels use ground-truth or counterfactual answers, whereas irrelevant-context labels use the pretrained model’s closed-book answer.
  • Datasets: KAFT combines four QA datasets spanning multiple formats and includes relevant, counterfactual, irrelevant, and empty-context examples with corresponding labels.The datasets cover multiple-choice, cloze, extractive, and open-domain QA; irrelevant or empty contexts receive pretrained few-shot closed-book answers.
  • Models: The evaluation uses T5 encoder-decoder models and PaLM decoder-only models, including PaLM sizes from 8B to 540B and T5 XL and XXL.Smaller T5 models are excluded because they respond poorly to in-context few-shot prompts needed to measure pretrained world knowledge.
  • Context construction: Context relevance is defined by logical entailment of the answer, distinguishing hard irrelevant contexts from same-topic contexts that do not entail it.The authors note that topical or entity overlap alone does not establish relevance and can cause models to ignore relevant contexts when trained against ground-truth labels.
  • KAFT construction: Counterfactual training data swaps original answer entities for plausible alternatives, creating contexts whose implied answers conflict with pretrained knowledge.A pretrained T5 XXL generates plausible counterfactual answers, which replace the original answers throughout the context.
  • Metrics: The study measures controllability by testing whether relevant counterfactual contexts produce their implied answers, using thresholded unigram recall for evaluation.The metric is applied to head questions that all five pretrained models answer correctly in closed-book few-shot settings, with positivity requiring over 80% answer-unigram recall.

4 Results

Experiments show that larger LLMs increasingly ignore conflicting context, while standard finetuning can reduce robustness to irrelevant context. KAFT improves controllability and robustness through targeted augmentations while preserving regular-task performance.

  • Larger models are more likely to ignore contexts: Larger models increasingly ignore contexts containing counterfactual answer entities, across pretrained and QA-finetuned models.This trend appears despite improving language understanding and entity knowledge with scale.
  • KAFT and Controllability: 24X better controllability is achieved by KAFT PaLM 540B than noisy finetuning on conflicting contexts, with similar performance on regular contexts.KAFT is also the only finetuning approach reported to consistently exceed pretrained models in controllability.
  • KAFT and Controllability: Most controllability gains originate from counterfactual augmentation, which teaches the model the intended priority order during conflicts.Relevant-only finetuning and KAFT without counterfactual augmentation also outperform noisy finetuning by reducing context-noise cases.
  • KAFT and Robustness: Robustness shows no clear improvement with model size and decreases slightly from T5 XL to XXL and from PaLM 8B to 62B.Standard finetuning severely reduces robustness, while relevant-only finetuning suffers the largest loss.
  • KAFT and Robustness: 6X better robustness than noisy finetuning and 1.6X better robustness than pretrained models are achieved by KAFT PaLM 540B.Adding counterfactual augmentation slightly reduces robustness, but the difference is comparably small.
  • Analysis and Ablation studies: Reducing or removing counterfactual and irrelevant augmentations severely harms controllability and robustness, respectively.KAFT also outperforms Unified QA V2 on both metrics, indicating that simply adding supervised data does not replace the method.
  • Analysis and Ablation studies: KAFT shows little unwanted memorization of counterfactual answers and instead learns the correlation between context and output.The evaluation probes closed-book answers on TriviaQA training questions with counterfactual augmentations.
  • Analysis and Ablation studies: Context noise reduces controllability, although good controllability can be achieved without explicit counterfactual augmentation when training noise is reduced.TriviaQA retrieval contexts can fail to logically imply answers, even when they contain exact answer matches.

5 Conclusion

The paper finds that LLMs can ignore conflicting context and be swayed by irrelevant context. KAFT substantially improves both behaviors without significantly affecting regular metrics, producing a clearer priority among information sources.

  • Conclusion: LLMs tend to ignore context that conflicts with internal knowledge and can be swayed by irrelevant context lacking a logical task link.These behaviors are characterized as controllability and robustness problems when noisy context is used to control working memory.
  • Conclusion: Figure 2 evaluates how controllability, robustness, and standard metrics change as KAFT augmentation sampling ratios vary.The ablations include reduced irrelevant augmentation, removed counterfactual augmentation, Unified QA V2, and noisy finetuning baselines.
  • Conclusion: KAFT shows little unwanted memorization of counterfactual answers in TriviaQA training questions with counterfactual augmentations.Table 6 measures the match rate between closed-book answers and counterfactual answers.
  • Conclusion: KAFT uses data augmentations to substantially boost controllability and robustness without significantly affecting regular performance metrics.The method aims to establish a clear priority order among pretrained knowledge and contextual information.
  • Conclusion: Context noise leads to lower controllability, with noise increasing from the first to the last row of Table 7.The table compares finetuning methods on head counterfactual questions.

6 Future work

The paper outlines priority-based handling of multiple information sources and proposes a dynamic training alternative, while identifying computational and effectiveness trade-offs.

  • 6.1 Multiple Sources: The current work trains models to use two information sources with a predefined priority order, and future work could expand this to multiple sources of varying quality or trustworthiness.One of the two sources is the model’s own parametric knowledge.
  • 6.1 Multiple Sources: The priority order places relevant contexts above the model’s parametric knowledge, then lower-priority contexts and all irrelevant context.This ordering determines how conflicts are handled, while irrelevant context should not influence the output.
  • 6.2 Dynamically enforce "learning to ignore": The proposed dynamic method would run the model twice for irrelevant-context examples and compute a loss against the current finetuned model’s closed-book answer.One pass uses the provided context and another omits it; stop-gradient is applied to the context-free prediction.
  • 6.2 Dynamically enforce "learning to ignore": The dynamic approach may improve simplicity but its robustness is unresolved and its two-forward-pass training procedure is more expensive.It would apply KAFT directly to datasets with ground-truth labels and context-relevance labels with minimal preprocessing.
  • 6.2 Dynamically enforce "learning to ignore": Classifier-free guidance is presented as a related approach whose tunable context-conditioning strength suggests a future direction.The cited benefit is adjusting context-conditioning strength after training.

A.1 Training Details

Training uses a shared learning rate, model-family-specific batch sizes, and checkpoints selected mainly to avoid overfitting under compute constraints.

  • A.1 Training Details: The learning rate is 0.0002 for all models.
  • A.1 Training Details: Batch size is 32 for PaLM models and 16 for T5 models.
  • A.1 Training Details: Checkpoints are selected at different finetuning steps, generally to avoid overfitting, with larger models additionally constrained by compute resources.The listed checkpoints are 100000 steps for T5 XL, 90000 for T5 XXL, 40000 for PaLM 8B and 62B, and 15000 for PaLM 540B.

A.2 Knowledge Probing Prompts

Knowledge-probing prompts are designed to identify questions the model can answer reliably and to mark questions for which it does not know the answer.

  • A.2 Knowledge Probing Prompts: The prompts contain regular QA pairs when the model answers a question correctly across multiple few-shot settings.
  • A.2 Knowledge Probing Prompts: They contain abstention-labeled QA pairs when the model cannot answer the question.T5 uses “I don’t know,” while PaLM uses “?”.
  • A.2 Knowledge Probing Prompts: The construction goal is to make the model answer only questions whose answers it knows.

A.3 Postprocessing

Pretrained-model predictions are postprocessed by truncating unwanted continuation markers and normalizing abstention responses before constructing irrelevant-context labels.

  • A.3 Postprocessing: The predictions are postprocessed in two ways: truncation and abstention normalization.
  • A.3 Postprocessing: Outputs are truncated at special tokens, punctuation, line changes, and question or context initialization symbols indicating a transition to another question.
  • A.3 Postprocessing: “I don’t know,” “unsure,” and “?” are normalized to “unsure” when constructing labels for irrelevant KAFT slices.

A.4 Mixture weights

KAFT combines multiple datasets and augmentation slices, sampling their training examples round-robin according to predefined mixture weights. Slice-level sampling rates are computed from dataset- and slice-level rates, with task slices distinguished by context relevance and answer type.

  • KAFT mixes multiple datasets, each containing multiple augmentation slices, during training.The mixture includes task slices such as relevant, irrelevant, empty, and counterfactual contexts.
  • Training data from the different sources is sampled round-robin according to predefined mixture weights.Table 11 lists the mixture weights and corresponding dataset statistics.
  • Knowledge-probing prompts are used with standard QA and Cloze-style QA datasets to probe pretrained-model answers.The listed prompts cover datasets including SQuAD 2.0, TriviaQA, and TReX.
  • Slice-level sampling ratios combine normalized dataset-level and slice-level rates.The ratio is defined for datasets d, d′ and slices s, s′, with an example given for the QASC relevant slice.
Loading 2211.05110v1…