(How) Do Language Models Track State?
Authors: Belinda Z. Li
Zifan Carl Guo
Jacob Andreas
Venue: ICML 2025 poster (2025)
Keywords: state tracking
simulation
interpretability
Primary area: deep_learning->large_language_models

Abstract: Transformer language models (LMs) exhibit behaviors—from storytelling to code generation—that seem to require tracking the unobserved state of an evolving world. How do they do this? We study state tracking in LMs trained or fine-tuned to compose permutations (i.e., to compute the order of a set of objects after a sequence of swaps). Despite the simple algebraic structure of this problem, many other tasks (e.g., simulation of finite automata and evaluation of boolean expressions) can be reduced to permutation composition, making it a natural model for state tracking in general. We show that LMs consistently learn one of two state tracking mechanisms for this task. The first closely resembles the “associative scan” construction used in recent theoretical work by Liu et al. (2023) and Merrill et al. (2024). The second uses an easy-to-compute feature (permutation parity) to partially prune the space of outputs, and then refines this with an associative scan. LMs that learn the former algorithm tend to generalize better and converge faster, and we show how to steer LMs toward one or the other with intermediate training tasks that encourage or suppress the heuristics. Our results demonstrate that transformer LMs, whether pre-trained or fine-tuned, can learn to implement efficient and interpretable state-tracking mechanisms, and the emergence of these mechanisms can be predicted and controlled. Code and data are available at https://github.com/belindal/state-tracking

(How) Do Language Models Track State?
Belinda Z. Li'
Zifan Carl Guo
Jacob Andreas
Abstract
execution to support prediction (Li et al., 2021; Nanda et al., 2023; Li et al., 2023). However, the mechanisms that LMs use to construct these representations are not un- derstood. Do LMs simulate state evolution step by step across successive hidden layers or token representations (Yang al., 2024)? Are states approximated through com- plex collection of heuristics (jylin04 et al., 2024)? Is state tracking an illusion (Bender & Koller, 2020)?
Transformer language models (LMs) exhibit behaviors- from storytelling to code generation- that seem to require tracking the unobserved state of an evolving world. How do they do this? We study state tracking in LMs trained or fine-tuned to compose permutations (i.e., to compute the or- der of a set of objects after a sequence of swaps). Despite the simple algebraic structure of this prob- lem, many other tasks (e.g., simulation of finite automata and evaluation of boolean expressions) can be reduced to permutation composition, mak- ing it natural model for state tracking in gen- eral. We show that LMs consistently learn one of two state tracking mechanisms for this task. The first closely resembles the "associative scan" con- struction used in recent theoretical work by Liu et al. (2023) and Merrill et al. (2024). The second uses an casy-to-compute feature (permutation par- ity) to partially prune the space of outputs, and then refines this with associative scan. LMs that learn the former algorithm tend to general- ize better and converge faster, and we show how to steer LMs toward one or the other with inter- mediate training tasks that encourage or suppress the heuristies. Our results demonstrate that trans- former LMs, whether pre-trained or fine-tuned, can learn implement efficient and interpretable state trateetracking mechanisms, and the emergence of these mechanisms can be predicted and con- trolled. Code and data are available at https: //gi tig.. wot//belinnall/sttteeeeeeeaae http:/stttte-ttrcking....
This paper studies the implementation and emergence of state tracking mechanisms in language models using per- mutation composition as a model system: given a fixed set of objects, we train or fine-tune LMs to predict the final position of each object after sequence of rearrangements. Previous work has used versions of this task to evaluate LMS' empirical state tracking abilities (Li et al., 2021; Kim & Schuster, 2023; Li et al., 2023). Additionally, as shown by Barrington (1989), many complex, natural, state-tracking tasks-including simulation of finite automata and evalua- tion of Boolean expressions-can be reduced permutation tracking with five or more objects. This makes a natural model for studying state tracking in general.
Our analysis proceeds in several steps. $2 provides technical preliminaries: $2.1 and $2.2 introduce state tracking prob- lems and the permutation composition task we use to model them (Figure 1A), and $2.3 reviews the set of interpretabil- ity tools we use analyze LM computations. Next, $3 lays out family of algorithms that past work has suggested LMs might, in principle, use solve the state tracking task (Figure ID), and describes the signatures lignatures-expecteed expected readouts from different interpretability methods-that we would expect to find if a given algorithm is implemented (Figure IB-C).
Finally, $4 and $5 present experimental findings. Across a range of sizes, architectures, and pretraining schemes, find that LMs consistently learn one of two state tracking mechanisms. The first mechanism, which we call the "as- sociative algorithm" (AA), resembles the associative sean construction used by Liu et al. (2023) and Merrill & Sabhar- wal (2024) to establish theoretical lower bounds on the ex- pressive capacity of Transformers. The second mechanism, which we call the "parity- algorithm" (PAA), first rules out a subset of final states using an easy-to-compute permutation parity heuristic, then uses an associative scan
1. Introduction
Language models (LMs) are trained to model the surface form of text. A growing body of work suggests that model internals contain a latent, decodable state of the world-e.g.. situations described by language and results of program
'MIT EECS and CSAIL. Correspondence to: Belinda Z. Li t.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

(How) Do Language Models Track State?
Permutation Composition
B
Prefix Patching Signature
Probing Signature
How much prefix must be modified at each layer to change
What percentage the state/ state parity sequence can be accurately probed from LMs' intermediate layers?
$ even
42315
40
outputs?
odd
partyy
Parch prefix up this token on this layer
Action
1253
even
Action perry
Train linear probe decode state sequence from wrth position in each layer's representations
0
odd 32514
42315
4
"Corupt" ourput 52
"Clean"
play)
output.
Legend
Input
%p 8,
Acc.
First layer
Correctly predicts final when prefix up to this coll is patchod.
Last layer-
Rules subsct uniquely identify the correct state, when prefix up this patched.
not
State Probe
Output states
State Parity Probe (chance accuracy is 0.5)
Transformer-mmpementaae Algorithms
Predicts final state with chance accurucy.
possible
State
Probe Accuracy
Sequential Algorithm
Layens
Parallel Algorithm [St only]
Parity
Probe
Accuracy
Layors
possible
Associative Algorithm
Probe
Accuracy
Layers
Parity
Purity-Associative
Probe
Algorithm (PAA)
Accuracy
Complement"
Figure 1. We use permutation word problems as simple model of state tracking. Actions permutations, and the products of those permutations; the current state can be tracked by taking the cumulative product from left to right ($2). We identify several possible algorithms that Transformers may Above, we depict the "signatures" of each algorithm under two types of interpretability analysis: prefix patching, where we create pairs of prompts differing only on the first token, where train linear probe to map from last-token representations across the layers to either the final state the final state parity ($2.3), Note: dotted lines indicate two different probing signatures consistent with this algorithm (see Appendix for more details).
solve permutation word problems: sequential, parallel, associative, and parity-associative ($3).
substitute all activation except the prefix up token at particular layer, and probing.
to obtain a final prediction. Notably, we fail to find evidence for either step-by-step simulation or fully parallel composi- tion, despite their being theoretically implementable by LMs. We support our findings with evidence from representation interventions (Meng al., 2022; Zhang & Nanda, 2024; $4.2), probes (Shi al., 2016; $4.3), patterns in prediction errors (Zhong et al., 2024; $4.4), attention maps (Clark et al., 2019; $4.5), and training dynamics (McCoy et al., 2020;
Olsson et al., 2022; Hu et al., 2023; $5.1).
The sean operation for PAA appears difficult for LMs im- plement robustly, and the choice of mechanism sometimes significantly impacts model performance long sequences ($5.1). Whether given LM learns AA PAA is highly stochastic ($5.2). However, each is associated with a char- acteristic set of phase transitions in the training loss (Chen et al., 2024), and LMs can be steered toward one solution or
2

---

