The Buffer Mechanism for Multi-Step Information Reasoning in Language Models

Thank you for your constructive comments on our paper. Below are our detailed responses to each of your points:

[R4W1] [Insufficient Explanations of Equations and Figures]

We have thoroughly revised the captions of every figure and table in the paper, adding detailed explanations to the flowcharts and supplementing the experimental figures with conclusions. Additionally, we have rewritten some paragraphs that might have been ambiguous to enhance readability. Due to the extensive nature of these changes, please refer to our revised manuscript for specific modifications.

[R4Q1] [75% Time Saving]

We apologize again for the lack of detailed explanation. As shown in Figure 8(b), after approximately one epoch, the accuracy of the baseline model rises above 95%, while the RMBA0.05 model reaches nearly the same accuracy after only about 0.25 epochs. Therefore, it saves about 75% of the training time. We have added the corresponding explanation in the new manuscript.

[R4Q2] [RMBA Stability]

According to our current experimental results, RMBA does not appear to have significant stability issues. Firstly, in Figure 8(c), we show the results after adding perturbations to different models; the experiments indicate that the RMBA algorithm has stability comparable to the baseline. We have also added the training loss curves of the PrOntoQA task in the appendix of the new manuscript, which also support our viewpoint.

From a theoretical perspective, RMBA can be understood as a special form of residual connection (just multiplied by a constant matrix), and residual connections are known to promote stability. We speculate that this may be one reason why RMBA does not lead to decreased stability.

[R4Q3] [Recommendations]

We will comprehensively revise the figure captions according to your suggestions.

[Other Modifications]

In addition, we have incorporated suggestions from other reviewers, adding visualization experiments of the matching score in the Phi-3 model, experiments using a 2-layer Transformer combined with CoT to achieve multi-step reasoning, and conducted detailed parameter sweeps for the RMBA experiments. We have also rewritten content in the paper that might cause ambiguity, enhancing the rigor of the manuscript. Specific changes can be found in the common response. Your comments, along with those of other reviewers, have greatly helped improve the quality of our paper.

We look forward to your reply!

Thank you for your constructive comments on our paper. Below are our responses to the issues you raised:

[R3W1] [Overclaim] We thank you and the other reviewers for highlighting the overclaim issues. In the revised manuscript, we have addressed these concerns by adding the qualifier "symbolic multi-step reasoning tasks" to the relevant statements.

[R3Q1] [Empirical Evidence]

1. In the revised manuscript, we add experiments in the appendix G on how to train a 2-layer model to achieve symbolic multi-step reasoning tasks by performing CoT. In this experiment, we trained a Transformer using only single-step reasoning data. During the testing phase, we allowed the model to use CoT for reasoning by re-inputting the output information from each step back into the model. This enabled us to achieve 2-step and 3-step reasoning. In the appendix, we present the actual information flow diagrams rather than the schematic diagrams shown in the main text. We hope this experiment can supports our viewpoint.

2. The core idea of this section is that, unlike vertical thinking where the model needs to form multiple buffers, when using CoT, the model repeatedly overwrites the same buffer to achieve multi-step reasoning. Therefore, during the training phase, only a few weight matrices are needed to form the reasoning circuit, which significantly reduces the training difficulty of the model. We have added these points to the main text.

[Other Revisions] In addition, we have incorporated suggestions from other reviewers, adding visualizations of the matching scores in the Phi-3 model and conducting a detailed parameter sweep for the RMBA experiments. We have also enriched the explanations of some figures and formulas in the paper, and rewritten content that might cause ambiguity, enhancing the rigor of the manuscript. The specific changes can be found in the common response. Your comments, along with those of other reviewers, have greatly helped improve the quality of our paper.

We look forward to your reply!

Thank you for your constructive comments on our paper. Below are our detailed responses to each of your points:

[R1W1.1] [Synthetic Tasks]

1. [Why start from synthetic tasks?] Due to the complexity of linguistic data and large language models, it is challenging to identify key mechanisms through direct research on real language data, limiting us to relatively macro-level statistical analyses. We chose to conduct research on simplified symbolic datasets, which allows us to perform sufficient numerical experiments with minimal computational resources and to closely observe the internal workings of the Transformer model. The buffer mechanism we proposed aids in understanding the expressive capabilities of Transformer models.

2. Applicability to real large language models. In the revised manuscript, we have added experiments calculating the matching score and kernel score in the Phi-3 model. We provided methods for computing the matching score and kernel score in multi-head models. By calculating the kernel score, we assessed the strength of inter-layer information interactions in Phi-3 and computed the matching score for each head in each layer. We found:

   • (a) When $l_1 \leq l_2$, $W^{qk(l_1)} W^{vo(l_2),T}$ exhibits almost no significant features (kernel score ≈ 0).
   • (b) When $l_1 > l_2$, the kernel score of $W^{qk(l_1)} W^{vo(l_2),T}$ decreases as $l_1 - l_2$ increases. These observations indicate that information stored in earlier layers (by buffer $W^{vo(l_2)}$) is extracted by later layers (via $W^{qk(l_1)}$), and the model tends to extract the most recently processed information.
   • (c) Heads with high matching scores are mostly concentrated in the middle layers, which aligns with previous findings that reasoning in large models predominantly occurs in the middle layers.

