Understanding the Language Model to Solve the Symbolic Multi-Step Reasoning Problem from the Perspective of Buffer Mechanism (2025)

Abstract: Large language models have consistently struggled with complex reasoning tasks, such as mathematical problem-solving. Investigating the internal reasoning mechanisms of these models can help us design better model architectures and training strategies, ultimately enhancing their reasoning capability. In this study, we constructed a symbolic multi-step reasoning task to investigate the information propagation mechanisms in Transformer models when solving the task through direct answering and Chain-of-Thought (CoT) reasoning. We introduced the concept of buffer mechanism: the model stores various information in distinct buffers and selectively extracts it through the query-key matrix. We proposed a random matrix-based algorithm to enhance the model’s reasoning ability. This algorithm introduces only 132 trainable parameters, yet leads to significant performance improvements on 7 multi-step reasoning datasets, including PrOntoQA, LogicAsker, and LogicInference. These findings provide new insights into understanding the large language models.

## introduction in recent years, llms have emerged and demonstrated remarkable capability across various tasks @xcite @xcite @xcite @xcite @xcite . these models have shown impressive in-context learning abilities @xcite @xcite and have been applied to logical reasoning problems, such as matching top human contestants at the international mathematical olympiad (imo) level @xcite and solving math problems @xcite . however, even the most advanced models still struggle with complex reasoning tasks. to truly enhance the reasoning capability of llms, it is crucial to investigate their intrinsic mechanisms. multi-step reasoning tasks encompass a broad concept, typically referring to the ability of a model to synthesize numerous complex conditions to answer questions. here, we consider a representative class of syntactic structures within multi-step reasoning tasks: a sentence that includes both the question and sufficient known information to answer it. for example, "given: [a]→ @xcite ... @xcite ..., question: 2-step reasoning starting from [a]", where "..." represents other textual content unrelated to logical reasoning. the answer to this question is "[c]". when a sentence contains only one logical reasoning step, it is often handled by the so-called induction head in transformer @xcite . however, multistep reasoning is not merely a linear accumulation of multiple induction heads but involves more complex mechanisms. large models employ two primary strategies for logical reasoning. the first, known as the vertical thinking strategy (vts), outputs reasoning results in a single forward based on their inherent structure. this approach is efficient but demands a high level of intelligence. as shown in fig. 1 , current large models exhibit significant limitations in their vertical thinking capability. another relatively less efficient approach is the (hts), such as chain of thought (cot) @xcite and tree of thought (tot) @xcite , diagram of thought (dot) @xcite . this strategy substantially enhances the model's reasoning performance. all models can produce the correct answers for tasks depicted figure 1 : the interaction results of multi-step reasoning tasks with large models. we tested each model 25 times, and when the number of reasoning steps exceeded three, these models exhibited random guessing. however, when we allowed the models to use cot prompting (by removing "directly" and "only return the answer"), all models achieved 100% accuracy. detailed interaction results are provided in the appendix i. in fig. 1 with the help of cot. considering the strengths and weaknesses of these two strategies, the ideal approach should combine both: decomposing problems into several coarse-grained subproblems (hts) and applying the vts to each subproblem. thus, researching how to improve vertical thinking capability and understanding why the horizontal thinking strategy can significantly enhance model reasoning ability are both crucial. in this work, we investigate the performance of transformer models on a symbolic multi-step reasoning dataset. our work aims to uncover these mechanisms and provide insights into how transformers process and integrate logical information across multiple layers to perform multi-step reasoning, which can help develop more effective strategies for improving their multi-step reasoning abilities. specifically, we found that transformers utilize a buffer mechanism when engaging in symbolic multi-step reasoning tasks. the model stores different intermediate results in separate buffers, allowing for quick retrieval as needed. we elaborate on how the model leverages this buffer mechanism for vertical thinking, and we explain why the horizontal thinking strategy can significantly enhance the model's multi-step reasoning capability from the perspective of the buffer mechanism. finally, based on this understanding, we propose a method to improve the model's reasoning abilities, leading to significant performance improvements for the gpt-2 model on 7 multi-step reasoning datasets, including clutrr @xcite , ruletaker @xcite , stepgame @xcite , logicinference @xcite , logicasker @xcite , pararuleplus @xcite , and prontoqa @xcite . the concept of "buffer" or similar concepts has also been mentioned in other works @xcite @xcite . our work provides a detailed description of the concept of buffer and applies this mechanism to enhance model performance. the main contributions of this work are as follows: ## reasoning dataset and transformer model dataset. to understand the mechanism of multistep reasoning in transformers, we design an abstract symbolic multi-step reasoning task. as shown in fig. 2 , reasoning chains are serialized into a sequence. every two tokens in the sentence represent a reasoning relation. the last token is the reasoning start token, and the label is the result with a fixed-step reasoning starting from the starting point. transformer model. we employ a decoder-only transformer. given an input sequence x @xmath0, where n is the sequence length and d is the dictionary size, the model first applies an embedding layer to obtain the input representation x (1) = x tgt + x @xmath1, where x tgt and x pos denote the token and positional embeddings, respectively, and d m is the hidden dimension. the single-head attention in each layer is computed as follows: ## vertical and horizontal thinking strategy we illustrate how the transformer model employs vertical or horizontal thinking strategies for multistep reasoning. fig. 3 depicts the schematic mechanisms by which these two strategies execute multi-step inference. the critical logic circuits highlighted in this figure were identified through causal intervention experiments (feng and steinhardt, 2023; meng et al., 2022; vig et al., 2020; wang et al., 2024a) (appendix h). from the figure, one can discern the core distinction between vts and hts. under vts, in every layer beyond the first-termed a "reasoning layer"-the final token attends to the token carrying the next inference result, conditional on the existing result, thereby effecting a single reasoning step. in contrast, hts accomplishes all logical inference within a single layer. this observation naturally prompts a key question: irrespective of whether vts or hts is employed, each node in the transformer stores multiple pieces of information. for instance, the red node in layer 1 encodes not only [a] but also [b]. how, then, does the model consolidate diverse information within a single node and accurately extract and utilize the relevant subset? addressing this question lies at the heart of our investigation into multi-step reasoning. ## buffer mechanism to address the above problem, we propose a mechanism employed by the transformer, which we refer to as the "buffer mechanism". specifically, we observe that the information transmitted via the attention module differs from that conveyed through the residual connections. in the attention module, information [a] with dimension d m is first mapped to a lower embedding dimension d v through the matrix w v , and then it is projected back to d m dimensions via the matrix w o . consequently, the information [a] after passing through the attention layer becomes [a]w vo . in contrast, information transmitted through residual connections does not require additional processing; thus, after the first layer, the information at even indices transforms to below, we introduce the buffer mechanism within the vts framework. all theoretical results are supported by experimental validation (section 5). we observe that each token in fig. 3 (a) stores various information in various forms (projected by different matrices w vo(l) ), which leads to a natural question: how does the attention layer effectively retrieve useful information while avoiding interference from others? to address this question, we introduce the following lemma (the proof can be found in appendix c): it can be observed that the matrices {w i } n @xmath2. each element of {a i } n @xmath3, and is almost unaffected by others. this property also applies for {b i } n @xmath4. by selecting the matrices associated with the relevant buffer, specific information contained in token x and token y can be extracted. the concept of the buffer mechanism is useful for understanding the internal logic mechanisms of transformer models. in each reasoning layer, the token at the last position contains the current intermediate result. additionally, there exists a token that contains both the current intermediate result and the next step's reasoning result, with these stored in different buffers. the role of w q and w k is to extract the current intermediate results from these two tokens, enabling them to attend to each other due to having the same token. we refer to this feature as "same-token matching". to quantitatively characterize this property, we define the following match matrix: we temporarily ignore the effects introduced by the feedforward layers only in our analysis; a detailed version that includes the feedforward layers can be found in appendix d. eq.( 4 ) means that the attention of all tokens is focused on themselves. furthermore, a much more remarkable observation is that the same-token matching is independent of the specific value of x tgt . for example, for x tgt,ood sampled from the untrained random vectors token ood , h (l) (x tgt,ood ) ≈ i still holds. therefore, when ker (l) ≈ i holds, the model exhibits ood generalization capability. understanding hts with buffer. numerous studies have shown that cot, a canonical representative of hts, can significantly enhance logical reasoning capability @xcite . @xcite in the original buffer w vo(1) , w vo(0) and i (identity matrix, which can also be treated as a buffer). thus, cot performs multi-step reasoning tasks through buffer reuse. unlike the vertical thinking strategy, which requires alignment across multiple weight matrices, the horizontal thinking strategy only requires a few layers to satisfy ker (l) ≈ i. this significantly reduces the difficulty of model training. these architectural efficiencies potentially underlie the empirically observed enhancement of reasoning performance in large language models employing cot methodologies. ## experiments this section presents the experimental validation of the buffer theory. to demonstrate multi-step reasoning without introducing additional complexity, we utilized a 3-layer, single-head transformer model to learn from a 2-step reasoning dataset. specific experimental settings can be found in appendix b. after training, the transformer exhibits generalization capability in both in-distribution and out-ofdistribution data (fig. 4 ). fig. 5 presents the experimental validation of eq.( 3 )(4). we visualize the computed values of h (i) (x tgt ) and ker (i) , @xmath5, 2 with the model parameters after training. notably, in the token ood region, token ood and h (2) (x tgt ) exhibit a diagonal structure similar to an identity matrix i, ensuring the model's ood generalization capability. to establish a strong correlation between the model's ood generalization and its ability to use the buffer mechanism for same-token matching, we define a metric for this capability, the matching score (ms): and ker (2) . according to (c)(d), and eq.( 4 ), the diagonal structure of the kernel matrix induces a diagonal structure in the matching matrix, even when x tgt is sampled from the token ood . fig. 4 and fig. 6 illustrates the accuracy curve and the dynamic changes in the matching score and kernel score during training. it is observed that the increase in the model's id and ood generalization coincide with the increases in the model's matching score and kernel score. by computing the cosine similarity between different buffers (w vo(0) , w vo(1) , w vo(2) and i), we observe that these buffers are nearly pairwise orthogonal(appendix d). based on the above experimental analysis, we conclude that the transformer model indeed employs the buffer mechanism for vertical thinking. according to the buffer mechanism, in theory, by setting the model's weights in the following way, an l-layer model can achieve (l -1)-step reasoning without the extra training process: ## improving transformer's data-efficiency in this section, we discuss how to improve data efficiency when employing the vertical thinking strategy, specifically, by enabling the model to learn the essence of logical reasoning with less data or training epoch. as mentioned in section 3, this can be achieved by setting w vo(l-1), t w qk(l) w vo(0),t ≈ i. a more intuitive approach is to replace w qk(l) and w vo(l) with w qk(l) + α (l) i and w vo(l) + β (l) i, where {α (l) } l @xmath6. this identity matrix-based algorithm (denoted as imba) was first proposed in boix-adsera et al. (2023) and has been validated from both theoretical and empirical perspectives for its role in facilitating the model's learning of one-step reasoning data. however, in multi-step reasoning tasks, imba may lead to information interference. for instance, if w vo(l) is replaced with w vo(l) + β (l) i, the storage representation of the two pieces of information transitions from [a]w vo +[b] to [a]w vo +(β (l) [a]+[b]), which introduces interference among the different pieces of information. therefore, this approach may not be effective in enhancing the model's ability to perform multi-step reasoning. based on the understanding of buffer mechanism, we propose a random matrix-based algorithm (rmba), specifically by substituting w qk(l) and w vo(l) with w qk(l) + α (l) z (l-1) and w vo(l) + β (l) z (l) , where {α (l) } l @xmath7, 1/d m ). in this case, the information storage representation changes from [a]w vo +[b] to [a]w vo +[a]z+[b], effectively creating a new buffer. in rmba, when α (l) • β (l) ≥ 0, the diagonal elements of (w vo(l-1),t + β (l-1) z (l-1),t )(w qk(l) + α (l) z (l-1) ) will exhibit positive increments, thereby enhancing the model's reasoning generalization capability. fig. 8 supports this perspective, while also showing that imba, based on induction heads, fails to improve multi-step reasoning performance. to strengthen the credibility of our results, we conducted a comprehensive sweep of hyperparameters such as weight decay, learning rate, and hidden dimensions (d m and d k ). the findings indicate that rmba parameterization can more robustly reach a high accuracy by a certain number of training steps across a wider range of hyperparameters. detailed experimental results can be found in appendix e. these results also demonstrate that multi-step reasoning is not achieved by simply "stacking" multiple single-step reasonings. an algorithm that enhances single-step reasoning may not be applicable to multi-step reasoning. therefore, investigating the multi-step reasoning is an important and meaningful topic. ## discussion in this study, we investigated the vertical and horizontal thinking strategy employed by transformer in a symbolic multi-step reasoning task from the perspective of the buffer mechanism. when utilizing the vertical thinking strategy, the model stores different intermediate results in separate buffers and transfers the reasoning results with the sametoken matching. in contrast, when applying the horizontal thinking strategy, the model reuses the existing buffers to store intermediate results. we validated that the buffer mechanism is a key factor in enabling the model's id and ood generalization capabilities. based on the buffer mechanism, we proposed a tailored approach, rmba, to enhance the model's multi-step reasoning ability, significantly improving data efficiency when training gpt-2 on 7 reasoning datasets. ## limitation our current work lacks an in-depth theoretical analysis. in future work, we aim to conduct deeper theoretical modeling for multi-step reasoning problems and extend the buffer mechanism to other types of multi-step reasoning tasks.