Hopping Too Late: Exploring the Limitations of Large Language Models on Multi-Hop Queries (2024)

Abstract: Large language models (LLMs) can solve complex multi-step problems, but little is known about how these computations are implemented internally. Motivated by this, we study how LLMs answer multi-hop queries such as “The spouse of the performer of Imagine is”. These queries require two information extraction steps: a latent one for resolving the first hop (“the performer of Imagine”) into the bridge entity (John Lennon), and another for resolving the second hop (“the spouse of John Lennon”) into the target entity (Yoko Ono). Understanding how the latent step is computed internally is key to understanding the overall computation. By carefully analyzing the internal computations of transformer-based LLMs, we discover that the bridge entity is resolved in the early layers of the model. Then, only after this resolution, the two-hop query is solved in the later layers. Because the second hop commences in later layers, there could be cases where these layers no longer encode the necessary knowledge for correctly predicting the answer. Motivated by this, we propose a novel “back-patching” analysis method whereby a hidden representation from a later layer is patched back to an earlier layer. We find that in up to 66% of previously incorrect cases there exists a back-patch that results in the correct generation of the answer, showing that the later layers indeed sometimes lack the needed functionality. Overall our methods and findings open further opportunities for understanding and improving latent reasoning in transformer-based LLMs.

## introduction despite groundbreaking performance in a multitude of tasks, large language models (llms) still struggle with complex knowledge queries @xcite . for example, llms often err when asked to complete multi-hop queries such as "the spouse of the performer of imagine is". answering such queries requires composition the spouse of the performer of imagine is yoko 1) first hop is resolved into bridge entity 2) information propagates to last token 3) second hop is resolved into target entity representation encoding bridge entity "john lennon" representation encoding target entity "yoko ono" back-patched representation backpatching hidden representation ⋮ ⋮ ⋮ ⋮ figure 1 : an illustration of our findings: we observe evidence of latent reasoning in two-hop queries where 1) during the early layers the first hop is resolved and the source entity imagine now encodes the bridge entity john lennon 2) during the middle layers critical information propagates to the last position 3) during the later layers the second hop is resolved and the last token now encodes the target entity yoko ono. we additionally illustrate back-patching: patching a hidden representation from a later layer back into an earlier layer in order to fix cases where the pathway fails. and reasoning skills, which have been the focus of many recent works @xcite @xcite @xcite . besides the general interest in compositional skills, the investigation of latent multi-hop abilities of llms would also have significant implications for areas such as generalization @xcite and model editing @xcite . to strengthen the latent reasoning ability of llms, it is essential to first understand the internal mechanism of how llms successfully complete two-hop queries using latent multi-hop reasoning. while there have been several works that investigate such a mechanism, a precise latent reasoning pathway has only been found in relatively small language models trained on synthetic datasets @xcite @xcite and has not been thoroughly investigated in large pretrained models. throughout this work we consider and compare two complementing settings, one where the model correctly answers the two-hop query and the other when it does not. to this end, in §2, we start by creating and publishing a dataset containing 82,020 two-hop queries based on data from wikidata @xcite . our approach to analyzing multi-hop reasoning begins with the hypothesis that it involves activating the same "knowledge-extraction module" twice: once for the first hop ("the performer of imagine is john lennon") and once for the second ("the spouse of john lennon is yoko ono"). the key question we ask in this work is where are these two procedures implemented in the llm. a natural way to answer this question is to seek the first point in the network where the outputs of these two steps appear. the typical approach to this would employ vocabulary projections @xcite @xcite @xcite . however, we instead opt for a more recent method named patchscopes @xcite which has clear advantages over the commonly used vocabulary projection. using this approach, in §3 we find that the outcome of the first hop is clearly evident in the hidden representation of the end-token of the first hop. we further find that resolving the first hop happens during the early layers, as depicted in figure 1 (stage 1). given the localization of the first hop, in §4 we then ask where the second hop is resolved. we find that it typically appears in the last token of the prompt only in layers after the resolution of the first hop, as further illustrated in figure 1 (stage 3 ). these observations suggest a mechanism where the first hop is resolved with the answer, that then propagates to the last token to resolve the second hop (as depicted by stage 2 in figure 1 ). we explore this propagation in §5. moreover, they suggest why this process may fail -if the two resolutions are done by different groups of layers, it must be that these two groups are both able to perform knowledge extraction, be it via their attention or mlp sublayers. however, since transformers have a limited amount of layers, and their functionality is different @xcite , it is quite likely that the layer at which the second hop is performed will not have the desired functional-ity. to verify this hypothesis, in §6, we propose an analysis method named back-patching. the crux of the method involves taking a hidden representation from a later layer, patching it back into the same position of the same prompt in an earlier layer, and then continuing the forward pass. in a way, this allows the model more layers to finish the computation without the need for additional parameters or training @xcite . unlike previous methods @xcite , backpatching makes no assumption on the structure of the prompt and does not require any knowledge of the participating entities. testing back-patching on queries initially predicted incorrectly, we find that up to 66% of incorrect cases have the potential to be correctly predicted using back-patching when choosing the optimal source and target layers. overall, our contributions can be summarized as follows: we release our code and dataset at https:// github.com/edenbiran/hoppingtoolate . ## localizing first hop resolution one could imagine several possible ways for a model to solve multi-hop queries. one such way would be to effectively treat the two-hop query as a single relation, enabling the extraction of the final answer without the need to recall e 2 . despite this, if the model truly performs latent reasoning, the most intuitive way to do so would start by resolving the first hop at some point within the computation. in order to localize where this happens, we use the patchscopes analysis method. our results show that the model indeed resolves the first hop during the early layers, at the position of t 1 . ## second hop is resolved at last position in §3, we observe that the first hop is often resolved in position t 1 . intuitively, one would expect the computation to use this result towards resolving the second hop (i.e., the final answer). in this section, we study this proposition by finding where e 3 first emerges and how it relates to the first emergence of e 2 . we find that e 3 indeed emerges predominantly after e 2 at the position of t 2 . furthermore, we search for the final part of the pathway, the predicted token. by projecting the residual updates made by the attention and mlp sublayers on the hidden representation of t 2 , we find that the mlp sublayers play the larger part in conjuring the predicted token. ## information propagates to last token as the model's prediction is emitted from the last position of the prompt, critical information regarding the bridge entity must propagate to this token. by using attention knockout, projection, and patchscopes, we verify this claim. we find that in a large majority of cases relevant information can indeed be detected passing from t 1 to t 2 . ## back-patching improves two-hop performance our results consistently display a sequential nature of computation, where the first hop is resolved into the bridge entity by the lower layers, which is then used to answer the second hop in the upper layers. these results suggest that when answering two-hop queries fail, it could be because the firsthop is resolved at layers that no longer contain the information needed to resolve the second hop. in order to verify this hypothesis, we test the following intriguing conjecture. we know that the earlier layers certainly contain knowledge, since they are regularly used to resolve the first hop. then, we ask what will happen if we take a hidden representation from a later layer and "feed" it into these earlier layers. will the correct information now be extracted? as our results next show, this is often the case. this "back-patching" analysis thus provides strong support that the success in two hop queries require the later layers to have the same functionality as earlier ones. overall, this could point to an inherent limitation of the transformer architecture, where the model "runs out" of layers, propagates and extracts information before it is fully resolved, or is forced to use only the upper layers for knowledge extraction. ## related work tracing and locating factual knowledge has been the focus of many recent works aiming to undercover and control the way transformers make predictions. several of these works are motivated by knowledge editing @xcite @xcite while others focus on understanding the inner workings of transformers @xcite @xcite . our work follows and expands on this second group by extending them to the more complex case of two-hop queries, which require multiple factual knowledge extractions and compositions. the question whether models can implicitly reason over stored parametric knowledge has also gained recent popularity. one avenue of related research includes circuit discovery @xcite @xcite , but these works focus mainly on synthetic tasks. more closely related to this paper, recent works @xcite @xcite have found existential evidence of latent reasoning in multi-hop queries, albeit without accounting for the flow of information throughout layers and positions. in a concurrent work by @xcite a multi-hop reasoning circuit similar to ours was discovered in a small transformer language model trained on synthetic data until grokking occurs. our work finds further proof of such latent reasoning in a more practical setting that includes large language models, real world knowledge and a regular pretraining setup. we additionally use different mechanistic interpretability methods and account for the flow of information during the forward pass. even without fully identifying the mechanism behind latent reasoning, several works have tried to address model's shortcomings in this area. one such work by @xcite assumes prior knowledge of the bridge entity, whose hidden representation is then strategically injected directly into the computation. likewise, recent work by @xcite also assumes knowledge of the bridge entity and then employs knowledge editing techniques to improve performance, which is known to have potentially disruptive effects on models @xcite @xcite . both these previously proposed methods have drawbacks, which a potential method based on back-patching would not share. ## conclusion our work uses mechanistic tools to study the latent reasoning abilities of llms on two hop queries. we observe strong evidence of a sequential latent reasoning pathway in which the second hop is answered in the mid-upper layers only after the first hop is answered by the lower layers. we find that this sequential nature could point to a possible limitation of transformers, where the second hop must be answered using only the knowledge encoded in the upper layers. to further validate this, we introduce and evaluate an analysis method named back-patching, based on extracting hidden representations from later layers and patching them back into earlier layers. finally, we show that backpatching significantly improves performance on previously incorrect queries. overall, our methods and findings open opportunities for understanding and improving latent reasoning in llms.