3. Algorithmic insights from synthetic tasks. The buffer mechanism we discovered based on synthetic tasks can inspire algorithm design. For example, as shown in Figure 8, the RMBA algorithm proposed—motivated by the buffer mechanism—can also aid in training real large language models.

[R1W1.2] [Prior Work]

According to our literature review, most of the current work in the field on multi-hop reasoning are evaluations of existing large models or use causal intervention starting from results to identify key pathways. In the Related Work section of the revised manuscript, we have added references to several recent articles related to multi-hop reasoning and causal intervention and discussed the novelty of our work. If the reviewer could provide more relevant references, we can more specifically discuss the connections and differences between our work and these studies.

[R1W2] [Causal Intervention]

1. Complementarity with our approach. First, there is no contradiction between these two studies, as causal intervention is one method to find key information transmission paths, while the buffer mechanism attempts to explain how this information transmission occurs. Moreover, causal intervention can only identify key paths but cannot fundamentally understand how these paths are generated, offering limited guidance for algorithm design. Our paper aims to study what kinds of model weights cause these paths to emerge, providing a more detailed modeling of the implementation mechanisms of causal paths. Accordingly, the proposed RMBA can also promote the generation of reasoning logic circuits.

2. Progress through synthetic tasks and aligned weight matrices. Some existing works have made significant progress via "synthetic tasks" and "weight alignment." For example, Boix-Adsera E, et al. [1] proposed the IMBA by studying single-step reasoning, which has been proven to aid in learning simple reasoning data. However, from the buffer perspective, we found that this algorithm causes the buffer to degenerate, making it difficult to learn multi-step reasoning. Therefore, we propose the RMBA to further improve the model's accuracy in multi-step reasoning tasks.

   [1] Boix-Adsera E, et al. When can transformers reason with abstract symbols

3. In the field of biological neural networks, performing causal intervention is common because the methods for recognizing and processing biological signals are limited. However, in artificial neural networks, we have richer means to recognize and process signals, allowing us to obtain deeper insights than those in biological neural network research.

Thanks for providing the additional experiments and clarifications. I have increased my score to 5 to reflect my greater confidence in the soundness of the experiments.

However, after the authors' clarification that they do not in fact have solid evidence for the buffer mechanism in Phi-3, I'm slightly concerned about the applicability of the results.

Specifically, it's possible that the buffer mechanism is a product of the particular training set up you used, and they may not emerge in pretrained language models. One reason why this might be true is that pretrained LMs tend to find directions in activation space that correspond to semantic concepts. This view is fundamentally at odds with the buffer mechanism, which says instead that there are multiple replicas of a semantic space.

Are there more concrete ways of determining if the buffer mechanism is used in Phi-3? This is the crux for me that will raise my score to a 6 or an 8.

A minor point: Sec. J is very sparse on details. In particular, Fig. 24 should probably report quantitative metrics for various predictions, and also some sense of the statistical significance. For example, you might report how frequently are the interventions working as your claim, and how many data points did you test with.

Thanks for adding the experiments with the cosine similarity on $W^{ov}$ matrices for Phi-3.

I still believe that this line of weight alignment argument is insufficient for establishing that Phi-3 exhibits the buffer mechanism. The weight alignment phenomena you're highlighting in Sec. H seems extremely general because it's a property of the Phi-3 weights, and not a property of any particular symbolic reasoning task you're studying for the main paper. This makes the hypothesis space much larger, because there could be many training dynamics considerations that could lead the formation of this weight alignment pattern.

If you can propose a concrete task on which Phi-3 exhibits the buffer mechanism, it will greatly improve the relevance and impact of the paper. Such evidence can come in the form of analyzing the representational role of activations, or analyzing behaviors of attention heads in this concrete task. I believe that if this is rigorously shown in an LLM, this paper can be really important.

Overall, while I appreciate the authors' efforts at improving the experimental rigor, I still believe that the the paper's contributions are limited because the authors made the choice to focus their analysis on synthetically trained models on toy synthetic datasets, rather than an LLM. I will therefore choose not to increase my score.

Thank you for the constructive feedback on our manuscript. Below are our responses to each of your comments:

[R2W1] [Overclaim] We appreciate the reviewer pointing out the overclaim issues in our writing. We have revised all instances of overclaim identified in our manuscript. Specifically, we have added the qualifier "symbolic multi-step reasoning tasks" to each relevant claim.

