Understanding and Patching Compositional Reasoning in LLM s (2024)

Abstract: LLMs have marked a revolutonary shift, yet they falter when faced with compositional reasoning tasks. Our research embarks on a quest to uncover the root causes of compositional reasoning failures of LLMs, uncovering that most of them stem from the improperly generated or leveraged implicit reasoning results. Inspired by our empirical findings, we resort to Logit Lens and an intervention experiment to dissect the inner hidden states of LLMs. This deep dive reveals that implicit reasoning results indeed surface within middle layers and play a causative role in shaping the final explicit reasoning results. Our exploration further locates multi-head self-attention (MHSA) modules within these layers, which emerge as the linchpins in accurate generation and leveraing of implicit reasoning results. Grounded on the above findings, we develop CREME, a lightweight method to patch errors in compositional reasoning via editing the located MHSA modules. Our empirical evidence stands testament to CREME’s effectiveness, paving the way for autonomously and continuously enhancing compositional reasoning capabilities in language models.

## introduction compositional reasoning stands as a pivotal mechanism, unlocking the ability of learning systems to decompose complex tasks into manageable subtasks and tackle them step-by-step @xcite . despite the revolutionary impact of large language models (llms) on the nlp landscape, they struggle at basic compositional reasoning tasks @xcite . this shortcoming is specifically highlighted by @xcite , who brought attention to the concerning "compositionality gap" in the realm of question-answering tasks. it was observed that there is a substantial failure rate of ∼ 40% in figure 1 : an example of a typical compositional reasoning error pattern :short-cut(section 3). before patching, llms take short-cut reasoning by directly binding "the nationality" and "c. auguste dupin" (a fictional french detective) to incorrectly predict "france". after patching, llms tend to firstly bind "the creator of" and "c. auguste dupin" to generate "edgar allen poe" (implicit reasoning result) and then correctly predict "u.s.a." (explicit reasoning result). this work, therefore, sets out to firstly delve into the specific failures to understand (rq1) what accounts for these compositional reasoning failures and (rq2) which parts of the llms are responsible for them, and secondly develop strategies for patching these failures. our initial step involves an analysis of a very recent dataset comprising compositional two-hop knowledge queries @xcite , selectively examining the cases where llms fail despite successfully answering the constituent single-hop queries. to ensure our findings and methodologies offer broad applicability, our analyses utilize two widely-used open-sourced llms: openalpaca-3b @xcite and llama-2-7b @xcite . through meticulous examination of the failure instances, we identify three prevalent types of errors. utilizing the logit lens tool (nostalgebraist, 2020), each error type highlights a critical shortfall in generating or leveraging the implicit reasoning result necessary for the explicit reasoning result foot_0 . this gap is particularly concerning as it contrasts sharply with the intuitive two-hop reasoning process inherent to human cognition. an illustrative example of compositional reasoning error is depicted in figure 1 , where the model incorrectly concludes its reasoning without properly generating and incorporating the implicit reasoning result. the above observations motivate our further empirical inquiry to answer the first question of what accounts for these failures, from the perspective of whether llms are indeed aware of implicit reasoning results during compositional reasoning. we inspect inner hidden states of llms via logit lens, from which we observe that implicit reasoning results not only manifest within the llms' intermediate layers but also tend to precede the generation of explicit reasoning results, often emerging statistically earlier. building on this, we further explore the relationship between implicit and explicit reasoning results through an intervention @xcite experiment, providing compelling evidence that the emergence of implicit reasoning results within llms plays a causative role in the generation of explicit reasoning results. the next question is, regarding rq2, in which modules llms generate implicit reasoning results? leveraging causal mediation analysis @xcite , we present both a compositional query and its corresponding second-hop query to the llm, resulting in the generation of two distinct computation graphs. we then intervene the computation graph g 1 , associated with the compositional query, by replacing the output of a single module with its counterpart from the second-hop computation graph g 2 . by identifying the modules whose replacement results in a significant enhancement in the predictive probability of the explicit reasoning result, we are able to locate several specific outputs from the multi-head self-attention (mhsa). intriguingly, the layers pinpointed through this approach show a strong correlation with those identified in preceding intervention experiments. this congruence reinforces the hypothesis that implicit reasoning results are not only present but are actively consolidated and utilized within these specific layers of the llm. grounded on our findings into rq1 and rq2, we develop creme (correcting compositional reasoning via model editing), a light-weight model-editing method to patch errors in compositional reasoning. creme follows @xcite ; @xcite by regarding the output matrix of the located mhsa, w l o , as a linear associative memory. to implement creme, we designate the input to w l o in the computation graph g 1 as k * and the output from w l o in g 2 as v * . we then proceed to insert the pair (k * , v * ) into w l o , ensuring that this insertion disrupts existing memories within w l o as minimally as possible. this objective is achieved by solving a convex optimization problem, which strikes a nuanced balance between the integration of new corrective information and the preservation of existing knowledge. our main contributions and takeaways are summarized below: (1) successful compositional reasoning within llms hinges on its awareness of generating and leveraging implicit reasoning results; ## analyzing compositional reasoning errors grounded on the observation of @xcite , we dive into the compositional reasoning failures: we identify three types of common errors among such failures and attribute the cause of these common errors to the failure of generating implicit reasoning result properly via inspecting hidden states. three types of common errors we query llms with all of compositional queries and the corresponding single-hop queries in d. we filter out two subsets of d: d single and d gap . for each datum (s 1 , r 1 , o 1 ) ⊕ (s 2 , r 2 , o 2 ) in d, d single contains the datum where the both of (s 1 , r 1 , o 1 ) and (s 2 , r 2 , o 2 ) are successfully answered. among d single , d gap contains the datum where the answer for the compositional queries (s 1 , r 1 , o 1 ) ⊕ (s 2 , r 2 , o 2 ) are mis-predicted. 3 in our analysis of d gap , we have discerned a few common patterns shared among a substantial portion of the failures. consequently, we have delineated three predominant types of errors, each characterized by distinct features, as outlined below. distortion: llms fail to effectively generate implicit reasoning results in the reasoning process. the predicted answer for the first example in in the second example of ## analyzing the inner hidden states of llms for compositional reasoning providing that llms are capable to perform compositional step-by-step reasoning @xcite , we hypothesize that they generate the implicit reasoning result o 1 (the notation is aligned with section 2.2) in the process of compositional reasoning, before finally obtaining the explicit reasoning result o 2 . we inspect inner hidden states of llms via logit lens (section 4.1) and observe that implicit reasoning results emerge in middle layers, implying that they may play a role in the compositional reasoning process (section 4.1). to verify this hypothesis, we design an intervention experiment (section 4.2) and demonstrate the emerging of o 1 has causal effect on predicting o 2 in the output layer (section 4.2). ## locating important modules in previous analysis, we attribute compositional reasoning errors to improperly generating implicit reasoning results. in this section, we aim to investigate if there sparsely exist some "key" modules (i.e., mhsa or mlp) 5 in llms that are responsible for properly generating implicit reasoning results in hidden states of middle layers. ## patching compositional reasoning grounded on the empirical insights in section 4 and section 5, we are poised to introduce the creme approach, designed to correct compositional reasoning failures via editing the parameters of mhsa at the located positions. we demonstrate its superiority through comparative analyses with two recent baselines for correcting compositional reasoning @xcite and a a widely recognized model editing baseline @xcite . specifically, our edit objective is the mhsa output matrix at the l-th layer w l o (for detailed description, please refer to eqn. 7). following santurkar et al. (2021), we view w l o as a linear associative memory (kohonen, 1972): w l @xmath0 for a given compositional query and its corresponding second-hop query, we run the llm twice: once with the compositional query and once with the second-hop query. in the first pass with the compositional query, the input of w l o at the last token position is k * ∈ r d×1 ; in the second pass with the corresponding second-hop query, the output of w l o at the last token position is v * ∈ r d×1 . we aim to edit w l o to ŵ l o such that: what is worthy noting is that applying creme only requires a single case of (compositional query, referenced second-hop query) while the patching effect can generalize to many other related cases. we showcase the effect of creme in ## conclusion in this paper we study and patch the compositional reasoning of llms. through examining failure instances and conducting diverse analysis experiments, we demonstrate successful compositional reasoning within llms hinges on its awareness of generating and leveraging implicit reasoning results. moreover, we locate few important mhsa modules in llms that are responsible for properly generating and leveraging implicit reasoning results via causal mediation analysis. to this end, we propose creme, to compositional reasoning failures via editing the located mhsa parameters and empirically demonstrate its superiority.