(How) Do Language Models Track State?
the other by training on an intermediate task that encourages or discourages LMs from learning a parity heuristic ($5.3).
This, combined with its simple structure, makes it a natural model system for studying state tracking in general.
As pretrained LMs sometimes re-use circuits when fine- tuned on related tasks (Prakash et al., 2024; Merullo et al., 2024), our results suggest possible mechanism by which real-world LMs might perform state tracking when mod- eling language, code, and games. We show preliminary evidence of these algorithms on a version of our permuta- tion composition tasks expressed in natural language (Ap- pendix E). Looking beyond state tracking. these findings underscore both the complexity and variability of LM so- lutions to complex tasks, which may involve both heuristic features and structured solutions.
More formally, the finite symmetric group S,, comprises the set of permutations of n objects equipped with a com- position operation. For example, 42315 denotes the permu- tation of objects (i.e. in St)) that moves the first object to the fourth position, the second object to the second position, etc. Importandly for our findings in this paper, every per- mutation can be expressed swaps (in Figure AA, do. but not (1, is an example of a swap). The parity of permutation (even or odd) is the parity of the number of swaps needed to create it.
composition of two-element
The composition of two permutations, standardly denoted a1 o dg. is the result of applying an after ag. Inputs to sequence models in machine learning are typically written with earlier inputs before later imputs (i.e., left-to-right), so for consistency with this convention, we will write 0g01 to denote the application of og then a1. Figure IA shows the result of composing 42315 and 12534 in sequence.
2. Background and Preliminaries
2.1. State Tracking
Inferring common ground in discourse (Li et al., 2021), navigating the environment (Vafa et al., 2024), reasoning about code (Merrill et al.. 2024), and playing games (Li et al., 2023; Karvonen, 2024) all require being able to track the evolving state of a real or abstract world. There has been significant interest in understanding whether (and how) LMS can perform these tasks. In theoretical work, researchers have observed that many natural state- trate-tracking problems (including the ones listed above) are associated with the complexity class NC1, but Transformers cannot track the state of arbitrarily long inputs (Merrill et al., 2022: Huet et al., 2025; Bhattamishra et al., 2020; Delétang et al., 2023; Strobl et al., 2024). However, prior work has shown that Transformers with Olog n)) depth can model inputs of up to length n (Liu et al., 2023; Merrill & Sabharwal, 2024). Empirical work, meanwhile, has found that large LMs learn to solve state tracking problems (Kim& & Schuster, 2023) and encode state information in their representations (Li et al., 2021; Li et al., 2023). But a mechanistic understanding of how trained LMs infer these states has remained elusive.
Finally, the word problem on S, the problem of comput- ing the product of a sequence of permutations. This product itself corresponds to a single permutation (32514 in Fig- ure 1). But, following the intuition given at the beginning of the section, it may equivalently be interpreted as the final or- dering of the objects being rearranged (DBAEC in Figure 1 IA). Following this intuitive explanation (and by analogy to other state tracking problems), we will use de to denote a single permutation ("action") in a sequence, and St = ag ae to denote the result of sequence of permutations (a "state").
Given sequence of permutations, we use e(a:) to denote the parity of the fth permutation, so:
E(st) = =((ao--... ..)) = Σ c((a:) mod? 2
(I)
(where is O for even permutations and 1 for odd ones).
All experiments in this paper train transformer language models to solve the word problem: they take as input a sequence of actions (ao. ., ar). and output a sequence of state predictions [sp..... (sp. sy). We also validate our findings on a natural language version of this task in Appendix E, where permutations are expressed instructions like swap positions 2 and 3.
2.2. Permutation Group Word Problems
Toward this understanding, the experiments in this paper focus on one specific state tracking problem, permutation composition. At a high level, this problem presents LMs with set of objects and sequence of reshuffling operations; LMs must then compute the final order of the objects after all reshufflings have been applied (Figure IA). Though less familiar than discourse tracking or program evaluation, Kim & Schuster (2023) used a version of this task to evaluate LM state tracking. More importantly, as shown by Bar- rington (1989) and recently highlighted by Merrill et al. (2024), permutation tracking (with five or more objects) is NC -comllee, meaning any other state tracking task in this family can be converted into permutation tracking class.
2.3. Interpretability Methods
Our experiments employ several interpretability techniques to understand how LMs solve permutation word problems, which we briefly describe below. Throughout this paper, we use h to denote the internal LM representation at token po- sition t after Transformer layer I, with T and L denoting the maximum input lenghh and number of layers respectively.
3

---

(How) Do Language Models Track State?
Probing In probing experiments (Shi al., 2016), we fix
s/- We then identify the hidden states that, when patched in, cause the model to output 8 rather than 8 with high probability. Our main experiments specifically perform prefix patching, where all hidden representations up to in- dex t. (h13.1 + 1 are patched at a particular layer I (Figure IB). Prefix patching allows us to localize how in- formation gets progressively transferred to the final token as we move deeper into the network. A value close to I means that some part of the prefix representation was used for prediction; a value close to 0 means that part was.
the target LM,, then train a smaller "probe" model (e-g.
linear classifier) to map LM hidden representations h to quantities z hypothesized to be encoded by the LM (Fig-
ure IC). Our experiments specifically evaluate whether (1) the state st, and (2) the final state parity is linearly encoded in intermediate-laye representations. For each layer we train (1) a state probe to predict p(s; he,s)and (2) a parity probe to predict p(c((s) hr.1). Given trained LM, we collect representations on one set of input sequences to train the probe, then evaluate probe accuracy on a held-out set.
We also experiment with other types of localization tech- niques (including suffix and window patching), as well as zero-ablating certain activations in Appendix B.
Activation Patching Probing experiments reveal what in- formation is present in an LM's representations, but not
that this information is used by the LM during prediction. Activation patching is a method for determining which rep- resentations play a causal role in prediction. Portions of the LM's internal representations are overwritten ("patched") with representations derived from alternative inputs; if pre- dictions change, we may conclude that the overwritten rep- resentations was used for prediction (Meng et al., 2022; Zhang& & Nanda, 2024; Heimersheim & Nanda, 2024).
3. What Algorithms Can Transformers Implement in Theory?
To use the methods described in $2.3 to interpret model be- havior, we must first establish a phenomenology for LM state tracking-iddeniffingg candidate state tracking algorithms that might be implemented by the model, along with the empirical probing and activation patching results we would expect to find if these algorithms are implemented. Below, we describe a set of state tracking mechanisms suggested by the existing literature.
Let p(y x; h + h!) denote the probability that an LM assigns to the output given an input æ. but with the repre- sentation ht replaced by some other representation h'. In a typical experiment, we first construct "clean" input that we wish to analyze, and "corrupted" input x' that alters or removes information from (e.g. by adding noise or chang- ing its semantics). Next, we compute the most probable outputs from clean and corrupted inputs:
For each mechanism, we first present a sketch of an im- plementation, in the form of rules for computing the value stored in the hidden state for each layer and timestep. We then describe the "signature" of each algorithm- -hhe result we would expect from the application of prefix patching and probing techniques described in the preceding section.
ÿ = arg max p(y æ))
y
V = arg max ply x²)
3.1. Sequential Algorithm
y
The sequential algorithm composes permutations one at a time from left to right (analogous to a mechanism some LMs use to solve multi-hop reasoning problems; Yang et al., 2024). Signatures of this algorithm would provide evidence that LMs implement step-by-step "simulation" in their hid- den states to solve state tracking tasks. In this algorithm, each hidden state hr stores the associated action a, until st can be computed, maintaining h,, st. As shown in the first row of Figure ID, this computation depends only on hidden states with S t.
We then re-run the LM on the corrupted input x', but sub- stitute a hidden representation from the clean input æ, and measure how much prediction shifts toward the clean output y using the normalized logit difference (Wang et al., 2023):
NLD LD(x'; hag +- helkan) LD(x')
LD(x) - LD(x")
where
LD(.) = p 1-) log log 1-)
and the representation helcam is taken from the clean run of the model. A value of NLD close to I indicates that we have restored part of the circuit that computes ý.
hr = 0: Vl
// initialize actions II by definition; see $2.2
(ho.0 = 8t)
for = 1..T.I = 1.. do
In this paper, we evaluate which representations are in- volved in prediction by presenting models with a clean sequence (og, .j...., a.) associated with a final state st- We then produce a corrupted sequence differing only in the first token, (a6: 01,,.............. ai), associated with final state
if 1 < then hy, = h1,1-1 = ag
propagate actions
if 1 = then h41 --1-11-1144p
=s:-104 = 8t
|| update states
if > then h4 =hhhi-11 St.
|| propagate states
end for
4

---

(How) Do Language Models Track State?
Patching Signature Because of this dependency, any patching experiment that replaces only hidden states with I >> will not affect the final model predictions, leading to the upper triangular patching signature shown in the first row of Figure IB.
1u.0 de 11 for = O..T, 1..L do
I/ initialize actions
if S log(t + 1) then
14,1 - 21 -1.111111
= -112'41 1.
I| compose actions
else her ht11111 = St end for
I/ propagate actions
Probing Signature Because 8, can only be predicted at layer/ 1 = f, we expect a state probe to show a linear de- pendence on depth: for sequences of maximum length T, a probe at layer will correctly label an 1/T fraction of states. If these state representations linearly encode parity, then the accuracy of the parity probe will also increase linearly; otherwise, it will remain constant."
(Defining hecoo,t = hoo for notational convenience.)
As seen in the third row of Figure ID, the model's prediction for s depends on the hidden representation hz/2 in the layer before the final state computed, the representation at h4/4 in the layer before that, etc.
3.2. Parallel Algorithm
As noted in $2.2, the word problem on Sz belongs NC² (and thus requires circuit depth that scales logarithmi- cally with sequence length). The word problem on S3, however, belongs to TC", the class of decision problems with constant-depth threshold circuits. See discussion in Ap- pendix A and Merrill & Sabharwal (2023). A constant-depth circuit will give rise to a set of hidden-state dependencies like the second row of Figure ID.
Patching Signature Consequently, for AA, the length of the prefix that must be modified to alter model behavior increases exponentially in depth, resulting in the signature in the third row of Figure 1B.
Probing Signature We similarly expect to see an exponen- tially increasing state probe accuracy (because 8, becomes predictable at layer ( = log t, a probe at layer I will correctly label 2 /T fraction of states). If state parity is encoded in state representations, then parity probe accuracy will also increase exponentially.
Patching Signature Let (p denote the number of layers needed implement the constant-depth circuit for this task. For patching interventions conducted
earlier than layer 1r, we expect the model's predictions change; deeper layers than (p, interventions will have no effect at all, re- sulting in the L-shaped pattern shown in the second row of Figure IB.
3.4. Parity- -astty-Assoiiative Algorithm
In this algorithm (PAA), LMs compute the final state in two stages: first computing the parity of the state (which can be performed in a constant number of layers using a subroutine from the Parallel algorithm); then separately computing the remaining information needed to identify the final state (the "parity complement") using a procedure analogous to AA. (Unlike the preceding algorithms, we are not aware of any previous proposals for solving permutation composition problems in this way; but as we will see, it is useful for understanding interactions between "heuristic" and "algorithmic" solutions in real LMs.)
Probing Signature We expect the probe to obtain perfect accuracy within a constant number of layers. Because the algorithm described in Appendix A computes state parity as an intermediate quantity, the parity probe will also obtain perfect accuracy within a constant number of layers.
3.3. Associative Algorithm
In the associative algorithm (AA), Transformers compose permutations hierarchically: in each layer, adjacent se- quences of permutations are grouped together and their product is computed. This is analogous to recursive sean in Liu et al. (2023) and flattened expression evaluation in Merrill et al. (2024). This algorithm takes advantage of the associative nature of the product of permutations, whereby agajaza, = (aza)) It ensures that h4 = 41-2'11 ... dr and thus that log((+1)+ = 0g- 0g--.-0p -r. Signa- tures of this algorithm would provide evidence that LMs perform state tracking not by encoding states, but rather by mapping between states, for prefixes of increasing length.
We model implementation of PAA with hidden states com- prising two "registers" and k (i.e. he, = (Er,) Kr,)) which store the parity and complement respectively.
Ko,r 0z Vt Co.t par(s:) Vt for = O.TT. = 1..I do
I/ initialize actions I| compute parities (App. A)
E,,1 = 6t,1-1
Il propagate parities
if < log(t + 1) then
K1,1 = comp 11112111111 2011-1111110 16,,11 II compose
else Keg Kt,1-1
// propagate complements
end for
See Appendix C.3 for representative model in which parity not linearly decodable.
In this algorithm, the hidden state at position i holds that
5

---

(How) Do Language Models Track State?
position's is state parity and parity complement (if computed at this point). Parity, like S3, may be computed with con- stant number of layers. The algorithm sketch given above is deliberately vague about the implementation of the par- ity complement composition operation (comp). In practice, different representations of this complement appear to be learned across different runs; see Figure 10 for evidence that these representations are computed using brittle (and perhaps heuristic- memorization-base -aadddd mechanism.
4.1. Experimental Setup
We generate million unique length-100 sequences of per- mutations in both S3 and S We split the data 90/10 for training/analysis, and fine-tune these models (using a cross- entropy loss) to predict the state corresponding to each prefix of each action sequence:
99
L=
log Pin(se| ag ap....)).
where PLM (8x [g)..aa) is the probability the language model places on state token 8n when conditioned on the length-n prefix of the document.
Patching Signature If the corrupted input has a different parity from the clean input, then in layers deeper than those used to compute parity, it is necessary to restore the entire prefix to cause the LM to assign full probability to the clean prediction. On these inputs, prefix patching will show a signature similar to the parallel algorithm (see Figure 8B). However, if the corrupted input has the same parity as the clean input, the portion of the hidden state computed in par- allel remains the same, while its complement is computed using the same mechanism as the associative algorithm (see Figure 8A). These inputs will thus exhibit an AA-like texponetiially-saaeed patching pattern. When averaged together, parity-matcihn and misttthis patching will produce a pattern with two regions, one shaped like the associative algorithm (associated with a 50% restoration in accuracy) and one shaped like the parallel algorithm (associ- ated with 100% restoration in accuracy). Again, this may be most easily understood graphically (Figure 1).
Except where noted, we begin with Pythia-160M, models pre-trained on the Pile dataset (Biderman et al., 2023). Re- gardless of initialization scheme, we fine-tune models for 20 epochs on Equation (3) using the AdamW optimizer with learning rate 5c-5 and batch size 128. For larger models (above 700M parameters), train using bfloat16.
4.2. Activation Patching
For both the S3 and S₃, tasks, across training runs, we find that activation patching results exhibit two broad clusters of behavior. For some trained models, they match the activa- tion patching signature associated with AA; in others, they match the signature of PAA- when the only source of variability across training runs is the order in which data is presented. Results for prototypical AA- and P-y--e models, on both S3 and S5. are shown in Figure 2. Addi- tional patching results in Appendix B confirm that patching intermediate representations of PAA-type models (the light-
Probing Signature We expect state probes to improve exponentially with depth, while parity probes converge to 100% at constant depth.
Pythia $3
Pythia $5
Inget Tniorss
Talons
R
X M R
4. What Mechanisms do Transformers Learn?
In this section, we compare these theoretical state track- ing mechanisms to empirical properties of LMs trained for permutation tasks. It is important emphasize that the vari- ous 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.
AA
PAA
0.0
0.25
0.5 0.3
10
Normalized Logit Difference
Nevertheless, our experiments successfully rule out some possible state tracking mechanisms and identify algorithmic features likely to be shared between the idealized mecha- nisms above and the true behavior learned by transformers. Specifically, our experiments yield evidence consistent with the associative algorithm (AA) in some models and the algorithm (PAA) in other models, across architectures, sizes, and initializations.
Figure 2. Activation patching the residual stream for various Pythia models trained on S3 and S.. Each cell at layer I and token represents the probability of the correct final state when the entire prefix up to t at layer I is restored. We find signatures matching the AA and PAA algorithms from Figure 1, with both models ignoring exponentially longer prefixes as we traverse down the layers, and PAA models containing intermediate representations that encode some information about the final state, but not its parity.
6

---

(How) Do Language Models Track State?
Pythia on S3 (PAA)
Pythia on (AA)
Pythia (PAA)
Pythia (AA)
1.00
1.00
1.00
1.00
State Probe
0.75
0.75
0.75
State Parity Probe
0,50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
eubee
10, (25)
12. res)
(25)
(4. (29)
(5. (25)
(6. res)
(7, (e5)
(8 reg)
(9, (08)
[10. res)
(11, (11,
(0) embed
me
(1. 199
15 r5)
109)
17 109
08 reS)
(9 res)
(10, res)
(1I. rog)
embor
(0. (0S)
(1, (es)
(2. res)
8 mg)
(4 (05)
55 reS)
16 re5)
18 res)
(9, (es)
[10, res)
(11, res)
embey
(0, res)
(1.15)
109
(3, 151
(4 res)
(5 (e))
(6. rg))
(7, (08)
(8 (p))
(9, res)
res)
(11, (os)
12
(4
12
Figure 3. Accuracy of state probe and state parity probe across layers on S3 and S; models sometimes match signatures for AA, and
sometimes PAA. In all models. the state probe accuracy increases roughly exponentially with model depth. We find that in PAA models. the parity of the state is linearly decodable from earlier intermediate layers, while in the AA models shown above, the parity is never linearly encoded in any layer of the model. (In other AA models, the parity can only be linearly decoded at the final layer.)
3D view
Top view (XY)
resentations of states can be cleanly decomposed into two orthogonal parts: the parity of the product and residual cluster identity, forming triangular prism. In Figure 4, we project representations from the PAA model for each of the six states onto these components. Even- peeettyrtty states (darker colors) and odd-parity states (lighter colors) are symmetric. The three cluster "spokes" are spaced 60 degrees apart.?
X (Cluster)
(Aggea Z
Y (Cluster)
Y (Cluster)
X (Cluster)
4.4. Generalization by Sequence Length
We next evaluate the state and parity accuracy of AA- and PAA-type models for held-out inputs of varying length. In general, we find that models learn to generalize perfectly to sequences of up to the length of their training data, then face steep accuracy dropoff after (which we refer to as the "cutoff length"", rather than generalizing uniformly across all sequence lengths.
Figure 4. In models that learn PAA on S3, representations of the final product can be geometrically decomposed into two orthogonal directions, corresponding to the parity of the product (represented as the Z-axis in the above graph) and cluster identity of the product (represented by the X-Y planc). Note that the clusters at 60 degrees to each other, and products of different parities within a cluster are equidistant from each other. with odd add-parity products in one plane, and even-pant even-parity products in another plane.
In Figure 5, where we plot the cutoff lengths at which each accuracy dips below 98%. We find that for models that learn PAA, the parity accuracy cutoff length is much longer than the state accuracy cutoff length, whereas, for models that learn an AA-type mechanism, the two cutoff lengths are equal. Furthermore, models that learn an AA-type mecha- nism tend to generalize better overall.
results specifically in predictions with incorrect parity. We find standard deviations be low in Figure 9, confirming the robustness of our signatures.
colored cells in Figure 2)
4.3. Probing
Test set accuracies of linear probes across LM layers/ are plotted in Figure 3. We report standard deviations of these accuracies in Table which are all less than 10-3 We again find empirical signatures consistent with those pre- dicted by AA and PAA, on both S3 and Sg. Models with AA-type probing signatures always have AA-type patching signatures, and vice versa. Throughout the rest of this paper, we refer to models (and state-tracking mechanisms) as "AA- type" or "PAA-type" based on which cluster of signatures they exhibit. Results in Appendix C break down probe curacies by sequence length, confirming that models solve sequences of exponentially longer length at deeper layers.
4.5. Attention Patterns
We look at attention patterns of LMs and check whether they can be used to differentiate between PAA models and AA models. Specifically, we find that in the carly layers, PAA models exhibit parity heads, heads that place attention to odd-parity actions. Recall that the parity of state can be determined by counting the number of odd-parit actions, and taking the parity of the count (Equation (1)). Examples of the parity head attention pattern are shown in Figure 15, We find no evidence of parity heads in any layer of AA models. (See Appendix for a formal metric measuring how much an attention head behaves like a parity head.)
What exactly is the "non- "non-parity residual" for PAA models? We visualize the linear components of representations near the final layer(s) of PAA models trained on Sz. The rep-
Furrher details, including an analysis of So. can be found in Appendix D.
7

---

(How) Do Language Models Track State?
Pythia-160M on S3 (PAA)
Pythin-160M on 53 (AA)
Pythin-160M on 55 (PAA)
Pythia-160M on S3 (AA)
State Purity
Suie citoo:: 104 Parity cuteft 100
State Purity
State Parity
Stale cotool:: 106 Parit cutell 106
State
Parity
0.8
0.3
0.8
0.8
Suee sutoff, 20.
Party ostoft 102
Accuncy
0.6
Stare cundi:
0.6
0.6
0.6
125 Purty cutll: 10p
0.2
02
0.2
0.2
0.0
0.0
0.0
50
150 Sequaence Length
75
100
175
200
50
75 100 125 150 Sequence Length
175
200
25
50
100 125 150 Sequence Length
175
200
25
50
75 100 125 150 175 200 Sequence Length
Figure 5. Generalization curves showing state and parity prediction accuracy as sequence lengths vary. Models are trained on length- sequences and asked to generalize to varying lengths of sequences. We plot generalization curves for AA and PAA models on Sy and Ss. In each plot, we show the 98% cutoff threshold, the sequence length at which accuracy dips below 98%. In the models that learned PAA, the parity cutoff is larger than the state cutoff, while in models that learned AA, the parity cutoff equals the state cutoff. Generally speaking, models that learned AA generalize better than ones that learned PAA.
Converge on state & parity simultaneously
Training Loss
Model
Pythis on 85 -AA Pythis 55 BAA
Phase 1: Converge on parity
20000
30000
50000
A30000
Trainine Stere
Phase 2: Converge on state
Figure 6. Annolated training curves for models that learn the AA and PAAA algorithms. In PAA models (blue), we find that convergence happens in two phases: in the first phase, they learn to generalize parities up to sequence length 100, and the second, they learn to generalize the states. In AA models (orange), parities and states are learned simultancously. Note that AA models also tend to converge faster to (ultimately) a lower loss than PAA models.
We also find evidence that attention patterns in AA models sparsify in later layers of the network, forming a tree-like pattern expected of AA, shown Figure 16.
determine which mechanism emerges during training.
5.1. When in Training Do Distinct Mechanisms Arise?
We find that an LM's eventual mechanism can be identified very early in training, based on the pattern of prediction er- rors. As in Figure 6, LMs that eventually learn AA improve the quality of their parity and state predictions in lockstep, while LMs that learn PAA learn in two phases: they first converge learning parity over the entire length of the
5. Why do Transformers Learn One Mechanism or Another?
Having determined that trained models consistently exhibit AA- or PiA-like signatures, we next study the factors that
8

---

(How) Do Language Models Track State?
training sequence; and only then do they learn to accurately predict the state itself.
mechanism LMs converge to, we generate length-100 doc- uments with only S3 elements as vocabulary items, and pre-train (randomly initialized) LMs with NTP on these documents, before training them on Sg. Specifically, the documents are generated from a topic model with parame- ters: # of topics = 4, a = 0.3, and = 0.1, where a is the density of topics each document and ,3 is density of words in each topic. As shown in Figure 7, when crom-scratch LMs are trained with our topic modeling NTP objective, they always learn an AA-type mechanism.
Because it is possible to identify these patterns early in train- ing, our subsequent experiments classify LMs as AA-type or PAA-type based on generalization curves (Section 4.4) after 10k training steps, rather than waiting for the full probing and patching signatures to emerge.
5.2. What Factors Affect Which Mechanism Learned?
Whether an LM learns AA or PAA is a deterministic func- tion of four factors: model architecture, size, initialization scheme, and fine-tuning data order. Our next experiments evaluate each of these factors in turn. We explore two dif- ferent model architecture families of various sizes (GPT-2, Radford et al., 2019, and Pythia, Biderman et al., 2023), several different model initializations (pre-trained on the Pile and trained from scratch with different random initial- izations), and up to 12 different data ordering seeds.
Parity Prediction We first train the entire model on pre- dicting the state parity of the sequence to output token if odd and 0 if even, before transitioning to training on the actual Sa objective. In Figure 7, we show that we can induce GPT-2 and Pythia models to learn PAA when trained from scratch on parity. Notably, from-scratch Pythia models already tend to learn PAA as baseline behavior. Therefore, we also apply this curriculum on Pythia models pre-trained on the Pile, and find that it consistently converts the mecha- nism learned from AA-type to PAA-typee
We find that model architecture and initialization, rather than model size, are the biggest determining factors of what mechanism the model chooses to learn. Figure 7 shows the ratio of LMs that learn each mechanism, aggregated by model architecture and initialization. The low variance indicates a minimal effect of model size.* 4 GPT-2 models, pre-trained or not, are split roughly evenly between the two mechanisms, while Pythia models tend to learn AA when pre-trained on the Pile, and PAA when not.
Control: Random Tee-TTTe--Peeiiie we train LMs on length- 100 documents of random S3 ele- ments sampled from a uniform distribution. We confirm that the control fine-tuning did not affect the ratio with which LMs learned each mechanism.
As control,
6. Conclusion
GPT Models
Pythia Models
n676
a-51
24
We have shown that LMs trained on permutation tracking tasks learn one of two distinct mechanisms: one consistent with an "associative algorithm" (AA) that composes action subsequences in parallel across successive layers; and an- other with a "parity-associativ algorithm" (PAA) which first computes shallow parity heuristic in early layers and then computes residual to the parity using an associative procedure. LMs that learn an AA-type mechanism tend to generalize better and converge faster; different choices of model architecture and training scheme encourage the discovery of one mechanism over another.
Proportion of Algoriitme Leamed
1.0
1.0
AA
0.8
47%
40%
0.8
PXA
0.6
0.6
83%
100%
100N
100%
100%
0.4
80%
0.2
49%
53%
0.2
12%
0.-
From scratch
Premainodd
Randon NTP
Topic Modsling
Pariyy
From- Acrotc
Prematood
Randon NTP
Topic Mede Medeling
Parily
Figure 7. Proportion of GPT-2 and Pythia models that learn an AA-type mechanism, PAA-type mechanism, or neither under different training regimes described in Section 5.
While large number of other state tracking tasks can be reduced to the more complex permutation task study (St), our experiments leave open the question of whether the specific mechanisms LMs use to solve S3 are also deployed for these other tasks.
5.3. How Does Pre-training Affect Mechanism Choice?
We show that appropriately designed intermediate tasks can encourage models to learn one mechanism or the other.
Impact Statement
Topic Modeling As a controlled way of studying how the next-tooee--redietion (NTP) objective affects which
The S3 and S₃ tasks we choose study in this paper be generalized to many different state tracking scenarios funda-
'Experimental details can be found in Appendix G. A finer-grained breakdown of the ratio over each model size can he found in Figure 17.
'we discuss another parity curriculum using an extra parity loss term in Appendix H.3.

---

(How) Do Language Models Track State?
mental many aspects of reasoning capabilities. Methods for identifying mechanisms that LMs implement, especially when these differ from human-designed algorithms, can provide crucial insights on how to build more robust LMs, control their behavior, and predict their failures. Our experi- ments focus on small-scale models, and we do not anticipate any immediate ethical considerations associated with our findings.
In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Chen, A., Shwartz-Ziv, R., Cho, K., Leavitt, M. I.., and Saphra, N. Sudden Drops in the Loss: Syntax Acquisi- tion, Phase Transitions, and Simplicity Bias in MLMs. In The Twelfth International Conference on Learning Rep- resentations, 2024. URL https: httpp:://ppnreview net/ forum?id:M05PiKKEER
Acknowledgments
Clark, K., Khandelwal, U., Levy, O., and Manning, C. D. What does BERT look at? an analysis of BERT's at- tention. In Linzen, T., Chrupata, G., Belinkov, Y., and Hupkes, D. (eds.), Proceedings of the 2019 ACL Work- shop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP,, pp. 276-286, Florence, Italy, August 2019, Association for Computational Linguisties. doi: N0.18653/v1/W199 org/w19-455 9--81//
This work was supported by the OpenPhilanthropy foun- dation, the MIT Quest for Intelligence, and the National Science Foundation under grant IIS-2238240. BZL is addi- tionally supported by Clare Boothe Luce fellowship, and JA is supported by Sloan fellowship. This work benefited from many conversations during the Simons Institute Pro- gram on Language Models and Transformers. The authors would also like to thank Reuben Stern, Sebastian Zhu, and Gabe Grand for feedback on drafts of the paper.
URL https://acllanthol.g.
Delétang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wen- liang, L. K., Catt, E., Cundy, C., Hutter, M., Legg, S., Veness, J., and Ortega, P. A. Neural Networks and the Chomsky Hierarchy. In 11h International Conference on Learning Representations, 2023.
References
Barrington, D. A. Bounded-width polynomial-size branching programs recognize exactly those lan- guages in NC1. Journal of Computer and Sys- tem Sciences, 38(1):150-164, 1989. ISSN 0022- 0000. doi:
Heimersheim, S. and Nanda, N. How use and interpret activation patching, arXiv preprint XXiv:2404.15255, 2024.
90037-8. URL httts:://ww..scienneerrree com/ ccinmee/ariccee//////200009999
Hu, M. Y., Chen, A., Saphra, N., and Cho, K. Latent state
models of training dynamics. Transactions on Machine Learning Research, 2023, ISSN 2835-8856. URL https:
Bender, E. M. and Koller, A. Climbing towards NLU: On meaning, form, and understanding in the age of data. In Jurafsky, D., Chai, J., Schluter, N., and Tetreault, J. (eds.), Proceedings of the 58th Annual Meeting of the As- sociation for Computational Linguistics, pp. 5185-511999 Online, July 2020. Association for Computational Lin- guistics. doi: 10. 10.18653/v1/2020..o acl- main.463 URL
/ //openrevie. wwt.:ee//foooo met:/forumidd-NExXXWoe
Huet, A., Houidi, Z. B., and Rossi, D. Episodic memo- ries generation and evaluation benchmark for large lan- guage models. In The Thirteenth International Confer- ence on Learning Representations, 2025, URL https:
https://aclanthologgy wrw.:/loooooogg///// www. acl 1aal-main..450 463/.
jylin04, JackS, Karvonen, A., and Can. OthelloGPT learned a bag of hcuristics, 2024. URL https:
Bhattamishra, S., Ahuja, K., and Goyal, N. On the Ability and Limitations of Transformers to Recognize Formal Languages. In Webber, B., Cohn, T., He, Y., and Liu, Y. Empiri-
//www. //www1leswwoo. www.wrowsoo.....o /ttps//rog...om/oss/////eeeeeee/ ctmm//posts//cpNuEEx/APPe Lewpllggptt-earred--abbbaee bag- of
othellogpt aaaarned-a-baaa
heuristics..
(eds.). Proceedings of the 2020 Conference cal Methods in Natural Language Processing (EMNLP), pp. 7096-7116, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/1//2200. emnlp- 576. URL 2020. wwt.emmmmlmm....0000
Karvonen, A. Emergent World Models and Latent Variable Chess-Playing Conference on Language Modeling, 2024. URL https: //openreview.n wet/fooummm doPummiddPPPrrreEnp
Estimation in Chess- Playing Language Models. In First
org/
Kim, N. and Schuster, S. Entity Tracking in Language Models. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 3835-3855, Toronto, Canada, July 2023. Association for Computational Linguistics.
Biderman, S., Schoelkopf, H., Anthony, Q. G., Bradley, H., O' Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., et al. Pythia: A suite for ana- lyzing large language models across training and scaling.
10

---

(How) Do Language Models Track State?
doi: 10. 10.18665/w11223..c44oogggg URL https:// aclanthology.orgg
Transactions of the Association for Computational Lin-
long. 213/.
guistics, 10:843-856, 2022. doi: 10.1 00493. URL hwtpthtts//ccaaaaeeeht
org. 2022. wtt..ttcc2.......0. .4%.
Li, B. Z., Nye, M., and Andreas, J. Implicit Represen- tations of Meaning in Neural Language Models. In Zong, C., Xia, F., Li, W., and Navigli, R. (eds.), Pro- ceedings of the 59th Annual Meeting of the Associa- tion for Computational Linguistics and the 11th Inter- national Joint Conference on Natural Language Pro- cessing (Volume Long Papers). pp. 1813-1827, On- line, August 2021. Association for Computational Lin- guistics. doi: 10.18653/v1/2011111 www.-oon....lloo..o......oo 143. URL https://aclanthologg org acl-long 143/.
Merrill, W., Petty, J., and Sabharwal, A. The Illusion of State in State-Space Models. In Forty-first International Conference on Machine Learning, 2024. URL https: //openreview. www.:ee///foooooe (one//fooumnddd2200337LLLL
Merullo, J., Eickhoff, C., and Pavlick, E. Circuit Component Reuse Across Tasks in Transformer Language Models. In The Twelfth International Conference Learning Rep- resentations, 2024. URL https:/ /ttts:://peereviex.. /ntte/ forum?id==fooYYVWskk
Li, K., Hopkins, A. K., Bau, D., Viégas, F., Pfister, H., and Wattenberg, M. Emergent World Representations: Ex- ploring a Sequence Model Trained on a Synthetic Task. In The Eleventh International Conference on Learning Rep- resentations, 2023. URL https://oppereviiee net/ ToT
Nanda, N., Lee, A., and Wattenberg, M. Emergent Linear Representations in World Models of Self supervupee Se- quence Models, 2023. URL attps://arxxxt http://aaivi.gg.aaaa/////// org/abs/ 2309.00944
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Transformers Learn Shortcuts to Automata. In The Eleventh International Conference on Learning Represen- tations, 2023. URL https: //tpprv///perreviewe.. wwwpwwweeeefee.... net/forum? HiDDe4FYqjFuel
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T.,, Drain, D., Ganguli, D., HatfieldDDodee
Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context Learning and Induction Heads. Transformer Circuits Thread, 2022, hutps:/rtansfooreeecccciiss http:rrrrruuiiisuuuuu///2iii pub/2022/in- context-learning wwwteniin--ann-iin wet--iidduttooohdddeeeee www.hhhhaddddhiniieeee htmm..
McCoy, R. T., Min, J., and Linzen, T. BERTs of feather do not generalize together: Large variability in general- ization across models with similar test set performance. In Alishahi, A., Belinkov, Y., Chrupala, G., Hupkes, D., Pinter, Y., and Sajjad, H. (eds.), Proceedings of the Third BlackboxNLP Workshop on Analyzing and Interpreting Neural Nenworks for NLP, pp. pp. 217-227, Online, Novem- ber 2020. Association for Computational Linguistics. doi: 10. 11111866// 1/2.2/1112222.2.// 1/000///b////.bb.booobooo// 1.21. URL https: //aclanthology org/ www..bboolllllllc.co.
Prakash, N., Shaham, T. R., Haklay, T., Belinkov, Y., and Bau, D. Fine-Tuuiig Enhances Existing Mechanisms: A Case Study on Entity Tracking, In Proceedings of the 2024 International Conference on Learning Representa- tions, 2024. arXiv:2402.1 14811.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Meng, K., Bau, D., Andonian, A., and Belinkov, Y. Locating and editing factual associations in GPT. Advances in Neu- ral Information Processing Systems, 15:1775--7737, 2022.
Shi, X., Padhi, I., and Knight, K. Does string-based neural
MT learn source syntax? In Su, J., Duh, K., and Carreras, X. (eds.), Proceedings of the 2016 Conference on Empiri- cal Methods in Natural Language Processing, pp. 1526- 1534, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/v1/D16.1 1159. URL https:::::/ htttps://aalanttholegs www.ey...rr//eoe orgge//-or//0111999900
Merrill, W. and Sabharwal, A. The Parallelism Tradeoff: Limitations of Log-Precision Transformers. Transactions of the Association for Computational Linguistics, 11:531- URL https:
545, 2023. doi: 1622ttcll_a00005 //aclanthology. . wrg./023. tacl-1.31/.
Merrill, W. and Sabharwal, A. A Little Depth Goes Long Way: The Expressive Power of Log-Depth Transformers. In NeurlPS 2024 Workshop on Mathematics of Modern Machine Learning, 2024.
Strobl, L., Merrill, W., Weiss, G., Chiang, D., and An- gluin, D. What Formal Languages Can Transformers Express? A Survey. Transactions of the Association for Computational Linguistics, 12:543-561 2024. doi: 10.1162/tacl lat.taclla0006666 org/2024. 2/g/2024.tacl-11 tacl-1.30.
Merrill, W., Sabharwal, A., and Smith, N. A. Saturated Transformers are Constant-Depth Threshold Circuits.
URL https: pektps://aclaan h/tpa//tlaaaaooo.....

---

(How) Do Language Models Track State?
Vafa, K., Chen, Y., Rambachan, A., Kleinberg, J., and Mullainathan, S. Evaluating the World Model Im- plicit in Generative Model. In The Tiihth-eighth An- mual Conference on Neural Information Processing Sys- tems, 2024. URL. https://openreviiee ner/forum? id=avK4JFpegy.
Wang, K. R., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Represen- tations, 2023, URL https: / /pwpr/vveeeeeeeeeeeeeeeeeeeeeeee net/forum? id=NpsVSN6o4ul.
Yang. S., Gribovskaya, E., Kassner, N., Geva, M., and Riedel, S. Do Large Language Models Latently Per- form Multi-Hop Reasoning? Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Compulational Linguistics (Volume 1: Long Papers), pp. 10210-10229, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10. 100118663/11/2222..allllooog 550. URL https://ac.antthooog...gg// http://oooooooorr///.....aaaaa.a acl- 550/.
Zhang. F. and Nanda, N. Towards best practices of activation patching in language models: Metrics and methods. In The Twelfth International Conference on Learning Rep- resentations, 2024. URL https://openreview.nnn net// forum?id-Hf11yyUBBC.
Zhong, 7., Liu, 7., Tegmark, M., and Andreas, J. The clock and the pizza: Two stories in mechanistic explanation of neural networks. Advances in Neural Information Processing Systems, 36, 2024.
12

---

(How) Do Language Models Track State?
A. A Constant-Deptl Algorithm Exists for Sy
S3 is the smallest non-solvable permutation group. Sa is isomorphic to D3, the symmetry group of an equilateral triangle, which can be generated by a transposition a (23) and 3-cycle b = (123). These generators satisfy the relation ab ba- which allows any word problem in S3 be reduced by tracking (1) the cumulative parity of transpositions and (2) the count of 3-cycles modulo 3. Since both parity checking and modular counting can be computed using constant-depth threshold circuits, the word problem for S3 belongs to TC".
B. Full Activation Patching Results
In the activation patching experiments, we overwrite ("patch") portions of the LM's internal representations and compute how much the resulting logits have changed. As discussed in Section 2.3, we perform prefix patching to localize important token positions. However, in addition to prefix patching, we also explore the following types of localization methods:
1. In suffix patching, all tokens starting from up to one before the last token of the sequence (he:]A)-1.1) are patched at a particular layer I.?
2. In window patching, all tokens in w-width window starting from (bot+--11,) are patched at layer 1.
For each of the above localization techniques, we patch the representation with several different types of content:
1. In representation deletion, we overwrite target representation(s) entirely with a zero vector,
h1 0
and measure the NL. as follows:
NLD = LD(a1... ... ..) LD(LD...... ac he; + 0)
LD(aj...a,
2. In representation substitution, we overwrite the representation(s) with those derived from running the LM on a minimally different (corrupted) representation Paompt. This is the setting described in Section 2.3.
Full results are shown in Figure 8. In general, we discover the following:
In PAA models, parities are computed in parallel in early-mid layers We use prefix substitution patching described in Section 2.3, but plot pairs that have same parity final states (c(ō)= = c()²)) separately from pairs that have opposte parity final states (c(T)= * (()"))
Results are shown in Figure 8A (for same parity) and 8B (for opposite parity). We find that in AA models, parities are computed with the state- with both the parity and opposite-pariee paaitte-ppttyy patching patterns displaying the same exponential curve. However, in PAA models, the patching patterns for same- and opposite-paritic differ drastically. When parities are the same, only the parity complement must be computed to infer the final state; the patching pattern in this case indicates that the parity complement is computed in an associative manner. When the parities are different, the patching signature has component that resembles a parallel patching signature, which is where the parity is computed. Restoring prefixes of layers before the parity is computed results in the entire prediction being of the correct parity, while restoring prefixes of layers after that results in the entire prediction being of the incorrect parity. We see that parities are computed roughly in parallel at early layers (around layers 3-5). Note that there is a middle region where restoring the prefixes shifts the logits towards the correct prediction, but not 100%: when prefixes in these regions are restored, the LM does not know the parity of the final answer, but does know some aspects of the parity complement, which was computed in an associative manner.
httpps: / http:////ooowwiii..ommo .org/ org/wwkii http///i///////iiiiieee ceatrricc_rrpopoooo @roup_on33__eett_ seLeeteersiie Isooorrhhicccc sichicctttt tttottttt_o___aaoe__e Diheeral wwt.edrallGrrou..D.. We do not patch the last token as the last token residual contains the current accumulated information necessary for computing the final product, and almost always will destroy the prediction when patched, making this method uninformative as localization tool. We wish to see what other tokens the final token uses when constructing its representation.
13

---

(How) Do Language Models Track State?
Prefix Substitution Patching
Prefix Substitution Patching
A
B
(Same Parity)
(Opposite Parity)
S3
S3
S5
Ingut Tokens
Inour Tokess
Inge Tokens
Expal Tidans X
Layss
AA
PAA
Layos
Laves
C
Suffix Deletion Patching
D
Window Deletion Patching
S3
SS
$3
Iepat
Iolons
R. R R 12
AA
Layos
Layos
PAA
Layox
0.0
0.25
0.5
0.75
1.0
Normalized Logit Difference
Figure& 8. Activation patching results across different types of localization (prefix, suffix, window patching) and different types of patching content (substitution, deletion). (A): Prefix-substiuutio results on only sequences with the same parity. We see the same exponential patching pattern in both AA and PAA models, showing that parity complements are computed in the same associative manner. (B): suddi-suustitttionn results on only sequences with opposite parities. in AA models, parities are computed with the state. In PAA models. however, the patching pattern is roughly parallel, meaning parities are computed roughly in parallel in early layers. (C): Suffix-deletinn results show that we can ignore progressively longer sequences of suffixes as we go down the layers of the network, consistent with how we believe AA and PAA work. (D): Window-deletion results show that important activations are arranged hicrarchically, again consistent with how we believe AA and PAA work.
see the exponential patching pattern in AA models, meaning that
Increasingly longer suffixes are ignored for AA and PAA models in later layers In Figure SC, we show suffix deletion patching results, finding that we can swap out exponentially longer both AA and PAA models without affecting the prediction. This is in line with how the associative algorithm in either model works: suffixes of progressively longer lengths are collected into the final token as we go down the layers.
Important activations are arranged hierarchically In Figure 8D, we show window deletion patching results, with a window size of 1. We find a patching pattern consistent with the associative algorithm: deleting any single token in the early layers is extremely important, but the spacing of important tokens gets sparser as we go down the layers, consistent with the depiction of AA/PAA in Figure 1. At the bottom layers, deleting any single token is unimportant for the final computation of the state.
Patching signatures are relatively consistent across examples In Figure 9, we plot the standard deviations across 200 pairs of S3 inputs for the following three sets of results: (A) prefix substitution patching of AA models, (B) prefix
14

---

(How) Do Language Models Track State?
All Inputs
Input Tokens 2
gmbed
0
8
1.00
(1
125
12x 1081
AA algorithm
Layens
FO5)
(05)
Pos)
-0.00
Same Parity
Opposite Parity
Inpul Takens
Inpul Tokens
9 4
X
2
8
8
8
R
X
suube
1.00
0.75
Layers
PAA algorithm
0.25
0000
Figure 9. Standard deviations of the prefix substitution activation patching results across 200 S3 input pairs, for (A) AA models, (B) PAA models on pairs with the same parity, and (C) PAA models on pairs with opposite parity. We find generally low standard deviations across examples.
substitution patching of PAA models (opposite parity inputs), (C) prefix substitution patching of PAA models (same parity inputs). We find relatively low standard deviations in all three cases, showing that these signatures hold across different examples.
C. Full Probing Results
C.1. Probing signatures are relatively consistent across examples
We investigate the sensitivity of our probing signatures to the data on which the probe was trained. We focus on Pythia models trained on S3. For each model type (PAA vs. AA) and each probe type (parity vs. state probe), we train 10 different probes across 10 different random subsets of the S3 dataset, and report the standard deviations of their accuracies across the 10 runs. Results can be found in Table 1. We find low standard deviations (less than 10 10-3 3) in all four cases, indicating that our probe signatures are robust to different subsets of the data and randomness in probe training.
C.2. Probe accuracies over sequence lengths
How do the probe accuracies in Figure 3 decompose over sequence lengths? We sweep over S3 sequences Q1 a, of lengths ranging from = 5 to 100, and train a linear prohe that takes in input /4,//a ... Gt the layer-l, position-I representation of and aims to predict the final state s, from the hidden representation.
the model on input sequence 01 ........ a; with
The mean probability the probe put the correct answer is plotted in Figure 10. We find that, generally speaking, both AA and PAA linearly encode states of exponentially longer sequences as they go down the layers. We find evidence that the PAA models use their intermediate layers to compute parity in parallel: at around the second residual layer, PAA models place! 3 probability on the correct answer (there are three actions of each parity in Sz).
C.3. Examples of Associative Algorithm Representations that Do or Do Not Linearly Encode Parity
As shown in Figure 1, models that learn AA sometimes encode parity linearly at the final layer but sometimes do not. The examples shown in Figure3 3 all do not linearly encode parity at the final layer. We show a 3D visualization of the S3 AA model's final hidden representations along the three principal components of the representation (which explain 41.8% the variance the data) in Figure 11. As we can see, parity not linearly encoded the final layer. Figure 12, we show
15

---

(How) Do Language Models Track State?
Layer
Pythia on S3 (PAA)
Pythia on S3 (AA)
State Probe
Parity Probe
State Probe
Parity Probe
0
4.16
10-6
3.03
10
6.33 10-6
2.71 10-6
1
3.36
10
3.31
10
2.62 x 10-6
2.82 x 10
2
4.52
10-6
5.45 x 10-6
1.46 x 10-6
8.97 x 10
3
4.31
10-5
5.59 x 10-4
5.00 x 10
1.38 x 10
4
1.58
x 10-3
3.72 10-6
4.16 x 10-6
1.11 x 10-6
5
1,70
10-3
5.38 x 10
1.97 x 10
1.11 x 10
1.87
10
7.02 10-6
4.73 x 10-6
1.59 x 10-6
7
1.29
10
1.34 10
7.79 x 10-6
3.05 10 6
8
3.18
10-5
3.08 10-6
1.18 x 10-5
2.64 10-6
9
2.96
10
6.85 10-5
2.36 x 10-5
3.53 x 10
10
3.10
10
7.98 x 10
2.60 x 10
2.96 x 10
11
1,86
10
8.22 x 10
8.83 x 10
2.77 X 10
12
7.05
X 10 10-5
4.01 x 10 1.20 x 10
2.75 x 10
Table Standard deviations across probe accuracies. We focus analysis S3 Pythia models and train 10 probes different subsets of the S3 dataset. Standard deviations are tiny in all cases, indicating robust signatures.
Probe Accuracy by Lover and Sequence Leneth
1.0
mg
AA (S3)
0.6
R33
0.4
0.0
Seqpesse Lengih
Probe Accuracy by Laver and Sequence Length
uubel
0.8
18: Rd
PAA (S3)
13
0.6
0.4
30830
0.2
51 Length
Figure We plot the average accuracy of linear probe trained predict final-token hidden representation of AA and PAA models on A. We find that both types of models can handle longer sequence lengths as we go down the network, and that PAA models compute the parities of sequenees at roughly layer 2, after which they can get the parity of the state correct but not the exact state.
state an action sequence A, given the corresponding
final-layer hidden representations from AA model that does linearly encode the parity (from- GPT2-base Sy). When projected onto three components that explain 49.9% of the variance in the data, we find clear linear separation between the odd and even parity representations.
D. Full Linear Decomposition Results
D.I. S3
We visualize the linear decomposition of the last-layer penultimate-layer representations across various PAA models. We find the triangular prism shape similar to Figure in all of them, but there no consistency in which states paired to form the clusters.
One interpretation is that PAA models may be learning various presentations of S3. with each clustering configuration corresponding to a different presentation. Generally speaking. S, can be generated by a 2-cycle and an n-cycle: any permutation of Sn can be created by composing these two permutations. For example, S3 can be generated by the 2-cycle 2 and 3-cycle 2 + 1, which corresponds to the clustering {(123, 213), (312, 132), (21,,21) the states
16

---

(How) Do Language Models Track State?
3D view
Top view (XY)
Component
(EV-21.4%)
Componen
(EV (EV-10.7%)
Compones
(BVV-2 (BV-21144 21,49)
Component
(EV-10.7%)
Side view (XZ)
Front view (Y2)
Component 45%)
(EV-10.7%0
Figure 11. Example activations from AA model that does not linearly encodes parity at the final layer, projected on three principal components with total explained variance 41.8%. Blue points have even parity, while orange points have odd parity.
3D view
Top view (XY)
Parity (LV-2113%)
Closer
Perty 30g
State Probe
(EV-12.75%
0.75
Parity Probe
Side view (X2)
Front view (YZ)
Oi Adl
Chater LEV 16.155
0.25
125
(1), res)
10.
5
(10,
Paricy (EV-21.3%)
Clister (EV-17.7%)
Figure 12. Example activations from an AA model that does linearly encode parity at final layer. Blue points have even parity, while orange points have odd parity. (Left) State and state parity probe signatures of this model. (Right) projection of hidden representations onto three components with a total explained variance of 49.9%.
within the cluster can be transformed into each other by applying ++ 2, while states between clusters are related to each other by + + 3 + 1. PAA models that cluster according to this pattern may have learned these generators.
D.2. S;
What happens models that learn PAA in Sg? We visualize the penultimate-layer representation of a Pothia-160Mg model that learned PAA on S3 in 3D space, with parity along one axis and two orthogonal directions along the other two. We find distinct clusters, corresponding to the position of the state (states having position and in position are clustered together).
E. Simulating State Tracking in Natural Language
To emulate a more practical scenario, we train pre-trained and from-scratch Pythia models on a version of the S3 permutation composition task expressed in natural language. For example, the permutation "132" would be expressed as "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, the following sequence:
17

---

(How) Do Language Models Track State?
3D view
Top view (XY)
view (XZ)
Front view (YZ)
X (Crroup)
(8 Z
(Group)
2 iPPanng
Z deesaog
Y (Group)
X (Group)
X (Group)
(Group)
Figure 13. Projecting models that learn PAA on Ss into 3D space. Unlike PAA models on S3 (Figure 4), the state cannot be fully represented by a clean decomposition into directions (notice the colors superimposed on each other). However, we do still find symmetry across the parity axis, similar Sg. Moreover, there are 4 neat clusters, one the center, and three outward protruding "prongs". We find that the clusters correspond to the position in the state.
Swap positions 2 and 3. Rotate the last item to the front Swap positions and 2.
would map to permutation sequence "132, 312, 213" and finally correspond to the state "123" after the swaps, We then conduct a similar style of probing and activation patching experiments on models trained on this task.
E.1. Probing experiments
We train probes to map from the activation of each layer at the position of the final "." token to the final state. Shown in Figure 14, as in our results on synthetic data, probing results are consistent with associative mechanisms - the state probe improves exponentially over layers. In pre-trained Pythia, we see the model represent the state parity much earlier (in terms of layer) than the actual state representation, a signature of the PAA algorithm. For ce--pr-ttriined Pythia, the accuracy of the state and state parity probes increases at similar rate, indicating that it more likely learning an associative algorithm: whether is the strict AA algorithm identified in the synthetic case is unclear.
E.2. Activation patching experiments
We patch prefixes up fixed token position. As shown in Figure 14, both models display distinctly associative signature with exponentially longer prefixes being disregarded for the final state prediction depth increases. Furthermore, the pre-trained Pythia model possesses a light blue middle section sign of the PAA algorithm. Interestingly, the pre-trained Pythia 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 L.M taking advantage of its innate natural language understanding (and perhaps pre-trained state tracking abilities!) to quickly solve the task in an early layer.
F. Full Attention Heads Analysis
F.1. Formalizing Parity Heads
To formalize metric for whether attention head behaves like parity head, define parity head score as the percentage of sequence lengths (ranging from 5 to 80) over which the head places significantly more attention on odd-parity permutations than even- veenttarity permutations, measuring significance using a 95% confidence interval.
Definition F.1. Let (H) a,E (x) be the attention weight of the /tth attention head in layer l at position for input where ar is the list of actions [a]...... J.).
Note that token positions may no longer he aligned in natural language: "swap" actions have 5 tokens, and "rotate" actions have 7. Thus, may not be replacing the activations the number of actions between prompts. Nonetheless, still believe the activation patching results serve as a good proxy for estimating how information gets propagated through the layers.
18

---

(How)) Do Language Models Track State?
Activation Patching
Probing
Probe Accuracy Across Layers
Input Tokens
25
50
75
100
123
150
117
200
225
220
275
300
1.0
1.00
Logi Differre Iffeeence
0.8
Pretrained Pythia-160M
0.75
Accuracy
Layers
0.6
0.50
0.25
0.4
0.00
0.2
state state_parity
0.0
2
Input Tokens
225 G
Layer
0
25
50
75
100
125
150
175
200
250
275
300
1.00
1.0
state
state_ state_parity
0.75
Logi Differeeeennneeeeeeeeeee
0.8
Non-pretrained
Layers
Pythia-160M
0.50
0.6
0.25
0.4
0.00
0.2
Layer
Figure Patching probing result Pythia models trained on natural language permutation composition task. plot signature of pre-trained Pythia model (top) and non-pre-trained Pythia model (bottom). In both cases, the signatures consistent with the state heing learned associatively (both the patching signature and state prohe have an exponential curve). The signature the pre-trained model is consistent with PAA signature, with the parity probe converging in early layers, and the activation patching signature containing light-blue middle section.
Define the sets attention weights on odd and even tokens of x as:
Ae, (x:) {a 1, (x) as is odd),
Ae, / (x) a, is even).
We find heads that respond to parity can be local: for example, attention head 5 in layer responds to odd-parity permutations in the midpoint of the sequence, while attention head 4 responds to ones late in the sequence.
Thus, we record the parity head score Jf.H.parity: which measures the proportion the sequence for which more attention placed on odd-parity actions compared even-parity actions:
Hariit
H.rt.parity.d
where
L
1=3
H = 1 HssLA,, ... a;l)] 0.95 Cle,H,oaulla...... ... ail) B[Ae, E[Ae, H.,vea
((aa...a)))))
Here, [ai...a]] a.] denotes the first elements of the sequence æ, and Hdd refers to the confidence interval around the
average attention weights on udd-parity tokens. We use sequence lengths of up to L = 80 for S3 and L = 50 for Sç.
We find no evidence of parity heads in any layer of AA models. However, we find at least two attention heads with o significantly exceeding 50% in the first few layers of PAA models, highlighted Table 2.
F.2. AA Attention Patterns
What sorts of attention patterns appear in AA models? Because attention is dense, we visualize only the top-K attention traces from and each position each layer. Specifically, we plot the attentions of LM on an input as a graph with:
19

---

(How) Do Language Models Track State?
Layer 1, Head 4
Layer 1, Head 5
Layer 1, Head 11
10
10
10
Targer Tokens
20
20
ML 11
20
0.8
30
1 MILL
0.6
MALL
40
40
50
50
50
0.4
60
60
60
0.2
70
41
80
80
80
0
20
40
20
40
20
40
60
Source Tokens
Source Tokens
Source Tokens
Figure 15. Examples of parity heads in PAA model. We plot heatmaps showing attention weights on each source token at each target
token location (we show up to only 60 source tokens for heads and for the sake of space). We draw arrows yellow lines at source tokens corresponding to odd-pan parity actions. Note that parity heads attend almost exclusively to odd-parik tokens.
Algorithm
Layer
Head
Parity Head Score
Sa (PAA)
90.1 90.1% 13.0%
86.4% 8..2%
0
67.1 14..1%
Sa (AA)
0
3.3%-+4455%
2
3.0%++..1%
11
2.0%
Ss (PAA)
3
83.6%
3
80.6%
2
50.3%+222.45
S: (AA)
0
3.9%
3.8% 15.4%
0
3. 3.2%
Table 2. Top-3 parity head scores across all attention heads in each type of model. We report average parity head scores (%) over 100 examples, as well their standard deviations. Informally, this metrie captures the proportion of sequence over which more attention placed on odd-parity actions than even-parity actions.
20

---

(How) Do Language Models Track State?
10
6
3
Embed
8
6
5
3
2
Embed
11
8
7
6
5
3
2
0
Embed
Figure Attention patterns AA models a tree-like pattern, with tokens successive layers attending to larger windows downstream tokens. This is in line with how expect the associative algorithm to function. We plot only the most salient attention weights by pruning odges for which the attention weight < 0.95, the target token not in the top-3 attended-to tokens from source token, the source token is not in the top-10 10 attended- ffom tokens for the target token. We expect that the attention patterns visualized here do not form a single clean tree, but the superimposition of multiple trees.
1. Nodes (), for each token position and layer I
2. Edges between two nodes (t1, / - 1) and (t2, () if position t2 at layer attends to position at layer 1. We define "attends to" as follows:
Let dii denote the maximum attention weight (across all attention heads at layer () from position at layer / to position to at layer I.
We say (12, !) attends (4,, 1-.1-11 1) all of the following conditions are met:
21

---

(How) Do Language Models Track State?
(a) !1)
-t 0.95
(b) (1)
11-22 top-3(1a :E [1,]]]): & is among the top-3 attended-to tokens for l2 at layer /
(c) t1 -2 E top-10(fa k E [1. })]): to is among the top-10 attended- frndedd-foom tokens for t at layer I
(1)
We show example attention patterns for an AA model on three sample prompts Figure 16. We only plot the attention subgraph directly connected the final token position at the final layer (which used to predict the state). We find that attention in AA models forms tree-like pattern where successive layers attend wider and wider context windows, with nodes that more and more spaced apart. This is in line with how we believe the associative algorithm works: adjacent pairs of actions are grouped together at each layer in hierarchical manner.
Note that the tree is not entirely clean: there are redundant edges and edges that cross over each other. We suspect that the subgraph we' ve picked up on is not a single tree, but rather the superimposition of multiple trees, each potentially contributing to not just the prediction for the final token, but also the predictions of the previous tokens.
G. How are these algorithms learned over the course of training?
We conduct a more detailed analysis of the training phases for two Pythia models trained on the S5 task: one that learned AA and one that learned PAA. The training curves are shown in Figure 6, and we investigate the generalization behavior at different points along these curves. Both models improve over training by progressively generalizing to longer sequence lengths, rather than making uniform gains across all lengths. In the case of the PAA model, convergence appears to occur in nogthh100~a two distinet phases: first, the model learns the parity of states across the entire length- 100 sequence. followed by learning how to predict the state. By contrast, the AA model learns generalize parity and state simultaneously.
H. Additional Factors Influencing Learned Algorithm
H.1. Model Size
We have investigated whether model size influences which algorithm the models learn to implement. As shown in Figure 17, model size empirically does not seem to have much effect on the choice of learned algorithm, with model architecture and initialization having much bigger effect. Some notable differences between the GPT-2 and Pythia architecture are the use of rotary embeddings, parallelized attention, and feedforward layers rather than sequential, and untied embedding and unembedding.
H.2. Topic Modeling
The topic model used in Section 5.3 is parameterized as follows: each document using random Dirichlet distribution with 0.3. The distribution p(token topic) for each token 123, 132, 213, 231, 312, 321 is:
We generate the distribution of the topics in
3.06 10
1.11
10-1
5.79-10-1
6.45 10-3
6.58
10
8.45-10-1
3.36 10-1
1.69 10-4
10-7
1.68
1068.10-
7.81
7.92 10
1.41
10-2
9.44
10-1
40-4
4.13 103-10- 10-2.
3.27
10
2.85 10-3
1.29 10-9.
7.06
10-1, 10-7,
2.58
10-3
2.89
10
We also trained LMs using topic model with a second token-topic distribution, aiming to distinguish the effect of this particular topic distribution from the effect topic modeling pretraining in general. On the second distribution, we also find that both randomly initialized GPT-2 and Pythia models learn AA Figure 18. The p(token topic) distribution for this model is listed below:
9.31 10-1
2.32
10-3
1.38
10-8
5.86
10-10
4.62
10-5
6.63
10-27
2.10-10-
10-7
9.07
2.53
10-1
4.95
10-1
1.18
10-3
4.55
10
2.17-
1.86
10
2.71
6.55
10
4.92
10
3.44
10
2.28 1
1.94 10-4
2.14-10-5
22

