Learning State-Tracking from Code: REPL Traces and Probabilistic Automata

Julien Siems, Riccardo Grazzi, Kirill Kalinin, Hitesh Ballani, Babak Rahmani

Keywords: state-tracking, linear RNNs, code execution, probabilistic automata, belief propagation, DeltaNet, next-token prediction, finite-state machines

TL;DR: Linear RNNs learn state-tracking from Python REPL traces with sparse supervision where Transformers fail, but probabilistic transitions in real code cause exponential norm decay.

Abstract: Over the last years, state-tracking tasks, particularly permutation composition, have become a testbed to understand the limits of sequence models architectures like Transformers and RNNs (linear and non-linear). However, current experiments use a sequence-to-sequence setup: learning to map actions (permutations) to states, that does not translate to the next-token prediction setting commonly used to train language models. We address this gap by converting permutation composition into code via REPL traces that interleave state-reveals through prints and variable transformations. We show that linear RNNs capable of state-tracking excel also in this setting, while Transformers still fail. Motivated by this representation, we investigate why tracking states in code is generally difficult: actions are not always fully observable. We frame this as tracking the state of a probabilistic finite-state automaton with deterministic state reveals and show that adversarial sequences exist where linear RNNs cannot guarantee stable probabilistic state-tracking.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 78