Constrained Belief Updates Explain Geometric Structures in Transformer Representations
Authors: Mateusz Piotrowski
Paul M. Riechers
Daniel Filan
Adam Shai
Venue: ICML 2025 poster (2025)
Keywords: transformer architectures
computational mechanics
mechanistic interpretability
optimal prediction
neural network representations
belief states
hidden markov models
Primary area: deep_learning->theory
TL;DR: Transformers implement parallelized Bayesian inference through constrained belief updating, creating predictable intermediate geometric structures that bridge optimal prediction with attention's architectural constraints.

Abstract: What computational structures emerge in transformers trained on next-token prediction? In this work, we provide evidence that transformers implement constrained Bayesian belief updating---a parallelized version of partial Bayesian inference shaped by architectural constraints. We integrate the model-agnostic theory of optimal prediction with mechanistic interpretability to analyze transformers trained on a tractable family of hidden Markov models that generate rich geometric patterns in neural activations. Our primary analysis focuses on single-layer transformers, revealing how the first attention layer implements these constrained updates, with extensions to multi-layer architectures demonstrating how subsequent layers refine these representations. We find that attention carries out an algorithm with a natural interpretation in the probability simplex, and create representations with distinctive geometric structure. We show how both the algorithmic behavior and the underlying geometry of these representations can be theoretically predicted in detail---including the attention pattern, OV-vectors, and embedding vectors---by modifying the equations for optimal future token predictions to account for the architectural constraints of attention. Our approach provides a principled lens on how architectural constraints shape the implementation of optimal prediction, revealing why transformers develop specific intermediate geometric structures.

Constrained Belief Updates Explain Geometric Structures in Transformer Representations
Mateusz Piotrowski
Paul M. Riechers
23 Daniel Filan
Adam S. Shai
Abstract
1. Introduction
Transformers excel at next-token prediction (Vaswani et al., 2017), but their success belics a fundamental tension: op- timal prediction requires Bayesian belief updating, a recur- sive process, while their architecture enforces parallelized, attention-driven computation (Fig. Al). How do transform- resolve this conflict? We show that they develop geomet- rically structured representations that approximate Bayesian inference under architectural constraints, revealing a precise interplay between theoretical necessity and implementation.
What computational structures emerge in trans- formers trained on next-token prediction? In this work, we provide evidence that transformers im- plement constrained Bayesian belief updating a parallelized version of partial Bayesian inference shaped by architectural constraints. We integrate the model-agnostic theory of optimal prediction with mechanistic interpretability to analyze trams- formers trained tractable family of hidden Markov models that generate rich geometric pat- terns in neural activations. Our primary analysis focuses on single-layer transformers, revealing how the first attention layer implements these con- strained updates, with extensions to multi-layer ar- chitectures demonstrating how subsequent layers refine these representations. We find that attention carries out algorithm with a natural interpre- tation in the probability simplex, and create rep- resentations with distinctive geometric structure. We show how both the algorithmic behavior and the underlying geometry of these representations can theoretically predicted in including the attention pattern, OV-vectors, and embedding vectors by modifying the equations for optimal future token predictions to account for the archi- tectural constraints of attention. Our approach provides a principled lens on how architectural constraints shape the implementation of optimal prediction, revealing why transformers develop specific intermediate geometric structures.
In this work, we combine insights from the theory of opti- mal prediction with neural network analysis. First, compu- tational mechanies (Shalizi & Crutchfield, 2001a; Marzen & Crutchfield, 2017: Riechers & Crutchfield, 2018b; Pepper, 2024; Shai et al., 2024) dictates what an optimal predic- tor must represent: belief states that encode distributions over futures. Second, mechanistic interpretability reveals how transformers approximate these states under architec- tural constraints, bending Bayesian updates into attention's parallelizable form (Elhage al., 2021; Nanda et al., 2023).
By combining these frameworks we reveal why transform- ers learn certain intermediate structures. We find that the geometry of a transformer's internal representations not an accident-it is mathematical signature of how architec- tural constraints warp otherwise optimal Bayesian inference. By interpreting learned weights and activations via stan- dard mechanistic interpretability. that is well-captured by the constrained belief updating equations. From first principles, we derive the constrained belief geometries, and reverse-engineer the transformer's computational blueprint, predicting attention patterns, value vectors, and residual stream geometries precisely. Thus, beyond verifying that transformers encode belief states, we show how the specific circuits that implement those states necessarily deviate from the unconstrained Bayesian ideal in predictable and theoretically tractable ways.
uncover algorithm
'MATS, Berkeley, CA. USA 'Simplex, Astera Institute, Emeryville, CA, USA Beyond Institute for Theoretical Sci- ence (BITS), San Francisco, CA, USA. Correspondence Paul M. Riechers <pmriechers@@e wrrechhrrs@@miil.co gma.l.aomm........mmmmmmmmm Adam S. Shai <adami- mos @ gmail.com>.
To concretize these ideas, we focus on transformers trained on data from the Mess3 class of hidden Markov models (HMMs) (Marzen & Crutchfield, 2017), which provides rich and visualizable belief-state geometries and also admits tractable optimal predictor. Our primary analysis exam-
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Constrained belief updates explain geometric structures in transformer representations
ines single-layer transformers to isolate how the first atten- tion layer implements constrained belief updating. though we also demonstrate that these principles extend to multi- layer architectures where subsequent layers refine the initial constrained representations. This allows us to rigorously compare the theoretically optimal geometry with the neural- activation geometry that transformers learn. More broadly, we anticipate that the same tension between architecture and optimal inference arises in large language models trained on natural text, and that our methodology would shed light on those more complex cases.
standing what information transformers represent, with geo- metric relationships between features revealing structured internal representations (Engels et al., 2024). Our work pro- vides mechanistic explanation for these non-orthogonal geometric structures, providing the theoretical why to com- plement the what of feature representations.
From features to circuits While feature directions re- veal what information is encoded, understanding how net- works process this information is done by identifying com- putational Prccuits-subeewwnn
that implement specific algorithmic operations. Circuits typically combine sim- pler features into more complex ones as information flows through the network. Examples include circuits that de- tect syntax (Elhage al., 2021), implement indirect object identification (Wang et al., 2022), or perform basic arith- metic (Nanda et al., 2023). However, identifying circuits remains largely a manual process, starting from observed behaviors and working backwards to discover relevant com- ponents (although active research is developing automated approaches; see Conmy al. (2023); Marks et al. (2024)).
Key contributions:
1. A Unified View of Optimal Prediction and Trans- former Computation: We bridge the model-agnostic theory of Bayesian belief states with the model-specific constraints of attention- bastddd---addddd- parallel processing. This synthesis explains why transformers trained next- token prediction discover distinct "constrained belief updating" geometry- optimal Bayesian in- ference with the functional form of attention.
Our work demonstrates that a principled, top-down theoret- ical framework, based on constrained belief updating, can guide the search for circuits and provide a deeper under- standing of their function within the larger network. We show how specific circuits in the attention mechanism di- rectly implement the computations predicted by our theory.
2. Spectral Theory of Constrained Belief' Updating: We develop a theoretical framework that analyzes how eigenvalues of the mata-gennrating transition matrices determine attention heads' behavior. By decomposing belief updates spectrally, we show that multi-head at- tention naturally implements these scalar updates in orthogonal modes-ever even handling oscillatory decay of influence influence- through a sum of specialized head outputs.
Belief state geometry and computational mechanics Our work draws inspiration from computational mechan- ics, a framework for studying information processing in dynamical systems (Shalizi & Crutchfield, 2001b; Crutch- field, 2012; Riechers & Crutchfield, 2018b). When ap- plied to sequential data, computational mechanics, in accor- dance with the POMDP framework (Kaelbling et al., 1998), shows that optimal prediction requires maintaining beliefs about the underlying latent states of the data-generating process (Upper, 1997). These belief states can be visual- ized as points on a probability simplex, evolving accord- ing to Bayesian updating rules, and forming characteristic geometric patterns (Crutchfield, 1994; Marzen & Crutch- field, 2017). Recent work shows that transformers naturally discover and encode these belief state geometries in their activations (Shai et al., 2024). This connection offers a prin- cipled way to analyze network representations: rather than reverse-engineering observed behaviors, we can study how architectural constraints shape the network's implementa- tion of theoretically optimal prediction strategies.
3. Predictive Experiments and Mechanistic Verifica- tion:
Our approach yields specific, testable predictions about attention patterns, value vectors, intermediate fractal representations, and final belief-state geometry. We confirm these predictions in trained transformers, demonstrating how the inherently recurrent next-token task is realized by an iiteentionbbsedd, parallelized im- plementation of Bayesian belief updates.
2. Background
2.1. Related Work
Features as directions in activation space Modern inter- pretability research views neural network representations through the lens of linear geometry, analyzing how acti- vation patterns align with specific directions that encode fundamental features (Park et al., 2024). This perspective is particularly useful given superposition (Elhage et al., 2022), where networks encode more features than available neurons using non-orthogonal vectors. Conceptualizing features as linear directions has been instrumental (Cunningham et al., 2023; Bricken et al., 2023; Templeton et al., 2024) in under-
This is the approach taken here. We move beyond prior work by proposing and validating a theory of constrained belief updating, demonstrating how specific architectural elements, like the attention mechanism, modify the idealized belief state dynamics. This perspective shifts the focus from
2

---

