SUSS · live demo · runs entirely in your browser
Watch a transformer suss out a system it has never seen
Each episode draws a new two-state linear system xt+1 = A xt + B ut + noise with random A and B. The 0.8M-parameter model from the paper sees only the history so far: no gradient steps, no hints. Next to it runs the Bayes-optimal estimator, computed exactly. Both predict the next state before it happens.
Next-state predictions
Dots: the state that actually happened. Each prediction for step t+1 is made at step t, before that state is seen.
What it thinks the system is
The hidden parameters [A B], the Bayes estimate from the context so far, and the transformer's own estimate: the matrix it applies to the newest input, ∂ŷ/∂z. In the first few steps the two disagree; with more context the transformer's estimate converges to the Bayes one.
Errors
Top: prediction error against the noise-free next state. Bottom: distance of each estimate from the true [A B]. Faint bands: average over 4096 episodes.
What you are looking at
The transformer was trained on many random systems, then frozen. It can only adapt through its context window, the way a language model picks up a pattern from a prompt. While data are scarce it is as good as the exact Bayes-optimal estimator; with more context it levels off slightly above it, 1.59× the Bayes excess risk on average.
The middle panel is the paper's functional test: the derivative of the prediction with respect to the newest input is the parameter estimate the network is effectively using. After a handful of steps it tracks the Bayes estimate; in the first steps it does not, even though the predictions themselves are already close to optimal, so early on the network is not yet applying a linear estimate. The paper then shows where the evidence for that estimate is stored: in the attention keys and values of the last two blocks, not in the residual stream. Swapping only those keys and values with another episode's moves 91–94% of the behaviour.
Project page · Paper (PDF) · Code · This model in 150 lines of JavaScript