A Formal Comparison Between Chain of Thought and Latent Thought — part 6/9

Cleaned-text code points [77818, 93498). Bibliography omitted here; source.md includes it.

In the proof, we assume that the computation graph contains at most n n output nodes. This assumption is without loss of generality: if the number of output nodes exceeds the number of input nodes, we can simply pad the input with dummy nodes (e.g., fixed zeros), thereby reducing the setting to the same case.

We construct a model in which (1) the attention layer aggregates the inputs, and (2) the looped feed-forward layer performs the computation of all nodes in parallel, as illustrated in Figure 9 . We first show that a feedforward layer followed by an attention layer can copy all input tokens to each position. Assume, given an input sequence x = ( x 1 , x 2 , … , x n ) ∈ Σ n x=(x_{1},x_{2},\ldots,x_{n})\in\Sigma^{n} and one-hot encoding 𝒆 : Σ → { 0 , 1 } | Σ | \bm{e}:\Sigma\to\{0,1\}^{|\Sigma|} . Assume each input at position i ∈ [ n ] i\in[n] is embedded as 𝒉 i = ( 0 , 𝒆 ⁡ ( x i ) , 𝒆 i ) ∈ { 0 , 1 } n ​ | Σ | + | Σ | + n , {\bm{h}}_{i}=\bigl(\,{\bm{0}},\;\bm{e}(x_{i}),\;{\bm{e}}_{i}\,\bigr)\in\{0,1\}^{n|\Sigma|+|\Sigma|+n}, (55) where 𝒆 i ∈ { 0 , 1 } n {\bm{e}}_{i}\in\{0,1\}^{n} . By Lemma B.6 , when substituting 𝒙 = ( 𝒆 ⁡ ( x 1 ) , … , 𝒆 ⁡ ( x n ) ) {\bm{x}}=(\bm{e}(x_{1}),\ldots,\bm{e}(x_{n})) into the lemma, there exists a feed-forward layer FF 1 \mathrm{FF}_{1} such that ( id + FF 1 ) ​ ( 𝒉 i ) = ( ( 𝒆 i ) 1 ⋅ 𝒆 ⁡ ( x i ) , ( 𝒆 i ) 2 ⋅ 𝒆 ⁡ ( x i ) , … , ( 𝒆 i ) n ⋅ 𝒆 ⁡ ( x i ) , 𝒆 ⁡ ( x i ) , 𝒆 i ) . (\id+\mathrm{FF}_{1})({\bm{h}}_{i})=\bigl(\,({\bm{e}}_{i})_{1}\cdot\bm{e}(x_{i}),\;({\bm{e}}_{i})_{2}\cdot\bm{e}(x_{i}),\;\ldots,\;({\bm{e}}_{i})_{n}\cdot\bm{e}(x_{i}),\;\bm{e}(x_{i}),\;{\bm{e}}_{i}\,\bigr). (56) To aggregate all positions via uniform attention, we use a single-head attention layer with: 𝐪 i = 𝒌 i = 𝟏 n ​ | Σ | , 𝐯 i = n ⁡ ( ( 𝒆 i ) 1 ⋅ 𝒆 ⁡ ( x i ) , ( 𝒆 i ) 2 ⋅ 𝒆 ⁡ ( x i ) , … , ( 𝒆 i ) n ⋅ 𝒆 ⁡ ( x i ) ) for all ​ i ∈ [ n ] , \mathbf{q}_{i}={\bm{k}}_{i}=\mathbf{1}_{n|\Sigma|},\quad\mathbf{v}_{i}=n\big(({\bm{e}}_{i})_{1}\cdot\bm{e}(x_{i}),\;({\bm{e}}_{i})_{2}\cdot\bm{e}(x_{i}),\;\ldots,\;({\bm{e}}_{i})_{n}\cdot\bm{e}(x_{i})\big)\quad\text{for all }i\in[n], (57) with an appropriate output projection, the output of the attention layer, at position i i , becomes 1 n ​ ∑ j = 1 n 1 ⋅ n ​ 𝒉 j \displaystyle\frac{1}{n}\sum^{n}_{j=1}1\cdot n{\bm{h}}_{j} = ( ∑ j = 1 n ( 𝒆 j ) 1 ​ 𝒆 ​ ( x j ) , ∑ j = 1 n ( 𝒆 j ) 2 ​ 𝒆 ​ ( x j ) , … , ∑ j = 1 n ( 𝒆 j ) n ​ 𝒆 ​ ( x j ) ) \displaystyle\;=\;\Bigl(\sum_{j=1}^{n}({\bm{e}}_{j})_{1}\,\bm{e}(x_{j}),\;\;\sum_{j=1}^{n}({\bm{e}}_{j})_{2}\,\bm{e}(x_{j}),\;\;\ldots,\;\;\sum_{j=1}^{n}({\bm{e}}_{j})_{n}\,\bm{e}(x_{j})\Bigr) (58) = ( 𝒆 ⁡ ( x 1 ) , 𝒆 ⁡ ( x 2 ) , … , 𝒆 ⁡ ( x n ) ) . \displaystyle\;=\;\bigl(\,\bm{e}(x_{1}),\;\bm{e}(x_{2}),\;\ldots,\;\bm{e}(x_{n})\,\bigr). (59)