Constrained belief updates explain geometric structures in transformer representations
A
M HMM
Theoretical Prediction
Fretction
Residual Stream Activations
a]/2
Intermediate
Atsention
2x
2x
Redioued
Training Data
Dobweddnng
218011811011222001199
Belief State Geometry
Transformer
Final
Intermediate
Figure 1: Transformers' internal representations exhibit complex geometric structure matching the belief-state geometry. (A) Mess3 HMM, vertices represent hidden states with their emission distributions. (B) Ground truth belief state geometry of Mess3. Each point represents a belief-state probability distribution over hidden states of the HMM, induced via Bayesian updates upon a sequence of observed emissions, with proximity to the vertices of the simplex corresponding to the probabilities of the three hidden states. (C) Schematic of a single-layer transformer with Intermediate activations after Attention, and Final activations after the subsequent MLP. (D) PCA projections of the model's 's final residual stream (left), before the unembedding, reveals geometric representation that closely matches the belief geometry shown in (B), whereas the PCA projection of the intermediate residual stream (right) after attention but before the MLP exhibits an intricate but different structure. In (B) and (D), points are colored according to the ground-trute belief states associated with the sequence of tokens that induces the point, taking the three constituent probabilities over hidden states of the HMM as RGB values.
reverse-engineering learned features understanding why particular geometric patterns emerge during training consequence of the interplay between optimal prediction and urchitectural constraints. Our work provides a concrete example of how this theoretical framework can be applied to understand the internal mechanisms of transformers.
via Bayes rule according the substochastic transition matrices of the HMM, (T(=) ZEZ' with matrix elements r( Pr(Z++1 = z, S s'|S, = s). In particular, the updated distribution, given context 21.d, is
no T
(313) full
77 H+
(2)
7 n (
2.2. Optimal Prediction and Belief State Geometry
where ((...)) ((.))....((............ Tr.x), and is the column vector of all ones. In this paper, we will make the simplifying as- sumption that the training data is sampled from stationary stochastic process, in which case the initial distribution over latent states is the stationary distribution Ve - . = TT' where T = Exezz T(*) is the row-stochastic transition ma- trix over hidden states.
Shai et al. (2024) showed that transformers minimizing next- token loss internally represent the context-induced proba- bility density over the entire future of possible token se- quences:
Pr(Za+1:00 Zi.a 21:a)
(1)
where Zu+1:0c Zu+1. Zu denotes the sequence of random variables for future tokens, Zi.a = Zi...., , Za denotes the sequence random variables for past tokens, which is realized by a particular sequence of tokens 21dd E zd known as the context up position
Thus, Eq. (2) embeds each token sequence into a probability simplex over the latent states of the HMM-a point in a real-valued vector space. The totality of these points forms particular geometry, called the belief state geometry, and is universally found in linear form within the activations of various deep neural networks, including RNNs (Pepper, 2024) and transformers (Shai et al., 2024).
When we conceptualize the training data being gener- ated by an edge-emitting hidden Markov model (Mealy HMM), we can derive a natural geometric embedding for these conditional probability distributions. HMMs generate training data by emitting tokens when moving among its hidden states S. from one hidden state S, time t to the next. The natural geometric embedding is then given by considering how an initial distribution over hidden states So no. as point in the vector space IR S) (with coor- dinates given by the probability elements), evolves upon seeing a particular sequence of tokens, 2jsd- This distri- bution over the hidden states, which uniquely induces a probability density over all possible futures, is updated
This precise framework for anticipating intermediate activa- tions in transformers provides a natural interpretation the attention mechanism in which it moves information in be- lief simplex for the purposes of building up the architecture- independent belief state geometry given in Eq. (2).
3. Methodology
Data Generation. Our study focuses on the Mess3 parametrized family of hidden Markov models (Marzen
3

---

Constrained belief updates explain geometric structures in transformer representations
& Crutchfield, 2017), which provide tractable yet rich set- ting for studying sequence prediction. As shown in Fig. 1A, these HMMs consist of three hidden states with observ- able emissions controlled by parameter a and transitions by parameter x. Higher values of a E [0, 1] mean each state more strongly prefers its unique emission symbol, providing clearer information about the generating state. The parame- ter (0. 2 controls state persistence- values create high inertia where states tend to persist, while high values increase transition probabilities between states. For each experimental run, we generate sequences by sampling from an HMM with specific (a, .z) values.
final representations align with the geometry of theoretical belief states, the intermediate fractals exhibit markedly different structure. The systematic difference between inter- mediate and final representations raises two key questions: (1) How does attention construct these intermediate fractals and (2) why do they take these particular geometric forms? The following results reveal the algorithmic process behind their construction and provides a theoretical explanation for their previously unexpected structure.
4.2. Intermediate representations are built by algorithms in the belief simplex
To determine how the intermediate representation is con- structed, we performed mechanistic interpretability on the attention heads. We find that attention performs an algo- rithm with a direct interpretation in the belief simplex.
Training Process. We train single-layer transformers on next-token prediction using gradient descent, with se- quences sampled from our parametrized HMMs as training data. The model learns to predict the next token in each sequence by minimizing cross-entropy loss (see App. for details). Our primary analysis focuses on single-layer ar- chitectures to clearly isolate how the attention mechanism implements constrained belief updating, though we also validate that these mechanisms persist as the foundational computation in deeper networks (Figs. A5;A6).
At every context position, the residual stream can be thought of as a Aimess--imeniionall skip connection communication channel streaming alongside all layers, carrying all working memory in transformer (Elhage et al., 2021). Attention and MLP modules read in linear transformations of the residual stream and then add their output to the local residual stream at each layer (Vaswani et al., 2017).
Analysis of Representations and Computations. To study how the model processes information, we analyze intermediate and final activations in the residual stream (Cig.IC). We apply principal component analysis (PCA) to these activations across all possible input sequences, finding that the representations are well well-captured by low- dimensional space. In some cases, we slightly rotate the PCA basis to align with theoretically meaningful directions. This dimensionality reduction enables us to visualize how representations evolve through the network-from input embeddings, through the intermediate state after attention, to the final output state after the MLP layer (Fig.ID). To understand how the network manipulates these representa- tions, we analyze the learned weights and attention patterns, examining how attention transforms input embeddings into intermediate representations and how the MLP layer trans- forms these into the final geometry. At each stage, we com- pare the learned representations to theoretical predictions derived from optimal Bayesian updates'.
Following (Elhage et al., 2021), we decompose the attention operation into two circuits: (i) the output-value (OV) circuit, which specifies what information is read from each position and how it transforms into a vector that can be broadcast to other positions, and (ii) the query-key (QK) circuit, which computes the similarity of a linearly transformed source and destination to determine how much to update the destina- tion't 's residual stream with that source's 's OV contribution.
For an attention head, the residual stream update ==(mid)) Z (pre) + Ca E Rddate at the destination position s is:
T
A
RSd
WoWWy (prc) represents the OV circuit's con- tribution from source position s, where Wo and Wy are the attention output and value weight matrices respectively, and T.(00) is the incoming residual stream vector at position s. Attention Aa. is determined by the QK circuit through query-key inner product and the causally masked softmax operations:
Here, Vo
4. Results
4.1. Intermediate representations are fractals, but not belief state geometry
A 8esd
(4)
ΣS= e.te /
Through PCA of the residual stream, we observe two distinct fractal structures in transformers trained on Mess3 HMM data: one after the attention mechanism but before the MLP, and another in the final layer output (Figs. 1, 4). While the
where Ta = Wo (prc)) is the query vector from position d. K = ( is the key vector from position s, dp is the head dimension, and Wo and Wx are each de x dmodel weight matrices. Recall that attention is negaaivee S Aa, S 1
"Code for analysis of can be found here.
4

---

Constrained belief updates explain geometric structures in transformer representations
A
Embedding and OV Vectors
B
Theoretical Prediction
C
Attention Pattern over All Inputs
9,
Embedding,
0.5
1
2nd Priciccal Component
7-
Attintron Value
simplex
Dr
Theoey
Token
0.1
(1-3x) 3x) n-1
Seg Length
-7.0
10
-8
-6
-5
-4
-3
-2
0
1st Principal Component
Source Position Relative to Destination
D
Attention Mechanism Computation for Input Sequence 01120
Subsequence:
Subsequence:
Subsequance:
Subsequenca:
Subsequence: 011120
1
Figure 2: Intermediate Representation Construction by Attention. A transformer trained on Mess3 with x = 0.15 and = 0.6 exhibits intermediate representations constructed through specific attention mechanism. (A) The OV vectors (arrows) form three distinet clusters, each corresponding
token and positioned the vertices of triangle, while token embeddings (circles) are clustered near the origin. (B) Our theoretical predictions for the OV vectors (shown for all (position, token) pairs) and embeddings (for positions > 2) align closely to those found in the trained transformer. (C) Attention patterns are primarily determined by the positional distance between the destination and source tokens, following an exponential decay described by (1 3z)|n-11 They are largely independent of specific token sequences. (C, inset) The theoretical (Eq. (13)) and actual values in the attention pattern align closely. (D) Construction of intermediate representations for five input subsequences of increasing length (from the example sequence 01120, shown left to right). The attention mechanism builds the fractal by taking linear combinations of the three U, vectors. The colored vectors illustrate the components of the sum for each example subsequence, while the gray dots represent all possible vector sums for all sequences that position.
and, for each destination position d, the attention to all sources sums to one: Exso Aa.s = 1. Eq. (3) shows how each attention head computes its update by weighting the transformed values (7%) from all previous positions accord- ing their relevance (Aa,s) to the current position.
integrated. As the attention weight decays with distance, the impact of past tokens on the current belief state dimin- ishes over time. Through this process of weighted vector addition within the belief simplex, the attention mechanism constructs the intermediate representations, resulting in the observed fractal structure shown in Fig. 2D. Incredibly, the computation the attention head performs is completely interpretable as a dynamic process in the belief simplex.
Our analysis yields several key insights into how the atten- tion mechanism constructs the intermediate representations. First, we find that projecting token embeddings (the inputs into the attention head) onto PCA space reveals three clus- ters that lie close to the origin, as shown in FAg. Mean- while, the OV projections form update vectors o that cluster in three directions pointing toward the vertices of triangle, naturally interpreted as the vertices of the belief simplex in Fig. 2. The model combines these directions through weights Aas determined by the QK circuit described by Eq. (3). For Mess3, these attention weights are invariant to token identity and decay exponentially with distance from the current position, controlling how past information is
4.3. Relating Intermediate Representations to Belief Updating Equations
The interpretation of attention operating in the belief sim- plex suggests connection to the theory of belief updating, Since the OV circuit is only able to access information from the source token that is attended to, we can write a con- strained belief updating equation that sums contributions from the value of the token value of n, assuming the initial belief is the stationary distri-
places back for each
5

---

