We are glad the reviewer found this paper to be a “really solid piece of work” and “as a whole [...] very well written”! We are happy to clarify your questions below:

Pretraining vs training distinction, were pre-trained models pre-trained explicitly on state-tracking problems or the entire internet?

When we talk about pre-trained models, we always mean models pre-trained on the entire internet (not state-tracking problems), unless otherwise specified. We will clarify this in future versions!

What is $n$? Is it the number of tokens?

Good catch! Yes, we meant $t$ instead of $n$ for the number of tokens in the algorithmic blocks. Apologies for the confusion! We will fix this in future versions of our paper.

Figure 1: In the prefix patching signature for AA, there are four rows of squares, where the first row is entirely blue, and the bottom row only has a blue square in the lower right-hand side.

This is merely a high-level extrapolation of the empirical result shown in Figure 2, which might make more sense with more details.

I am understanding the top row to be the input layer, and the bottom row to be a deepest layer, the token index is increasing from left to right.

This is correct!

Why does just patching the first token result in an incorrect state?

Great question! We create pairs of prompts by corrupting only the first token of the prompt. This is shown visually in section B of Fig 1 (the two prompts differ on action $a_1$), but we realize how this may be easy to miss and we will make it clearer in the final submission. Corrupting only the first token would result in two prompts with a differing final state. Restoring the activation of the first token at the embedding layer is equivalent to prompting the LM with the correct prompt, and thus is sufficient for producing the final correct state.

Having the two prompts differ only at the first token allows us to track the impact of that first token through the network. In early experiments where we corrupt a token near the middle or end of the prompt, we found similar patching patterns, only more vertically compressed (as would be expected).

We will add this clarification to Section 2.3 and Section 4.2 in the paper!

Is claim “A growing body of work suggests that these models learn to represent the latent state of the world” true in light of other work?

This is a fair point. We want to emphasize the difference between (1) "state representations are decodable" and (2) "models use state tracking procedures that generalize to arbitrarily long inputs." The first has been shown by prior work [2,3], while the second remains underexplored [4] but also is not a necessary condition for (1). We will add the aforementioned references for (1) to the second sentence, and soften the claim of the sentence to the following:

A growing body of work suggests that the latent state of the world can be decoded from model internals—e.g. situations described by language and results of program execution—to support prediction.

We will also add the suggested sources – thanks for providing them!

The extra dotted lines in Figure 1 (for the state parity probe) are a bit confusing until one reads the footnote on the bottom of page 5

Thanks for noting this! We will add a line in the caption in Figure 1: “Note the dotted lines indicate two different probing signatures that would both be consistent with this algorithm (see Appendix C.1 for more details).”

[2] Li et al., 2022. Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task. https://arxiv.org/pdf/2210.13382

[3] Li et al., 2021. Implicit Representations of Meaning in Neural Language Models. https://arxiv.org/abs/2106.00737

Thank you so much for your time and feedback. Below, we address some of your general critiques and specific questions.

How could these mechanisms be applied to practical scenarios?

To emulate a more practical scenario, we train Pythia models on a version of our task with permutations expressed in natural language, e.g. “132” would be “swap positions 2 and 3.” while “312” would be “rotate the last item to the front.” We train LLMs to predict the final state (e.g. 231) from the final period token of the sequence. For example,

Swap positions 2 and 3. Rotate the last item to the front. Swap positions 1 and 2.

would map to state “123”.

We then conduct probing and activation patching experiments.

Probing experiments

We train probes to map from the activation of each layer at the position of the final “.” token to the final state.

• Pretrained Pythia-160M (PAA signature): https://ibb.co/5XwfF5h7
• Non-pretrained Pythia-160M (AA signature): https://ibb.co/R4HfbXhn

As with our results on synthetic data, probing results are consistent with the AA and PAA mechanisms respectively.

Activation patching experiments