Then, we show that the feed-forward layer can encode the entire computation graph into its weights and simulate all nodes simultaneously. Let the flag vector for each node be ( t 1 , … , t N ) ∈ { 0 , 1 } N (t_{1},\ldots,t_{N})\in\{0,1\}^{N} , where N ≔ | V n | = size ⁡ ( G n ) N\coloneqq|V_{n}|=\mathrm{size}(G_{n}) . By Lemmas B.7 and B.8 , there exist feed-forward layers FF 2 , FF 3 : 𝔽 s N ⁡ ( | Σ | + 1 ) → 𝔽 s N ⁡ ( | Σ | + 1 ) \mathrm{FF}_{2},\mathrm{FF}_{3}:\mathbb{F}_{s}^{N(|\Sigma|+1)}\to\mathbb{F}_{s}^{N(|\Sigma|+1)} such that, for the input vector ( z 1 , … , z N ) ∈ Σ N . (z_{1},\ldots,z_{N})\in\Sigma^{N}. , ℱ ​ ℱ ​ ( 𝒆 ⁡ ( z 1 ) , t 1 , … , 𝒆 ⁡ ( z N ​ ( x ) ) , t N ) \displaystyle\mathcal{FF}\Bigl(\,\bm{e}(z_{1}),t_{1},\ldots,\bm{e}(z_{N}(x)),t_{N}\Bigr) = ∥ i = 1 N ( t i ⋅ 𝒆 ( f v i ( 𝒛 ( i ) ) ) , 1 m i ∑ j = 1 m i t p i , j ) , \displaystyle=\big\|_{i=1}^{N}\,\Bigl(t_{i}\cdot\bm{e}\bigl(f_{v_{i}}({\bm{z}}^{(i)})\bigr),\;\frac{1}{m_{i}}\sum_{j=1}^{m_{i}}t_{p_{i,j}}\Bigr), (60) where ℱ ​ ℱ ≔ ( id + FF 3 ) ∘ ( id + FF 2 ) \mathcal{FF}\coloneq(\id+\mathrm{FF}_{3})\circ(\id+\mathrm{FF}_{2}) . Here, f v i f_{v_{i}} denotes the function associated with node v i v_{i} , and p i , 1 , … , p i , m i p_{i,1},\ldots,p_{i,m_{i}} denote the indices of the predecessor nodes of v i v_{i} , and 𝒛 ( i ) = ( z p i , 1 , … , z p i , m i ) {\bm{z}}^{(i)}=(z_{p_{i,1}},\ldots,z_{p_{i,m_{i}}}) denotes their values. The last term 1 m i ​ ∑ j = 1 m i t p i , j \frac{1}{m_{i}}\sum_{j=1}^{m_{i}}t_{p_{i,j}} can be obtained using a linear layer.

