Transformers Use Causal World Models in Maze-Solving Tasks

Dear Reviewer 1yTd,

Thank you for taking the time to understand our work and help us to improve its soundness. We have incorporated your suggestions into our work, in particular adding more details in several places and adding a new figure as well as more detail to improve the narrative flow and make our contributions clearer. In addition, we discuss the points you raised below:

1. No clear overall understanding of the internal structured of the learned weights: Existing works in the field do not target a complete understanding of internal weight structure and it is not clear what would constitute a clear overall understanding. We attempt to thoroughly analyse the features implicated in representing the maze’s connectivity structure within both Terry and Stan through our SAE and attention analyses. Showing a sensible agreement between the circuit (weight level) analysis and the SAE feature analysis is one of our main contributions, and brings us closer to an overall understanding of our world models than prior works - though we would not claim that is it yet a complete understanding.
2. Regarding the organization of the paper: We thank the reviewer for highlighting concerns around clarity as a key shortcoming of the initial draft. We have made numerous changes to improve the flow and clarity of our work. Regarding, the contributions - our paper is an empirical investigation into the world models learned by our transformers, which yields a number of interesting findings on the properties of these World Models (along with their existence, which might be considered the implicit central claim). As such, there are multiple contributions, some of which are methodological (efficacy of decision trees and SAEs) and most of which are empirical.
3. “Many claims are not grounded”: If the reviewer could list the claims which he finds insufficiently justified we will be happy to respond in more detail. It is worth noting, as we do below, that some of our claims are observations which apply in our context and we do not assert their generality; as such, it is not clear what might be considered “conclusive” evidence beyond what we have shown.
4. “What is the formal definition of the world model in this paper?”: A mathematical definition of world models is not adopted (and we do not know of any such definition which is considered canonical in the field). However, the paper states what is canonically meant by “world model” and, as such, how the features we discover can be said to constitute a world model. We have attempted to clarify this further in the paper.
   1. “How does such a world model come into play in predicting shortest path”: This is an excellent question and one we begin to explore in Section 4 - here we show that perturbing the internal representations ****corresponding to the world model causes the transformer to predict paths in-line with having observed a corresponding changed maze in its input. This establishes the causal role of these features and shows that the model uses these internal representations to construct the shortest path. It is, however, beyond the scope of our work to perform a full mechanistic analysis on the question of “how” these perturbations affect representations further downstream (at subsequent layers) in our transformers.
5. How the decoupling effects of SAE come into play for WM? It is not entirely clear what is meant by “decoupling effects”. We leverage SAEs as a way to recover features of interest from the model’s residual stream without presuming that any of these will correspond to a world model. After training such SAEs with a generic reconstruction objective over the internal residual stream of the model, we then attempt to isolate features relevant to representing the maze (which we dub the “world model”) through the use of decision trees - it need not have been the case that such features were formed by the model.
6. How is "Isolating irrelevant features" helpful for WM? We assume the reviewer means “isolating relevant features”. As the “World Model” refers to a structured internal representation in our transform models, there is no sense in which features which we identify (relevant or not) are “helpful for WM”. The WM either exists, or does not, in the internal representations of our trained models - we merely recover these and attempt to understand their structure and role within the transformers. It is possible that transformers which acquire WMs are better able to perform tasks, but this is not a claim we make in this work.

Dear Reviewer 2SJe,

Thank you for providing such a thorough review of our work and offering many helpful comments and critiques. We have incorporated as many of these as we were able to into the paper, and provide some further discussion below:

• Terry and Stan have different numbers of heads: We chose to analyze terry and stan in spite of this difference simply as they were the most performant models given their position embedding scheme. We understand that this makes it less clear that all differences can be attributed solely to the position embedding scheme, but would make the following two observations which gave us sufficient confidence in analysing these models nonetheless:
  • The models do in fact have the same number of parameters, up to the difference caused by Stan’s having learnable position embeddings. In our implementation, the dimensionality of the residual stream is fixed, and e.g. decreasing the number of heads means each head has fewer parameters. (we have added the parameters counts to the paper for completeness, and commented on this fact)
  • Although Stan has more attention heads than terry, in stan only 3 heads are actually implicated in constructing the world model (connection) features, whilst the others seem to play less clear roles.
• Stan’s failure under the addition of a 6x6 intervention: This was perhaps not sufficiently clearly described - but Stan never saw sequences of length greater than those of a 6x6 fully-connected maze during training and thus fails as soon as even one connection is added (and in Appendix Figure 1 we show also fails for 7x7 fully-connected mazes). It is indeed the case that Stan saw 7x7 sparsely connected mazes at training, but these had adjacency lists shorter than those required by 6x6 fully-connected with an additional connection. As these experiments attempt to remove a connection when the input is a 6x6 maze + 1 connection, Stan failed to generalize.
• Addition of more than one feature: This is a good suggestion, but as we wish to guarantee that interventions affect the shortest path over a balanced and large number of samples, this requires non-trivial dataset generation. We agree that such an experiment with n>1 wall additions would be highly interesting to advance the findings relating to latent interventions which are not possible in the input space, but believe a systematic investigation of the phenomenon of simultaneous intervenability of SAE features, and the asymmetry in addition / removal of features, should be left to future work.

Thank you again for finding our work interesting and providing a thorough, helpful review. We hope that with the above considerations, and the changes made to the paper in mind, you might consider revising your score.

The Authors

Dear Reviewer gA1f,

Thank you for your thoughtful and detailed review. We appreciate the constructive feedback, particularly regarding the clarity of our writing and methodology. We have made several improvements to address your concerns, which we discuss below.

