Memory Injections: Correcting Multi-Hop Reasoning Failures During Inference in Transformer-Based Language Models (2023)

Abstract: Answering multi-hop reasoning questions requires retrieving and synthesizing information from diverse sources. Large Language Models (LLMs) struggle to perform such reasoning consistently. Here we propose an approach to pinpoint and rectify multi-hop reasoning failures through targeted memory injections on LLM attention heads. First, we analyze the per-layer activations of GPT-2 models in response to single and multi-hop prompts. We then propose a mechanism that allows users to inject pertinent prompt-specific information, which we refer to as “memories,” at critical LLM locations during inference. By thus enabling the LLM to incorporate additional relevant information during inference, we enhance the quality of multi-hop prompt completions. We show empirically that a simple, efficient, and targeted memory injection into a key attention layer can often increase the probability of the desired next token in multi-hop tasks, by up to 424%.

## introduction transformer-based large language models (llms) @xcite have shown exceptional promise for basic knowledge retrieval and language generation; however, they often lack the ability to perform basic reasoning tasks @xcite @xcite . in this work, we focus on the simple task of answering multi-hop prompts (i.e., prompts in which the subject is not stated explicitly), which humans handle easily but with which llms often struggle (see fig. 1 ). researchers have attempted to rectify multihop reasoning failures by using various prompting methods such as chain-of-thought (cot), treeof-thought (tot), and graph-of-thought (got) reasoning @xcite @xcite @xcite the largest coral reef in the world is located off the coast of the philippines llm (a) multi-hop prompt. ## background & notation we define single-vs. multi-hop prompts and provide a formal definition of the transformer model. ## experimental overview our central aim is to better understand how the outputs of the attention heads affect model performance with respect to predicting the correct next token in prompts requiring single-hop reasoning versus in prompts requiring multi-hop reasoning. ## proposed methods recent work suggests that attention heads are knowledge retrievers during a model's inference pass @xcite . extending this result to multi-hop prompts, we hypothesize that attention layers play an important role in retrieving memories relevant to the "hop" in a given prompt. therefore we define two algorithms below: one for analyzing attention head outputs in embedding space and the other for injecting a targeted memory into a model's hidden activations in order to correct faulty/incomplete reasoning. ## results and discussion we report, in turn, on our curated memory, random memory, and part-of-speech injection experiments. ## related work much recent work has focused on the inner workings of transformers @xcite @xcite . @xcite explore how the emergent properties of llms form during training. recent interpretability research has focused on the mechanisms by which linear layers in llms retrieve information, characterizing them as keyvalue stores of information @xcite and showing that tokens can be characterized by their distribution in the output vocabulary @xcite . others have also examined the intermediate activations of llms in order to uncover underlying reasoning mechanisms. nostalgebraist (2021) applied gpt-2's unembedding matrix to intermediate layers to interpret how the model arrives at its final answer. @xcite employed a learned transformation to mitigate the effect of any bias introduced by using the unembedding matrix. there has been much recent interest in whether llms are reliable stores of information for attempting to both identify where knowledge exists and how to edit stored factual knowledge effectively @xcite @xcite . recent approaches to knowledge editing make use of learned hyper-models to edit weights, additional trained parameters, or direct interventions on model weights @xcite @xcite . however, these approaches raise another issue: dealing with knowledge retention and preventing catastrophic forgetting @xcite @xcite . additionally, it is not clear that the mechanisms by which model predictions are constructed is fully understood, limiting our ability to improve model performance @xcite . some approaches propose to use external knowledge stores such as knowledge graphs to augment the factual capabilities of llms @xcite @xcite . ## conclusions and future directions we demonstrate that a key reason llms perform worse on multi-hop prompts is because they fail to recall intermediary information that is relevant to a hop. we find that attention heads play an important role in this factual recall process, and that in the case of multi-hop reasoning, certain attention layers fail to recall relevant information. to rectify this shortcoming, we establish an algorithm for injecting "memories" directly into the model's hidden activations during inference. through experimentation, we find that injecting relevant memories into the hidden activations of the attention heads during inference is an efficient way to boost model performance on multi-hop prompts. we anticipate that our memory injection scheme can extend a model's longevity by enabling less frequent retraining/fine-tuning. we also hope in future work to demonstrate the use of memory injections to correct stale or incorrect information, remove private or harmful information, and combat bias during llm inference. there is also a tremendous opportunity to scale online-memory injections to enhance the quality of thousands/millions of model inferences, if we can automate the process of memory selection via unsupervised algorithms, for instance by connecting llms with knowledge bases.