For the k k -th loop, assume by induction that the hidden state is 𝒉 ⁡ ( k ) ≔ ( t 1 , k − 1 ⋅ 𝒆 ⁡ ( v 1 ​ ( x ) ) , t 1 , k , t 2 , k − 1 ⋅ 𝒆 ⁡ ( v 2 ​ ( x ) ) , t 2 , k , … , t N , k − 1 ⋅ 𝒆 ⁡ ( v N ​ ( x ) ) , t N , k ) , {\bm{h}}(k)\coloneq\bigl(t_{1,k-1}\cdot\bm{e}(v_{1}(x)),\,t_{1,k},\,t_{2,k-1}\cdot\bm{e}(v_{2}(x)),\,t_{2,k},\,\ldots,\,t_{N,k-1}\cdot\bm{e}(v_{N}(x)),\,t_{N,k}\bigr), (61) where v i ​ ( x ) ∈ Σ v_{i}(x)\in\Sigma denotes the value computed by node v i v_{i} given the input x x , and t i , k ∈ { 0 , 1 } t_{i,k}\in\{0,1\} indicates whether node v i v_{i} lies within depth at most k k . Under this assumption, it holds that ℱ ​ ℱ ​ ( 𝒉 ⁡ ( k ) ) \displaystyle\mathcal{FF}({\bm{h}}(k)) = ∥ i = 1 N ( t i , k ⋅ 𝒆 ( f v i ( v p i , 1 ( x ) , v p i , 2 ( x ) , … , v p i , m i ( x ) ) ) , 1 m i ∑ j = 1 m i t ( p i , j , k ) ) , \displaystyle\;=\;\big\|_{i=1}^{N}\,\Bigl(t_{i,k}\cdot\bm{e}\bigl(f_{v_{i}}(v_{p_{i,1}}(x),v_{p_{i,2}}(x),\ldots,v_{p_{i,m_{i}}}(x))\bigr),\;\frac{1}{m_{i}}\sum_{j=1}^{m_{i}}t_{(p_{i,j},k)}\Bigr), (62) = ∥ i = 1 N ( t i , k ⋅ 𝒆 ( v i ( x ) ) ) , t i , k + 1 ) = 𝒉 ( k + 1 ) . \displaystyle\;=\;\big\|_{i=1}^{N}\,\Bigl(t_{i,k}\cdot\bm{e}\bigl(v_{i}(x))\bigr),\;t_{i,k+1}\Bigr)\;=\;{\bm{h}}(k+1). (63)

To extract the output node corresponding to each position denoted by o i o_{i} , in the final loop iteration, by Lemma B.6 , there exists a feedforward layer FF 4 \mathrm{FF}_{4} such that ( id + FF 4 ) ​ ( 𝒉 ⁡ ( k ) , 𝒆 o i ) = ( t o i , k ⋅ 𝒆 ⁡ ( v o i ​ ( x ) ) , ∗ ) . (\id+\mathrm{FF}_{4})({\bm{h}}(k),{\bm{e}}_{o_{i}})=\bigl(t_{o_{i},k}\cdot\bm{e}(v_{o_{i}}(x)),\;*\bigr)\,. (64)

##### Summary

We construct the looped model as follows. Each input token x i ∈ Σ x_{i}\in\Sigma at position i ∈ [ n ] i\in[n] is embedded as 𝒉 i ( 0 ) = ( 𝒆 i , 𝒆 o i , 𝒆 ⁡ ( x i ) , 0 , 𝒆 ⁡ ( x i ) , 0 , … , 𝒆 ⁡ ( x i ) , 0 , 0 ( 1 + | Σ | ) ​ ( N − n ) + | Σ | ) ∈ { 0 , 1 } 2 ​ n + ( 1 + | Σ | ) ​ N + | Σ | . {\bm{h}}^{(0)}_{i}=\bigl(\,{\bm{e}}_{i},\;{\bm{e}}_{o_{i}},\;\bm{e}(x_{i}),\;0,\;\bm{e}(x_{i}),\;0,\;\ldots,\;\bm{e}(x_{i}),\;0,\;{\bm{0}}_{(1+|\Sigma|)(N-n)+|\Sigma|}\bigr)\in\{0,1\}^{2n+(1+|\Sigma|)N+|\Sigma|}. (65)

