(How) Do Language Models Track State?

• URL: http://arxiv.org/abs/2503.02854v2
• Date: Tue, 11 Mar 2025 15:36:40 GMT
• Title: (How) Do Language Models Track State?
• Authors: Belinda Z. Li, Zifan Carl Guo, Jacob Andreas,
• Abstract summary: Transformer language models (LMs) exhibit behaviors that appear to require tracking the unobserved state of an evolving world.<n>We study state tracking in LMs trained or fine-tuned to compose permutations.<n>We show that LMs consistently learn one of two state tracking mechanisms for this task.
• Score: 50.516691979518164
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Transformer language models (LMs) exhibit behaviors -- from storytelling to code generation -- that appear to require tracking the unobserved state of an evolving world. How do they do so? 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, then refines this with an associative scan. The two mechanisms exhibit markedly different robustness properties, 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 pretrained or fine-tuned, can learn to implement efficient and interpretable state tracking mechanisms, and the emergence of these mechanisms can be predicted and controlled.

Related papers

• Self-Steering Language Models [113.96916935955842]
  DisCIPL is a method for "self-steering" language models. DisCIPL uses a Planner model to generate a task-specific inference program. Our work opens up a design space of highly-parallelized Monte Carlo inference strategies.
  arXiv  Detail & Related papers  (2025-04-09T17:54:22Z)
• Finite State Automata Inside Transformers with Chain-of-Thought: A Mechanistic Study on State Tracking [41.3496135369579]
  Chain-of-Thought (CoT) significantly enhances the performance of large language models (LLMs) across a wide range of tasks.<n>In this work, we evaluate the state tracking capabilities of Transformer+CoT and its variants, confirming the effectiveness of CoT.<n>We propose two metrics, compression and distinction, and show that the neuron sets for each state achieve nearly 100% accuracy.
  arXiv  Detail & Related papers  (2025-02-27T14:24:51Z)
• Algorithmic Capabilities of Random Transformers [49.73113518329544]
  We investigate what functions can be learned by randomly transformers in which only the embedding layers are optimized. We find that these random transformers can perform a wide range of meaningful algorithmic tasks. Our results indicate that some algorithmic capabilities are present in transformers even before these models are trained.
  arXiv  Detail & Related papers  (2024-10-06T06:04:23Z)
• Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers [54.20763128054692]
  We study how a two-attention-layer transformer is trained to perform ICL on $n$-gram Markov chain data. We prove that the gradient flow with respect to a cross-entropy ICL loss converges to a limiting model.
  arXiv  Detail & Related papers  (2024-09-09T18:10:26Z)
• Talking Heads: Understanding Inter-layer Communication in Transformer Language Models [32.2976613483151]
  We analyze a mechanism used in two LMs to selectively inhibit items in a context in one task.<n>We find that models write into low-rank subspaces of the residual stream to represent features which are then read out by later layers.
  arXiv  Detail & Related papers  (2024-06-13T18:12:01Z)
• Are you still on track!? Catching LLM Task Drift with Activations [55.75645403965326]
  Task drift allows attackers to exfiltrate data or influence the LLM's output for other users. We show that a simple linear classifier can detect drift with near-perfect ROC AUC on an out-of-distribution test set. We observe that this approach generalizes surprisingly well to unseen task domains, such as prompt injections, jailbreaks, and malicious instructions.
  arXiv  Detail & Related papers  (2024-06-02T16:53:21Z)
• In-Context Convergence of Transformers [63.04956160537308]
  We study the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent. For data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process.
  arXiv  Detail & Related papers  (2023-10-08T17:55:33Z)
• Transformers as Algorithms: Generalization and Implicit Model Selection in In-context Learning [23.677503557659705]
  In-context learning (ICL) is a type of prompting where a transformer model operates on a sequence of examples and performs inference on-the-fly. We treat the transformer model as a learning algorithm that can be specialized via training to implement-at inference-time-another target algorithm. We show that transformers can act as an adaptive learning algorithm and perform model selection across different hypothesis classes.
  arXiv  Detail & Related papers  (2023-01-17T18:31:12Z)
• Unsupervised Learning of Equivariant Structure from Sequences [30.974508897223124]
  We present an unsupervised framework to learn the symmetry from the time sequence of length at least three. We will demonstrate that, with our framework, the hidden disentangled structure of the dataset naturally emerges as a by-product.
  arXiv  Detail & Related papers  (2022-10-12T07:29:18Z)
• Inducing Transformer's Compositional Generalization Ability via Auxiliary Sequence Prediction Tasks [86.10875837475783]
  Systematic compositionality is an essential mechanism in human language, allowing the recombination of known parts to create novel expressions. Existing neural models have been shown to lack this basic ability in learning symbolic structures. We propose two auxiliary sequence prediction tasks that track the progress of function and argument semantics.
  arXiv  Detail & Related papers  (2021-09-30T16:41:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.