State tracking abilities vary across language model architectures
A Fine-Grained Hierarchy and Benchmark for State Tracking in Language Models
Artificial Intelligence
Summary
Tracking the state of information, like keeping track of variables over time, is hard for some language models such as transformers. The authors break down state tracking into four levels of difficulty to better understand which tasks different models can handle. They show theoretically and with tests that new model designs like Gated DeltaNet do better on complex tracking tasks than standard transformers. However, even large pretrained models struggle with many types of state tracking, pointing to a clear challenge for improving language models.
What this means in practice
- •For language model developers: Assess and improve model designs for tracking variables and entity states over long texts using the StateBench benchmark.
- •For software engineers building ai tools: Choose or develop language models with better state tracking capabilities for code understanding or evaluation tasks involving variable assignments.
Authors
Jackson Petty, William Merrill, Ashish Sabharwal
Abstract
Theoretical limitations of transformers on state tracking---the ability to track information about entities or variables over long contexts---have motivated new recurrent language modeling (LM) architectures such as Gated DeltaNet (GDN). These architectures have, in particular, been benchmarked against the $S_5$ word problem, a theoretically hard state-tracking task that captures abilities such as code evaluation under arbitrary variable assignment. It, however, remains open whether other, simpler forms of state tracking might both be more relevant in practice and already explain observed differences between these architectures. We propose evaluating state tracking using a hierarchy of word problems with four increasing levels of difficulty: r-trivial, corresponding to immutable assignment to variables in programming; aperiodic, where assignments are mutable; soluble, where assignments are self-referential; and general finite state tracking. Theoretically, we show that the class of problems transformers are thought to length-generalize on does not even contain all aperiodic state tracking. On the other hand, we prove that GDN with positive eigenvalues expresses all aperiodic problems, and the regular languages recognized by certain transformers are precisely all soluble ones, thus providing a fine-grained picture of state tracking capabilities of various architectures. Taking these ideas to practice, we construct StateBench, a 4-tier long-horizon state-tracking benchmark. Our systematic direct training experiments across architectures largely align with theoretical expectation, albeit with interesting observations at very large scale length-generalization. Additionally, we also evaluate pretrained models---both open hybrid and closed frontier ones---on StateBench, finding that none of them robustly solves even soluble state tracking, suggesting a direction for future improvement.