The first attention layer is an identity map, while the first feedforward layers compute 𝒉 i ( 1 ) = ( 𝒆 i , 𝒆 o i , ( 𝒆 i ) 1 ⋅ 𝒆 ⁡ ( x i ) , 0 , ( 𝒆 i ) 2 ⋅ 𝒆 ⁡ ( x i ) , 0 , … , ( 𝒆 i ) n ⋅ 𝒆 ⁡ ( x i ) , 0 , 0 ( 1 + | Σ | ) ​ ( N − n ) + | Σ | ) . {\bm{h}}^{(1)}_{i}\;=\;\bigl(\,{\bm{e}}_{i},\;{\bm{e}}_{o_{i}},\;({\bm{e}}_{i})_{1}\cdot\bm{e}(x_{i}),\;0,\;({\bm{e}}_{i})_{2}\cdot\bm{e}(x_{i}),\;0,\;\ldots,\;({\bm{e}}_{i})_{n}\cdot\bm{e}(x_{i}),\;0,\;{\bm{0}}_{(1+|\Sigma|)(N-n)+|\Sigma|}\,\bigr). (66)

The second attention layer uniformly gathers all positions and appends a constant 1 1 to each block: 𝒉 i ( 1.5 ) \displaystyle{\bm{h}}^{(1.5)}_{i} = ( 𝒆 i , 𝒆 o i , 𝒆 ⁡ ( x 1 ) , 1 , 𝒆 ⁡ ( x 2 ) , 1 , … , 𝒆 ⁡ ( x n ) , 1 , 0 ( 1 + | Σ | ) ​ ( N − n ) + | Σ | ) \displaystyle\;=\;\bigl(\,{\bm{e}}_{i},\;{\bm{e}}_{o_{i}},\;\bm{e}(x_{1}),\;1,\;\bm{e}(x_{2}),\;1,\;\ldots,\;\bm{e}(x_{n}),\;1,\;{\bm{0}}_{(1+|\Sigma|)(N-n)+|\Sigma|}\,\bigr) (67) = ( 𝒆 i , 𝒆 o i , 𝒉 ⁡ ( 0 ) , 0 | Σ | ) . \displaystyle\;=\;\bigl(\,{\bm{e}}_{i},\;{\bm{e}}_{o_{i}},\;{\bm{h}}(0),\;{\bm{0}}_{|\Sigma|}\,\bigr). (68)

We now proceed by induction. Assume that at the k k -th iteration, the output after the second attention layer at position i i is 𝒉 k , i ( 1.5 ) = ( 𝒆 i , 𝒆 o i , 𝒉 ⁡ ( k ) , t o i , k − 1 ⋅ 𝒆 ⁡ ( v o i ​ ( x ) ) ) , {\bm{h}}^{(1.5)}_{k,i}\;=\;\bigl(\,{\bm{e}}_{i},\;{\bm{e}}_{o_{i}},\;{\bm{h}}(k),\;t_{o_{i},k-1}\cdot\bm{e}(v_{o_{i}}(x))\bigr), (69) where t o i , k − 1 ≔ ( 𝒉 k − 1 , i ( 1.5 ) ) 2 ​ n + ( 1 + | Σ | ) ​ ( o i − 1 ) t_{o_{i},k-1}\coloneq\bigl({\bm{h}}^{(1.5)}_{k-1,i}\bigr)_{\,2n+(1+|\Sigma|)(o_{i}-1)} denotes the indicator flag showing whether the o i o_{i} -th node has been reached, i.e., whether it lies at depth k − 1 k-1 .