We patch prefixes up to a fixed token position N. Note that token positions may no longer be aligned in natural language: e.g. “swap” actions have 5 tokens, while “rotate” actions have 7. Thus, we may not be replacing the activations of the same number of actions between prompts.

Nonetheless, we still believe the activation patching results serve as a good proxy for estimating how information gets propagated through the layers.

• Pretrained Pythia-160M: https://ibb.co/Tqdcfb4N
• Non-pretrained Pythia-160M: ​​https://ibb.co/0VXSNNxP

As above, these resemble AA and PAA signatures from the paper. Interestingly, the pretrained results are significantly more “compressed” over the layers – the LLM computes the state very early on. We suspect this may be due to the pre-trained LLM taking advantage of its innate natural language understanding (and perhaps pre-trained state tracking abilities!) to quickly solve the task in an early layer.

Root causes for why models are biased towards one or the other algorithm

The goal of this work wasn’t to provide a definitive account of how these algorithms emerged, but rather identify them, provide evidence for how they contrast, and factors that influence how they emerge over training. We believe a comprehensive causal account of the emergence of these algorithms is out of the scope of this paper, and may require novel interpretability tools that have not yet been invented. We hope future research can be conducted on this topic, and we will add a point on this in our conclusion!

Exclusive focus on final-state accuracy neglects intermediate state accuracy, robustness to long-range dependencies, and error tolerance under adversarial perturbations

We respectfully disagree about the first two.

1. We train and evaluate models to predict all intermediate states up to the sequence length 100 (Section 4.1). We also use probing techniques to verify whether the model's international representation aligns with the ground truth (4.3).

2. We also examine long-range dependencies by performing activation patching early in the sequence (the very first token!) and seeing how it affects the state prediction at the last token (Section 4.2).

We would appreciate clarification on what kind of experiments would best elucidate “error tolerance under adversarial perturbations,” as we're uncertain which adversarial perturbations would be most relevant in the context of our setup.

Overreliance on parity shortcuts risks overestimating models’ true tracking capabilities.

We’d appreciate some clarification on how our measurements of the degree to which models rely on parity heuristics affects our estimation of LMs’ state tracking capabilities! For context, our analysis is grounded in whether the model can complete the permutation group task, i.e. their “true tracking capabilities”. We find that models implement an associative scan in both algorithms, and the parity shortcut is simply a heuristic that eliminates incorrect answers in earlier layers for some models (see Sec. 4.2 and Appendix B). In this sense, the parity shortcut cannot implement state tracking on its own and is only complementary to the “true tracking” algorithm. Note that in all our Section 4 analyses (probing, activation patching, generalization, etc.), models are able to achieve 100% accuracy on exponentially longer sequences as we go down the layers, which is impossible if they were relying on a parity heuristic alone.

Do deeper/wider networks inherently favor AA?

We do this analysis in Section 5.2 (please check out Figure 6). We found that model size isn’t correlated with the algorithm the model chooses to learn; rather, model architectural and initialization differences matter more.

We are thankful for your positive feedback and thoughtful comments. Below, we address some specific questions:

Are permutations provided as individual tokens (i.e., the size of the vocabulary is the number of permutations)?