• Whether the features we find constitute a “world model”: Whilst it is true that the sense in which we claim that connection features constitute a “world model” is not as abstracted from the inputs as those of other related works, this was what motivated our choice of domain - as we wished to form a comprehensive understanding of both how well SAEs are able to capture world-model features and the extent to which identified features agree with those identified from a circuit perspective. We still believe that these features may be meaningfully considered a “world model” in that they form an intermediate representation which is not completely trivial; as we demonstrate, the transformer must learn to combine the tokens which constitute an edge into a single feature - something required to consistently perform well on the task (as our intuitions would have told us) in the absence of naive heuristics or mere path following (which would not require the model to internally represent the connections in full).

• Whether our interventions are straight-forward / uninteresting: It is indeed the case that our interventions can be reduced to multi-token perturbations, but we do not believe this makes them uninteresting - for two reasons. Firstly, in Othello one could also perturb the input sequence, albeit less trivially, to achieve the desired outcome by using e.g. search over an othello solver. Secondly, even though these interventions may be straight-forward there are still interesting behaviours which we observe. I.e.:

  • There is an asymmetry in adding / removing features (which would not be the case with input perturbations).
  • It is possible to e.g. add features (toggle additional connections on) to the Stan model and coherently effect its maze solving behaviour in the presence of these. However, the model fails to generalize to longer sequences, so attempting the same intervention in the input space, by e.g. adding four tokens specifying the edge, would lead the model to fail 100% of the time (as we demonstrate).

  As such, whilst the content of the interventions is by the choice of environment straightforward, we observe behaviours which we consider non-obvious.

• Our SAEs have too many features / the features are not atomic: This is a valid concern, especially given the lack of detail we provided on our SAE training and evaluation in the initial submission. We have added our validation results and representative examples of SAE features to the appendices, and make the following observations:

  • The features occur after one transformer block on the semicolon position, and are not mere semicolon-demarkers as they have a causal role in model behaviour and capture the maze structure.
  • The SAEs intentionally form an overcomplete basis but through sparsity penalties are decentivised from encoding redundant information. Furthermore, metrics measuring sparsity and feature density distributions give reasonable values.
  • We observe other SAE features which display clear structure, as we have shown with the added figures in the appendices.

• We have also improved the writing in places to hopefully improve the clarity of the work. We have also provided further details on the use of decision trees in the appendices to answer your question regarding the target - they are trained in a fashion similar to the linear probes of prior work, with a supervised classification target for the presence or absence of an edge (thus isolating the edge features).

Questions:

• On the significance of patterns switching from 1/3 back to 5/7 back: You are right to point out that this is not well-explained, and indeed we have no strong explanation as to why our transformers decided to varying circuits, besides high-level arguments around symmetry breaking and the likely fact that these may be learned at different times during training. Pending additional experiments, we leave further investigation to future work as this does not significantly concern the focus of our paper.
• Regarding even/odd parity attention heads: This is a very interesting question which we do not have a good answer to! This was a finding which we did not particularly expect, but which is not too surprising after-the-fact. At a high-level we believe this related to subspaces of differing attention-heads at the same layer being maximally orthogonal to reduce interference, but leave a complete investigation to future work.

We hope that in light of the changes to the paper, and the clarifications we have attempted to provide in the discussion above, you might reconsider your scores. Once again, we are grateful for your help in improving our work.

The Authors

Dear Reviewer SDpz,

Thank you for taking the time to understand our work and to provide us with such valuable suggestions. We’ve made a number of changes to the paper on the basis of these, and provide detailed responses to the points you raised below:

• Circuit analysis limited to attention visualizations and magnitudes of vectors from applying OV matrices: Whilst we do observe even/odd parity we were mainly interested in understaning consistency vetween the sae features and the rresults of looking at heads directly. to this end the most interesting part of the circuit analysis was the results of ablating attention heads and seeing which SAE features they affected. The description of these experiments and their corresponding plots were not very clear in the original submission, so we have added detail to make the purpose and findings of these experiments more clear.
• No results on probes: This was an omission as previous work using probes was unable to perform causal interventions, and performing comparable interventions with probes would require knowing exactly what we were looking for; i.e. training probes on latents collected after the first block, only on “;” tokens corresponding to a particular connection, and using other “;”s for negative examples. Whilst is it feasible that this would result in causally effective probes (though less likely so for the transformer’s leveraging compositional codes to encode connections), it is clear that probing here is not an effective way of discovering the features that constitute the world model.
• SAE Training details and Features: This was indeed an omission in the first submission. We have now added thorough training details, as well as information about the sweep we performed over SAEs and representative examples of other features discovered in the models.
• The paper does not add much to Igorevich et.al 2023.: Igorevich et al.’s workshop paper demonstrated that the information about mazes was linearly decodable with probing after layer 2 at a single token position. However, there was no demonstration of the causal role of the probes they used. Alongside the use of decision trees, our findings show that transformers in such tasks do acquire a causal world model, but this model is constructed in the first layers and can be understood clearly in terms of subsets of the transformer’s attention heads which specialize in constructing representations of connections.

Questions:

1. In Figure 9, it is mentioned that "Stan removal interventions fail as the inputs in these cases have more connections than the model is able to handle"; however, if I understand correctly, removing the connections reduce the length so the Stan model should not fail in this case?
2. What is the rough statistics of SAE features and what features other than edge connections are encoded: This is an excellent question which we overlooked due to space-constraints. We have added details in the appendices to expand on these findings and show that there are at least four-types of features which SAEs on both models of interest discovered: 1) The connection features which form the WM, 2) specific-token features in the adjacency list, 3) specific-token features in the path (these differ from those in the adjacency list in both models, even though only one model is found to utilize a compositional code), 4) semi-specific features over various coordinates in the adjacency list.

We hope that with these clarifications, and the changes we have made to the paper at your suggestion, you may reconsider your score. We thank you again for helping us to improve the quality of our work.

The Authors