ICML 2025

(How) Do Language Models Track State?
Belinda Z. Li 1 Zifan Carl Guo 1 Jacob Andreas 1
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 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 a 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 easy-to-compute feature (permutation par-
ity) to partially prune the space of outputs, and
then refines this with an 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 heuristics. Our results demonstrate that trans-
former 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 con-
trolled. Code and data are available at https:
//github.com/belindal/state-tracking.
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
1MIT EECS and CSAIL. Correspondence to: Belinda Z. Li
<bzl@mit.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouve...