Paper: Toward a Mechanistic Understanding of Stepwise Inference in Transformers: A Synthetic Graph Navigation Model (ICLR-2024)

Reviewer comment (Theory):
Emergent abilities of large language lodels were thought to be a consequence of unpredictable scaling, therefore, I doubt if a simpilified DAG substitute can replicate this behaviour.

Author response:
>**Comment:** Emergent abilities of large language models were thought to be a consequence of unpredictable scaling, therefore, I doubt if a simpilified DAG substitute can replicate this behaviour.


**Response:** We *respectfully* disagree with the claim that emergent abilities only occur in large scale models. Abilities like in-context learning (ICL) and stepwise inference protocols like chain-of-thought (CoT) have been shown and extensively studied in small toy models by several recent works by designing well defined synthetic tasks [1,2,3,4,5,6,7,8,9,10]. The goal of such studies (and ours well) is to study interesting phenomenology of LLMs and design abstract setups that enable development of precise mechanistic hypotheses for how the phenomenon works. For example [4,5] use a synthetically designed setup to develop hypotheses for how ``knowledge” about an entity is stored in a pretrained model, showing that such knowledge can be manipulated via relatively simple linear transformations. This development of mechanistic hypotheses is precisely what we have done with our experiments. 

We also note that to ensure the point about relation to LLMs is addressed in the paper, we have substantially revised our manuscript to further motivate the explicit connection between graph navigation and stepwise inference in more realistic settings. Specifically, we have rewritten parts of the introduction (we have elaborated our exposition on the relation between stepwise inference and graph navigation, added a discussion of the limitations of a synthetic data approach), rewritten our abstract, contextualized our results more thoroughly with existing literature (added Appendix figure 8 which shows examples of graph navigation in a LLM), and also edited the formatting and arrangement of figures and results in paper (including promoting the training dynamics figure to the main text since it shows that global planning capability arises *after* models learn to take correct steps). As an example, our edits now highlight several recent studies that either use or model stepwise inference with LLMs to study tasks that are specific instantiations of our proposed graph navigation problem (e.g., see [11,12, 13, 14]). We have included a figure in Appendix (Fig. 8) that describes example prompts and model outputs from each of these papers to exemplify the relation of our proposed synthetic benchmark to realistic setups and tasks considered by prior work in showing the value of stepwise inference. Given that these studies exist in the (very recent) literature, we argue the connection between stepwise inference, planning, and graph navigation is explicit and our updated draft ensures this connection is emphasized upon. 

With our detailed responses and changes to the paper, we hope that we have addressed the reviewer's concerns and the reviewer increases their score.

[1] Bai et al. Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection (NeurIPS 2023)

[2] Akyurek et al. What learning algorithm is in-context learning? Investigations with linear models (2022) https://arxiv.org/abs/2211.15661

[3] Ahn et al. Transformers learn to implement preconditioned gradient descent for in-context learning (NeurIPS 2023)

[4] Zeyuan Allen-Zhu and Yuanzhi Li, (2023) Physics of Language Models: Part 3.1, Knowledge Storage (2023)

[5] Zeyuan Allen-Zhu and Yuanzhi Li, Physics of Language Models: Part 3.2, Knowledge Manipulation (2023)

[6] Maya Okawa et al. Compositional Abilities Emerge Multiplicatively: Exploring Diffusion Models on a Synthetic Task (NeurIPS 2023)

[7] Lindner et al (2023) Tracr: Compiled Transformers as a Laboratory for Interpretability https://arxiv.org/abs/2301.05062

[8] Feng et al, Towards demystifying the mystery behind chain of thought: A theoretical perspective (2023)

[9] Li et al. Dissecting Chain-of-Thought: Compositionality through In-Context Filtering and Learning (NeurIPS 2023)

[10] Prystawski et al. Why think step by step? Reasoning emerges from the locality of experience (NeurIPS 2023)

[11] Mommenejad et al. Evaluating Cognitive Maps in Large Language Models with CogEval: No Emergent Planning (NeurIPS 2023)

[12] Saparov et al. Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thought (ICLR 2023)

[13] Dziri et al. Faith and Fate: Limits of Transformers on Compositionality (NeurIPS 2023)

[14]  Hou et al (2023) Towards a Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Models  https://arxiv.org/pdf/2310.14491.pdf

