Distributional reasoning in LLMs: Parallel Reasoning Processes in Multi-hop Reasoning

Authors: Yuval Shalev, Amir Feder, Ariel Goldstein
Primary area: interpretability and explainable AI
Keywords: LLM, Multihop reasoning, Interpretability

Abstract:
Large language models (LLMs) have shown an impressive ability to perform tasks believed to require "thought processes”. When the model does not document an explicit thought process, it's difficult to understand the processes occurring within its hidden layers, and to determine if this process can be referred to as reasoning. We introduce a novel and interpretable analysis of internal multi-hop reasoning processes in LLMs. We demonstrate that the prediction process for compositional reasoning questions can be modeled using a simple linear transformation between two semantic category spaces. We show that during inference, the middle layers of the network generate highly interpretable embeddings that represent a set of potential intermediate answers for the multi-hop question. We use statistical analyses to show that a corresponding subset of tokens is activated in the model's output, implying the existence of parallel reasoning paths. These observations hold true even when the model lacks the necessary knowledge to solve the task. Our findings can help uncover the strategies that LLMs use to solve reasoning tasks, offering insights into the types of thought processes that can emerge from artificial intelligence. Finally, we also discuss the implication of cognitive modeling of these results.

Decision: reject
Reviewer scores: 5, 6, 3, 6

Reviews and discussion:
<thread>
creation date:
1730268283621

thread_type:
official_review

<thread_object>
creator: reviewer_p5is

<content>
summary:
The paper investigates the inner processes of LLMs on two-hop reasoning tasks. The authors identify the categories of the intermediate and final answers and study their activation in the LLMs. They find that these categories can respectively be found in the intermediate and final layers and that the connection between the intermediate and final activations can be mapped using a linear transformation. Moreover, they find that LLMs consider multiple hypotheses in parallel, with the middle layers representing the potential intermediate answers and the final layers the potential final answers.

strengths:
The paper is well written and mostly easy to read and understand.

The proposed methodology is an interesting way for studying the challenging problem of LLM interpretability. The resulting findings provide intriguing insights into the reasoning of LLMs in two-hop composition tasks and can be very useful to the research community.

weaknesses:
Although the analysis provided in the paper is quite interesting, the main issue of the paper is its limitation to a single dataset and two-hop questions, which restricts the possibility to extrapolate the findings to larger settings. More specifically:

1. The activation of a word is independent from the rest of the logits, therefore this activation may not be the main reasoning mechanism happening inside the LLM but a byproduct or a minor function only marginally used in the total prediction. The paper does not clearly establish if the activation is the main reason for the final model prediction. It might be interesting to compute relative activations with respect to the other dimensions of the layer representation with (e.g. with softmax). 

2. The example provided in Figure 1 is a composition of two simple queries. The category mapping found could be an artifact coming from the task instead of a property of the model on two-hop reasoning tasks. The evidence provided in the experiments does not seem to eliminate this possibility.

3. Figure 4 does not include legends for the plots, making them hard to interpret. It would be good to include what each data point/plot corresponds to.

4. The experiments focus on two-hop reasoning. To assess the robustness of the findings, it would be interesting to see if they hold to n-hop queries with n higher than two. Queries recalling names could easily be modified into three-hop by asking to return the first letter of the name instead of the full name as done in Figure 1.

questions:
1. Have you conducted experiments to verify if the identified mechanism is the main function responsible for the final model prediction or how much does it contribute to it?

2. How diverse are the questions asked in the experiments? How did you assess that the dataset used contain sufficiently diverse queries for your findings to be robust?

3. Have you conducted experiment on n-hop reasoning tasks with $n>2$ to verify if your findings scale to more complex multi-hop reasoning tasks?

4. What are the average sizes of the semantic categories $c_1$ and $c_2$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730408868750

thread_type:
official_review

<thread_object>
creator: reviewer_ndbn

<content>
summary:
This paper investigates the mechanism of multi-hop reasoning of LLM. Specifically, the paper uses a compositional two-hop question answering setup as a test bed for the investigation. It proposes the hypothesis, referred to as distributional reasoning, that the first attribute extraction creates a distribution of possible attributes on which the second extraction can operate simultaneously to derive the final answer. To test the proposed hypothesis, the paper constructs a linear model to predict the activation of the final layer based on the activations of intermediate layers. The results indicate that once the model depth reaches a certain threshold, the activations of intermediate layers can linearly predict the activations of the final layer. Additionally, the paper demonstrates that the intermediate layers activate a small subset of tokens representing potential intermediate answers.

strengths:
1. The paper proposes a hypothesis to interpret the multi-hop reasoning ability of LLM and designs sound experiments to test the proposed hypothesis. The findings provide interesting insights on neural cognition and reasoning.
2. The paper is clearly written. The experiments are well-designed with main results (Section 4.1 & 4.2) providing direct support to the hypothesis and additional results (Section 4.3) ruling out the confounding factor of the possibility of LLM simply memorizing knowledge.

weaknesses:
1. While I can follow the proposed hypothesis and think the experiments are well-designed for testing them, I am unsure whether we can draw a conclusion that LLMs use distributional reasoning to answer multi-hop questions. As the authors also point out in the Limitations section, the paper only experiments on the two-hop question answering setup and all the questions are synthesized from one data source, the results may be different on questions with very different structures. For example, I am wondering whether the proposed hypothesis still hold on answering math problems that require multiple steps to solve. Or, still within the question answering setup, can the concept of distributional reasoning be identified on answering n-hop questions when n > 2?
2. Another major concern is the implication of the proposed hypothesis. Although the experimental results support the proposed hypothesis, it does not exclude other possibilities - therefore, I personally think the contributions of "computational framework demonstrating the role of associations in structured reasoning" is a bit oversold. A way to strengthen the contribution is to demonstrate the potential usage of the discovered hypothesis/pattern. Can the finding on distributional reasoning help us design better reasoning algorithm (e.g., discover self-inconsistency) or help us train models better at multi-hop reasoning? Even though the paper mentions some relevant aspects in Section 5, it lacks a proof-of-concept to showcase this is a promising way to go.

questions:
See the questions raised in "Weaknesses".
</content>
</thread_object>

</thread>
<thread>
creation date:
1730504822517

thread_type:
official_review

<thread_object>
creator: reviewer_b2q5

<content>
summary:
This paper provides a study of the internal reasoning processes of LLMs when answering compositional questions. 
They find that one can apply linear models on the activations for specific task-related tokens at intermediate transformer layers to approximate the reasoning of the model at the "outcome" final prediction layer. This supports a hypothesis that LLMs consider a distribution of possible answers at intermediate layers then narrow down the final answer at later ones, in a process akin to the parts of the "spread of activation" and propositional reasoning theories of human cognition.

The authors consider a synthetic dataset (the Compositional Celebrities dataset from Press et al 2023, as well as carefully crafted variants to test for hallucinatory scenarios) and provide strong evidence for their hypothesis. The authors argue that transformer-based LLMs can be used to test questions in cognitive modeling.

strengths:
* The paper is well-written and well-motivated. 
* The study design is clear and effective, and the results are interesting and well-supported. The secondary experiments (4.2 on the interpretability of the intermediate layers, and 4.3 on whether this type of reasoning still holds when the answer is hallucinatory) do a good job of preemptively answering follow-up questions I had on the original study. 
* The finding will be of interest to people concerned with whether current LLMs can shed light on questions in cognitive modeling of reasoning processes.

weaknesses:
**W1** The primary weakness of the paper is that it only considers one clean, synthetic dataset with one type of question. It is unclear whether such a clean result would hold on a more realistic reasoning dataset. The authors acknowledge this in their limitations section, but it is still the case that this paper shows __one__ scenario in which a clean form of reasoning can be approximated using linear modeling for __one__ kind of problem. The paper would be stronger with at least one other dataset that is more complex than 2-hop compositional questions and/or noisier than the templatic questions in Compositional Celebrities. Otherwise, this analysis should really be considered exploratory rather than comprehensive.

**W2** l501 "This bridges a gap that was observed by cognitive sciences decades ago and emphasis the role of AI
research in cognitive modeling." -- it would be helpful to describe more specifically how the result contributes the debate surrounding this gap (i.e., why they are incongruous) more specifically in the context of existing work. The paper seems to argue that both associative and propositional reasoning can take place at the same time in a single framework -- whether this is possible has been the subject of debate for many years, and various other cognitive models have been proposed that integrate both paradigms to differing degrees, e.g. [ACT-R](https://wires.onlinelibrary.wiley.com/doi/full/10.1002/wcs.1488?casa_token=9nMiql_9dRMAAAAA%3A1kKURVyse66RR4wdR3wxa1DDWQI_HO0D-oUAPfHepYuJzEhYj_0Jo77bV5qSxGHQq9O4Sz3VPpYa) and [CLARION](https://en.wikipedia.org/wiki/CLARION_(cognitive_architecture)). Comparing your consideration of transformers to these other frameworks and discussing how the contributions differ would benefit the paper.

