---
cluster: "C"
day: "C.3"
title: "Computational Mechanics"
contributors:
  - Xavier Poncini (Simplex)
  - Adam Shai (Simplex)
  - Paul Riechers (Simplex)
summary: >-
  What must a next-token predictor remember about the past? Causal states,
  belief states and their geometry, from HMMs to generalised HMMs, ending
  with the evidence that transformers learn this structure.
---

## 1. Prerequisites

Familiarity with the goals and basic methodology of mechanistic interpretability, including the notions of features and circuits; the linear representation hypothesis; linear probes as a method for reading off internal representations; and a basic understanding of the transformer architecture.

Students should be comfortable with the idea that one can train a linear map from model activations to some target structure and evaluate its quality (e.g. via MSE or $R^{2}$).

**Mathematical background.** The module is mathematically self-contained: all formal definitions (HMMs, GHMMs, belief states, MSPs) are introduced from scratch. However, students will engage with the material more fluently if they are comfortable with the following:

- **Linear algebra:** row-stochastic matrices, rank and invertibility, null spaces (kernels).
- **Probability:** conditional probability, Bayes' rule, probability distributions over finite sets, the probability simplex.
- **Basic machine learning:** next-token prediction, loss functions (cross-entropy), the concept of model activations.

<LearningOutcomes>

- Construct their own hidden Markov models.
- Explain the motivation for introducing generalised hidden Markov models (GHMMs).
- Explain why the belief state is a useful object for making predictions about GHMM data.
- Compute the belief state of a GHMM.
- Explain the difference between the belief state of a HMM and the belief state of a GHMM (that is itself not an HMM).
- Explain the relationship between belief states and next-token observation probabilities for GHMMs.
- Compute the mixed state presentation (MSP) for GHMMs.
- Explain why generating data is easier than predicting data.
- Explain the evidence for the claim: transformers represent belief state geometry in their residual stream.
- Develop hypotheses for where transformers represent: belief states, observation probabilities, log observation probabilities.

</LearningOutcomes>

## 2. Content

### 2.1 Fast-track