Reviewer comment (Writing):
The definition of hierarchical graph and random graph is confusing to me. It seems possible to convert the random graph in fig. 10 into a hierarchical graph by simply re-grouping the nodes.

Author response:
**Response**: 
In a hierarchical graph, nodes are pre-partitioned into layers. The only connections allowed are between successive layers and these are chosen randomly. Whereas in a random graph, all connections are chosen at random. It is not possible to simply regroup nodes in a random graph and convert it to a hierarchical graph. To provide evidence for this, we now refer the reviewer to Appendix Figure 9 from our paper, where we highlight and discuss the differences between a hierarchical and random graph. In that figure, the *bottom panels* show the distribution of path lengths in the two families of graphs: hierarchical graphs have an exponential distribution owing to their structure and random graphs have a gaussian distribution. For example, given any pair of nodes, in a hierarchical graph every path between them is of the same length (determined by the difference of their layer numbers), whereas in a random graph there is a diversity of path lengths (for instance, there might be an edge between the nodes as well as a path of length more than 10, as is the case with our synthetic random graphs). We wish to emphasize that our specific approach to using a hierarchical graph is designed to enable more precise control over the data distribution.

To further ensure this point is clear, we have added pseudocode for the algorithms for generating both the random and hierarchical graphs—**Please refer to Algorithms 1 and 2 in the Appendix A.3**. We hope that this,along  with the path characterization provided in Figure 9, addresses concerns about possible similarities between the two setups.

Reviewer comment (Experiments):
The highlighted finding, i.e., the step-wise inference gap is influenced by (i) the underlying DAG structure, and (ii) the length of the training samples is only studied in a very shallow level. The DAG structure in the context seems to be not well-defined and well-categorized. The discussion about the length of the training samples mainly focuses on how it affects the length of the model output.

Author response:
**Response:**
Our highlighted findings in a single graph scenario include: (i) the inference gap consistently persists across diverse data distributions in the training data (Figure 3); (ii) a diversity versus accuracy trade-off exists in finite temperature stepwise inference for transformers (Figure 4); and (iii) demonstrating that there is a bias towards shorter paths and the observation that over the course of training the model first learns to take correct steps and then acquires the capability of global planning (Figure 5). We recognize that our approach serves as a first step towards formulating a set of hypotheses for further validation in large-scale, realistic settings, and there is considerable scope for extending our research to more realistic settings such as different data distributions, noisy or incomplete information, and different graph structures. As part of this endeavor, We have now repeated our analysis on a setup where a ground-truth graph is “noisily observed”. Specifically, every time a path is sampled between two nodes, we intentionally corrupt a random fraction of the tokens on the path by 5-20%. The model is now trained on these corrupted paths—this randomization further ensures that overfitting is infeasible, since a given path is never seen twice (due to corruptions). *In this setup, we again see that our claim on the performance gap induced by stepwise inference continues to persist.* We have also defined our training protocol and description of the graphs used in training in more detail (in Appendix).

Reviewer comment (Experiments):
While the paper claims to address stepwise inference in transformers and introduces a graph navigation model, the experiments seem to focus narrowly on synthetic tasks without sufficient evidence of the model's generalizability to a more realistic or applicable datasets.

Author response:
Thank you for raising this comment. We have substantially revised our manuscript to address your concern by further motivating the explicit connection between graph navigation and stepwise inference in more realistic settings. Specifically, we have rewritten parts of the introduction (we have elaborated our exposition on the relation between stepwise inference and graph navigation, added a discussion of the limitations of a synthetic data approach), contextualized our results more thoroughly with existing literature, and also edited the formatting and arrangement of figures and results in paper (including promoting the training dynamics figure to the main text since it shows that global planning capability arises *after* models learn to take correct steps). As an example, our edits now highlight several recent studies that either use or model stepwise inference with LLMs to study tasks that are specific instantiations of our proposed graph navigation problem (e.g., see [1, 2, 3, 4]). We have also included a figure in Appendix (Fig. 9) that describes example prompts and model outputs from each of these papers to exemplify the relation of our proposed synthetic benchmark to realistic setups and tasks considered by prior work in showing the value of stepwise inference. Given these (very recent) studies, we argue the connection between stepwise inference, planning, and graph navigation is explicit and our updated draft ensures this connection is emphasized upon. 

