F ast M em: Fast Memorization of Prompt Improves Context Awareness of Large Language Models (2024)

Abstract: Large language models (LLMs) excel in generating coherent text, but they often struggle with context awareness, leading to inaccuracies in tasks requiring faithful adherence to provided information. We introduce FastMem, a novel method designed to enhance instruction fine-tuned LLMs’ context awareness through fast memorization of the prompt. FastMem maximizes the likelihood of the prompt before inference by updating only the last Feed-Forward Network (FFN) module. This targeted approach ensures efficient optimization without overfitting, significantly improving the model’s ability to comprehend and accurately follow the context. Our experiments demonstrate substantial gains in reading comprehension, text summarization and adherence to output structures. For instance, FastMem improves the accuracy of Llama 3-8B-Inst on the NQ-SWAP dataset from 59.1% to 71.6%, and reduces the output structure failure rate of Qwen 1.5-4B-Chat from 34.9% to 25.5%. Extensive experimental results highlight FastMem’s potential to offer a robust solution to enhance the reliability and accuracy of LLMs in various applications. Our code is available at: https://github.com/IAAR-Shanghai/FastMem.

## introduction large language models (llms) function within a versatile input space defined by prompts. these prompts set the task context, provide background information for answering questions, and offer examples for in-context learning @xcite @xcite @xcite . despite their remarkable utility, llms have difficulty in fully understanding and following the provided context, which we refer as the ability of context awareness. in this paper, we consider both the instruction in the prompt, such as query or predefined output format, and the reference provided by the user or retrieval system as the context. a major issue with context awareness is maintaining faithfulness to the contextual information. for instance, in summarization tasks, llms may fabricate entities or relationships that are not present in the original reference @xcite @xcite . this issue also reduces the effectiveness of retrieval-augmented generation (rag) @xcite , as the llms may falsify the information in references found by rag when answering the question. another aspect is adhering to a predefined output structure, which is necessary for many applications @xcite . llms undergo pre-training and fine-tuning. the pre-training stage involves learning on a large corpus of text to parameterize extensive knowledge into the network @xcite . after pre-training, various abilities of the model are further elicited by finetuning on more specific tasks or via reinforcement learning with human feedback @xcite . although fine-tuning improves llms context awareness, previous work has shown that llms may not always faithfully follow the information provided in the context when answering questions @xcite @xcite . this issue becomes particularly pronounced when there is a conflict between the ingrained parametric knowledge and the contextual information provided during inference. for instance, such conflicts can happen when the pre-training dataset is outdated or lacks recent information. additionally, the diversity of pre-training corpus can result in a mass-seeking effect, where the model assigns nonzero probabilities to many tokens @xcite . mass-seeking behavior of llms can exacerbate the challenge of context awareness. these findings imply that llms perform worse when they have high perplexity regarding the context. we also verify that models perform worse in extractive question answering when the perplexity of the answer span in the context is high (see sec. 3.2). inspired by these insights, we propose a novel approach named fastmem to improve context awareness via fast memorization of prompts, as depicted in fig. 1 . specifically, we adapt an instruction fine-tuned model by maximizing the prompt likelihood before generation. to ensure the efficiency of the optimization and prevent overfitting to small text snippets, we update only the last ffn module, achieving significant improvements with minimal computational cost. our contributions are summarized as follows: ## related work fine tuning. instruction tuning typically enhances a model's ability to understand and respond to contextual information @xcite @xcite . many studies have explored tuning language models on dedicated datasets to further enhance their context awareness and improve aspects such as faithfulness and factual accuracy @xcite @xcite . unlike these methods, fastmem optimizes on individual data examples and does not require curating a specific dataset. additionally, fastmem is applied upon instruction fine-tuned models to further enhance their context awareness. decoding strategy. many decoding strategies have been proposed to improve context awareness. one mainstream approach leverages the evaluative capabilities of llms to enhance the quality and factuality of their generated text @xcite @xcite . another approach applies the concept of contrast to strengthen context awareness @xcite @xcite . ours fastmem and these decoding strategies are both inference-time methods, while fast-mem is compatible with existing decoding strategies. in this work, we also investigate incorporating two decoding strategies to further enhance the effectiveness of our method. knowledge editing. knowledge editing methods update and correct factual information within language models by modifying the model's parametric knowledge @xcite @xcite . both knowledge editing and fastmem involve training on individual data examples, but knowledge editing requires input-output pairs (e.g., question and answer), while fastmem uses only the input, making it more widely applicable. additionally, knowledge editing is not as efficient as fastmem to be applied at inference time. test-time training. test-time training (ttt) was initially proposed for vision tasks to adapt models for out-of-distribution data using psuedolabels @xcite . in contrast, fastmem can improve model performance when the train and test data distributions are inconsistent or consistent. our experiments reveal that a negative correlation between the performance and perplexity with respect to the context exits in llms, as discussed in sec. 3.2. this negative correlation potentially leads to a more general applicability of fastmem. recently, @xcite studied ttt for language models. however, they only consider pretrained models, in which raw text for ttt can be directly fed to the network. if instruction fine-tuned models follow the same method without modification, they risk being corrupted based on our experiments. in fastmem we address this problem with dedicated control tokens. furthermore, we explore the reduction of latency and achieve real-time application of our method. thus, fastmem distinguishes itself from existing ttt methods. ## method in this section, we elaborate on our method, fast-mem. in sec. 3.1, we clarify the notations. in sec. 3.2, we present the motivation behind our method. in sec. 3.3, we explain how we develop our method to be robust and efficient. in sec. 3.4, we explore combining fastmem with existing decoding strategies. finally, in sec. 3.5, we demonstrate the use cases. an illustration of our approach and its variants is provided in fig. 2 . ## experiments in this section, we evaluate fastmem on q&a and summarization tasks that require llms to respond based on the given context. ## conclusion in this work, we introduced a novel approach, fast-mem, to enhance context awareness in llms. by optimizing an instruction fine-tuned model to memorize the prompt, we effectively improve its performance across a variety of tasks. furthermore, fast-mem's efficiency in optimizing only the last ffn module ensures minimal computational cost while achieving substantial gains. overall, our method provides a robust solution to the challenges of context awareness, paving the way for more reliable and contextually faithful llm applications.