---

(How) Do Language Models Track State?
GPT Models (From Scratch)
GPT Models (Pretrained)
n-11
n-11
n-13
1h-3
1.0
In-12
n-32
a-20
n-12
45%
2756
31%
No convengeice
42%
30%
No contergence
59%
18%
0.6
0.4
27%
33%
62%
20%
70%
0.2
41%
27%
wwworrooooflggrriiiie. Leamndd
0.0
0.0
GPT
GPTZ-lerge
GPT2 mediumm
OPPT2-maa --snall
GPT2-1
GPT2-tagge
GPTZ matlinn
GPT2-amall
Preporice of Algoriihms
Pythia Models (From Scratch)
Pythia Models (Pretrained)
e-11
m-11
9-11
n-36
1.0-
myl1
m-14
n=13
AA
convengetce
converpence
PAA
0.8
36%
PAA
0.6
8256
8216
89%
0.6
100%
92%
100%
0.4
64%
0.2
0.2
18%
18%
Pythia:2.80 2,89
Pytha-10 1B
Pythia 10M
Pythia. 1e0kt
Pythls-2.90 80
Pythi- BB
Pythia 410M
Pythia- 160M
Model
Model
Figure 17. Proportion of S3 algorithms learned by models of various sizes from the GPT-2 and Pythia families. Model architecture and initialization are much bigger factors in influencing the algorithms learned than model size.
H.3. Parity Loss Curriculum
We explored an additional procedure that encourages models to learn parity via an extra loss term. We train an extra linear classifier that takes in the residual activations of an early layer (e.g., layer 3) in which we find parity to be typically computed through probing. We train with this additional loss term induce the representation to linearly encode parity early on. The classifier is trained to output (1) parity or (2) (parity, action) tuple, ensure that the residual also encodes the original action, and not just the parity. After training, we evaluate the model the original S3 task. Both procedures induce the model to learn PAA consistently, though the (parity, action) classifier typically allows the model to generalize better.
H.4. Length Curriculum
We implemented curriculum training approach where the model progressively exposed documents of increasing length. We first trained only the initial 10 tokens, then expanded 25 tokens, 50 tokens, and finally the complete 100-token sequences. Each stage of the curriculum was trained for a fixed number of epochs (data). The goal of such curriculum to push the model to learn the associative algorithm (AA), as the parity heuristic, we hypothesize, might be less useful for shorter sequences. However, empirically, such a curriculum has no obvious effect on the kind of algorithm the model learns. Of the trials, trials of GPT-2 learn AA, and the other trials learn PAA.
We discovered that all five models can perfectly generalize both parity and state when trained on the first 25 tokens. It only when trained to generalize to sequence length 50 that the distinction between PAA and AA emerges. Models learning the PAA algorithm do not generalize, while models learning AA can generalize from length 25 to 50, illustrated by Figure 19, This further confirms our finding that the model learns the S3 algorithm early on.
23