questions:
**Q1** Does this analysis generalize to any sort of task that isn't 2-step compositional reasoning? What classes of task/problem does this framework make sense for, and what kind would it be unrealistic?

**Q2** This analysis suggests that the suggest-then-narrow-down reasoning approach can be __induced__ from the model's intermediates, but that is not the same thing as the model necessarily __using__ that approach. How can we have confidence that the model is actually using this reasoning chain and not rely upon other forms of reasoning instead/ in addition?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730708783074

thread_type:
official_review

<thread_object>
creator: reviewer_f2fn

<content>
summary:
This paper introduces a method to analyze multi-hop reasoning in LLMs by modeling their prediction process as a linear transformation between semantic spaces. The authors find that intermediate model layers create interpretable embeddings representing potential answers, activating specific token subsets and suggesting parallel reasoning paths. This framework offers insights into how LLMs mimic thought processes, even without direct task knowledge, with implications for understanding AI reasoning and cognitive modeling.

strengths:
- This paper addresses a critical challenge in understanding LLM reasoning by probing the multi-hop reasoning processes. The discussion on human cognitive processes and LLM reasoning mechanisms is also inspiring.

- The study presents a clear illustration of distributional reasoning in LLMs (esp. with the example in Figure 2), which is a novel finding to my knowledge.

weaknesses:
- While the illustration of distributional reasoning is interesting, the remaining analysis seems confusing to me. For example, what's the motivation of the analysis that uses the matrix to approximate the second-hop reasoning? The paper mentioned "This method reveals the extent to which the second-hop operation is invariant to the prompt specifics." why is this an important question to explore? and what's the takeaway from this analysis?

- The writing needs to be improved to make the logic clearer. Especially, there should be strong and clear motivation before going into any notations and technical details of the analysis. The current version makes readers very easy to get lost.

questions:
See the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732047395219

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We sincerely appreciate the comprehensive feedback provided by the reviewers. We are diligently working to enhance our paper in accordance with your comments, and we aim to address each point thoroughly and effectively.

Thanks to the reviewers' insightful comments, we've expanded the paper with three new sections in the appendix. For the camera-ready submission, these sections will be integrated into the main body of the paper. Please refer to sections C.4, C.5, and D in the latest version for these additions.

In response to questions about our dataset selection, we would like to provide further explanation for our choices:

In this work, we are trying to uncover a general mechanism of LLMs implicit reasoning. This mechanism contains two main characteristics. One, LLMs activate parallel reasoning paths, and second, the process works as a pipeline of two sequential stages in which the second relies only on the output of the first (’y’ will be the first letter of ‘yellow’ no matter if the question was about ‘banana’ or ‘lemon’). When trying to demonstrate this mechanism, there are several constraints that we need to keep in mind:

First, LLMs are not very good at solving multi-step questions implicitly. This is the reason why the Chain-of-Thought method has become such a popular approach. This fact limits us to keeping the structure of the prompts relatively simple, because otherwise they just won't be able to reach the correct answer, and no reasoning process could be traced. We hope our paper will promote directions for addressing this problem in future models.

Second, as in other scientific disciplines, our proposed approach for demonstrating the proposed mechanism requires its own limitation of data curation. This doesn't mean that our results don't apply to other structures and relations, but other relations will need a different approach to examine them. Same as brain mechanisms discovered in fMRI experiments might not apply exactly to any other situation outside of the lab, the experimental design constrains its setting, yet it is still highly valuable.

Now in more details about our choice of dataset:

Our framework describes the reasoning process as a transformation between two semantic categories. In order to conduct the analyses we proposed, we need a dataset of two-step questions with unique characteristics: each question should contain sub-questions which relate to two distinctive semantic groups (e.g., colors, letters, cities, etc.). In addition, the questions should be built in a way that the final answer shouldn't be related to the question subject directly ('y' is not related directly to 'banana' without the context of its color). Lastly, the questions should be constructed in a way that the answer could be generated in a single token.

The Compositional Celebrities dataset is ideal for demonstrating our claims because the alternative intermediate and final answers for its questions come from relatively distinct categories. The dataset's semantic categories (countries, digits, domains, etc.) are well defined, making it easy to identify possible alternatives answers in the model’s logits. The dataset has 6547 prompts in 14 relation types. Each type has 468 prompts, allowing us to effectively fit the linear models.