We’d also like to contextualize how our methodology of proposing and studying a synthetic problem fits into the literature. Specifically, the goal of our work was to analyze an interesting phenomenon witnessed in LLMs, i.e., stepwise inference, by designing an abstraction of the problem. This enables development of precise mechanistic hypotheses. In LLM and generative model literature, this is a fairly common approach. For example, [5,6] use a synthetically designed setup to develop hypotheses for how ``knowledge” about an entity is stored in pretrained model, show that such knowledge can be manipulated via relatively simple linear transformations. [7] used a procedurally defined synthetic dataset with a relatively small model to demonstrate that emergent abilities seen in neural networks are partially driven by the compositional nature of real world data. In other work [8], use Tracr-compiled transformers to show that if primitive operations involved in a formal algorithm can be implemented in a model, stepwise inference is sufficient for length generalization. Similarly, [9] used context-free grammars to demonstrate that stepwise inference is sufficient to solve problems that require dynamic programming. [10] use CoT to show transformers can learn arbitrary compositions of MLPs in-context. Our work lies in the same regime as these studies and has the goal of defining a well-designed task for studying stepwise inference protocols. 


[1] Mommenejad et al. Evaluating Cognitive Maps in Large Language Models with CogEval: No Emergent Planning (NeurIPS 2023)

[2] Saparov et al. Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thought (ICLR 2023)

[3] Dziri et al. Faith and Fate: Limits of Transformers on Compositionality (NeurIPS 2023)

[4] Hou et al. Towards a Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Models. (arXiv, 2023)

[5] Zeyuan Allen-Zhu and Yuanzhi Li, (2023) Physics of Language Models: Part 3.1, Knowledge Storage (2023)

[6] Zeyuan Allen-Zhu and Yuanzhi Li, Physics of Language Models: Part 3.2, Knowledge Manipulation (2023)

[7] Maya Okawa et al. Compositional Abilities Emerge Multiplicatively: Exploring Diffusion Models on a Synthetic Task (NeurIPS 2023)

[8] Zhou et al. What Algorithms can Transformers Learn? A Study in Length Generalization. (arXiv, 2023; https://arxiv.org/abs/2310.16028)

[9] Feng et al, Towards demystifying the mystery behind chain of thought: A theoretical perspective (arXiv, 2023)

[10] Li et al. Dissecting Chain-of-Thought: Compositionality through In-Context Filtering and Learning (NeurIPS 2023)

Reviewer comment (Evaluation):
The approach primarily involves modeling the decision-making process using directed acyclic graphs (DAGs). However, there are concerns that the model may overfit these synthetic graph structures. The paper does not adequately address how the model handles noisy real-world graphs, which may exhibit cycles, incomplete information, or random behavior that are common in real-world applications. In this context, a deeper analysis of the robustness of the proposed method is crucial.

Author response:
These are valid concerns and we have both highlighted our existing results and added a new experiment to address them, as discussed below.

**We have added new experiments with corrupted graphs.** Thanks to your feedback, we have added a new experiment to better model noisy graphs. In the experiment, we have now repeated our analysis on a setup where a ground-truth graph is "noisily observed". Specifically, every time a path is sampled between two nodes, we intentionally corrupt a random fraction of the tokens on the path by 5-20%. The model is now trained on these corrupted paths—this randomization further ensures that overfitting is infeasible, since a given path is never seen twice (due to corruptions). *In this setup, we again see that our claim on the performance gap induced by stepwise inference continues to persist.* 

**Test-train split and randomness by sampling from diverse paths:** We’d also like to clarify and emphasize that **overfitting cannot occur in our setup**. In our current setup, please note that given a start and goal node pair, there is a huge diversity of paths between them: for example, please see Appendix Figure 9, which shows that there are up to $10^3$ valid paths between a pair of nodes. We divide all node pairs into test and train.  For the train node pairs, in each context window receives 1 random path:  this itself is a source of randomness that we now highlight better in the revised manuscript. This randomness corresponds to how a particular stepwise path is sampled in “noisy” real-world settings. (e.g., a particular sequence of logics out of several possible solutions is sampled in a solution set of math word problems.) Importantly, this randomness also acts as a strong regularizer against overfitting on the training set. Further, all our results are on the test set node pairs that have never been seen in a train set, i.e., overfitting to this data is not possible. We note we have updated discussion of our training protocol in Appendix A.3 to summarize these details better.

Summary: Thank you very much for your constructive feedback! We hope that we have sufficiently answered your questions with the above new experimental results, arguments and justifications, and hope that you will consider increasing your score to support the acceptance of our work.