Statement
Let a sequence-generating process have, at each step , a state such that the pair is a sufficient statistic for the entire future of the process, while only a smaller, derived quantity is sufficient for predicting the immediate next step alone. The hypothesis claims that a Transformer trained end-to-end on the next-step objective nonetheless recovers the larger, task-over-complete statistic — not merely — and does so at a specific, architecturally-identified location: the residual stream immediately after the first attention block linearly encodes via a coordinate system shared across otherwise-different downstream prediction targets, while the representation after the MLP/LayerNorm sublayers have acted linearly encodes only the smaller, target-specific . Recurrent architectures (RNNs) trained on the identical objective reach the same Bayes-optimal loss without this staged separation — the two statistics remain entangled across the hidden state rather than occupying two cleanly separated processing stages.
Intuition
Most representation-geometry claims on this map ask “does the network encode more than a human labeler would think it needs to?” This hypothesis asks a sharper, structural question: if the network builds more than the objective needs, does that extra information have a specific address — a sublayer where it appears in full, and a later sublayer where it gets thrown away? The claim is that attention behaves like a general-purpose “state estimator” (it builds the whole world model, the same way regardless of which of several different next-step tasks is eventually asked of it), while the MLP behaves like a task-specific “decision layer” (it discards whatever part of the world model the current objective doesn’t need). An RNN, updating one hidden state with no such two-sublayer split, has no comparable architectural seam at which to perform this separation — so even if it reaches the same accuracy, there’s no obvious place forcing it to keep the full state and the compressed state apart.
Properties
- A staged compression tied to a specific architectural boundary, not just “somewhere in the middle layers.” Intrinsic-dimension profile across depth documents a generic expansion-then-compression trajectory across depth in several networks; this hypothesis is a sharper, falsifiable claim about where exactly the compression step happens (the attention→MLP boundary within a single block) and what is being compressed (a known, analytically-derivable sufficient statistic, not just “intrinsic dimension” in the abstract).
- Extends the belief-state-geometry lineage to a non-stationary, non-ergodic setting — a different shape, not a restatement of Belief State Geometry Hypothesis (Mixed-State Presentation). Both hypotheses trace back to the same computational-mechanics tradition (Shai et al. 2024) and both concern a network’s internal representation of a process’s state, but Belief State Geometry Hypothesis (Mixed-State Presentation) claims a curved Mixed-State-Presentation simplex geometry for belief over a stationary, ergodic HMM’s hidden states; this hypothesis instead claims a two-stage linear-subspace separation (world state vs. task state) for a non-stationary, non-ergodic process (a grid walk with a finite horizon), and locates the claim at a specific sublayer boundary that Belief State Geometry Hypothesis (Mixed-State Presentation) does not address.
- The staging is architecture-specific, not task-necessitated. The distinguishing, falsifiable content is that reaching the same Bayes-optimal loss does not force a network to keep the full world state and the compressed task state geometrically apart — an RNN control that matches the transformer’s loss while failing to show the same staged separation is direct evidence the separation is attention’s doing, not an automatic consequence of optimal prediction.
- Operationalized via cross-task representational alignment (Centered Kernel Alignment (CKA)), not just single-task probe accuracy. The strongest version of the claim is not “the world state is linearly decodable post-attention” alone (which a single high probe could already suggest) but that the same post-attention coordinate system is reused, with high pairwise CKA, across networks trained on genuinely different next-step objectives built from the same underlying state space — while the post-MLP representations of those same networks are comparatively far apart.
- Predicts that attention-only ablation should remove exactly the task-irrelevant part of the world state. A causal-validation test distinguishing this from a purely correlational reading: patching or ablating only the attention-block output (not the full residual stream) should remove information about the over-complete state beyond what the current objective’s needs, while leaving next-step task performance itself unaffected until the point where the MLP would have recompressed it — a prediction not yet tested by the hypothesis’s originating paper, which reports only correlational (probing, alignment) evidence.
Exercises
Base
- The hypothesis distinguishes a statistic sufficient for “the entire future” from one sufficient for “the immediate next step alone.” Give a one-sentence reason the first must contain at least as much information as the second.
Solution
Predicting the entire future trajectory requires being able to predict every individual next step along it (in particular the very next one), so any statistic sufficient for the whole future must in particular support next-step prediction — the next-step-sufficient statistic is derivable from the full-future-sufficient statistic (as ), but the reverse need not hold, since can carry information relevant only to later, not immediate, steps.
- Why is an RNN control (rather than just showing high probe accuracy for in a Transformer alone) necessary to support the claim that the staging is architecture-specific?
Solution
High probe accuracy for in the Transformer alone would only show that the information is present somewhere and linearly decodable — it says nothing about whether building and then discarding the full state in two separate geometric stages is a special property of the attention mechanism, or simply an inevitable by-product of reaching optimal prediction on this task, which any sufficiently powerful architecture would exhibit. An RNN that reaches the same optimal loss without the same staged, two-location separation directly rules out the second explanation, isolating the finding to something about attention specifically.
Middle
- Explain concretely why Belief State Geometry Hypothesis (Mixed-State Presentation)‘s barycentric-decomposition test (Exercise set on that page) is not a test this hypothesis could reuse directly, even though both hypotheses descend from the same computational-mechanics tradition.
Solution
Belief State Geometry Hypothesis (Mixed-State Presentation)‘s barycentric test presupposes a simplex geometry — a convex combination of per-vertex predictive distributions over a finite set of hidden causal states — which only makes sense for a stationary process with a fixed, small set of discrete causal states whose Mixed-State Presentation can be computed in advance. This hypothesis instead concerns a non-stationary process (a grid walk with a time-varying horizon ) whose relevant state is not naturally a point in a fixed probability simplex over discrete causal states at all, but a continuous, time-indexed geometric coordinate — so there is no vertex set to take a barycentric combination over, and the operational test needed is a different one (cross-task CKA and staged-location probing) suited to a continuous, non-stationary target.
- A study finds high post-attention CKA across three task variants, but also finds high post-MLP CKA across the same three variants (i.e. no drop after the MLP). Does this refute the hypothesis, weaken it, or leave it untested? Justify your answer.
Solution
This weakens (though does not by itself fully refute) the hypothesis: the hypothesis’s specific, falsifiable content is that a sharp drop in cross-task alignment occurs specifically after the MLP sublayer — that the MLP is doing task-specific compression work the attention sublayer’s output does not yet reflect. Finding that both stages remain equally well-aligned across tasks would suggest either that no meaningful task-specific compression is happening at all in this network (the tasks were too similar to require it) or that if it is happening, it isn’t cleanly localized to the MLP sublayer the way the hypothesis claims — either reading undercuts the “staged, sublayer-specific separation” claim, even though the earlier, separate claim (attention builds a shared world-state coordinate system) could still hold on its own.
Pro
- Sketch, in terms of an ablation/patching experiment on the residual stream, the specific causal test proposed in the hypothesis’s fifth Property, and explain what result pattern (across a “task-irrelevant information” probe and a “next-step task performance” metric) would most cleanly distinguish this hypothesis from a purely correlational reading of the same probing and CKA evidence.
Solution
Take the residual stream immediately after the attention sublayer, and construct an intervention that selectively removes (e.g. via directional ablation, or patching in a value with the task-irrelevant component of scrambled or replaced by its value from a different context) only the part of not needed to compute for the current objective, leaving the -relevant component intact. The cleanest distinguishing result pattern: (a) a probe for the ablated, task-irrelevant part of at the post-attention site should now fail (confirming the intervention worked), (b) next-step task performance (the model’s actual loss on its trained objective) should be unaffected, since by hypothesis the MLP was going to discard that information anyway, and (c) if a different, counterfactual next-step objective requiring exactly the ablated information were tested on the same (unablated-elsewhere) network, performance on that counterfactual objective should degrade. This three-part pattern — selective probe failure, no effect on the trained objective, but a causal effect on a hypothetical objective needing the removed information — is what distinguishes “attention genuinely builds a reusable, task-general world state that the MLP discards the unneeded parts of” from the weaker, purely correlational claim that the two statistics merely happen to be separately decodable at two different sites without either one playing an actual causal role in what’s ultimately used downstream.