Interactive illustration
Length-Independent State Tracking Under a Parallel Scan
To track a finite state at every sequence length, a recurrent network must, in one update, erase its own numerical error and execute the next transition.
The code of the paper is available, during the review period included: github.com/dd8wgy467d-hash/reviewanonymous.
How to read this page
Read it next to the PDF: the tag Paper §3, Thm 3.1 names the part of
the paper each section follows. </> Code flips a figure to the code it runs, with your
settings highlighted.
| Paper | Notebook | What you do |
|---|---|---|
| §2, Defs 2.1–2.2 | N1, N2 | run an automaton; watch an affine RNN realize ℤ5, then lose it |
| §3, Thms 3.1–3.3, Fig. 1 | N3, N4 | switch restoration on and off; compare one rate with two |
| §4, App. E, Fig. 3 | N5, N6 | watch a perfectly trained rotation drift; see why definite words survive |
| §5, App. D, Fig. 2 | N7 | scan over tables, then over affine maps in fp16 |
| §6, App. G | N8, N9 | edit the logits of an NFSM head; split S3 over two heads |
| §7, App. H | N10–N12 | browse the results; read a TSO story; look at trained states |
Paper §1, Introduction
1. Why ask for length independence?
The recurrences that scale (state space models, gated linear RNNs, linear attention) are affine in the hidden state, so a parallel scan evaluates them in \(O(\log L)\) depth. Trained on state tracking, they work up to the training length and then collapse. The paper isolates the part of that gap that training cannot remove, finite precision, and asks: what must an RNN satisfy to track a state at every length? Can an affine recurrence satisfy it? If not, can a scan-compatible one?
Paper §2, Formalization, State tracking and automata
2. Automata, RNNs, and finite precision
A task is a semiautomaton \((Q,\Sigma,\delta)\). A word \(w\) acts on the states through the composite map \(\delta_w\), and these maps form the transition monoid \(M_\delta\).
Paper §2, eqs. (1)–(3), Definitions 2.1–2.2, App. D, Lemma D.1
Realization at finite precision
An RNN realizes the automaton through an encoding \(\iota\), an update \(\Phi_x\) and a readout \(\pi\) (eq. 3). Finite precision is modelled as a perturbation of size at most \(\eta\) at every step; every machine run is such an η-trajectory (Lemma D.1). Length independence asks for (T1) a correct readout on every η-trajectory and (T2) a bounded reachable set \(U^\star\).
N2 is the smallest affine realization of \(\mathbb Z_5\): \(A_x = \rho R(2\pi x/5)\). The dashed disk holds the endpoints of all η-trajectories; its radius follows \(r_t = \rho\, r_{t-1} + \eta\).
Try Set \(\rho = 1\), tick worst case, then add samples: the disk crosses a cell boundary at \(L^\star\).
Paper §3, Theorems 3.1–3.2, eq. (4), Figure 1
3. Three conditions on one step, two rates in one map
Theorem 3.1 reduces length independence to three conditions on one step: (R1) a decodable code, (R2) restoration, \(\Phi_x(U)\oplus\bar B_\eta \subseteq U\), and (R3) a correct transition under perturbation. Equivalently, every cell must land inside its successor cell with a margin \(\eta\) (eq. 4). One map must be both an executive and a restoring organ.
Try Raise \(\lambda\) to 1: the reachable set widens every step until it straddles two cells.
Paper §3, Theorem 3.3, Lemma F.11
Two rates in one map
Along any word, eq. (4) needs local contraction inside cells and global separation between them
(Thm 3.3). The hardest word holds two states, like id in the flip-flop; every
non-definite automaton has one (Lemma F.11). N4 iterates such a map from both code points \(\pm1\). There,
\(\beta\) sets how far the map is from affine: at \(\beta = 2\) the rate is below 1 inside the cells and
above 1 between them, so two cells are held, while at \(\beta = 0.3\) the recurrence is almost affine.
Try Move \(\beta\) from 0.3 to 2: the transition from almost affine to two held cells.
Paper §4.1, eq. (5), Theorem 4.1, App. C.2, App. E, Figure 3
4. An affine recurrence supplies a single rate
In an affine recurrence, \(\Phi_w(h) - \Phi_w(h') = A_w (h-h')\) (eq. 5): one operator moves both the signal and the perturbation. By Theorem 4.1, if \(\rho(A_w)\ge 1\) for some word, perturbations accumulate; if \(\rho(A_w) < 1\) for every word, old inputs are forgotten and only definite automata are realized. Most scan-compatible architectures sit on one of the two branches (App. C.2).
Even perfect training does not help (App. E). N5 takes the exact \(\rho = 1\) realization of \(\mathbb Z_7\), a rotation by \(2\pi/7\), rounds its entries once to a number format, and runs it in fp64. The stored rotation has a slightly wrong angle and radius, so the phase error grows until the readout is wrong.
Try Play in bf16, then in fp32: a finer format moves the failure, it does not remove it.
Paper §4.1, Def. F.10, Prop. F.13, §7, Table 1
Forgetting is harmless only for definite targets
A definite automaton depends only on its last \(k\) inputs, and a contracting affine map realizes every such target (Prop. F.13). The flip-flop shows the boundary: DFF5 never draws more than four identities in a row, while FF draws runs of any length.
Paper §4.2, Lemma F.17, Prop. F.22, Theorem F.25
Two rates require multistability
An affine step that satisfies (R2) has a single fixed point, so it holds at most one cell, at any width (Prop. F.22), and a stack of affine layers does no better (Thm F.25). The nonlinearity must act on \(h_{t-1}\) itself (Lemma F.17). PD-SSM's hardmax acts on the input, so its step stays affine.
Paper §5, Definition 5.1, Proposition 5.2
5. Scan compatibility is a budget, not affinity
A parallel scan needs every composite \(\Phi_w\) stored in \(b\) bits and merged in \(c\) operations (Def. 5.1). Then only finitely many hidden states are reachable, and the recurrence acts on them as a table (Prop. 5.2). An exactly restoring recurrence sends each basin to one attractor, so each symbol is a table on \(K\) indices: \(K\lceil\log_2 K\rceil\) bits, \(K\) lookups per merge, and every merge is exact.
Paper §5, App. D, Figure 2
Affine composites are rounded instead, and floating-point arithmetic is not associative, so a scan and a sequential loop return different trajectories. N7 (b) reruns a small version of Figure 2.
Paper §6, eq. (6), Props. G.1–G.3
6. The neural finite-state machine
An NFSM head updates \(\Phi_x(h) = \mathrm{rd}(\theta_x\, \mathrm{rd}(h))\) (eq. 6), where \(\mathrm{rd}\) keeps the one-hot vector of the largest coordinate and \(\theta_x\) is a logit matrix computed from the input. The product with \(\theta_x\) executes, sending index \(k\) to the row of the largest entry of column \(k\); the outer \(\mathrm{rd}\) restores, putting the result back on a code point. The inner one is the identity on a trajectory, where \(h\) is already a basis vector. The head is correct when its tables match the target (Prop. G.2), and a logit perturbation below half the column margin \(\gamma\) changes nothing, at any length (Prop. G.3).
Try Click the dashed cell of each wrong column until every column shows ✓.
Paper §6, Prop. G.4, Example G.7, App. H.6.2, Table 5
Heads as memory channels
A block of heads reads the same input and updates independently (Prop. G.4). The trained \(S_3\) model uses one 3-index head for the image of point 1 and one 2-index head for the sign.
Paper §7, Table 1, App. H.4, Table 4, App. H.6
7. Experiments
One NFSM layer (two on TSO) is compared with four-layer Mamba\(^-\) (\(\rho < 1\)), AUSSM (\(\rho = 1\)) and PD-SSM (\(\rho < 1\)). A † means the tables read off the trained NFSM match the target, which certifies correctness at every length.
Paper §7, App. H.1, App. H.6.3, Figure 5
Tracking Shuffled Objects
TSON gives the same algebra as text: people swap items, then a question asks who holds one. In the trained two-layer NFSM, each second-layer head holds the holder of one item.
Try Drag tokens read through the story and watch the state (p, g) follow.
Paper App. H.5, Figure 4
What the trained states look like
The recurrent state on \(\mathbb Z_5\) under \((+1)^T\), projected on two principal components. The NFSM visits five fixed points. Mamba\(^-\) collapses onto one point (\(\rho<1\)), AUSSM's clusters interleave (\(\rho=1\)), and PD-SSM loses the separation too.