After passing through the second feedforward layer, the third attention layer with all weights set to zero, and the third feedforward layer, the hidden state updates to 𝒉 k , i ( 3 ) = ( 𝒆 i , 𝒆 o i , 𝒉 ⁡ ( k + 1 ) , t o i , k − 1 ⋅ 𝒆 ⁡ ( v o i ​ ( x ) ) ) . {\bm{h}}^{(3)}_{k,i}\;=\;\bigl(\,{\bm{e}}_{i},\;{\bm{e}}_{o_{i}},\;{\bm{h}}(k+1),\;t_{o_{i},k-1}\cdot\bm{e}(v_{o_{i}}(x))\,\bigr). (70)

After the fourth feedforward layer, the hidden state becomes 𝒉 k , i ( 4 ) = ( 𝒆 i , 𝒆 o i , 𝒉 ⁡ ( k + 1 ) , t o i , k ⋅ 𝒆 ⁡ ( v o i ​ ( x ) ) ) . {\bm{h}}^{(4)}_{k,i}\;=\;\bigl(\,{\bm{e}}_{i},\;{\bm{e}}_{o_{i}},\;{\bm{h}}(k+1),\;t_{o_{i},k}\cdot\bm{e}(v_{o_{i}}(x))\bigr). (71)

By induction, after the final iteration of depth depth ⁡ ( G n ) \mathrm{depth}(G_{n}) of the computation graph G n G_{n} , we obtain 𝒉 depth ⁡ ( G n ) , i ( 4 ) = ( ∗ , t o i , depth ⁡ ( G n ) ⋅ 𝒆 ⁡ ( v o i ​ ( x ) ) ) = ( ∗ , 𝒆 ⁡ ( v o i ​ ( x ) ) ) . {\bm{h}}^{(4)}_{\mathrm{depth}(G_{n}),i}\;=\;\bigl(\,*,\;t_{o_{i},\mathrm{depth}(G_{n})}\cdot\bm{e}(v_{o_{i}}(x))\bigr)\;=\;\bigl(\,*,\;\bm{e}(v_{o_{i}}(x))\bigr). (72)

The final output is given by z i = 𝐎𝐔𝐓 ⁡ ( 𝒉 depth ⁡ ( G n ) , i ( 4 ) ) = [ 𝟎 𝑰 | Σ | ] ⁡ ( 𝒉 depth ⁡ ( G n ) , i ( 4 ) ) = 𝒆 ⁡ ( v o i ​ ( x ) ) . z_{i}=\mathbf{OUT}({\bm{h}}^{(4)}_{\mathrm{depth}(G_{n}),i})=\begin{bmatrix}{\bm{0}}&{\bm{I}}_{|\Sigma|}\end{bmatrix}({\bm{h}}^{(4)}_{\mathrm{depth}(G_{n}),i})=\bm{e}(v_{o_{i}}(x)). (73)

The decoding function then selects the symbol corresponding to the maximum score: y i = Dec ⁡ ( z i ) = arg ⁡ max j ∈ [ | Σ | ] ⁡ ( 𝒆 ⁡ ( v o i ​ ( x ) ) ) j = v o i ​ ( x ) . y_{i}\;=\;\mathrm{Dec}(z_{i})=\arg\max_{j\in[|\Sigma|]}\,\bigl(\bm{e}(v_{o_{i}}(x))\bigr)_{j}=v_{o_{i}}(x). (74)

The parameter size grows proportionally with the input dimension size ⁡ ( G n ) \mathrm{size}(G_{n}) , since the function must be approximated along each dimension. Therefore, it can be bounded by O ⁡ ( ff ​ _ ​ param ​ ( G n ) ⋅ size ⁡ ( G n ) ) . O\!\left(\mathrm{ff\_param}(G_{n})\cdot\mathrm{size}(G_{n})\right). ∎