---

(How) Do Language Models Track State?
GPT Models
Pythia Models
D-
n-60
1-15
D-69
n-
n-24
1.0
1.0
Proportion of Algorithms Learned
AA
0.8
44%
0.8
convergence
PAA
0.6
0.6
88%
100%
100%
100%
100%
0.4
49%
0.2
0.2
12%
0.0
0.0
From-scratch
Topic (Primary)
Topic (Secondary)
From-scratch
Topic (Primary)
Topic (Secondary)
Figure Proportion S3 algorithms learned by models first pre-trained on the two topic models compared to randomly initialized baselines. Topic modeling, regardless of the specific distributions, pushes the model to learn AA.
Truined Tokens (AA)
Trained Takens (AA)
Trained Tokens (AA)
Truined 100 Tokens (AA)
1.0
State
1.0
State
Parily
Partty CURATT
Parily
Parily
Party Parily sotul
Parity
State sanut 43. Partyy
Accunicy
0.6
0.6
hypp
0.4
0.4
Stale Sulent 25
0.4
0.4
State sulell 65
0.2
0.2
0.2
My Www.haa
0.2
0.0
0.0
100
125
150 175
25
50
75
125
150
175
200
200
200
Sequence Length
Sequence Length
Sequence Length
Sequence Length
Trained (PA.)
Trained
(PAA)
Trained (PNA)
(PAA)
1.0
State Parity
1.0
State
Parity
Parity
0.8
0.8
Parily cunoit
0.8
0.8
Seaee coset 28
Stone suutott 30
Cilyy cutoll 97
Accuray
0.6
0.6
0.6
cubbll:
0.6
M. Mww.hiiiy
0.4
0.4
0.4
0.4
Stite qutatt
0.2
0.2
www.
25
50
125 Sequence Length
75
100
150
175
200
25
50
125 Sequence Length
75
100
150
175
300
100 125 150 175 200 Sequence Lenath
125 150 175 300
Sequence Length
Figure 19. Model generalization curves when trained using length curriculum after training 10, 25, 50, and 100 tokens, respectively. While the length curriculum doesn't it push the model to learn one algorithm or the other, it shows that the model learns these algorithms early indicated by whether the model generalize well from training 25 tokens 50 tokens.
24