Yes, they are individual tokens! We append them as special tokens to the tokenizer, but since none of the original tokens are used, the vocabulary is effectively the number of permutations. See our response to qUEj (point #1) about results with natural language text rather than these tokens.

Clarify training supervision: Is this a causal LM trained to perform a seq2seq task? Is this NTP? Or an encoder-only seq2seq model? What is the equivalent NTP variant?

You can think of this as a causal decoder-only LM trained to perform a sequence classification task – for each prefix of a sequence, it is trained to predict the state from the last token of the prefix. Because any action is valid from any state in the permutation composition task, there would be no way to distinguish between the various states through NTP alone. Thus, we did not run an NTP version of the task, though we concur that this would be an interesting avenue of future work.

Would training end-to-end to output final state change results/conclusions?

Great question! We chose our original method for a denser supervision signal from each example. We suspect that training end-to-end won’t make a big difference, but would slow training down.

Is the claim that the model always learns the most compressed (fewer number of layers) version of the algorithm starting from the first layer?

These Section 3 descriptions are simply meant as high-level guides and illustrations of general algorithms; in practice, LMs may skip layers or tokens, or even use a combination of algorithms. As shown in Figure 14, LMs appear to implement something much more redundant and complex than the illustrated associative algorithm in Figure 1.

Why not train on fresh data? Why data repetition? Does this change results?

This is purely for memory reasons. It is hard to check each new data point is unique beyond a certain point, due to needing to store all previously-generated data points in memory. (We ensure all our training examples are distinct from each other, and distinct from the evaluation set.) We use a sufficiently big dataset that we don’t think it’ll make a real difference realistically.

Figure 5: seq length the model is trained on?

The model is trained up to sequence length 100, see L280 in the main text. We’ll also include this information in the figure caption in future versions of the paper!

Sec 5.3: Intuitively why does topic modelling steer towards associative algorithms?

We don’t have a definitive answer to this question, but here’s a hypothesis: when we're training with topic modeling, the model learns a summary representation of the action tokens that are in the place of the parity scheme the model would normally learn otherwise. Training downstream on state tracking, It would be hard for the model to unlearn the topic modeling representation and pick up on the parity heuristic representation, rather than adapting the existing representation to solve the state tracking task.

Adding an experiment with a more realistic setting

Great question! See our response to qUEj (point #1).

Minor typos

We will fix, thank you!

We appreciate your thoughtful review, especially the time you took to examine the full submission, including the Appendix. We’re also grateful for your affirmation of our work’s contribution to the literature. Below, we address some of your specific questions:

Unclear if signatures are sufficient conditions. Would an alternate scheme give similar signatures?

You raise an important point: the signatures we describe are necessary, but not sufficient, conditions for the proposed algorithms. We tried to make a distinction between theoretical “algorithms” and empirical “mechanisms” in the paper, but to make this distinction clearer, we will add the following clarification at the beginning of Section 4 in the final version:

It is important to emphasize that the various signatures described above provide necessary, but not sufficient, conditions for implementation of the associated algorithm; the exact mechanism that LMs use in practice is likely complex and dependent on other input features not captured by the algorithms described above.

Sensitivity analysis for linear probe: how would varying the probe affect results

We trained the probe across 10 different random data subsets on the S3 AA and PAA models, and measured the standard deviations across these runs.

The standard deviations are all quite small, indicating that our results are highly insensitive to randomness in training.

Please let us know if this aligns with what you had in mind for “sensitivity analysis,” or something else!

Sensitivity of activation patching to varying $x$ or $x’$

We performed activation patching experiments across 200 different input pairs $(x,x’)$. We plot a heatmap of the standard deviations across various (layer, position) of these pairs in:

• S3 AA: https://ibb.co/5gmqB2Q2
• S3 PAA: https://ibb.co/R4YGwC4z

Note that for the PAA models, the standard deviations are highest in the middle chunk — for half the prompt pairs (when parity is equal), replacing the middle chunk results in the corrected answer, while for the other half, replacing the middle chunk has little to no effect.

In Figure 1, what is represented by the two bifurcating dashed lines for state parity probe?

The dashed lines refer to the interesting fact that for models implementing AA, some encode parity linearly, and some do not. We go into more detail on this in Appendix C.1. We will add the following to the caption of Figure 1 to avoid confusion:

“Note the dotted lines indicate two different probing signatures that would both be consistent with this algorithm (see Appendix C.1 for more details).”

jylin04 seems incorrect

We were referencing a blog post [1] and weren’t sure about the citing format ourselves, but we think the reference is correct given that we don’t know the authors’ full names.

Other typos / inconsistencies

We will correct, thanks!

[1] OthelloGPT learned a bag of heuristics https://www.lesswrong.com/posts/gcpNuEZnxAPayaKBY/othellogpt-learned-a-bag-of-heuristics-1