Discussion: Our proof compresses the entire input into a single position before applying an arbitrary feed-forward computation, which may appear to deviate from the standard Transformer architecture. An alternative approach is to distribute computation across multiple positions, as shown by ( Sanford et al., 2024b ) , which can reduce the required embedding dimension. We nevertheless adopt the single-position construction to isolate the core characteristics of looped models: attention is used purely as an information aggregation mechanism, a single Transformer layer suffices, and all computation is carried out by the feed-forward network in latent space. This latent-space computation is strictly more expressive than computation in the language space. This is supported by recent work for looped ReLUs ( Liang et al., 2024 ) .

### B.6 Proof of Theorem 3.6 for Continuous Thought

###### Proof.

The proofs are based on the construction for CoT and looped TF. Let the vocabulary be 𝒱 = Σ {\mathcal{V}}=\Sigma and each node v ∈ V n v\in V_{n} is labeled by a one-hot vector 𝒆 ⁡ ( v ) ∈ { 0 , 1 } | ℱ | \bm{e}(v)\in\{0,1\}^{|\mathcal{F}|} . At decoding step k k , the model has access to the concatenated sequence ( x 1 , x 2 , … , x n , h 1 , h 2 , … , h k ) (x_{1},x_{2},\ldots,x_{n},\,h_{1},h_{2},\ldots,h_{k}) (75) where x = ( x 1 , … , x n ) ∈ Σ n x=(x_{1},\ldots,x_{n})\in\Sigma^{n} denotes the input, and h i ∈ 𝔽 s m h_{i}\in\mathbb{F}_{s}^{m} is the hidden state generated at the i i -th Coconut step. For each node v j v_{j} , let v j ​ ( x ) v_{j}(x) denote its value on input x x . We assume that h k ≔ ( t 1 , k − 1 ⋅ 𝒆 ⁡ ( v 1 ​ ( x ) ) , t 1 , k , t 2 , k − 1 ⋅ 𝒆 ⁡ ( v 2 ​ ( x ) ) , t 2 , k , … , t N , k − 1 ⋅ 𝒆 ⁡ ( v N ​ ( x ) ) , t N , k , 𝒆 k ′ , 0 | Σ | ) , h_{k}\coloneq\bigl(t_{1,k-1}\cdot\bm{e}(v_{1}(x)),\,t_{1,k},\,t_{2,k-1}\cdot\bm{e}(v_{2}(x)),\,t_{2,k},\,\ldots,\,t_{N,k-1}\cdot\bm{e}(v_{N}(x)),\,t_{N,k},\,{\bm{e}}^{\prime}_{k},\,{\bm{0}}_{|\Sigma|}\bigr), where N ≔ size ⁡ ( G n ) N\coloneqq\mathrm{size}(G_{n}) , t i , k ∈ { 0 , 1 } t_{i,k}\in\{0,1\} indicates whether node v i v_{i} lies within depth at most k k , and 𝒆 : Σ → { 0 , 1 } | Σ | \bm{e}:\Sigma\to\{0,1\}^{|\Sigma|} denotes one-hot encoding, and the vector 𝒆 k ′ ∈ { 0 , 1 } N {\bm{e}}^{\prime}_{k}\in\{0,1\}^{N} is defined by 𝒆 k ′ = { 𝟎 k < depth ⁡ ( G n ) , 𝒆 ⁡ ( v o i ) k = depth ⁡ ( G n ) + i , {\bm{e}}^{\prime}_{k}=\begin{cases}\mathbf{0}&k<\mathrm{depth}(G_{n}),\\[4.0pt] \bm{e}(v_{o_{i}})&k=\mathrm{depth}(G_{n})+i,\end{cases} (76) where v o i v_{o_{i}} denotes the i i -th output node, and 𝒆 k ′ {\bm{e}}^{\prime}_{k} can be encoded by positional embeddings. Under this assumption, we prove by induction that the model generates h k + 1 ​ ( x ) h_{k+1}(x) .

We first consider the base case k = 1 k=1 , in which the model receives only 𝒙 {\bm{x}} . We focus on the final token. Using the same construction as in the looped TF, we can show that a single feed-forward layer followed by an attention layer suffices to copy all input tokens to every position. Specifically, the hidden representation at the last position becomes h 1 = ( 𝒆 ( v 1 ( x ) ) , 1 , … , 𝒆 ( v n ( x ) ) , 1 , 0 , 1 , 0 , 𝒆 0 ′ , 0 ) , h_{1}=\bigl(\bm{e}(v_{1}(x)),\,1,\,\ldots,\,\bm{e}(v_{n}(x)),\,1,\,{\bm{0}},\ 1,\,{\bm{0}},\,{\bm{e}}^{\prime}_{0},\,{\bm{0}}\bigl), (77) where v i ​ ( x ) = x i v_{i}(x)=x_{i} for all i ∈ [ n ] i\in[n] , corresponding to the input nodes. In what follows, we consider the case k > 1 k>1 . By the same argument as for Looped TF, based on Equation 62 , we show that the feed-forward layers can encode the entire computation graph into their weights and simulate all nodes simultaneously. That is, there exist two feed-forward layers whose composition, denoted by ℱ ​ ℱ \mathcal{FF} , satisfies ℱ ​ ℱ ​ ( h k ) = h k + 1 . \mathcal{FF}(h_{k})=h_{k+1}.

Consider the last feed-forward layer applied after the zero-weight attention layer. By substituting 𝒆 i = 𝒆 k ′ {\bm{e}}_{i}={\bm{e}}_{k}^{\prime} and 𝒙 = ( h k ) 1 : m − | Σ | {\bm{x}}={(h_{k})}_{1:m-|\Sigma|} into Lemma B.6 , there exists a feed-forward network FF \mathrm{FF} such that ( id + FF ) ( h k ) = { h k + 1 , if ​ k < depth ⁡ ( G n ) , ( ( h k + 1 ) 1 : m − | Σ | , 𝒆 ( v o i ( x ) ) ) , if ​ k = depth ⁡ ( G n ) + i . (\id+\mathrm{FF})(h_{k})=\begin{cases}h_{k+1},&\text{if }k<\mathrm{depth}(G_{n}),\\[6.0pt] \Bigl({(h_{k+1})}_{1:m-|\Sigma|},\,\bm{e}(v_{o_{i}}(x))\Bigr),&\text{if }k=\mathrm{depth}(G_{n})+i.\end{cases} (78) In particular, for all k < depth ⁡ ( G n ) k<\mathrm{depth}(G_{n}) , the output of the last token, for the input ( x 1 , x 2 , … , x n , h 1 , h 2 , … , h k ) (x_{1},x_{2},\ldots,x_{n},\;h_{1},h_{2},\ldots,h_{k}) is h k + 1 h_{k+1} . For the final positions corresponding to output nodes, namely k = depth ⁡ ( G n ) + i k=\mathrm{depth}(G_{n})+i , the hidden state 𝒉 k {\bm{h}}_{k} contains the embedding of the output symbol v o i ​ ( x ) v_{o_{i}}(x) . Applying the output projection yields 𝐎𝐔𝐓 ⁡ ( 𝒉 k ) = [ 𝟎 𝑰 | Σ | ] ​ 𝒉 k = 𝒆 ⁡ ( v o i ​ ( x ) ) . \mathbf{OUT}({\bm{h}}_{k})=\begin{bmatrix}\mathbf{0}&{\bm{I}}_{|\Sigma|}\end{bmatrix}{\bm{h}}_{k}=\bm{e}\!\left(v_{o_{i}}(x)\right). (79)

Finally, the decoding function outputs the symbol in Σ \Sigma corresponding to the maximum coordinate of 𝐎𝐔𝐓 ⁡ ( 𝒉 k ) \mathbf{OUT}({\bm{h}}_{k}) . ∎

### B.7 Proof for Theorem 3.12