Read [Transformers represent belief state geometry in their residual stream](https://arxiv.org/pdf/2405.15943). Focus on being able to answer the following questions:

- What is a hidden Markov model (HMM)?
- What is a belief state?
- What is the mixed state presentation (MSP)?
- About the map learnt from activations to belief states:
  - What is the source?
  - What is the target?
  - How is the quality of the map evaluated?

### 2.2 Schedule

- **10:00–10:30** — [Lecture: Overview and scope](https://iliad-intensive.org/downloads/computational-mechanics/computational-mechanics-slides.pdf)
- **10:30–11:15** — [Lecture: Predicting the future](https://iliad-intensive.org/downloads/computational-mechanics/computational-mechanics-slides-2-predicting-the-future.pdf)
- **11:15–11:30** — *Break*
- **11:30–12:30** — [Exercises: Predicting the future](#3-predicting-the-future)
- **12:30–13:30** — *Lunch*
- **13:30–14:15** — [Lecture: Representing the past](https://iliad-intensive.org/downloads/computational-mechanics/computational-mechanics-slides-3-representing-the-past.pdf)
- **14:15–15:45** — [Exercises: Representing the past](#4-representing-the-past)
- **15:45–16:00** — *Break*
- **16:00–17:45** — [Reading & Discussion: Transformers represent belief geometry](#5-identifying-belief-geometry-in-transformers)
- **17:45–18:00** — Feedback

## 3. Predicting the future

<Callout type="note" title="Instructions">

Give a clear justification for each answer. You may use results proved in the lecture unless an exercise asks you to establish them. Questions marked **Extension** are optional.

</Callout>

<Exercise id="ex-p2-1">
**Exercise 3.1 (From statistics to partitions).** Let $u_{1},\ldots,u_{6}\in\mathcal{X}_{+}^{*}$ be six possible histories. Suppose that they induce the following conditional distributions over the entire future:

$$
\begin{array}{c|cccccc}\text{history} & u_1 & u_2 & u_3 & u_4 & u_5 & u_6\\ \hline \text{future distribution} & P & P & P & Q & Q & S\end{array}
$$

where $P$, $Q$, and $S$ are pairwise distinct.

Three statistics $R_{A}$, $R_{B}$, and $R_{C}$ induce the partitions

$$
\begin{aligned}\Pi_{A}&=\bigl\{\{u_{1},u_{2}\},\{u_{3}\},\{u_{4}\},\{u_{5}\},\{u_{6}\}\bigr\},\\ \Pi_{B}&=\bigl\{\{u_{1},u_{2},u_{3}\},\{u_{4},u_{5}\},\{u_{6}\}\bigr\},\\ \Pi_{C}&=\bigl\{\{u_{1},u_{2},u_{3},u_{4}\},\{u_{5}\},\{u_{6}\}\bigr\}.\end{aligned}
$$

Recall that a statistic is sufficient for prediction when histories assigned the same value induce the same distribution over the future.

**(a)** What does it mean for two histories to belong to the same cell of the partition induced by a statistic? Illustrate your answer using $u_{1}$ and $u_{2}$ in $\Pi_{A}$.

**(b)** Determine which of $R_{A}$, $R_{B}$, and $R_{C}$ are sufficient. If a statistic is not sufficient, give two histories witnessing the failure.

**(c)** Which cells of $\Pi_{A}$ can be merged without losing sufficiency? Perform all such merges and identify the resulting partition.

**(d)** Explain why no two cells of $\Pi_{B}$ can be merged while preserving sufficiency. What does this show about $R_{B}$?

**(e)** Prove that every refinement of a sufficient partition is also sufficient. Is every coarsening of a sufficient partition necessarily sufficient? Use $\Pi_{A}$, $\Pi_{B}$, and $\Pi_{C}$ to illustrate your answer.

**(f)** **Extension.** Consider the partition

$$
\Pi_{D} =\bigl\{\{u_{1}\},\{u_{2},u_{3}\},\{u_{4},u_{5}\},\{u_{6}\}\bigr\}.
$$

Show that $\Pi_{A}$ and $\Pi_{D}$ are both sufficient but that neither refines the other. How can both nevertheless refine the same minimal sufficient partition?
</Exercise>


<Exercise id="ex-p2-2">
**Exercise 3.2 (Causal states are minimal sufficient statistics).** For $x,x'\in\mathcal{X}_{+}^{*}$, define the relation $\sim$ by

$$
\begin{aligned}x\sim x' \quad\Longleftrightarrow\quad&\Pr\!\left( X_{|x|+1:|x|+|y|}=y\mid X_{1:|x|}=x \right)\\&\qquad= \Pr\!\left( X_{|x'|+1:|x'|+|y|}=y\mid X_{1:|x'|}=x' \right)\end{aligned}
$$

for every $y\in\mathcal{X}^{*}$. Define the causal state of $x$ by

$$
\epsilon(x):=[x]_{\sim},
$$

and write $\mathcal{S}:=\epsilon(\mathcal{X}_{+}^{*})$ for the set of causal states.

**(a)** Prove that $\sim$ is an equivalence relation on $\mathcal{X}_{+}^{*}$ by establishing each of the following:
1. reflexivity: $x\sim x$ for every $x\in\mathcal{X}_{+}^{*}$;
2. symmetry: if $x\sim x'$, then $x'\sim x$;
3. transitivity: if $x\sim x'$ and $x'\sim x''$, then $x\sim x''$.

**(b)** Prove that $\epsilon$ is sufficient for prediction.

**(c)** Let $R:\mathcal{X}_{+}^{*}\to\mathcal{R}$ be any sufficient statistic. Prove that

$$
R(x)=R(x')\quad\Longrightarrow\quad \epsilon(x)=\epsilon(x').
$$

Deduce that every sufficient partition refines the causal-state partition.

**(d)** Using parts (b) and (c), conclude that $\epsilon$ is a minimal sufficient statistic.

**(e)** **Extension.** Let $T$ be another minimal sufficient statistic. Prove that $T$ and $\epsilon$ induce the same partition of $\mathcal{X}_{+}^{*}$. Deduce that there is a bijection

$$
\psi:\mathcal{S}\to\operatorname{im}(T)
$$

satisfying $T=\psi\circ\epsilon$.
</Exercise>


<Exercise id="ex-p2-3">
**Exercise 3.3 (From finite generators to causal states).** In this exercise, $S_{t}$ denotes the hidden state immediately before the transition that emits $X_{t+1}$. All parameters lie strictly between $0$ and $1$.

**(a)** **The Zero–One–Random Process.** Let $S_{0}$ be uniformly distributed over $\{s_{0},s_{1},s_{R}\}$. The process has transitions

$$
s_{0}\xrightarrow{\,0:1\,}s_{1}, \qquad s_{1}\xrightarrow{\,1:1\,}s_{R}, \qquad s_{R} \begin{cases}\xrightarrow{\,0:1-p\,}s_{0},\\[-1mm] \xrightarrow{\,1:p\,}s_{0}.\end{cases}
$$

An edge label $x:q$ means that the process emits $x$ and follows the edge with probability $q$.

1. List all possible histories of length two. For each history, determine whether it uniquely identifies the current hidden phase. When it does, identify that phase.
2. For each ambiguous length-two history, determine which additional observations resolve the ambiguity. Identify the resulting phase and explain why it remains known after every subsequent observation.
3. Assume that $\epsilon(\emptyset)$, $\epsilon(0)$, $\epsilon(1)$, and $\epsilon(10)$ are pairwise distinct. Show that the causal states have shortest-length representatives

$$
\emptyset,\quad 0,\quad 1,\quad 10,\quad 00,\quad 01,\quad 11.
$$

Explain why no additional causal states arise.
4. Under the assumption in part (iii), draw the automaton associated with the causal-state dynamics.

**(b)** **The Even Process.** Let $S_{0}$ be uniformly distributed over $\{A,B\}$ and consider the generator

$$
A\xrightarrow{\,0:q\,}A, \qquad A\xrightarrow{\,1:1-q\,}B, \qquad B\xrightarrow{\,1:1\,}A.
$$

Thus every run of $1$s between successive $0$s has even length.

1. For each of the length-one histories $0$ and $1$, determine whether the current hidden state can be identified exactly.
2. Starting from each ambiguous history in part (i), consider its possible one-symbol extensions. Which extensions resolve the ambiguity, and which remain ambiguous? You may use the fact that $\epsilon(\emptyset)=\epsilon(11)$. What does this imply for histories of the form $1^{n}$?
3. Assume additionally that $\epsilon(\emptyset)\neq\epsilon(1)$. Show that the causal states have shortest representatives

$$
\emptyset,\qquad 1,\qquad 0,\qquad 01,
$$

and explain why there are no additional causal states. You may distinguish causal states using possible and impossible future continuations.
4. Using $\epsilon(\emptyset)=\epsilon(11)$ and assuming $\epsilon(\emptyset)\neq\epsilon(1)$, draw the automaton associated with the causal-state dynamics.

**(c)** **The Simple Nonunifilar Source.** Let $S_{0}$ be uniformly distributed over $\{A,B\}$ and consider

$$
A\xrightarrow{\,0:1/2\,}A, \qquad A\xrightarrow{\,0:1/2\,}B, \qquad B\xrightarrow{\,0:1/2\,}B, \qquad B\xrightarrow{\,1:1/2\,}A.
$$

1. Show that observing a $1$ identifies the new hidden state as $A$.
2. Starting from $A$, show that there are $n+1$ hidden paths that emit $0^{n}$, each with probability $2^{-n}$. How many of these paths can subsequently emit $1$? Hence calculate

$$
\Pr(X_{n+2}=1\mid X_{1:n+1}=10^{n}).
$$
3. Show that this probability is different for every $n\geq0$. Conclude that the histories

$$
1,10,100,\ldots
$$

belong to distinct causal states and that the process has infinitely many causal states.
</Exercise>


<Exercise id="ex-p2-4">
**Exercise 3.4 (Why next-token prediction is enough).** Recall that

$$
\mathcal{X}_{+}^{*} :=\left\{x\in\mathcal{X}^{*}: \Pr\!\left(X_{1:|x|}=x\right)>0\right\}
$$

is the set of allowable histories.

Consider a model with parameters $\theta$ satisfying

$$
h_{n}=f_{\theta}(h_{n-1},x_{n}), \qquad y_{n}=g_{\theta}(h_{n-1},x_{n}).
$$

Fix the initial hidden state $h_{0}$ and suppose that the model predicts the next-token distribution exactly: for every $n\geq1$,

$$
y_{n} =\Pr\!\left( X_{n+1}\mid X_{1:n}=x_{1:n}\right), \qquad \forall\,x_{1:n}\in\mathcal{X}_{+}^{*}.
$$

For every nonempty allowable history, define

$$
R(x_{1:n}):=(h_{n-1},x_{n}).
$$

For $a\in\mathcal{X}$, define the deterministic update

$$
\Phi_{a}(h,x):=\bigl(f_{\theta}(h,x),a\bigr).
$$

For a word $z_{1:j}$, write

$$
\Phi_{z_{1:j}}:=\Phi_{z_j}\circ\cdots\circ\Phi_{z_1}, \qquad \Phi_{z_{1:0}}:=\operatorname{id}.
$$

**(a)** Show that, whenever $x_{1:n}z_{1:j}$ is allowable,

$$
R(x_{1:n}z_{1:j}) =\Phi_{z_{1:j}}\!\left(R(x_{1:n})\right).
$$

Hence show that, if $R(x_{1:n})=R(x'_{1:m})$ and both extended histories are allowable, then

$$
R(x_{1:n}z_{1:j})=R(x'_{1:m}z_{1:j}).
$$

**(b)** Let $R(x_{1:n})=R(x'_{1:m})$ and fix $z_{1:k}\in\mathcal{X}^{k}$. As the symbols of $z_{1:k}$ are read from left to right, use part (a) and exact next-token prediction to show that for each $0\leq j<k$ such that the prefixes through $z_{1:j}$ are allowable,

$$
\begin{aligned}&\Pr\!\left( X_{n+j+1}=z_{j+1}\mid X_{1:n+j}=x_{1:n}z_{1:j}\right)\\&\qquad= \Pr\!\left( X_{m+j+1}=z_{j+1}\mid X_{1:m+j}=x'_{1:m}z_{1:j}\right).\end{aligned}
$$

Deduce that either the same first zero-probability extension is encountered after both histories, or every prefix of $z_{1:k}$ is allowable after both histories.

**(c)** In the case where every prefix is allowable, use the conditional chain rule to write

$$
\begin{aligned}&\Pr\!\left( X_{n+1:n+k}=z_{1:k}\mid X_{1:n}=x_{1:n}\right)\\&\qquad= \prod_{j=0}^{k-1}\Pr\!\left( X_{n+j+1}=z_{j+1}\mid X_{1:n+j}=x_{1:n}z_{1:j}\right).\end{aligned}
$$

Write the analogous product after $x'_{1:m}$ and compare its factors. Deal separately with the zero-probability case from part (b), and conclude directly that

$$
\begin{gathered}R(x_{1:n})=R(x'_{1:m}) \quad\Longrightarrow\\[1mm] \Pr\!\left( X_{n+1:n+k}=z_{1:k}\mid X_{1:n}=x_{1:n} \right) = \Pr\!\left( X_{m+1:m+k}=z_{1:k}\mid X_{1:m}=x'_{1:m} \right).\end{gathered}
$$

Hence $R(x_{1:n})=(h_{n-1},x_{n})$ is a sufficient statistic for predicting the entire future.
</Exercise>


## 4. Representing the past

<Callout type="note" title="Instructions">

Give a clear justification for each answer. You may use results proved in the lecture unless an exercise asks you to establish them. Questions marked **Extension** are optional.

</Callout>

<Exercise id="ex-p3-1">
**Exercise 4.1 (Translating between HMM diagrams and matrices).** For an edge-emitting hidden Markov model, use the convention

$$
T^{(a)}_{ij}:=\Pr(X_{t+1}=a,S_{t+1}=j\mid S_{t}=i).
$$

Thus rows index the current hidden state, columns index the next hidden state, and the superscript records the emitted symbol. Throughout this exercise, $0<p,q<1$.

**(a)** Suppose that $T^{(a)}_{ij}=r$. Explain each component of the following diagram:

$$
i\xrightarrow{\,a:r\,}j.
$$

What does a zero entry mean? Why is the matrix $T:=\sum_{a\in\mathcal{X}}T^{(a)}$ row-stochastic?

**(b)** **From matrices to a diagram: Zero–One–Random.** Take the state order to be $(0,1,R)$ and suppose that

$$
\eta^{(\emptyset)}=\begin{bmatrix}\tfrac{1}{3}&\tfrac{1}{3}&\tfrac{1}{3}\end{bmatrix}, \qquad T^{(0)}= \begin{bmatrix}0&1&0\\ 0&0&0\\ 1-p&0&0\end{bmatrix}, \qquad T^{(1)}= \begin{bmatrix}0&0&0\\ 0&0&1\\ p&0&0\end{bmatrix}.
$$

1. Draw the corresponding edge-emitting HMM. Describe informally the three phases represented by $0$, $1$, and $R$.
2. Calculate $\Pr(X_{1:3}=011)$ and $\Pr(X_{1:3}=000)$ in two ways: first by tracing paths through your diagram, and then from

$$
\Pr(X_{1:3}=x_{1:3}) =\eta^{(\emptyset)}T^{(x_1)}T^{(x_2)}T^{(x_3)}\mathbf{1}.
$$

**(c)** **From a diagram to matrices: the Even Process.** Consider the edge-emitting HMM

$$
A\xrightarrow{\,0:q\,}A, \qquad A\xrightarrow{\,1:1-q\,}B, \qquad B\xrightarrow{\,1:1\,}A,
$$

with no other edges. Take the state order to be $(A,B)$ and let $\eta^{(\emptyset)}=\begin{bmatrix}\tfrac{1}{2}&\tfrac{1}{2}\end{bmatrix}$.

1. Construct $T^{(0)}$ and $T^{(1)}$. Account explicitly for every zero entry in the two matrices.
2. Form $T=T^{(0)}+T^{(1)}$ and verify that it is row-stochastic.
3. Calculate $\Pr(X_{1:3}=011)$ and $\Pr(X_{1:3}=010)$ both by tracing paths through the diagram and by multiplying the appropriate symbol matrices. Explain how the second calculation reflects the defining constraint of the Even Process.

**(d)** **Extension.** Let $\mathcal{X}$ be the observation alphabet of an arbitrary finite-state HMM, let $T=\sum_{a\in\mathcal{X}}T^{(a)}$, and let $\eta^{(\emptyset)}$ be a probability distribution. Prove that, for every $n\geq 0$,

$$
\sum_{w\in\mathcal{X^n}}\Pr(X_{1:n}=w) =\eta^{(\emptyset)}T^{n}\mathbf{1} =1.
$$
</Exercise>


<Exercise id="ex-p3-2">
**Exercise 4.2 (Belief states and a lossy next-token readout).** Continue with the Zero–One–Random HMM from [Exercise 4.1](#ex-p3-1), using state order $(0,1,R)$. For an allowable history $x_{1:n}=x_{1}\cdots x_{n}$, compute its belief from scratch as

$$
\eta^{(x_{1:n})}:=\frac{\eta^{(\emptyset)}T^{(x_1)}\cdots T^{(x_n)}}{\eta^{(\emptyset)}T^{(x_1)}\cdots T^{(x_n)}\mathbf{1}}.
$$

Here $\eta_{i}^{(x_{1:n})}=\Pr(S_{n}=i\mid X_{1:n}=x_{1:n})$. For an allowable extension $x_{1:n+1}=x_{1:n}a$, update recursively as

$$
\eta^{(x_{1:n+1})}=\frac{\eta^{(x_{1:n})}T^{(a)}}{\eta^{(x_{1:n})}T^{(a)}\mathbf{1}}.
$$

**(a)** Compute $\eta^{(0)}$ and $\eta^{(1)}$. Interpret their components as posterior probabilities over the three hidden phases.

**(b)** Show that the seven beliefs

$$
\mathcal{B}= \bigl\{ \eta^{(\emptyset)},\eta^{(0)},\eta^{(1)},\eta^{(00)}, \eta^{(01)},\eta^{(10)},\eta^{(11)}\bigr\}
$$

are distinct and that $\mathcal{B}$ is closed under every allowable symbol update. Draw the symbol-labelled transition diagram on $\mathcal{B}$.

**(c)** Compare your belief-state diagram with both the causal-state automaton constructed in Part 2 and the three-state generative HMM from [Exercise 4.1](#ex-p3-1).
1. Match each belief to the causal state with the same shortest-length representative. Do the symbol-labelled transitions agree?
2. Which beliefs correspond to knowing the hidden phase exactly, and which represent uncertainty over several phases? Show that the hidden phase becomes known after at most three observations and remains known thereafter.
3. Observe that the finite-history belief-state presentation has seven states while the generative HMM has only three. What does this suggest about the possible relative difficulty of generation and inference?

**(d)** For each belief $\eta^{(x_{1:n})}\in\mathcal{B}$ from part (b), compute its associated next-token probability vector

$$
\begin{gathered}p^{(x_{1:n})} :=\left( \Pr(X_{n+1}=a\mid X_{1:n}=x_{1:n}) \right)_{a\in\{0,1\}} =\eta^{(x_{1:n})}A,\\[1mm] A_{ia}:=\sum_jT^{(a)}_{ij}.\end{gathered}
$$

Visualise the belief geometry on the $2$-simplex over hidden states $(0,1,R)$ and the next-token geometry on the $1$-simplex over symbols $(0,1)$.

**(e)** Calculate

$$
\Pr(X_{3:4}=00\mid X_{1:2}=10) \qquad\text{and}\qquad \Pr(X_{3:4}=00\mid X_{1:2}=01).
$$

Deduce that the next-token probability vector need not be sufficient for predicting the entire future.

**(f)** **Extension.** Let $\delta:=\eta^{(10)}-\eta^{(01)}$. Show that

$$
\delta A=0 \qquad\text{but}\qquad \delta T^{(0)}T^{(0)}\mathbf{1}\neq0.
$$

Interpret these two statements in terms of one-step and longer-horizon prediction.
</Exercise>


<Exercise id="ex-p3-3">
**Exercise 4.3 (Nonunifilar generators and large belief-state automata).** For a belief $\eta$ and a symbol $a$ satisfying $\eta T^{(a)}\mathbf{1}>0$, write

$$
F_{a}(\eta):=\frac{\eta T^{(a)}}{\eta T^{(a)}\mathbf{1}}.
$$

Thus the belief-state automaton contains the transition

$$
\eta\xrightarrow{\,a:\eta T^{(a)}\mathbf{1}\,}F_{a}(\eta).
$$

**(a)** **The Simple Nonunifilar Source.** Revisit the source from Part 2. With state order $(A,B)$, its symbol matrices and initial belief are

$$
T^{(0)}= \begin{bmatrix}\tfrac{1}{2}&\tfrac{1}{2}\\0&\tfrac{1}{2}\end{bmatrix}, \qquad T^{(1)}= \begin{bmatrix}0&0\\\tfrac{1}{2}&0\end{bmatrix}, \qquad \eta^{(\emptyset)}=\begin{bmatrix}\tfrac{1}{2}&\tfrac{1}{2}\end{bmatrix}.
$$

1. For $n\geq0$, show that

$$
\eta^{(10^n)}= \begin{bmatrix}\dfrac{1}{n+1}&\dfrac{n}{n+1}\end{bmatrix}, \qquad \eta^{(0^n)}= \begin{bmatrix}\dfrac{1}{n+2}&\dfrac{n+1}{n+2}\end{bmatrix}.
$$

Deduce that $\eta^{(0^n)}=\eta^{(10^{n+1})}$ and observe that $\eta^{(\emptyset)}=\eta^{(10)}$.
2. Show that, for every $n\geq1$,

$$
\eta^{(10^n1)}=\eta^{(1)}.
$$
3. Draw the first four states of the belief-state automaton and indicate its continuing structure with an ellipsis. Show that the beliefs $\eta^{(10^n)}$ are all distinct, identify their limiting point in the $1$-simplex, and explain why the two-state generator produces a countably infinite belief-state automaton.

**(b)** **Mess3.** Now take hidden-state and symbol sets both equal to $\{0,1,2\}$, with

$$
\alpha=\frac{1}{5},\qquad b=\frac{1-\alpha}{2}=\frac{2}{5}, \qquad x=\frac{3}{20},\qquad y=1-2x=\frac{7}{10},
$$

and

$$
\begin{aligned}T^{(0)}&= \begin{bmatrix}\alpha y&bx&bx\\ \alpha x&by&bx\\ \alpha x&bx&by\end{bmatrix},&T^{(1)}&= \begin{bmatrix}by&\alpha x&bx\\ bx&\alpha y&bx\\ bx&\alpha x&by\end{bmatrix},\\[1mm] T^{(2)}&= \begin{bmatrix}by&bx&\alpha x\\ bx&by&\alpha x\\ bx&bx&\alpha y\end{bmatrix},&\eta^{(\emptyset)}&= \begin{bmatrix}\tfrac{1}{3}&\tfrac{1}{3}&\tfrac{1}{3}\end{bmatrix}.\end{aligned}
$$

1. Verify that $T=\sum_{a=0}^{2}T^{(a)}$ is row-stochastic and explain why this generator is nonunifilar. Compute $\eta^{(0)},\eta^{(1)}$, and $\eta^{(2)}$.
2. Compute $\eta^{(00)},\eta^{(01)}$, and $\eta^{(02)}$. Use the cyclic symmetry of the process to sketch the first two generations of the belief-state automaton, including the symbol labels.
3. Compute the next-token matrix $A$, with $A_{ia}=\sum_{j}T^{(a)}_{ij}$, and show that it is invertible. What does this imply about whether distinct beliefs can be merged by the map $p^{(x_{1:n})}=\eta^{(x_{1:n})}A$? Contrast this with [Exercise 4.2](#ex-p3-2).
4. Using a short program, form

$$
\mathcal{B}_{d}:= \bigl\{\eta^{(x_{1:d})}:x_{1:d}\in\{0,1,2\}^{d}\bigr\}, \qquad 0\leq d\leq7.
$$

Record the number of distinct beliefs at each depth. Visualise $\mathcal{B}_{7}$ on the hidden-state $2$-simplex, colouring each point by its final symbol. Also visualise the associated next-token vectors on the symbol $2$-simplex. Describe the geometric structure that emerges and how the three maps $F_{0},F_{1},F_{2}$ generate the outgoing transitions of the belief-state automaton.
5. Explain why the belief-state automata constructed in this exercise are unifilar even though their generators are not. Compare the ladder-with-resets structure of the Simple Nonunifilar Source with the branching structure of Mess3. Finally, explain why the set of beliefs reached by finite histories is countable even though its closure can have a much richer fractal geometry.
</Exercise>


<Exercise id="ex-p3-4">
**Exercise 4.4 (Generalised hidden Markov models and predictive vectors).** A $d$-dimensional generalised hidden Markov model (GHMM) consists of a finite alphabet $\mathcal{X}$, real $d\times d$ matrices $(T^{(a)})_{a\in\mathcal{X}}$, an initial row vector $\eta^{(\emptyset)}$, and a column vector $\phi$. Writing

$$
T:=\sum_{a\in\mathcal{X}}T^{(a)},
$$

these objects satisfy

$$
T\phi=\phi,\qquad \eta^{(\emptyset)}\phi=1,\qquad \eta^{(\emptyset)}T^{(x_1)}\cdots T^{(x_n)}\phi\geq0
$$

for every $x_{1:n}\in\mathcal{X}^{n}$ and every $n\geq0$.

**(a)** **A probability law from matrix products.** Consider

$$
\Pr(X_{1:n}=x_{1:n}) =\eta^{(\emptyset)}T^{(x_1)}\cdots T^{(x_n)}\phi.
$$

1. Show that these quantities are nonnegative and that

$$
\sum_{x_{1:n}\in\mathcal{X^n}}\Pr(X_{1:n}=x_{1:n})=1
$$

for every $n\geq0$.
2. Show that

$$
\sum_{a\in\mathcal{X}}\Pr(X_{1:n+1}=x_{1:n}a) =\Pr(X_{1:n}=x_{1:n}).
$$

Explain why this is the consistency condition needed to define a one-sided stochastic process.
3. The GHMM conditions do not require $\eta^{(\emptyset)}T=\eta^{(\emptyset)}$. Explain why the process need not therefore be stationary. Show that this additional equality is sufficient for stationarity.

**(b)** **Predictive vectors.** For any $x_{1:n}$ satisfying $\Pr(X_{1:n}=x_{1:n})>0$, define its predictive vector by

$$
\eta^{(x_{1:n})}:=\frac{\eta^{(\emptyset)}T^{(x_1)}\cdots T^{(x_n)}}{\eta^{(\emptyset)}T^{(x_1)}\cdots T^{(x_n)}\phi}.
$$

1. Show that

$$
\eta^{(x_{1:n})}\phi=1
$$

and, for every continuation $y_{1:m}$,

$$
\Pr(X_{n+1:n+m}=y_{1:m}\mid X_{1:n}=x_{1:n}) =\eta^{(x_{1:n})}T^{(y_1)}\cdots T^{(y_m)}\phi.
$$
2. Derive the recursive update and identify its normalising factor:

$$
\eta^{(x_{1:n}a)}=\frac{\eta^{(x_{1:n})}T^{(a)}}{\eta^{(x_{1:n})}T^{(a)}\phi}, \quad \Pr(X_{n+1}=a\mid X_{1:n}=x_{1:n}) =\eta^{(x_{1:n})}T^{(a)}\phi.
$$
3. Conceptually contrast an HMM belief state with a GHMM predictive vector. Compare the meaning of their coordinates, their positivity and normalisation properties, and how they encode predictions of future observations.

**(c)** **Predictive vectors need not be probability vectors.** Consider the two-state HMM

$$
M^{(0)}= \begin{bmatrix}\tfrac{4}{5}&0\\0&\tfrac{1}{5}\end{bmatrix}, \qquad M^{(1)}= \begin{bmatrix}\tfrac{1}{5}&0\\0&\tfrac{4}{5}\end{bmatrix}, \qquad \alpha=\begin{bmatrix}\tfrac{1}{2}&\tfrac{1}{2}\end{bmatrix}.
$$

Thus a hidden state is selected initially and thereafter emits a biased coin without changing state. Let

$$
S=\begin{bmatrix}2&-1\\0&1\end{bmatrix}, \qquad T^{(a)}:=S^{-1}M^{(a)}S,\qquad \eta^{(\emptyset)}:=\alpha S,\qquad \phi:=S^{-1}\mathbf{1}.
$$

1. Compute $T^{(0)},T^{(1)},\eta^{(\emptyset)}$, and $\phi$. Show that, for every $x_{1:n}$,

$$
\eta^{(\emptyset)}T^{(x_1)}\cdots T^{(x_n)}\phi = \alpha M^{(x_1)}\cdots M^{(x_n)}\mathbf{1}.
$$

Deduce that this change of coordinates gives a valid GHMM for the same process.
2. Compute $\eta^{(0)}$ and $\eta^{(1)}$. Verify directly from $\eta^{(0)}$ that

$$
\Pr(X_{2}=0\mid X_{1}=0)=\frac{17}{25}, \qquad \Pr(X_{2}=1\mid X_{1}=0)=\frac{8}{25}.
$$

Explain why the entries of a GHMM predictive vector need not themselves be probabilities.

**(d)** **The probability matrix and minimal dimension.** For finite sets $\mathcal{U},\mathcal{V}\subseteq\mathcal{X}^{*}$, define the probability matrix by

$$
\mathsf{P}_{\mathcal{U},\mathcal{V}}(u,v) :=\Pr(X_{1:|u|+|v|}=uv), \qquad u\in\mathcal{U},\quad v\in\mathcal{V}.
$$

1. Use the GHMM matrix product to factor $\mathsf{P}_{\mathcal{U},\mathcal{V}}$ as a matrix of prefix row vectors times a matrix of continuation column vectors. Deduce that

$$
\operatorname{rank}\mathsf{P}_{\mathcal{U},\mathcal{V}}\leq d
$$

for every $d$-dimensional GHMM realisation.
2. For the process in part (c), take $\mathcal{U}=\mathcal{V}=\{\emptyset,0,1\}$. Compute $\mathsf{P}_{\mathcal{U},\mathcal{V}}$ and use an SVD to find its singular values. Deduce that the two-dimensional GHMM is minimal-dimensional.
</Exercise>


## 5. Identifying belief geometry in transformers

So far, we have developed the conceptual framework of computational mechanics. We now ask whether neural networks trained for next-token prediction—including transformers and recurrent architectures—learn the predictive structures identified by that framework. You will read selected sections of two papers, then discuss their evidence and limitations in small groups.

<Callout type="note" title="Session at a glance">

1. **Readings (75 minutes).** Begin with the main reading and use the route below. Continue to the extension if time permits.
2. **Discussion (30 minutes).** Form groups of three or four. Compare what you found convincing, confusing, or incomplete; use the prompts on page 2 if useful.

</Callout>

### 5.1 Readings

Prioritise the main reading. If you can answer its evaluation question, you have probably understood one of its central results. If time is short, skim the extension overview before deciding whether to continue.

#### 5.1.1 [Main] [Transformers represent belief state geometry in their residual stream](https://arxiv.org/pdf/2405.15943)

**Reading route**

- **Skim.** Sections 1, 2.1, and 2.2.
- **Read.** Sections 2.3 and 3–5.
- **Evaluation.** Explain Figure 6D. What do the comparison and shuffle control establish?

**Overview.** This paper presents evidence that transformers trained on HMM-generated sequences encode belief-state geometry in their residual streams. The claim is tested on the Mess3 and Random–Random–XOR (RRXOR) processes. For each process, the authors fit an affine map from residual activations to the exact HMM belief state and evaluate it using mean-squared error (MSE). The error decreases over training checkpoints. A shuffle control tests whether low MSE could arise merely because a high-dimensional activation space is being projected into a low-dimensional belief space.

#### 5.1.2 [Extension] [Neural networks leverage nominally quantum and post-quantum representations](https://arxiv.org/pdf/2507.07432)

**Reading route**

- **Skim.** Sections 1–3.
- **Read.** Sections 4–8.
- **Evaluation.** Explain Figure 3B. How are the light-blue and light-orange plots related to the corresponding dark-blue and dark-orange plots?

**Overview.** This paper extends the question from HMMs to generalised hidden Markov models (GHMMs). As in Part 3, a GHMM retains the matrix-product expression for output probabilities while relaxing the entrywise probabilistic constraints on its internal matrices. This larger model class can realise some processes in finite dimension even when a corresponding minimal HMM requires exponentially many more hidden states.

Using techniques related to those in the main reading, the authors study transformers, LSTMs, RNNs, and GRUs trained on processes with distinct GHMM and HMM realisations. Their results indicate that the networks learn belief-state representations associated with the lower-dimensional GHMM realisation. The finding suggests that neural networks can select among representational model classes in ways that capture dimensionality savings.

### 5.2 Discussion

Here are some prompts to consider:

- What is the most convincing evidence that transformers represent belief-state geometry? What is the least convincing?
- What additional experiment would most strengthen or weaken the claim?
- Which alternative hypotheses could explain the reported results?
- Can you explain each element of the diagram below—including the generative HMM, Bayesian updating, transformer activations, fitted affine map, and belief-state geometry?

<div id="fig-shai-belief-state-geometry">
<Figure src="/uploads/computational-mechanics/shai-belief-state-geometry.png" alt="figure" />
</div>

## 6. Python exercises: HMM from scratch

The Python exercises are ARENA-style notebooks in two parts. Each part has an **exercises** notebook (with `# YOUR CODE HERE` stubs, inline tests, and collapsible solutions) and a fully worked **solutions** notebook.

- **Part 1 — Sequence probabilities & next-token distributions** (exercises) — [Open in Colab](https://colab.research.google.com/github/iliad-team/iliad-intensive-C.3/blob/main/python-exercises/colab/part1_sequence_probabilities_exercises_colab.ipynb)
- Part 1 (solutions) — [Open in Colab](https://colab.research.google.com/github/iliad-team/iliad-intensive-C.3/blob/main/python-exercises/colab/part1_sequence_probabilities_solutions_colab.ipynb)
- **Part 2 — Belief states** (exercises) — [Open in Colab](https://colab.research.google.com/github/iliad-team/iliad-intensive-C.3/blob/main/python-exercises/colab/part2_belief_states_exercises_colab.ipynb)
- Part 2 (solutions) — [Open in Colab](https://colab.research.google.com/github/iliad-team/iliad-intensive-C.3/blob/main/python-exercises/colab/part2_belief_states_solutions_colab.ipynb)

The Colab notebooks are self-contained: their setup cells write the helper modules into the session, so there is nothing else to install or download.

## 7. Learn more

- Beyond toy architectures – in-context learning
  - LLMs have an amazing ability to adapt internal representations in-context – [Read: ICLR: [In-Context Learning of Representations](https://arxiv.org/pdf/2501.00070)]
  - LLMs can predict HMM data in-context – [Read: [Pre-trained Large Language Models Learn Hidden Markov Models In-context](https://arxiv.org/pdf/2506.07298)]
  - This suggests that internal representations of LLMs predicting GHMM data in-context, may resemble belief states.
- Processes consisting of GHMMs with richer structure [Read: [Transformers learn factored representations](https://arxiv.org/pdf/2602.02385v1)]
- Review article of computational mechanics [Read: [Between order and chaos](https://www.nature.com/articles/nphys2190)]
- [Simplex](https://www.simplexaisafety.com/) is a non-profit research organisation working on applying computational mechanics to AI safety.