These data properties are rare in other datasets, often lacking diverse relation types or clear semantic categories. For instance, the dataset in [1] has only 1000 prompts split into 19 relation types, but with only 3 having over 100 prompts. All relations are attributes of a person (e.g., "The place of birth of the father of John Penn"). "Father of", for example, isn't a well-defined category for our purposes, as it will be hard to obtain a vector representing all fathers. It doesn't mean the model doesn't use distributional reasoning. It may still use the activation of alternative options for John Penn's father to create the output result.

To expand our dataset and empirical results, we curated a synthetic dataset that allows us to track, in a controlled setup, the effects of intermediate activations on the output tokens. The hallucination experiments, which demonstrate our hypothesis in a much harder setting, use a novel type of dataset to assess internal semantic processes in a neutral manner.

Lastly, following the comments from this rebuttal, we added to the appendix results for 1593 new prompts, which are not questions regarding birthplaces of celebrities. These results will be integrated into the main paper by the camera-ready version. See C.4 section in the last revision.

[1] Memory Injections. https://aclanthology.org/2023.blackboxnlp-1.26/
</content>
</thread_object>

</thread>
<thread>
creation date:
1733239354631

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We appreciate the detailed feedback provided by the reviewers. Your comments greatly improved our work.

We would like to highlight several novel contributions of this work, which we will clarify in the camera-ready version:

1. An analysis of multi-hop reasoning processes reveals the presence of parallel reasoning paths. This insight of multiple paths suggest a more precise understanding of the internal workings of LLMs.

2. A linear model fitted to each type of second-hop operation is used to predict unseen data, revealing the extent to which the second-hop operation is invariant. This indicates that the linear transformation performing the reasoning process remains the same across different prompts. For instance, a linear model fitted on “the first letter of the name of the color of a banana is” will predict the outcome of the same reasoning operation on an apple's color. This separates reasoning from the prompt specifics, which is crucial since questions about color names' first letters aren't inherently related to the fruits.

3. Our "hallucinations experiments" provide a new method for evaluating internal processes in LLMs. These experiments highlight the difference between expecting accurate answers from LLMs and expecting a valid reasoning process, which can sometimes produce incorrect answers. Moreover, the synthetic dataset we used, consisting of fictitious items, demonstrates a novel way of creating datasets for assessing the inner workings of LLMs.

4. Our analysis indicates that this parallel process is interpretable. Projecting the embedding to the vocabulary space, it is as if the model "thinks out loud" allowing us to trace its simultaneous reasoning processes. This finding is significant for interpretability and can help trace associative errors that lead to incorrect predictions.

5. This work examines the internal thinking processes of LLMs from a cognitive science perspective, raising fundamental questions about the nature of intelligence and cognition. Our observational findings reveal mechanisms that can be viewed as a mix of two main and competing approaches to cognitive modeling. This highlights the significance of exploring LLMs' inner workings for cognitive literature.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper investigates the internal multi-hop reasoning processes in large language models (LLMs). The authors propose that LLMs utilize a "distributional reasoning mechanism", whereby intermediate layers generate embeddings representing potential intermediate answers, which are then refined in subsequent layers. The study uses a synthetic dataset to demonstrate that these processes can be modeled using linear transformations between semantic spaces, suggesting the presence of parallel reasoning paths.

Strengths:
* Addresses a core problem in LLM research: understanding the nature of "reasoning" in LLMs. 
* A novel methodological approach to tackle this problem: an interpretable analysis of internal multi-hop reasoning processes.
* A solid, clear, well-motivated experimental design.

Weaknesses:
* The paper relies on a single synthetic dataset and only two-hop questions, which raises questions about the generalizability of the findings to more complex or natural datasets. In its current form, the findings are overly narrow.
* The paper might suffer from confirmation bias / non-falsifiable hypothesis testing. The experimental results are consistent with but stop short of proving the main hypothesis (distributional reasoning). As noted by reviewer nbdn, corroborating evidence could be provided by experiments that rule out other alternative hypotheses or that show that certain interventions motivated by the hypothesis result in changes consistent with it. But it its current form, saying that the distributional reasoning is "demonstrated" in this paper does feel overblown.

Overall, I think this is a good but borderline paper, that provides findings that are likely to be of interest to the community but whose validity and robustness is still unclear.