[R2W2] [Circular Logic]
The ambiguity in this section was due to our inadequate phrasing; in fact, there is no circular reasoning involved. We have revised the relevant section in the new manuscript as follows:
To achieve the same-token matching, it is sufficient for $\text{Ker}^{(l)} \approx I$, in which case
$h^{(l)}(X_{tgt}) \approx X_{tgt}X_{tgt}^T = I + O(1/\sqrt{d_m}).$ This equation indicates that the attention of all tokens is focused on themselves. Furthermore, we observe that same-token matching is independent of the specific value of $X_{{tgt}}$. For example, for $X_{{tgt, OOD}}$ sampled from untrained random vectors $token_{OOD}$, $h^{(l)}(X_{{tgt, OOD}}) \approx I$ still holds. Therefore, when the model weights satisfy $Ker^{(l)} \approx I$, the model demonstrates out-of-distribution generalization capability.

[R2W3] [Minor Things]
We conducted more detailed interactions with various large models. Specifically, we used four GPT and Claude models to perform 2-, 3-, and 4-step reasoning tasks, examining their performance with and without Chain-of-Thought (CoT) prompting. Each setup was repeated 25 times, totaling 600 interactions, to ensure statistically robust results. The findings show that current large models struggle to perform 2(or higher)-step reasoning tasks directly. However, with CoT prompting, all these models excel at multi-step reasoning. Based on these results, we have revised Figure 1 and included detailed interaction results in the Appendix. Additionally, we have corrected the typographical error in Figure 2.

[R2Q1] [Reasonable]
There is an intuitive explanation for the diagonal structure observed between the $W^{vo}$ matrix of the previous layer and the $W^{qk}$ matrix of the next layer: information derived in one layer is directly utilized by the next layer, aligning with logical reasoning.

1. Diagonal Structure in Real Models:
   In the revised manuscript, we have supplemented experiments using the Phi-3 model to compute matching scores and kernel scores, and defined methods for calculating these metrics in multi-head models. Key observations include:

   • (a) When $l_1 \leq l_2$, $W^{qk(l_1)} W^{vo(l_2),T}$ exhibits almost no significant features (kernel score ≈ 0).
   • (b) When $l_1 > l_2$, the kernel score of $W^{qk(l_1)} W^{vo(l_2),T}$ decreases as $l_1 - l_2$ increases. These observations indicate that information stored in earlier layers (by buffer $W^{vo(l_2)}$) is extracted by later layers (via $W^{qk(l_1)}$), and the model tends to extract the most recently processed information.
   • (c) Heads with high matching scores are mostly concentrated in the middle layers, which aligns with previous findings that reasoning in large models predominantly occurs in the middle layers.

2. Improving Reasoning:
   Both the RMBA algorithm proposed in our paper and the IMBA algorithm by Boix-Adsera E, et al.[1] improve the model's reasoning capability. This supports the conclusion that the diagonal structure of $W^{qk(l_1)}W^{vo(l_2),T}$ facilitates information transfer.
   Reference: [1] Boix-Adsera E, et al., When can transformers reason with abstract symbols

3. Future Directions:
   The real-world data have different characteristics (e.g., memory-centric vs. reasoning-centric). In the future, the MOE (Mixture of Experts) architectures could potentially be designed to incorporate diagonal structures in certain layers and heads, improving both reasoning and memory capabilities.

[R2Q2] [Accuracy Suddenly Improved]
This is a good question. Understanding the dynamics of how models learn multi-step reasoning is a key direction of our future work. Currently, we lack a rigorous explanation but have observed that this phenomenon is consistent, regardless of changes to learning rate, weight decay, $d_m$, and $d_k$. When the number of epochs is fewer than 50, the model has yet to establish a stable information flow (even on the training data) and appears to rely more on memorization. The cause of this phenomenon remains unclear.

[R2Q3] [OOD Embedding]
In fact, Transformers do not need to learn embeddings for OOD tokens. Based on our mechanism analysis and random matrix theory, as long as $Ker = I$, it follows that $h(X_{{OOD}}) \approx X_{{OOD}}X_{{OOD}}^T \approx I$, ensuring same-token matching. This is one of the most intriguing aspects of the paper, as the OOD generalization capability effectively characterizes whether the model has truly “understood” the rules.

[R2Q4] [Parallel Reasoning]
The parallel reasoning described in Appendix also utilizes the buffer mechanism for information transfer. Here, the buffer mechanism is represented as {$\{\prod_{l \in J} \mathbf{W}^{vo(l)} | J \subset \{0, 1, \ldots, L-1\}\}$}. Figure 12(c) shows that the model still stores different information in different buffers. This section complements the vertical reasoning section and applies to cases where the reasoning steps $n \geq L$ (model depth). It does not conflict with the main text, which focuses on $n < L$. Generally, the depth of a large language model greatly exceeds the reasoning steps required for single predictions, especially with the prevalent use of CoT prompting in modern large models.

[Additional Revisions]
In addition, we have incorporated feedback from other reviewers by adding experimental studies on CoT and conducting a thorough parameter sweep for the RMBA experiments. We also enriched the explanations of some figures and formulas and clarified ambiguous content to improve the rigor of the paper. Detailed revisions can be found in the common response. Your comments, along with those of other reviewers, have greatly helped improve the quality of our paper.

We look forward to your reply!