Constrained belief updates explain geometric structures in transformer representations
bution of the HMM, This gives the following equation for the constrained belief at position d in the sequence:
increasing deviations between predicted and actual repre- sentations, though the overall structure remains similar. A complete characterization of how these deviations scale with a remains for future work.
(zsa)
Σ (TI x)
Fig. A3 demonstrates that transformers reliably discover these theoretical geometries during training, with MSE to the constrained belief theory decreasing rapidly for post- attention activations while MSE to the full belief geometry simultaneously decreases for post- post-MLP activations. The final converged representations show strong quantitative agreement with our theoretical predictions (Fig. A4), with MSE values orders of magnitude lower than random initial- ization for both the constrained and full belief geometries.
where T is the HMM's hidden state transition matrix (marginalizing out the emissions), and T is the HMM transition matrix conditioned on seeing token (see App. A for details). Eq. (5), interpreted as a context-induced point in a vector space, is the natural geometric embedding of
Pr((a) + Σ Pr(Sa|2, 24) Pr(Sa)
(6)
4.4. Attention Implements a Spectral Algorithm to Build the Constrained Beliefs
This equation describes the best possible embedding if you haven't seen any context, Pr(Sa) - ग,, followed by inde- pendent corrections to that prediction from the token each preceding context position. (r(5a11 7.g 23) Pr(Sa) TT
Eq. (5) shows how the attention pattern in our model must relate to powers of the Markov transition matrix of the underlying hidden states, 7mM where is the relative token distance.
Notably, since Eq. (6) is a distribution over latent states Sa rather than merely the next token Za+ I, this constrained-update equation naturally implemented by attention implies a probability density over all extended futures Z6+1:20 rather than just the next timestep.
To understand how this works, we turn to spectral analysis. The main goal is to understand how information or influence from a past token (at source position s) propagates to affect the belief state at the current destination position d. This influence mathematically depends on the sequence of hidden state transitions between and captured by powers of the HMM transition matrix, T
It is useful to take a step back and get handle on the intu- ition for Eq. (5) and Eq. (6). Bayesian inference (Eq. (2)) requires multiplying tokenfspecific transition matrices, fundamentally recursive process where updates depend on the full history integrated up to the previous step. In con- trast, an attention head computes its output at position d via a parallel, feedforward weighted sum (Eq. 3) of value vectors, Us from source positions S d. Crucially, each v, contains only local information from position s. It cannot directly access or depend on the specific tokens between and d due to the parallel nature of the value computation and attention weighting. Therefore, the most information that token zs can independently contribute to the belief at position d within this single-layer constraint is the correc- tion derived from knowing 2, occurred d s steps prior, assuming default starting belief π and knowledge of intervening tokens. These independent displacements from the stationary distribution over latent states is the difference of probability distributions: Pr(Sa] Z.) - Pr(St). Linear algebraically, this contribution is precisely -.Tl,.dd= - n Summing these over all past sources naturally yields the constrained belief update form in Eq. (5) (see Fig. AI for a conceptual diagram). It represents the best possible parallel approximation achievable by single attention layer given its architectural limitations.
Spectral decomposition (using eigenvalues À and associated projectors 72) is standard mathematical tool to analyze matrix powers because it simplifies T" into a sum Σx A"T,. The cigenvalues (A) are crucial because they tell us the rate at which the influence of past information decays (if [|] < 1) or even oscillates (if À is negative complex) as the distance = d increases.
Our key finding is that the learned attention weights (Ad,. in Eq. 3) directly implement this propagation effect, effec- tively learning to approximate the A decay predicted by the theory. This explains why attention patterns often show exponential decay, and why multiple heads are needed to capture oscillatory patterns arising from negative eigen- values. Furthermore, this spectral perspective allows us to make precise, verifiable predictions about the learned OV vectors and token embeddings, directly connecting the dynamics of the data (via T's eigenvalues) to the specific parameters learned by the transformer.
When T is diagonalizable with a set of eigenvalues Ar, then has simple spectral decomposition such that we can
Eq. (5)'s constrained belief geometry closely matches the intermediate structure observed in the central range of a E [0.2, 0.6), as shown in the left two columns in Fig. 4. As a moves further from this range, we observe gradually
6

---

Constrained belief updates explain geometric structures in transformer representations
rewrite Eq. (5) as
span ( ((()) A (A) - (N) and X + 1)) and, in particular, (()))1 0 such that adding any of the OV vectors to any stochastic vector (whose elements by definition add to one) keeps you in the hyperplane of the probability simplex.
p TT
For the Mess3 family of processes, T has single eigenvalue C - 1 - 3x with multiplicity aç = 2 besides its eigenvalue of 1. Accordingly, Eq. (11) simplifies to
where 73 is the spectral projection operator associated with eigenvalue (Riechers & Crutchfield, 2018a). In this di- agonalizable case, TA Ex=
k=1 ( where ax is the algebraic multiplicity of the eigenvalue X, with right eigen- states satisfying T' (Ax} = A (AR), left eigenstates satisfying (A.|T= = X (Ax): I. all satisfying the orthonormality condition (AylAe)= = 6 Notably in Eq. (7), all dependence on inter- token distance now lies solely in the exponentiation of the eigenvalues, which all live on or within the unit circle in the complex plane for stochastic transition matrix like T.
(as = - (TTTTTIITT
for d > s,
(12)
which forces f( = cT To for some e E IR indepen- dent of d. from which we obtain
Aatm. = 5 - Ad,
for d > s.
(13)
So, for example, A2, implies Ad,1 for all destinations d
2 2; and A3,2 implies Aa,2 for all destinations d 2 3.
For the Mess3 process, the stochastic matrix T has eigenval- ues Ar = (1,(), {). where Ç = 1 - 3x is a degenerate eigen- value with multiplicity ag- = 2. We observed that the atten- tion weight n tokens back is approximately <" = (1 (1-3x) 3x"", which suggests a strong connection between the theoreti- cally motivated Eq. (7) and the catiiieettrallimpleeeennaa
[1) = I 1M, since all projection operators must sum the identity. Combining this insight with Eq. (12) tells us about the OV-vector for all positions:
For Mess3, T3
C (IT - T
f(ūm) = Am+1.m
(14)
Eq. (3). Encouraged by this correspondence and further ev- idence of similarity, we make the ansatz that the role of attention in the first layer is to implement the constrained belief update of Eq. (6) via Eq. (7)'s 's spectral mechanism" 2. Taking this ansatz seriously allows us to precisely anticipate the analytic form of the learned attention pattern.
Notably, Eq. (14) tells us that all OV-vectors associated with the same token must be (T)) x f(Tg)) if 2, which is consistent with what observe in our experiments (Fig. 2A). Moreover, the magnitude of the mh OV-vector is inversely proportional to the attention element Am+1.m which is again consistent with our experiments (Fig. 2AB). In our experiments, we find A2, to be signif- icantly larger than all the other Am+1.m elements, while the latter all cluster together; the magnitude of 51 is corre- spondingly smaller than all of the other strongly clustered o. magnitudes.
To derive the analytic form of the attention pattern, we assume that there is linear map ʃ IR dgdmaan + R S1
from the residual stream to the hyperplane containing the probability simplex over the hidden states of a minimal generative model of the data (the 2-simplex in this case). Let a = - Ti = I - 1 be the projection from IR S to the hyperplane IR S 11 containing the simplex. Our full ansatz is thus f(
Combining Eqs. (10) and (14) constrains the embedding
o ((21)) IT or, more explicitly:
f (pre)) m = (1
CA m Amtiim
(T Am
π
(15)
f(7 (ma))
Σ
X TTx
(8) to be parallel to the OV-vectors, as we indeed observe.
s=1 \EAri(1)
Eqs. (13), (14) and (15) make strong predictions about the form of attention pattern and how it relates OV-vectors and token embeddings, which must be true if the first layer of attention is indeed implementing the constrained belief updates over latent states of a generative model of the train- ing data. These relationships are all borne out in our ex- periments (Fig. 2ABC), except for some scalar discrepancy in the first two embedding vectors (see Appendix D for quantification), which is strong validation of the predictive power of our framework.
=
) +
(1.).
(9)
From this, we group source- sperci-specific terms to infer that
f( + Aa,af(Ta) = TT - T
(10)
and
Ad,f(T))=
Σ
p TTSSTT ford (11)
AEAr((1)
From Eq. (11), we notice that f(7)) is in the linear span of the non-stationary left eigenstates of T. Le., f())) E
4.4.1. 1. NEGATIVE FIGENVALUES REQUIRE MORE ATTENTION HEADS
"The details of this correspondence break down if there are many attention heads in the first layer.
For the transition matrix T to be row stochastic (a require- ment for a valid HMM), x must be in the range [0, 0.5).
7

---

Constrained belief updates explain geometric structures in transformer representations
Interestingly, when < 0 (which occurs when x 1/3), the predicted patterm oscillates and cannot be captured by single attention head, since attention pattern entries must be non-negative. In these cases, we observe that a single-heus heage- transformer captures an incomplete representation of the belief state geometry, and the transformer performs corre- spondingly worse (App. F). However, upon adding a second attention head, the model converges to the solution pre- dicted by the belief updating equation, even in the presence of oscillatory dynamics, as shown in Fig. 3.
A
PCA of Input Space and OV Vectors
B
Attention Patterns
15-
vectors, V,
0.6
Heze
Head 2
Second Principal Component
Embedding. R
0.4
0.2
5-
0.0
-5
-0.4
-0.6
-0.0
The anticipated need for second attention head when the data-gener transition matrix has a negative eigenvalue further demonstrates how our analysis provides a handle to relate the architectural constraints of the attention mech- anism to the structure of the training data. In fact, our framework provides more specific predictions for the atten- tion pattern and its relation to embedding and OV-vectors in this case too.
-15
-10
6.6 4 2
First Principal Component
Relative Position
Figure 3: Attention heads combine to capture oscillatory dynamics belief updating. (A) the token embedding space, the model uses each attention head embed tokens opposite poles of the simplex.
(B) The attention patterns of the two heads (shown here averaged over all sequences) act as positive and negative components. When combined, they produce the oscillatory pattern predicted by the expo- nentiated eigenvalue ç" = (1 - -zz)" = (-1)" 1)". (1-3z)"
With two attention heads, the update to the residual stream at the destination position s becomes
(16)
more complex. Through purely local computations at each position, the MLP learns continuous nonlinear warping that transforms the intermediate fractal structure into the final belief geometry.
where each head now has its own QK and OV matrices. With the negative cigenvalue Ç < 0 and two attention heads, we can relate the constrained belief update to the details of attention and embedding via
Fig. 4 shows the close match between theoretical predic- tions and observed representations across different Mess3 parameter settings, confirming our theoretical understanding. The transformation involves stretching and compressing dif- ferent regions of the space while maintaining topological structure, though a full characterization of its mathematical properties remains for future work.
[( (mid)
TTI
(17)
[ (pto)
A ((v
A (T(P))
=
This is naturally accommodated by
To verify that our single-layer analysis captures the funda- mental computation even in deeper networks, we analyzed 4-layer transformers trained on the same Mess3 data. We find that the first attention layer consistently implements the constrained belief update mechanism, Eq. (5), with subse- quent layers progressively transforming the representation toward the full Bayesian belief geometry, Eq. (2). Quantita- tive MSE analysis (Fig, A6) shows that first-layer attention outputs align well with constrained beliefs while later lay- ers systematically improve alignment with full Bayesian beliefs. This progression shown in Fig. A5, confirming that our theoretical framework for the first layer remains the foundational computation in multi-layer networks.
1 + + -6 TT and (18)
( - - ----1---,1----11 1
' TT
(19)
for d > which implies that the Ve-tors point in op-
posite directions, f(o!")) 2)( with f( x
(TT|== T) and
A d42m, = 2m (h)
for > 8,
(20)
consistent with our experiments as shown in Fig. 3. We note that the magnitudes of OV vectors are tied to attention magnitudes via cça-s = A f( - l (x(()) with c = | v π| E R. which is also observed in Fig. 3.
5. Discussion and Conclusion
4.5. Post-MLP geometries
We have shown how combining computational mechanics with mechanistic interpretability yields a principled un- derstanding of why transformers trained on Mess3 HMM
While the intermediate geometry is well characterized by our constrained belief equations, the MI. transformation is
8

---

Constrained belief updates explain geometric structures in transformer representations
Intermediate
Final
general observations that attention heads specialize, our analysis reveals precisely why and how multiple heads must coordinate: the non-negativity constraints of attention, com- bined with oscillatory patterns in optimal belief updates, necessitate specific decompositions across heads, providing concrete mechanistic understanding of their functional roles. This demonstrates how combining theoretical understanding with architectural constraints can yield precise, verifiable interpretations of neural network components.
Theoretical
Activations
Theoretical
Activations
x-0.15
0-0.6
0-0.
0=0.2
Limitations and future work. Our experiments used small transformers (1 and 4 layers), with analysis focused on the first attention layer. For training data we used the specialized Mess3 family of HMMs with full support over the space of all possible sequences of tokens. We discovered how transformers implement belief updates when attention patterns depend primarily on positional distances, while token-specific information is handled through value vectors. While we validated that the first layer of multi-layer trans- formers implements the same constrained belief updating mechanism (Figs. A5;A6), with subsequent layers refining toward full Bayesian beliefs, techniques must be adapted to both larger transformer architectures and data- getargenng processes that capture the complexities of real-world data. While this setting offers clear insights, it does not capture many aspects of natural language. Future work could apply these techniques to processes that better reflect properties of natural language--hierarchical, with sparse support over sequences. Morcover, the interplay between multi-heat at- tention and deeper layer stacks likely exhibits additional nuances that our primary single-layer analyses only begins to uncover. Finally, while we showed that the final MLP layer refines partial updates to approximate full Bayes, the deeper question of why gradient descent converges on these circuits remains ripe for further investigation.
x=0.5
o=0.6
x=0.5
0=0.2
Figure 4: Comparison of model representations and theo- retical predictions for different Mess3 hyperparameters in each row. Each subfigure shows four columns: (i) Interme- diate representation from Eq. (5). (ii) PCA projection of the model activations in the intermediate layer. (iii) Ground
truth belief state geometry from Eq. (2). (iv) PCA projection of the final activations after the MIP.
data learn intermediate fractal-like structures, and how these structures systematically transition into final belief-state representations. By focusing primarily on single-layer transformers, we isolated the fundamental computation per- formed by the attention mechanism, while also demonstrat- ing that these principles extend to the first layer of deeper networks. This approach provides top-down theoretical ex- planation grounded in the tension between optimal Bayesian belief updates and the parallel, attention-basee constraints of transformer computation, developing geometric observa- tions of activation space into mechanistic understanding of the underlying computational principles.
Conclusion. By combining computational mechanics with mechanistic interpretability, we have shown how trans- formers implement inherently recursive Bayesian updates through parallel computations via the attention mechanism, and how these intermediate representations are refined into the final form. This reconciles model-agnostic theories of next-token prediction with the reality of architecture- specific constraints. We hope our results not only advance interpretability for HMM-like toy tasks but also inspire deeper theoretical insights into how large-scale transformers produce- sellitt-bllief-tee belief-like structures in applications.
Implications for interpretability. Our work demonstrates an alternative to bottom-up up architectural analysis. Knowing the structure of optimal predictors allows us to predict and verify the specific intermediate computations that are imple- mented by the attention mechanism. Our analysis reveals the computational role of specific directions in activation space- showing how the geometry of belief updates shapes the learned representations. Additionally, by focusing on a small, tractable HMM, we see how specific properties of its transition matrix lead to oscillatory patterns that require specialized multi-head head solutions due to the non non-negatyvity constraints of attention mechanisms. Rather than relying
A Acknowledgments
The authors are grateful for the community and financial support from MATS, PIBBSS, FAR Labs, BITS, and Astera Institute, and MP's further financial support from Open

---

Constrained belief updates explain geometric structures in transformer representations
Philanthropy during the MATS extension program, which made this project possible.
McCandlish, S., and Olah, C. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. URL /ttp://////asa/mmeeerr/ee -otrrannsfommerriirttttt circuits. wub//22l//rraewwzk////
Author Contributions
html.
MP discovered the attention-based constrained belief updat- ing algorithm in the simplex, and performed the bulk of the experiments with mentorship from ASS. PMR developed the mathematical theory together with MP and ASS. ASS su- pervised the project, and DF provided project management. MP. PMR, and ASS wrote the manuscript, with helpful guidance from DF. MP, PMR, and ASS performed analysis, and established the correspondence between transformer behavior and theoretical predictions.
Elhage, N., Hume, T., Olsson, C.,, Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds. Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D., Wattenberg, M., and Olah, C. Toy models of superposition. Transformer Circuits Thread, 2022. URL https: htttss:/transfoomm: pub hubb/202//oyy mwde /aeeeliinndde html.
ci circuits.
Engels, J., Michaud, E. J., Liao, I, Gurnee, W., and Tegmark, M. Not all language model features lin- ear. arXiv preprint arXiv:2405.
Impact Statement
2024.
This paper presents work whose goal is to advance the inter- pretability of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Kaelbling, L. P., Littman, M. L., and Cassandra, A. R. Plan- ning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-33, 1998.
Kingma, D. P. and Ba, J. Adam: A method for stochastic op- timization, 2017. URL https: /ttps///axvvvvooo abs/ 1412. 6980.
References
Bricken, I., Templeton, A., Batson, J., Chen, B., Jermyn, A., Conerly, T., Turner, N., Anil, C., Denison, C., Askell, A., Lasenby, R., Wu, Y., Kravec, S., Schiefer, N., Maxwell, T., Joseph, N., Haffield-Dodds, Z., Tamkin, A., Nguyen, K., McLean, B., Burke, J. E., Hume, T., Carter, S., Henighan, T., and Olah, C. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. URL https: // monoeemantic-featueess. html..
Marks, S., Rager, C., Michaud, E. J., Belinkov, Y., Bau, D., and Mueller, A. Sparse feature circuits: Discovering and editing interpretable causal graphs in language mod- els, 2024. URL https: /tap:::///rxiiiii /ttpiv..ivvvv..///////.../aaaa org/ abs/ abs/2433 19647.
Marzen, S. E. and Crutchfield, J. P. Nearly maximally predictive features and their dimensions. Physical Re- view E, 95(5), May 2017. ISSN 2470-0053. doi: 10. 10.1103/phyatrv..55555551.. URL http: / doi. o
Conmy, A., Mavor-Parker, A. N., Lynch, A., Heimersheim, S., and Garriga-Alonso, A. Towards automated circuit discovery for mechanistic interpretability, 2023. URL https:/ / //taxiivv../////////////////aav/ www.wi.ww./aaaaaaaeee..oo btgg/bs//333 14997.
95. 051301.
Nanda, N., Chan, L., Lieberum, T., Smith, J., and Steinhardt, J. Progress measures for grokking via mechanistic inter- pretability, 2023. URL https: /atxxaaxvvv/v////////ea///aaaa org/ abs// 05217.
Crutchfield, J. P. The calculi of emergence: Computation, dynamics, and induction. Physica D, 75:11-54, 1994.
Park, K., Choe, Y. J., and Veitch, V. The linear representa- tion hypothesis and the geometry of large language mod- els, 2024. URL https: /ariiiiiivv.////ii//.aaa/aaaiiiv org/ abs 2311. 03658.
Crutchfield, P. Between order and chaos. Nature Physics, 8(1):17-24, 2012.
Cunningham, H., Ewart, A., Riggs, L., Huben, R., and Sharkey, L. Sparse autoencoders find highly interpretable features language models, 2023. URL https:// arxiv. org. abs 2000/2339. 08600.
Pepper, K. RNNs represent belief state geometry in their hidden states. https:/appartessarcc.com Research submission to the Computational Mechanics Hackathon research sprint co-hosted by Apart, PIBBSS, and Simplex.
June 2024.
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D.,
Riechers, P. M. and Crutchfield, J. Beyond the spectral theorem: Decomposing arbitrary functions of nondiago- nalizable operators. AIP Advances, 8:065305, 2018a.
Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J.,
10

---

Constrained belief updates explain geometric structures in transformer representations
Riechers, P. M. and Crutchfield, J. P. Spectral simplicity of apparent complexity, Part I: The nondiagonalizable metadynamies of prediction. Chaos, 28:033115, 2018 8b. doi: 10. 1063/1031100 14985,,,
Shai, A. S., Marzen, S. E., Teixeira, L., Oldenziel, A. G., and Riechers, M. Transformers represent belief state geometry in their residual stream. accepted to Ad- vances in Neural Information Processing Systems 37 (NeurIPS 2024), 2024. URL https: /ttp::::///xxxooooo abs/2405.. 15943.
Shalizi, C. R. and Crutchfield, J. P. Computational mechan- ies: Pattern and prediction, structure and simplicity. J.
Stat. Phys., 104:817-879, 2001a.
Shalizi, C. R. and Crutchfield, J. P. Computational me- chanics: Pattern and prediction, structure and simplicity. Journal of statistical physics, 104:817-879, 2001b.
Templeton, A., Conerly, T., Marcus, J., Lindsey. J., Bricken, T., Chen, B., Pearce, A., Citro, C., Ameisen, E., Jones, A., Cunningham, H., Turner, N. L., McDougall, C., MacDiarmid, M., Freeman, C. D., Sumers, T. R., Rees, E., Batson, J., Jermyn, A., Carter, S., Olah, C., and Henighan, T. Scaling monosemanticity: Ex- tracting interpretable features from claude sonnet. Transformer Circuits Thread, 2024. URL https: //transformer- circuits. secdlingmmoooemmttiittt/nnee...
Upper, D. R. Theory and algorithms for hidden Markov models and generalized hidden Markov models. Univer- sity of California, Berkeley, 1997.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A., Kaiser, L.,, and Polosukhin, I. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Wang, K., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J. Interpretability in the wild: a circuit for indirect object identification gpt-2 small, 2022. URL https: /arxivv... owg/abs///aaeeoooooo..ooo..e.bee abs /21122211..0 00593.

---

Constrained belief updates explain geometric structures in transformer representations
A. Mathematical Details of HMMs and Belief State Geometry
In this work we created training data from a class of Hidden Markov Models (HMMs) called Mess3. The HMMs have three hidden states S {1, 3) and emit from a vocabulary of three tokens = {0. 1. 2.
The HMMs in this class are parameterized by and x, with dependent quantities 3 (1 a)/ and y =1- - 2x.
The labeled transition matrices define the probability of moving to state j (indexing columns) and emitting the token on the label, 2, conditioned on being in state i (indexing rows), P(33, 2|s;) and are:
ory
3r
3x
T
a
By
3
(21)
002
fx
By
By 0.
3x
T(1)
Br
oxy
Bar
(22)
Br 060
By
By
3x
0.E
(
Bx
By
OCE
(23)
Br
Ax
ozy
Note that even though the dynamics amongst the emissions are infinite-Markon order, the dynamics amongst the hidden states are Markov, with a transition matrix given by marginalizing out the token emissions: T T.e)T(=,
Since Mess3 has non-zero row sums for each labeled transition matrix, we can also define a conditional transition matrix, Tls, with elements T = P(sj (2,88.)) 8s), which is given by normalizing each labeled transition matrix such that every row sums to
A.1. Full belief updates
An important part of the work presented here is about how an optimal observer of token emissions from the HMM would update their belicfs over which the hidden states the HMM is in, given a token sequence. the observer a belief state given by a probability distribution 77 (a row vector) over the hidden states of the data-generating process, then the update rule for the new belief state ny' given that the observer sees new token is:
T) TT)
r''
(24)
where 1 is column vector of ones of appropriate dimension, with the denominator ensuring proper normalization of the updated belief state. In general, starting from the initial belief state na. we can find the belief state after observing a sequence tokens 20: 21,...,,N. 2N:
No ((o)) ... T(Ex)) no ((1) ... T
(=1a))
(25)
For stationary processes, the optimal initial belief state is given by the stationary distribution n = π over hidden states of the HMM (the left-eigenvector of the transition matrix T Σ. T()) associated with the eigenvalue of 1).
The beliefs have geometry associated with them, called the belief-state geometry. The belief-state geometry is given by plotting the belief distribution over the HMM's hidden states induced from each possible sequence of tokens as a point in the probability simplex over these hidden states.
A.2. Constrained belief updates
Incorporating past contributions to belief updates in parallel, as the attention mechanism suggests, we instead obtain
Tm
π
π
(26)
# (=
12

---

Constrained belief updates explain geometric structures in transformer representations
For processes like Mess3 that have non-zero sums for each labeled transition matrix, this can be written more simply as:
1
daia)
k π
(27)
which is the form that appears in the main text. For other processes that don't satisfy this condition, slight modifications of the equations in the main text follow straightforwardly from Eq. (26).
B. Model architecture and training procedure
We employ standard single-layer transformer model with learned positional embeddings. The model architecture follows the conventional transformer design, with dmodel 64 and dæ 256. Depending on the Mess3 parameters, we use either single-head or double-head attention mechanism. We conduct systematic sweep over the HMM parameters a and r. training separate model for each pair. Models trained next-token prediction using cross-entropy loss, with batch size 128. We use Adam optimizer (Kingma & Ba, 2017) with 10 learning rate and no weight decay. Each model is trained for approximately 15 million tokens.
We generate all possible input sequences up to length 10, recording hidden activations from the transformer's residual stream. These activations are organized into a dataset capturing the model's response to all input patterns.
Input sequences consist of three symbols, embedded with positional information, without beginning-of-sequemne (BOS)
token.
13

---

Constrained belief updates explain geometric structures in transformer representations
C. Recurrent vs. Parallel Belief Updating
Optimal Bayesian Inference Fundamentally Recurrent
Input Sequence
Z+
ZZ
Z
Z4
Zs
Full Bayesian Inference
Belief at position i depends directly belief
Belief, P(S, Th
Belief, P(S,(z.2g) Bellef, Tie
Belief; P(S.(2,2,2,) Bellef, T
Belief, ((S.12,2,2,2..)) Bellefa T
Beliefs
Bellef States
2(,,12,2,2,,225))) Bellef, T
Sequential multiplication token specific transition matrices
Belief = n7 77-7% T T 17.1 7
Each value vector, v, contains local information from position
The Attention Mechanism Fundamentally Parallel
The residual attention weighted sum of vectors
Residual Stream
2t
Zz
Zz
Z4
Zs
contributions cannot directly access information between the source target positions
Value Vectors
V4
V22
V/3
V4
Vs
ZZ
attention weights
015
825
B
ass
Attn Updates agvi
Resid, After 2, aivi
Resids Σ a
Constrained Belief Updating Recapitlates the Parallel Nature of Attention
Input Sequence:
Each source token contributes correction based on its value and its position relative to the destination
Z₁
Zz
Z4
Zs
direct access to tokens between source and destination tokens
Contributions
Independent Contribs.
Contrib,
Contrib;
Contrib,
Contrib,
Contrib;
Contributions rañlect difference between knowing just source token's identity (with no detalls about tokens in between) having no knowledge of any tokens.
m-T+ T
T-PP-P--
T TP
Beliefs T Σ contrib,
This is parallel approximation of process
Figure A1: Diagrammatic visualization comparing the recurrent nature of Optimal Bayesian Inference (top), to the parallel attention mechanism (middle), and the parallel Constrained Belief Updating presented this paper.
14

---

Constrained belief updates explain geometric structures in transformer representations
D. Quantification of Theoretical Predictions
A
Position
Position
1.0
1.0
0.5
0.5
vectors,
Embedding, &4
0.0
0.0
0.0
-0.5
-0.5
-0.5
-1.0
-0.5
0.0
0.5
1.0
-1.0
0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
B
0207
3887
8500
220
0387
Normal Fit Cross valiatalottion
0.0
0.8
0,532
Soure
0.1
6.4
Nocoal
Uztsa valiclation
0002
0.0
0.8
Source Plu 2.2
Sourse Parkfoms at
Stare Pouition
Soure Posfituss
4225
OV Model
Embed Model
Nocual
2.0
LEENCE
Meen Squared Error
Squared
La otede
Bourne Pasllsst
Beuuroe Photioal
Braaten Prottact
Beuzte Pasiiion at
Figure A2: Embeddings for the first two positions are correctly predicted to parallel to the OV vectors, as with all of the embeddings; however the sign of the predicted embedding for these first two positions deviates from the observed embedding. We do not yet understand the reason for this discrepancy, but still find it remarkable that the bulk of the high-dimensional computation carried out by attennioo pattern, OV vectors, and all embeddings beyond the first two positions-co can be very precisely understood by sequence of operations in the two-dimensional simplex.
E. Multi-layer experiments
F. Minimal architectural requirements
To verify our theoretical understanding of the transformer's computational requirements, we conduct a systematic evaluation across different architectural configurations. Fig. A7 shows that the model achieves good performance with minimal architecture: a single layer with two attention heads sufficient to achieve low KI. divergence across different Mess3
15

---

Constrained belief updates explain geometric structures in transformer representations
Evolution of Normalized MSE During Training (with Zoomed Inset)
Mid (Post-AttnPPee-MLPP)
Post (Post-MLP)
Random Init
x15 a20
x15 a60
2.0
0.45
0.4
1.00
1.6
0.75
1.2
Normalized MSE (MSE Initial MSE)
0.50
0.8
0.4
0.25
0.0
0.00
x50 a20
X50 a60
4.0
1.00
3.2
0.75
2.4
10
15
0.50
1.6
0.25
0.8
0.0
0.00
4
12
16
20
4
8
12
16
20
Training Checkpoint Index (0 = Random Init)
Training Checkpoint Index (0 Random Init)
Figure A3: Evolution of Normalized Mean Squared Error (MSE) During Training Across Experimental Conditions. Each subplot shows the MSE for one experimental condition (labeled by title), normalized by the initial MSE value from the randomly initialized model (represented at Checkpoint Index 0). The Y-axis represents MSE relative to this random baseline (Lower is better), where the dashed line at Y=1 indicates performance equal to random initialization. The X-axis represents the training checkpoint index, starting from the random initialization at index 0. Lines: Show the normalized MSE for Mid (Post-Attenion/Pre-MLP blue circles) and Post (Post-MLP, orange squares) activations, reflecting their fit to their respective thcoretical geometries over training. For Mid activations we regress to the Constrained Belief geometry, while for Post activations we regress to the Full (unconstrained) Belief geometry. Insets: Provide a zoomed-in view of the Y-axis from -0.02 0.5, highlighting the convergence behavior low MSE values. Observations: Across all conditions, both Mid and Post activation representations show significantly improved geometric fits (normalized MSE drops well below compared to the random baseline training progresses. Some conditions (c.g., a=20) may show initial transient increase in normalized MSE before rapid improvement.
16

---

Constrained belief updates explain geometric structures in transformer representations
Final Trained Model Fit to the Theoretical Predictions
0.008
Mid (Random Init)
Absolute MSE (Lower is Better)
0.007
Mid (Final Trained)
Post (Random Init)
0.006
Post (Final Trained)
0.005
0.004
0.003
0.002
0.001
5.1e-04
2.1e-04
8.1e-05
1.9e-04
2.5p.04
0.000
x15 a60 Experimental Condition
x50 a20
x50_a60
Figure A4: Comparison of Final Trained Model Fit Random Initialization Across Experimental Conditions. The bar chart displays the Absolute Mean Squared Error (MSE, lower is better), comparing the geometric fit of activations at the final training checkpoint against the initial random model state. Results are shown for four different experimental conditions (X-axis). Within each condition, bars represent: Mid PPost-AttentionPPee-MLL Activations, MSE of the regression to the Constrained Belief Geometry: Blue bars. Post (Post-MLP) Activations, MSE of the regression to the Full (unconstrained) Belief Geometry: Orange bars. Solid bars indicate the MSE for the final trained model, while transparent bars of the same color show the MSE for the randomly initialized model (baseline). Numerical labels specify the precise MSE values achieved by the trained models. Observations: Across all conditions, the trained model (solid bars) achieves lower MSE, indicating much better fit to the underlying theoretical geometries compared to the random baseline (transparent bars).
17

---

Constrained belief updates explain geometric structures in transformer representations
x=50, a=20
x= a=
Layer
Layer 1
Layer 2
Layer 3
Layer
Layer
Layer
Layer
Resid Mid
Resid Post
x=15, a=20
x=15, a0
Resid Mid
Resid Post
Figure A5: Visualization of Activation Geometry via Regressions onto the Full Belief State. The figure displays 2D projections of transformer activations after linear regression onto the Full (unconstrained) Belief state geometry (Eq, (2)) for all stages, arranged in a 2x2 grid. Each quadrant corresponds to a different experimental condition (combinations of x and a, labeled above each quadrant). Within each quadrant, the top row shows visualizations for Mid (Post-Attention/Pre MLP) activations, and the bottom row shows visualizations for Post (Post-MIPP) activations, across Layers 0-3 (columns). Points in all plots colored according the ground truth Full Belief state coordinates (Eq. (2)), projected into RGB space. Importantly, all visualizations show regressions to the Full Belief state geometry (and not the Constrained Belief geometry). Top Row (Mid Activations): Shows Mid activations after being regressed onto the Full Belief geometry (Eq. (2)). Bottom Row (Post Activations): Shows Post activations after being regressed onto the Full Belief geometry (Eq. (2)). This visualization qualitatively shows how well activations linearly map to the Full Belief state throughout the network. Notably, the residual stream activations after the first attention layer but before the MLP (Layer 0 mid, top-left plot within each quadrant), despite being regressed onto the Full Belief target, visually retain structure resembling the Constrained Theory geometry (Eq. (5)). This indicates that the constrained structure is the dominant feature linearly recoverable from early Mid activations, even when seeking the best fit to the final target geometry. Comparing subsequent Mid and Post stages across layers (moving rightwards) reveals the accurate fit to the Full Belief state geometry.
18

---

Constrained belief updates explain geometric structures in transformer representations
Multilayer Transformer Fits Constrained and Full Belief Geometry
Trained: vs Constrained Theory
Trained: vs Full Belief Random: vs Full Belief
10A0: Random: vs Constrained Theory
--.-.
X= 50.0, a=20.0
x=50.0, a660.0
10
10-3
Absolute MSE (Log Scale)
10'4
x=15.0, 8=20.0
x=15.0, a=60.0
10-4
10a
10
Pest
sed
Host
Mid
Post
Mid
L3 Rool
Phar
Post
Mid
Past
Mid
Post
Residual Stream Stage
Figure A6: Quantitative Fit of Multi-Layer Transformer Activations to Theoretical Geometries Across Residual Stream Stages. Absolute Mean Squared Error (MSE, log scale) comparing activation representations to the Constrained Theory (Eq. (5)) and the Full (unconstrained) Belief state geometry (Eq. (2)) across interleaved residual stream stages (LO Mid, LO Post, L3 Post) of 4-layer transformer. The figure presents results for four experimental conditions (combinations of x=15, 50 and a=20, 60) in 2x2 grid. Lines: Show absolute MSE comparing activation fits to theoretical geometries. Solid lines represent the trained model; dashed lines represent the random baseline. Green Stars: Fits of the residual stream activations to the Constrained Theory (Eq. (5)). Orange Circles: Fit of the residual stream activations to the Full Belief geometry (Eq. (2)). At the residual stream after the first attention but before the first MLP (LO Mid), the fit to the Constrained Theory (green) is better (lower MSE) than the fit to the Full Belief geometry (orange) across all conditions, supporting the hypothesis that the initial layer's attention mechanism implements the constrained update. The fit to the Full Belief geometry (orange line) improves dramatically after the first MLP, and at that point the residual stream activations switch to fitting the Full Belief geometry better than the Constrained Belief geometry. The fit to the Constrained Theory (green line) does not show this convergence and may worsen in later layers. The MSE values for the trained transformer (solid lines) are consistently orders of magnitude lower than the corresponding random baselines (dashed lines), demonstrating that these geometric alignments are learned features resulting from training.
19

---

Constrained belief updates explain geometric structures in transformer representations
layar
layers
6:15, 26
10.3 1-t0
0.02
Voddetn
0.00
Figure A7: Validation KL divergence between model predictions and optimal probabilities across different architectural configurations. Results shown for various Mess3 parameter settings (æ and a:) and model architectures (number of heads and layers). The model achieves good performance with minimal architecture: single layer with two attention heads is sufficient across parameter settings.
parameter settings. This empirical finding aligns with our theoretical analysis when x >> 1/3, the belief update patterns contain oscillatory components that require two heads to implement due to the non-negativity constraint of attention. The necessity two heads visually demonstrated in Fig. A8. For x 0.5, where the optimal update pattern has significant oscillatory components, a single-head transformer fails to capture the correct belief geometry. With heads, the model can properly implement these updates through complementary attention patterns, resulting in representations that closely match the ground truth geometry.
G. Dimensionality of Residual Stream Activations
Table 1: Cumulative explained variance ratios for PCA components of the residual stream activations the intermediate position (after attention) and the final position (before unembedding). The table shows results for different settings of the Mess3 HMM parameters x and a.
Intermediate
Final
T 0.15
0.15
0.5
0.5
0.15
0.15
0.5
component
0.2
0.6
0.6
0.2
0.2
0.6
0.6
0.2
0.5408
0.4648
0.4074
0.5268
0.9618
0.4947
0.4596
0,6503
0.8768
0.8894
0,8028
0.8519
0.9825
0.7681
0.7096
0,8592
0.9673
0.9859
0.8913
0.9173
0.9943
0.9811
0.8855
0.9689
0.9749
0.9903
0.9455
0.9649
0.9960
0.9897
0.9189
0.9755
0.9815
0.9929
0.9848
0.9886
0.9969
0.9916
0.9428
0.9807
5
0.9870
0.9942
0.9978
0.9977
0.9976
0.9931
0.9586
0.9850
6
0.9914
0.9955
0.9986
0.9984
0.9981
0.9945
0.9723
0.9886
We perform PCA on the residual stream activations after the attention module (intermediate) and before the unembedding layer (final). The effective dimensionality of the residual stream is low, with the first few components capturing most of the variance (Sec Table 1). most cases, the first components explain over 90% of the variance. For r = 0.5, the effective dimensionality is higher, possibly due to the oscillatory dynamics the belief updating equation this regime. Further investigation is needed to fully understand this phenomenon.
20

---

Constrained belief updates explain geometric structures in transformer representations
Final geometry (1 head)
Ainal geometry heads)
Ground bruth
X-0.5 8-0.5
Figure A8: Comparison of learned belief geometry with one head (left) versus two heads (middle) against ground truth (right) for different Mess3 parameter settings. With 0.5, where the optimal update pattern requires both positive and negative components, single head fails capture the correct geometry due to the non-negativity constraint of attention. Two heads allow the model to properly implement these updates, resulting in geometry that closely matches the ground truth.
21


Reviews:
[{"review_id": "dICuplTPon", "reviewer": "Reviewer_rrWD", "summary": "This paper discussed the circuit implemented within trained Transformer to implement (partial) Bayesian inference on a type of Hidden Markov Model. The paper combines theoretical analysis with interpretability tool to showcase that the implemented circuit and corresponding representation is optimal when accounting the architectural constraint of 1-layer Attention. Overall, this paper produce a satisfactory explaination on the inner working mechanism of Transformer on the Mess-3 HMM.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your strong endorsement of our work. We appreciate your positive assessment of our paper's strengths, and we agree that this paper should make a nice contribution to both interpretability and the study of representations.\nWe also appreciate your constructive feedback, which nudged us to investigate and confirm that our results indeed generalize to the multi-layer case. \n\n**Multi-Layer Extension:**  To directly address your question about multi-layer extensions, we have trained and analyzed 4-layer transformers. Our key finding is that the first attention layer consistently implements the constrained belief update mechanism (Eq. 5) derived in the paper, producing the intermediate fractal geometry. Subsequent layers (specifically the MLPs and later attention blocks) then act to transform this intermediate representation progressively closer to the full Bayesian belief state geometry (Eq. 2).\n\n We plan to add quantitative evidence for this to the appendix (and summarize in the main text). This includes:\n\n\n1. PCA visualizations of activations after each layer in the 4-layer models, showing the transition from constrained to full belief geometry. See new figure [here](https://docs.google.com/document/d/1CDwiJp8qG4d4OqgntBkIMYQhxJ1vy-2j2BecubhrL5U/edit?tab=t.0#heading=h.7gjcfnihl09i)\n2. Mean Squared Error (MSE) analysis from linear regressions of residual stream activations onto both the theoretical constrained geometry (Eq. 5) and the full belief geometry (Eq. 2). This analysis, performed layer-by-layer, quantitatively shows the initial fit to Eq. 5 post-attention-1, and the increasing fit to Eq. 2 in deeper layers. See new figure  [here](https://docs.google.com/document/d/1CDwiJp8qG4d4OqgntBkIMYQhxJ1vy-2j2BecubhrL5U/edit?tab=t.0#heading=h.uauqgkf7zz7l)\n\nThese results demonstrate that our single-layer analysis captures the fundamental computation performed by the initial attention mechanism, even within deeper networks, while later layers perform refinement.", "claims_and_evidence": "Yes, the claims are well supported.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I checked the theoretical claims produced by this paper and am convinced that they are correct.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The authors made a very thorough discussion in Section 2.1 already. This paper should be of interest to the community of both mechanistic interpretability and representation theory.\n\nThe most related work is [1], which discuss the same setting and find that Transformers approximately learn to encode the belief state linearly in the hidden state. This paper produces a comprehensive analysis on how this belief state is calculated within Transformer in a 1-layer setting.\n\n\n[1] Transformers represent belief state geometry in their residual stream", "essential_references_not_discussed": "No that I can think of.", "other_strengths_and_weaknesses": "# Strength \n1. The visualization in Figure 2 is very intuitive.\n2. The theory in Section 4 is clean and captures the essence of the architectural constraint.\n\n# Weakness\n\n1. Only 1 layer case is studied here, while this keeps the paper compact, how the current result generalize to multi-layer is unclear.", "other_comments_or_suggestions": "In Section 2.2, it would be better to mention that the same deduction and notation is inherited from Shai et.al.", "questions_for_authors": "1. Following weakness 1, has the authors tried to ablate with number of layer and will the first layer implement similar algorithm? Is there any new phenomenon in the multi-layer case that can't be explained by the 1-layer analysis here?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VMUQPy7jLO", "reviewer": "Reviewer_bbxf", "summary": "This work bridges mechanistic interpretability and Bayesian definitions of optimal predictions (constrained belief updating) to study emerging geometric structures in small (1-layer) transformer models. This is done with a solid theoretical foundation, in a ‘controlled’ experiment: We know the generative model behind the dataset used for training (the Mess3), that is an HMM with 3 states, that depends on two parameters. \n\nSadly, this is quite far from my area of expertise, so please do not give much weight to my review, and forgive misunderstanding/mistakes.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer bbxf,\n\nThank you for your positive assessment and thoughtful feedback. We appreciate that you find the work well-motivated and structured. We care very much about making the paper understandable and clear, so your suggestions for clarification are very helpful!\n\n**Emergence of Fractal Geometry:**\n\nYou asked for clarification on how the fractal drawings emerge from the belief states. That's an excellent point. These structures (visualized in Figs 1D, 4) arise directly from the iterative application of Bayesian belief updating rules within the probability simplex representing beliefs over the HMM's hidden states.\n\n- Each time a new token is observed, the current belief state (a point in the simplex) is transformed according to the update equations (Eq. 2 for the full Bayesian update, Eq. 5 for our derived constrained approximation implemented by attention).\n- Because the underlying Mess3 HMM involves probabilistic transitions between states, applying these update rules repeatedly for different token sequences acts mathematically like an Iterated Function System (IFS). IFSs are known to generate self-similar, fractal geometries.\n- The specific fractal shapes shown in our figures represent the collection of all possible belief states reachable after observing sequences up to a certain length. The \"folding\" or self-similarity comes from the fact that different sequences can lead to nearby belief states, and the update rules map regions of the simplex onto smaller regions within itself. (This connection between HMMs, Bayesian updates, and fractals is explored theoretically, e.g., by Jurgens & Crutchfield, 2021).\n\nWe will add a clearer explanation of this link between the iterative belief updates and the resulting fractal geometry either in the main text (likely near Figure 1/Section 3, space permitting) and/or expand on it in Appendix A.\n\n**Intuition for the Spectral Algorithm (Section 4.3):**\n\nThank you for highlighting that Section 4.3 could benefit from more intuition, which we will add to the text. The main goal of the spectral analysis is to understand precisely how information or influence from a past token (at source position s) propagates to affect the belief state at the current destination position d.\n- This influence mathematically depends on the sequence of hidden state transitions between s and d, captured by powers of the HMM transition matrix, $T^{d-s}$.\n- Spectral decomposition (using eigenvalues $\\lambda$ and associated projectors $T_\\lambda$) is a standard mathematical tool to analyze matrix powers because it simplifies $T^n$ into a sum $\\sum \\lambda^n T_\\lambda$.\n- The eigenvalues ($\\lambda$) are crucial because they tell us the rate at which the influence of past information decays (if $|\\lambda|<1$) or even oscillates (if $\\lambda$ is negative or complex) as the distance n=d-s increases.\n- Our key finding here is that the learned attention weights ($A_{d,s}$ in Eq. 3) directly implement this propagation effect, effectively learning to approximate the $\\lambda^{d-s}$ decay predicted by the theory (Eq. 10). This explains why attention patterns often show exponential decay, and why multiple heads are needed (Fig 3) to capture oscillatory patterns arising from negative eigenvalues.\n- Furthermore, this spectral perspective allows us to make precise, verifiable predictions about the learned OV vectors ($\\vec{v}_s$ in Eq. 13) and token embeddings ($\\vec{x}_s$ in Eq. 14), directly connecting the dynamics of the data (via T's eigenvalues) to the specific parameters learned by the transformer.\n\nWe will revise Section 4.3 to include a paragraph explaining this intuition more clearly.\n\n**Notation Suggestion:**\n\nWe completely agree that using $Z_{d+1}:$ for a sequence could be confusing. Thank you for pointing this out! We will revise the notation throughout the paper to use $Z_{d+1:\\infty}$ to unambiguously denote the sequence of future tokens.\n\n**Conceptual Diagram:**\n\nTo further aid in understanding the core theoretical idea, we have also created a new diagram [link here](https://docs.google.com/document/d/1CDwiJp8qG4d4OqgntBkIMYQhxJ1vy-2j2BecubhrL5U/edit?tab=t.0#heading=h.pwspurd6saek), that visually explains how the architectural constraints of attention lead naturally from the ideal Bayesian update (Eq. 2) to our derived constrained belief update (Eq. 5).\n\nThank you again for your helpful suggestions, which will undoubtedly improve the clarity of our paper. We appreciate your support for our work.\n\nSincerely,\nThe Authors", "claims_and_evidence": "The claims seem well supported by the methodology. In this, it is a good work: research question, claim, method,  and experimental evidence are clearly explained.", "methods_and_evaluation_criteria": "The proposed method and evaluations are well motivated and defined, in a controlled setup.", "theoretical_claims": "No, they are outside my area of expertise", "experimental_designs_or_analyses": "They seem valid, the models used are justified.", "relation_to_broader_scientific_literature": "They allow us to understand transformer models a little better than before.", "essential_references_not_discussed": "I'm not familiar with the literature.", "other_strengths_and_weaknesses": "**Pros**: The paper seems very well motivated and structured, and the fact that it uses a ’toy’ controlled setup for training and testing makes it more reliable, and allows to produce testable predictions. I have also enjoyed the discussion on one and two head transformer models. For what I could understand, it was quite an interesting read.  \n\n\n**Cons**: I found some parts of the paper, especially the parts related to the spectral algorithm harder to understand. While this is not a cons overall (it could be given by my ignorance on the topic), I would suggest the authors to try to add a couple more intuitive explanations. \n\nAlso, I think it is not clear, how the fractal drawings emerges from the belief states. As it is a key point, maybe it could be worth to add a couple more lines of explanation to this, either in the main body, or in the first section of the supplementary material.", "other_comments_or_suggestions": "Minor: Z_{d+1} for a sequence is confusing, maybe use Z_{d:T}, or some other notation to express that it is a series and not a single token?", "questions_for_authors": "No questions", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U6NHG34eA3", "reviewer": "Reviewer_tAc4", "summary": "This paper studies internal representations in transformers trained with next-token prediction on sequences generated by Hidden Markov Models. It reveals that transformers (focusing on single-layer transformers) perform constrained Bayesian belief updates, implementing an approximate version of optimal Bayesian inference in their parallel architecture. The paper proposes a mechanistic algorithm allowing attention heads to implement this process, creating predictable geometrically structured representations. The authors also develop a spectral theory to explain how data-generating transition matrices influence attention behavior. Through experiments on Mess3 class Hidden Markov Models, the authors validate their predictions about attention patterns and representation geometry. In summary, this work investigates the connection between optimal prediction theory and mechanistic interpretability, providing a principled understanding of how transformers balance optimality and architectural design.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive assessment of our work, and also your constructive feedback, which improved the clarity and impact of our work.\nWe address your primary concerns regarding the single-layer focus and the justification for Eq. 5, as described below, with new figs and analysis:\n\n**Single Layer Focus and Multi-Layer Extension** You are correct that our analysis focuses on a single-layer transformer, and we acknowledge this could have been stated more explicitly throughout the main text. We will revise the Abstract, Introduction, Methodology, and Discussion sections to clearly delineate that the primary analysis investigates the computation within the first layer, while also discussing how this relates to deeper networks. Although we decided to focus on the first layer for this initial work, we plan on extending the theory out to the multilayer setting in future work.\n\nTo directly address your question about multi-layer extensions in experiment, we have trained and analyzed 4-layer transformers. Our key finding is that the first attention layer consistently implements the constrained belief update mechanism (Eq. 5) derived in the paper, producing the intermediate fractal geometry. Subsequent layers (specifically the MLPs and later attention blocks) then act to transform this intermediate representation progressively closer to the full Bayesian belief state geometry (Eq. 2).\n\n We will add quantitative evidence for this to the appendix the main text. This includes:\n\n- PCA visualizations of activations after each layer in the 4-layer models, showing the transition from constrained to full belief geometry. New fig [here](https://bit.ly/4ch5W1m)\n- Mean Squared Error (MSE) analysis from linear regressions of residual stream activations onto both the theoretical constrained geometry (Eq.5) and the full belief geometry (Eq.2). This analysis, performed layer-by-layer, quantitatively shows the initial fit to Eq.5 post-attention-1, and the increasing fit to Eq.2 in deeper layers. New fig [here](https://bit.ly/4iP8vdo)\n\nThese results demonstrate that our single-layer analysis captures the fundamental computation performed by the initial attention mechanism, even within deeper networks, while later layers perform refinement.\n\n**Intuition for Eq.5:** We appreciate you highlighting the need for a clearer explanation of Eq. 5. We will incorporate the following improved intuition into Section 4.3:\n- Bayesian inference (Eq. 2) requires multiplying token-specific transition matrices, a fundamentally recursive process where updates depend on the full history integrated up to the previous step.\n- In contrast, an attention head computes its output at position d via a parallel, feedforward weighted sum (Eq. 3) of value vectors ($v_s$) from source positions $s \\le d$.\n- Crucially, each $v_s$ contains only local information from position $s$. It cannot directly access or depend on the specific tokens between $s$ and $d$ due to the parallel nature of the value computation and attention weighting.\n- Therefore, the most information that token z_s can independently contribute to the belief at position d within this single-layer constraint is the correction derived from knowing $z_s$ occurred d-s steps prior, assuming a default starting belief ($\\pi$) and no knowledge of intervening tokens. These independent displacements from the stationary distribution over latent states is the difference of probability distributions: $Pr(S_d | Z_s) - Pr(S_d)$. Linear algebraically, this contribution is precisely $\\pi T^{|z_s} T^{d-s} - \\pi$.\n- Summing these over all past sources naturally yields the constrained belief update form in Eq. 5. It represents the best possible parallel approximation achievable by a single attention layer given its architectural limitations.\n- [link](https://bit.ly/3XGCU5g) depicts each of these above points. \nWe believe this revised explanation clarifies why Eq. (5) is the natural computation for the attention layer to implement when reconciling optimal prediction with its inherent structure.\n\nRegarding the minor scalar discrepancy for the first two embeddings, we agree it doesn't invalidate the core findings. While still under investigation (as noted in Appendix C), the remarkable agreement across the attention pattern, OV vectors, and subsequent embeddings strongly supports our theoretical framework.  And as the referee noted, the spectral perspective allows us to anticipate that multiple attention heads are needed when the transition matrix for the training data has negative eigenvalues.\n\nWe will also strengthen the narrative as suggested, ensuring the scope (initial focus on first layer) and the derivation logic are clearer. While exploring other HMMs is interesting future work, Mess3's tractability was essential for this detailed mechanistic study.\nWe hope these clarifications and planned revisions address your concerns and demonstrate the significance of our findings. Thank you again for your valuable feedback.", "claims_and_evidence": "I find the empirical evidence in support of all claims made in this publication sufficiently compelling. The authors convincingly demonstrate that the proposed belief state update algorithm is the primary mechanism implemented by the single-layer transformer. A separate treatment of the negative eigenvalue case and the role of multi-head attention is also quite notable. My only concern at the moment is the observation of a \"scalar discrepancy in the first two embedding vectors\". I do not see this as being critical for the main conclusions of the paper, but hope that this might be understood eventually.", "methods_and_evaluation_criteria": "Given the nature of this publication, the simplistic Mess3 class Hidden Markov Models appears to be an adequate choice that allowed the authors to probe the precise mechanisms implemented by the transformer model. The paper could potentially benefit from discussing more general families of HMMs and potential qualitative differences in the state dynamics, but in my opinion, this is not critically important.", "theoretical_claims": "While I have not exhaustively verified all of the claims (for instance, I have not tried computing the eigenvalues of T), I read all of the authors theoretical claims and could not identify any clear mistakes. Most theoretical statements and intermediate results appear to be correct.", "experimental_designs_or_analyses": "I reviewed the overall design of experiments aimed at exploring the properties of trained transformers and comparing them with theoretical predictions. The proposed analysis seems adequate and appears to capture remarkable similarity between theoretical expectations and properties of trained models.", "relation_to_broader_scientific_literature": "This work investigates the specific mechanisms by which transformers capture belief state geometry in their residual stream, building upon a setup similar to “Transformers Represent Belief State Geometry in their Residual Stream.” It contributes to the broader field of mechanistic interpretability of transformers, which includes studies on linear and nonlinear regression, arithmetic, and other tasks. The paper also relates to research exploring transformers' ability to learn Hidden Markov Models (HMMs), as well as studies on the limitations of transformers in this context (for example, “On Limitation of Transformer for Learning HMMs”). Finally, it connects to: (a) 'An Explanation of In-context Learning as Implicit Bayesian Inference,' which uses HMMs to understand in-context learning through Bayesian inference, and (b) research on the information-theoretic properties of in-context learning (the observed fractal-like structure of system attractors can be linked to the growth of mutual information between the observed token sequence and the current state representation).", "essential_references_not_discussed": "I am not aware of any essential references that were not discussed in the paper.", "other_strengths_and_weaknesses": "In my opinion, two most important strengths of this publication are: (a) its creative setup and approach to understanding how transformers gradually attain information about the hidden state underlying the observed system dynamics; (b) the depth of theoretical understanding of very particular mechanisms that a trained transformer appears to implement.\n\nI also identified two weaknesses that made understanding this paper a little difficult. One is that the proposed mechanism and experimental studies seem to be focused on a single-layer transformer. The publication hardly mentions this fact and there appear to be just a few places (one in Appendix) where this exceptionally important detail is stated explicitly. This makes the main contributions stated by the authors in the abstract and the introduction less general than they may appear to be. I would expect the learned mechanisms in multi-layer transformers to be much more nuanced and while this clearly falls outside the scope of the current publication, it needs to be stated very explicitly.\n\nSecondly, the narrative structure of the paper is a little difficult to follow at times. For example, an exceptionally important equation (5) is not explained sufficiently well. It is my understanding that it can be related to the lowest-order term in the decomposition of the operator product $(T+\\epsilon T_1)(T+\\epsilon T_2)\\dots (T+\\epsilon T_n) \\approx T^n + \\epsilon (T_1 T^{n-1} + T T_2 T^{n-2} + \\dots + T^{n-1} T_n) + \\dots$, but this needs to be explained. Otherwise, it is not entirely clear why this particular expression is “one and only” sensible way of approximating the constrained belief update using a single-layer attention mechanism.", "other_comments_or_suggestions": "My primary suggestion is to strengthen the paper's narrative. I believe it would benefit from more in-depth discussions, such as the derivation of equation (5) and highlighting potential state geometry differences in other HMMs. If space is limited, detailed derivations could be moved to the appendix.", "questions_for_authors": "Two primary questions I have are related to two potential weaknesses:\n\n1. Is my understanding correct and all stated results (both theoretical and experimental) are only in relation to single-layer transformers. If so, have authors considered a very natural multi-layer extension? What were the most notable differences in how multi-layer transformers learn the updates in their residual stream?\n\n2. Is my understanding of the origins of equation (5) correct? If not, what is the explanation and are there any formal statements explaining why we expect this exact form to be implemented by a self-attention layer?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Mxt1KraaFx", "reviewer": "Reviewer_Wjfr", "summary": "In this paper, the authors propose a theoretical framework suggesting that transformers implement constrained Bayesian belief updating, and explain observed geometric patterns in transformer representations. Using the tools of mechanistic interpretability, they derive and empirically validate precise predictions of attention patterns, intermediate fractal-like structures, and final belief states in the trained transformers. The study focuses on the Mess3 family of Hidden Markov Models (HMMs).\n\n## Update after rebuttal\nI believe the authors have addressed and clarified some of my concerns. I maintain some reservations on the generalizability of the present findings, especially given the ad hoc nature of the considered learning task (not clearly connected to realistic tasks where transformers are employed) and of the theoretical analysis, here made possible because of the simplicity of the task. I also think that giving more prominence to the requisites for converging to a good representation (in terms of amount of data and training) could clarify when such a \"constrained\" regime could appear in the training of a transformer in a more realistic setting. \nHowever, I believe the paper is well written, and the presented analysis is sound and original, so I will raise my score to a 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and insightful feedback. We appreciate your recognition of our works rigor and the empirical support in our controlled setting. We also acknowledge the validity of your concerns regarding quantitative evaluation, learning dynamics, scope, and unexplained deviations. We have performed new analyses and made new figures to address these points and will make changes to the text, described below.\n\n**1. Learning Dynamics & Quantitative Evaluation:** You rightly pointed out the lack of analysis of the learning procedure despite the abstract's phrasing, and the need for quantitative evaluation beyond visual inspection (e.g., Fig 4).\n\n* **Learning Dynamics:** We have tracked how well the network activations fit our theoretical predictions during training. New Figure S ([link](https://bit.ly/4ldbEVZ)) shows the Mean Squared Error (MSE) when regressing activations onto theoretical geometries (Constrained Belief for Post-Attention, Full Belief for Post-MLP) throughout training. It demonstrates that the model rapidly converges towards the predicted geometries (Normalized MSE \\<\\< 1), quantitatively showing *that* gradient descent finds these solutions.  \n* **Quantitative Fit:** New Figure R ([link](https://bit.ly/3QTfFkG)) provides the absolute final MSE values, confirming a quantitative fit to our theory in the trained model compared to random initialization, addressing the need for metrics beyond visual assessment.  \n* **Abstract Phrasing:** We agree the phrasing \"gradient descent resolves the tension\" is misleading. Our core point is how the transformer's architecture itself (parallel attention vs. recursive Bayes) necessitates a specific computational form (Constrained Belief Update) to approximate the optimal solution. While Fig S shows GD finds this solution, our analysis focuses on what solution the architecture predisposes. We will revise the abstract to emphasize the role of architectural constraints in shaping the emergent computation, and get rid of the reference to GD.\n\n**2. Scope (Single-Layer Focus & Generalizability):** Our core derivation indeed focuses on the computation within the first layer. However, we have now analyzed deeper models.\n\n* **Multi-Layer Analysis:** As detailed in our response to Reviewer tAc4, we analyzed 4-layer transformers. New Figures X (quantitative MSE across layers ([link](https://bit.ly/4iP8vdo)) and Y (qualitative PCA visuals across layers ([link](https://bit.ly/4ch5W1m)) show that the 1st attn. layer consistently implements the Constrained Belief update, while subsequent layers refine the representation towards the Full Bayesian belief state.  \n* **HMM Focus & Generalization:** We used Mess3 HMMs not just for tractability, but because such controlled systems are essential for isolating and rigorously testing fundamental principles – like how architectural constraints shape optimal prediction – before scaling hypotheses to the complexities of natural language, where similar principles likely apply but are harder to isolate. As for the MLP, it's role in refining might generalize, though its specific function will be task-dependent. \n\n**3. Discrepancies & Justification:** We acknowledge the discrepancies noted (e.g., Appendix C embeddings, deviations outside the central alpha range mentioned in Sec 4.3).\n\n* **Contextualizing Deviations:** While these deviations warrant further study (as noted in our limitations), the new quantitative results (Figs R, S, X) demonstrate that the overall fit to our theory is strong. Also, at large alpha Mess3 generates near IID data, which leads to a lack of learnable structure - we believe this explains why there are deviations at large alpha.  \n* **Quantitative Predictions:** Importantly, our spectral theory correctly predicts the necessity of multiple attention heads when the HMM transition matrix has negative eigenvalues (Sec 4.4.1, Fig 3), a non-trivial validation of the underlying principles.  \n* **Justification for Eq. 5:** We have clarified the derivation of the Constrained Belief Update (Eq. 5) in our response to tAc4, including a new diagram [link](https://bit.ly/3XGCU5g) explaining why it represents the natural parallel approximation achievable by a single attention layer due to its architectural inability to access intermediate tokens recursively. This clarification hopefully makes the theoretical motivation and assumptions behind our core mechanism clearer.\n\n**4. Relation to Literature:** Our work's unique contribution is in using computational mechanics not just to describe optimal prediction, but to theoretically predict how specific architectural constraints (attention's parallel structure, preventing recursive updates) lead to systematic, geometric deviations (the Constrained Belief Update) from the unconstrained Bayesian ideal. This predictive, theory-driven approach complements empirical circuit discovery or other mech interp. methods.\n\nThank you again for your valuable feedback!", "claims_and_evidence": "The paper makes clear theoretical claims regarding the correspondence between transformer mechanisms and Bayesian belief updating constrained by architectural elements. Empirical evidence from analyses of transformers trained on synthetic Mess3 datasets mostly validates these claims, although some discrepancies (e.g., in the analysis in the appendix, and when the alpha parameter exceeds a certain range) are not addressed in the present work, and in my opinion are slightly at odds with the strong claim of full interpretability of the model computations.\n\nOn the other hand, in the abstract the authors state that their approach \"provides a principled lens on how gradient descent resolves the tension between optimal prediction and architectural design\", but in the paper there is no analysis of the learning procedure and the \"algorithm discovery\" phase, since the analysis only takes into account the asymptotic solution after training (and does not even consider what happens in a data scarse regime).", "methods_and_evaluation_criteria": "The authors utilize mechanistic interpretability (analysis of attention circuits and residual streams) and principal component analysis (PCA) to demonstrate the emergence of predicted geometric structures within the transformers' internal representations, relying on the synthetic and tightly controlled scenario under study. However, in some cases the authors rely on visual evaluation for quantifying the degree of agreement between their predictions and the empirical results (see e.g. fig. 4, where there is a general qualitative agreement but also visible differences, and assessing how relevant they are is difficult for the reader).", "theoretical_claims": "The presented theoretical framework is mathematically rigorous and predicts intermediate and final representation geometries. However, assumptions behind this theoretical setup, particularly regarding optimal Bayesian inference in constrained attention mechanisms, require clearer justifications, especially concerning their applicability to natural tasks beyond simplified HMM settings. In particular, the authors claim that it is crucial to investigate how model-specific constraints can warp the optimal Bayesian inference, but it is not clear how this interpretative route can be generalized to more complex cases (e.g. in LLMs trained on natural text, where there are open questions about the computational order mismatch between the transformer computations and the correct inference procedure).", "experimental_designs_or_analyses": "The empirical evaluations, confined exclusively to synthetic datasets (Mess3), mostly validate the theoretical predictions of the authors. However, the authors hint at the possible discrepancies that can arise in some regimes that are not investigated in this paper (e.g., large alpha) and some predictions that are not matched by the transformer implementation are left without explanation. This undermines the claims about the explanatory framework presented by the authors.", "relation_to_broader_scientific_literature": "While the paper appropriately references computational mechanics and transformer interpretability literature, it does not contrast this study against other mechanistic or interpretability-focused frameworks (e.g., circuit-based methods, dictionary learning, causal interventions), making its novelty and unique advantages somewhat unclear. Many of these works also study how exact inference can be embedded within the architectural constraint of the transformer.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "*Strengths:*\n* Rigorous mathematical derivations with clear empirical support in a controlled environment.\n\n*Weaknesses:*\n* Empirical scope confined to toy synthetic scenarios.\n* Limited to no analysis of the deviations from the predicted computational behavior.\n* No analysis of the training process and how the optimization lands on the discovery of the optimal inference algorithm.\n* Limited discussion of how findings extend or generalize to practical transformer models on natural language data.", "other_comments_or_suggestions": "The authors should explore practical implications or architectural modifications inspired by their theory. Additionally, discussing the conditions under which the theory may fail or require adjustments in more realistic or complex tasks would provide valuable balance and depth.", "questions_for_authors": "* Can the authors discuss how gradient descent helps the transformer converge to the optimal representation? E.g., how much data is needed? \n* Can the authors discuss when there are deviations from the predicted behavior and why? Does something change qualitatively in the inference model (e.g. at large alpha)?\n* Is the MLP layer warping mechanism something the authors think will extend to different settings?\n* How might the theoretical framework be extended to transformers trained on natural language data?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "None"}}]

Author rebuttal:

Meta-review: This paper presents a well-executed and theoretically grounded analysis of how single-layer transformers trained on synthetic HMM data implement a constrained form of Bayesian belief updating. Using tools from mechanistic interpretability and computational mechanics, the authors derive and empirically validate predictions about attention patterns and geometric structures in internal representations.

The majority of the reviewers support acceptance from the beginning, citing strong technical contributions and a clear theoretical-experimental link. Some concerns were raised about scope and clarity, which were addressed through expanded analysis (including multi-layer extensions) and clearer explanations. 

There was a strong concern about generalisability and missing training dynamics, but the rebuttal resolved it and helped prompt the work’s rigour and contribution.

While the synthetic setting limits immediate applicability, the paper makes a novel and substantive contribution to mechanistic interpretability by offering a predictive theory of transformer computation under architectural constraints. I recommend acceptance.
Decision: Accept (poster)