Reviews:
[{"review_id": "0bLReDBHpi", "reviewer": "Reviewer_jeEg", "summary": "It is known that transformers are theoretically able to capture certain formal language tasks of length $n$ with depth $\\log(n)$. Empirically, large language models, which are primarily based on transformers, do appear to learn to state track. A full understanding of the mechanism that they learn for state-tracking is still missing. This paper analyses language models trained on permutation word problems (which for a dictionary of greater than 5 is $NC^1$-complete).\n\nFour possible mechanisms are hypothesised, along with what patterns would be expected from prefix patching and probing experiments. This paper shows that the language models tested consistently learn one of two state tracking mechanisms for this task. Generalisation properties of these mechanisms are explored and how to steer the network to converge from one mechanism to the other.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n## Unclear if signatures are sufficient conditions. Would an alternate scheme give similar signatures?\n\nYou 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:\n\n> 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.\n\n## Sensitivity analysis for linear probe: how would varying the probe affect results\n\nWe trained the probe across 10 different random data subsets on the S3 AA and PAA models, and measured the standard deviations across these runs.\n\n| Layer | S3 Pythia-160M (PAA) State Std | S3 Pythia-160M (PAA) Parity Std | Pythia-160M (AA) State Std | Pythia-160M (AA) Parity Std |\n|-------|-------------------------------|----------------------------------|----------------------------|-----------------------------|\n| 0 \t| 4.16e-06                  \t| 3.03e-06                     \t| 6.33e-06               \t| 2.71e-06                \t|\n| 1 \t| 3.36e-06                  \t| 3.31e-06                     \t| 2.62e-06               \t| 2.82e-06                \t|\n| 2 \t| 4.52e-06                  \t| 5.45e-06                     \t| 1.46e-06               \t| 8.97e-07                \t|\n| 3 \t| 4.31e-05                  \t| 5.59e-04                     \t| 5.00e-06               \t| 1.38e-06                \t|\n| 4 \t| 1.58e-05                  \t| 3.72e-05                     \t| 4.16e-06               \t| 1.11e-06                \t|\n| 5 \t| 1.70e-05                  \t| 5.38e-06                     \t| 1.97e-06               \t| 1.11e-06                \t|\n| 6 \t| 1.87e-05                  \t| 7.02e-06                     \t| 4.73e-06               \t| 1.59e-06                \t|\n| 7 \t| 1.29e-05                  \t| 1.34e-05                     \t| 7.79e-06               \t| 3.05e-06                \t|\n| 8 \t| 3.18e-05                  \t| 3.08e-05                     \t| 1.18e-05               \t| 2.64e-06                \t|\n| 9 \t| 2.96e-04                  \t| 6.85e-05                     \t| 2.36e-05               \t| 3.53e-06                \t|\n| 10\t| 3.10e-04                  \t| 7.98e-05                     \t| 2.60e-05               \t| 2.96e-06                \t|\n| 11\t| 1.86e-04                  \t| 8.22e-05                     \t| 8.83e-06               \t| 2.77e-06                \t|\n| 12\t| 7.05e-05                  \t| 4.01e-05                     \t| 1.20e-06               \t| 2.75e-06                \t|\n\n\nThe standard deviations are all quite small, indicating that our results are highly insensitive to randomness in training.\n\nPlease let us know if this aligns with what you had in mind for “sensitivity analysis,” or something else!\n\n## Sensitivity of activation patching to varying $x$ or $x’$\n\nWe 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:\n\n* S3 AA: https://ibb.co/5gmqB2Q2 \n* S3 PAA: https://ibb.co/R4YGwC4z  \n\nNote 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.\n\n## In Figure 1, what is represented by the two bifurcating dashed lines for state parity probe?\n\nThe 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: \n\n“Note the dotted lines indicate two different probing signatures that would both be consistent with this algorithm (see Appendix C.1 for more details).”\n\n\n## jylin04 seems incorrect\n\nWe 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. \n\n## Other typos / inconsistencies\n\nWe will correct, thanks!\n\n[1]  OthelloGPT learned a bag of heuristics  https://www.lesswrong.com/posts/gcpNuEZnxAPayaKBY/othellogpt-learned-a-bag-of-heuristics-1", "claims_and_evidence": "This paper proposes 4 possible schemes that language models can learn to state track and illustrates the expected signatures for the probing and patching experiments. The extensive experimental results match the signatures expected from two of these four schemes.\n\nThe signatures are necessary for the AA and PAA schemes proposed, however, it wasn't clear whether these are sufficient conditions. Would another similar scheme give similar signatures?", "methods_and_evaluation_criteria": "For interpretation of how the language models are able to track state, understanding how important weights/inputs to each layer is, is important. The probing and the patching experiment seems useful to test the sensitivity of the network outputs.", "theoretical_claims": "There are no theoretical claims in this paper.", "experimental_designs_or_analyses": "There seems to be extensive evaluation of the probing and patching experiments, with more results seen in the Appendix. This also enable readers to understand how early on in training the network seems to converge to these state tracking method.\n\nA linear probe is trained, and it would be good to perform a sensitivity analysis to see how varying this probe would affect the results.", "relation_to_broader_scientific_literature": "This paper sits well in the literature of mechanistic understanding and comprehending how transformers may achieve their state tracking abilities. It's interesting to see how one of the mechanisms learned is similar to the associative scan.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Overall this is a relatively well-written paper, which explores the mechanistic understanding of language models and has novel results.\n\nSome explanations can be make clearer (see below).", "other_comments_or_suggestions": "There is a lot going on in Figure 1, some of which I find difficult to interpret. For example, for the probing signature, what is represented by the two bifurcating dashed lines for state parity probe for the sequential and AA algorithm?\n\nMinor points:\n\nLine 12: jylin04 reference seems to have incorrect author name.\n\nLine 151: $h_{lt}$ -> $h_{t,l}$\n\nLine 291: lenght -> length\n\nThere is a mix of using $\\S$ and Section to refer to section, which can be unified.", "questions_for_authors": "How sensitive is activation patching across different data samples $x$ or alternative $x'$?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tSPdrMxllL", "reviewer": "Reviewer_4TFk", "summary": "The paper studies the mechanisms that Transformers learn for performing state tracking - predicting the state after a sequence of operations. Specifically, the model is provided with a sequence of permutations $a_1, \\dots, a_t$, and needs to compute the state, which is the composition of the permutations $s_t = a_t \\circ \\dots \\circ a_1$. The authors discuss possible algorithms for performing state tracking, and demonstrate that Transformers learn either an \"associative scan\" algorithm, which computes in parallel compositions of permutations, or a \"parity-associative scan\" which uses a computation of the parity of the permutations for ruling out some of the outputs before performing associative scan. The authors discuss the difference between the two mechanisms and how a model can be steered towards performing one mechanism and not the other.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful for your positive feedback and thoughtful comments. Below, we address some specific questions: \n\n## Are permutations provided as individual tokens  (i.e., the size of the vocabulary is the number of permutations)?\n\nYes, 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.\n\n## 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?\n\nYou 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.\n\n## Would training end-to-end to output final state change results/conclusions?\n\nGreat 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. \n\n## Is the claim that the model always learns the most compressed (fewer number of layers) version of the algorithm starting from the first layer?\n\nThese 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.\n\n## Why not train on fresh data? Why data repetition? Does this change results?\n\nThis 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.\n\n## Figure 5: seq length the model is trained on?\n\nThe 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!\n\n## Sec 5.3: Intuitively why does topic modelling steer towards associative algorithms?\n\nWe 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. \n\n## Adding an experiment with a more realistic setting\n\nGreat question! See our response to qUEj (point #1).\n\n## Minor typos\n\nWe will fix, thank you!", "claims_and_evidence": "I think the paper is very well-written, and studies a very central question regarding the computational abilities of Transformers and how they learn to solve state tracking problems. The results are novel and interesting, the authors demonstrate that Transformers learn a surprising algorithm for state tracking that leverages the permutation parity computation in order to find the final state. Additionally, the results on steering the model towards one solution instead of the other are very interesting.\n\nSome questions and comments:\n1. I think the description of the task is not clear enough:\n- Are the permutations provided as individual tokens (i.e., the size of the vocabulary is the number of permutations)?\n- \"the input to the model is a sequence of actions $[a_1, \\dots, a_t]$ and the output is a sequence of state predictions $[s_1, \\dots, s_t]$\": is this a causal language model trained to perform a sequence-to-sequence task (i.e., not next-token prediction)? Or is this an encoder-only sequence-to-sequence model? What would be the equivalent next-token prediction variant of this task, and do you think this changes the results/conclusions?\n2. Related to the previous point: if I understand correctly, the model is provided with the sequence of states after each step as a supervision. This can potentially change the function that the model learns, as it could potentially guide it to generate a particular state at a particular position. Did you try training this end-to-end (i.e., get a sequence of actions and output only the final state)? Would this change the results/conclusions?\n3. In Section 3, the description of how the different algorithms are implemented by the Transformer implicitly assumes that everything is implemented starting from the first layer. If the Transformer has more layers than the minimal number of layers required for solving the task, it could potentially pass the tokens through the first layers and start implementing the algorithm deeper in the network, or otherwise skip layers in the middle. Is the claim that the model always learns the most compressed (fewer number of layers) version of the algorithm starting from the first layer?\n4. The experiments are done with data repetition (generating 1 million sequences repeated for 20 epochs). Is there a reason for not training on fresh data? Does this change the results?\n5. In Figure 5 - what is the sequence length that the models are trained on?\n6. In Section 5.3 - do you have any intuitive explanation for why adding the topic modeling steers the model towards the associative scan algorithm?\n7. I think that adding an experiment that connects the results in the paper to a more realistic setting (e.g., word problems that require state tracking) would be a nice addition to the paper.\n\nMinor:\n- I think that putting labels on the axes of the matrices in Figure 1 (Layers/Sequence Length) will make it easier to parse.\n- Typo in line 302 left and 282 right: \"S3 and S3\" => \"S3 and S5\".", "methods_and_evaluation_criteria": "See above.", "theoretical_claims": "See above.", "experimental_designs_or_analyses": "See above.", "relation_to_broader_scientific_literature": "See above.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uTmmhSWMyP", "reviewer": "Reviewer_qUEj", "summary": "This paper investigates how language models track dynamic states through systematic analysis of permutation composition problems . The study reveals that both pretrained and fine-tuned Transformer models learn two distinct state-tracking mechanisms: an Associative Algorithm (AA) and a Parity-Augmented Associative Algorithm (PAA). The former solves tasks via hierarchical permutation composition, while the latter first prunes the state space using parity heuristics before refining results with AA. Experimental validation demonstrates the divergence between these mechanisms, and the authors propose intermediate tasks to guide model specialization.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your time and feedback. Below, we address some of your general critiques and specific questions. \n\n## How could these mechanisms be applied to practical scenarios?\n\nTo 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,\n\n> Swap positions 2 and 3. Rotate the last item to the front. Swap positions 1 and 2.\n\nwould map to state “123”.\n\nWe then conduct probing and activation patching experiments.\n\n**Probing experiments**\n\nWe train probes to map from the activation of each layer at the position of the final “.” token to the final state.\n\n* Pretrained Pythia-160M (PAA signature): https://ibb.co/5XwfF5h7\n* Non-pretrained Pythia-160M (AA signature): https://ibb.co/R4HfbXhn\n\nAs with our results on synthetic data, probing results are consistent with the AA and PAA mechanisms respectively.\n\n**Activation patching experiments**\n\nWe 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.\n\nNonetheless, we still believe the activation patching results serve as a good proxy for estimating how information gets propagated through the layers.\n\n* Pretrained Pythia-160M: https://ibb.co/Tqdcfb4N \n* Non-pretrained Pythia-160M: ​​https://ibb.co/0VXSNNxP \n\nAs 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.\n\n\n## Root causes for why models are biased towards one or the other algorithm\nThe 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!\n\n## Exclusive focus on final-state accuracy neglects intermediate state accuracy, robustness to long-range dependencies, and error tolerance under adversarial perturbations\nWe respectfully disagree about the first two. \n\n1. 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).\n\n2. 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).\n\nWe 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.\n\n## Overreliance on parity shortcuts risks overestimating models’ true tracking capabilities.\n\nWe’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.\n\n## Do deeper/wider networks inherently favor AA?\n\nWe 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.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "### Strengths\n1. The permutation composition framework offers theoretical significance and generalizability to practical tasks like finite automata simulation, establishing a concise yet versatile testbed for state tracking research.\n2. The work elucidates how architectural choices and initialization parameters influence algorithmic preference, while proposing actionable interventions to optimize state-tracking capabilities.\n\n### Weakness\n1. Exclusive reliance on artificially designed permutation tasks (S3/S5 groups) leaves unverified whether observed mechanisms generalize to real-world scenarios like natural language inference or code execution. Despite claims of generalizability, the absence of cross-task validation or concrete application cases substantially weakens practical relevance.\n2. While noting that initialization and architecture dictate algorithm preference, the paper fails to uncover root causes(e.g., how initial parameters bias computational pathways or how attention implement parity heuristics). This black-box attribution reduces findings to correlational observations rather than reproducible causal mechanisms.\n3. Exclusive focus on final-state accuracy and parity correctness neglects critical dimensions of state tracking, including intermediate state consistency, robustness to long-range dependencies, and error tolerance under adversarial perturbations. Overreliance on parity shortcuts risks overestimating models’ true tracking capabilities.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. The paper mentions that model architecture influences algorithmic preference but lacks concrete analysis of how depth (number of layers) or hidden layer dimensionality (width) modulates the AA/PAA propensity. Do deeper networks inherently favor AA?\n\n2. If applying the mechanisms discovered in permutation tasks to practical scenarios (e.g., tracking dialogue states or program execution traces), would structural modifications or specialized training strategies be required?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_flag": "Flag this paper for an ethics review."}}, {"review_id": "S031CuCISY", "reviewer": "Reviewer_nf2n", "summary": "This paper investigates state tracking mechanisms in transformer based language models. In particular, this paper applies interpretability analyses to Pythia and GPT-2 models as they solve the word problems in $S_3$ and $S_5$. The authors find signatures relating to two computational models, PAA and AA, and demonstrate that different training curricula can give rise to models using these two mechanisms.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n## Pretraining vs training distinction, were pre-trained models pre-trained explicitly on state-tracking problems or the entire internet?\n\nWhen 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!\n\n## What is $n$? Is it the number of tokens?\n\nGood 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.\n\n## 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.\n\nThis is merely a high-level extrapolation of the empirical result shown in Figure 2, which might make more sense with more details. \n\n## 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.\n\nThis is correct!\n\n## Why does just patching the first token result in an incorrect state?\n\nGreat 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.\n\nHaving 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).\n\nWe will add this clarification to Section 2.3 and Section 4.2 in the paper!\n\n\n## 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? \n\nThis 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: \n\nA 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.\n\nWe will also add the suggested sources – thanks for providing them!\n\n\n## 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\n\nThanks 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).”\n\n\n[2] Li et al., 2022. Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task. https://arxiv.org/pdf/2210.13382 \n\n[3] Li et al., 2021. Implicit Representations of Meaning in Neural Language Models. https://arxiv.org/abs/2106.00737", "claims_and_evidence": "The authors provide strong evidence that the signatures they propose are realized in the transformers they analyze.", "methods_and_evaluation_criteria": "The proposed methods make sense. \n\nHowever, I was a bit confused about the pretraining vs training distinction.\n\nWhen the models are said to be pretrained (cf line 282), does this means sometimes they were LLMs trained on the entire internet, and then other times (Section 5.3) the weights were initialized randomly, and they were pre-trained explicitly on state tracking problems?", "theoretical_claims": "I did not follow the rationale for the AA and PAA signatures. This is an important issue, because the empirical results hinge on the AA and PAA signatures actually corresponding to the AA and PAA algorithms. I think it's very important that the authors clarify this important point.\n\nHere are my questions and confusions in detail about the theoretical aspects of this paper.\n\n1. In the algorithm blocks, $n$ is used? What is $n$? Is it the number of tokens? In the description of the word problem (Section 2.2), $t$ is used for the number of tokens. I find this shift in notation very confusing.\n\n2. I do not understand the signatures for AA and PAA as described in this paper. For example, take Figure 1. In the prefix patching signature section for AA, there are four rows of squares, where the first row is entirely blue, and the bottom row only as a blue square in the lower right hand side.\n\nThis signature simply does not make sense to me with the definition of prefix patching as described in this paper.\n\nI am understanding the top row to be the input layer, and the bottom row to be a deepest layer, as would seem to be suggested by the diagrams of the algorithms in Figure 1. My understanding is that the token index is increasing from left to right.\n\nSo, I don't understand how merely patching the very first token in the first first layer (which seems to be the upper left) could result in an algorithm with a corrupted input predicting the correct state. I would have thought that this clean input would have been corrupted by running the AA algorithm with states corresponding to an entirely different input sequence.\n\nI've thought as hard as I can about this issue and it just doesn't make any sense to me. Unfortunately I can't recommend accepting this paper while this confusion persists, because everything else about this paper seems like a really solid piece of work. I hope that the authors can provide a crystal clear explanation of the PAA and AA signatures during the rebuttal period.", "experimental_designs_or_analyses": "I did not check the experiments.", "relation_to_broader_scientific_literature": "There is a lot of discussion about whether or not foundation models can track state, and this paper makes a nice contribution by showing concretely mechanisms by which transformers could track state and empirically verifying them on pretrained models.", "essential_references_not_discussed": "I was a bit surprised to read the second sentence of this paper\n> A growing body of work suggests that these models learn to represent the latent state of the world\n\nThese papers would seem to argue differently\n* Merrill and Sabharwal '23, \"The Parallelism Tradeoff: Limitations of Log-Precision Transformers\" and Merill et al '22 \"Saturated transformers are constant-depth threshold circuits\"\n* Huet et al '25, \"Episodic Memories Generation and Evaluation Benchmark for LLMs\"\n* Bhattamishra et al '20, \"On the Ability and Limitations of Transformers to Recognize Formal Languages.\"\n* Deletang et al 2023\n* Sarrof et al 2024\n* Strobl et al 2024\n\n\nA more nuanced discussion was provided in Section 2.1, however.\n\nIt would be nice to at least have some references accompanying the second sentence of the paper. And, the authors might consider a more balanced second sentence also acknowledging work that suggests that transformer based LMs can really struggle with tracking state.", "other_strengths_and_weaknesses": "The paper as a whole is very well written.", "other_comments_or_suggestions": "The extra dotted lines in Figure 1 (for the state parity probe) are a bit confusing until one reads the footnote non the bottom of page 5, but this takes quite a while to get to.", "questions_for_authors": "My essential question is about the correctness of the AA and PAA signatures, see my question under \"Theoretical Claims\"", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: ## Summary
The paper studies how transformer language models can keep track of hidden state. It uses permutation word problems—objects moved by a sequence of swaps—as a controlled test case. Models trained or fine‑tuned on this task learn one of two procedures. The first is an associative scan that matches prior theoretical work by Liu et al. (2023) and Merrill et al. (2024). The second first computes permutation parity to narrow the output space, then refines the result with an associative scan. These two procedures differ in robustness, and the authors show that intermediate training tasks can push a model toward either one. The findings show that transformer language models can learn efficient, interpretable, and broadly applicable state‑tracking routines.

## Decision

This paper investigates the mechanisms by which transformers perform state tracking, proposing four possible schemes and clearly delineating expected experimental signatures. The reviewers appreciated the paper's novel and compelling experimental results, particularly highlighting the identification of an intriguing algorithm akin to associative scan and permutation parity computation. The permutation composition framework is noted for its theoretical significance and practical generalizability, serving as an insightful testbed for further studies on state tracking in transformers. Additionally, reviewers praised the clarity and quality of writing, as well as the valuable insights regarding architectural choices and interventions that influence algorithmic preferences. Overall, the paper is considered a strong contribution to the mechanistic understanding literature, concretely demonstrating state-tracking capabilities in language models. The reviewers unanimously agreed to accept this paper. I believe this paper makes a meaningful contribution to the field. Thus, I recommend this paper for acceptance at ICML.
Decision: Accept (poster)