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

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

Set the first layer as 𝑾 1 ( 1 ) = [ 𝟎 𝑾 1 ( f ) 𝟎 w f × 1 𝟎 𝑰 L 𝟎 L × 1 ] , 𝒃 ( 1 ) = [ 𝒃 ( f ) 𝟎 L ] , 𝑾 2 ( 1 ) = [ 𝑾 2 ( f ) 0 𝟎 L × w f 0 𝟎 1 × w f 𝟎 ] , {\bm{W}}^{(1)}_{1}=\begin{bmatrix}{\bm{0}}&{\bm{W}}^{(f)}_{1}&\mathbf{0}_{\,w_{f}\times 1}\\ {\bm{0}}&{\bm{I}}_{L}&\mathbf{0}_{\,L\times 1}\end{bmatrix},\quad{\bm{b}}^{(1)}=\begin{bmatrix}{\bm{b}}^{(f)}\\ {\bm{0}}_{L}\end{bmatrix},\quad{\bm{W}}^{(1)}_{2}=\begin{bmatrix}{\bm{W}}^{(f)}_{2}&\,{\bm{0}}\\ \mathbf{0}_{\,L\times w_{f}}&\,{\bm{0}}\\ \mathbf{0}_{\,1\times w_{f}}&{\bm{0}}\end{bmatrix}, (26) and define FF 1 ​ ( 𝟎 , 𝒖 , t ) ≔ 𝑾 2 ( 1 ) ​ ReLU ​ ( 𝑾 1 ( 1 ) ​ ( 𝟎 , 𝒖 , t ) + 𝒃 ( 1 ) ) \mathrm{FF}_{1}({\bm{0}},{\bm{u}},t)\coloneqq{\bm{W}}^{(1)}_{2}\,\mathrm{ReLU}\big({\bm{W}}^{(1)}_{1}({\bm{0}},{\bm{u}},t)+{\bm{b}}^{(1)}\big) . Then, it holds that FF 1 ​ ( 𝟎 , 𝒖 , t ) \displaystyle\mathrm{FF}_{1}({\bm{0}},{\bm{u}},t) = 𝑾 2 ( 1 ) ​ ReLU ​ ( [ 𝑾 1 ( f ) ​ 𝒖 + 𝒃 ( f ) 𝒖 ] ) \displaystyle={\bm{W}}^{(1)}_{2}\,\mathrm{ReLU}\left(\begin{bmatrix}{\bm{W}}^{(f)}_{1}{\bm{u}}+{\bm{b}}^{(f)}\\ {\bm{u}}\end{bmatrix}\right) (27) = [ 𝑾 2 ( f ) ​ ReLU ​ ( 𝑾 1 ( f ) ​ 𝒖 + 𝒃 ( f ) ) 𝟎 ] \displaystyle=\begin{bmatrix}{\bm{W}}^{(f)}_{2}\mathrm{ReLU}({\bm{W}}^{(f)}_{1}{\bm{u}}+{\bm{b}}^{(f)})\\ {\bm{0}}\end{bmatrix} (28) = [ FF f ​ ( 𝒖 ) 𝟎 ] . \displaystyle=\begin{bmatrix}\mathrm{FF}_{f}({\bm{u}})\\ {\bm{0}}\end{bmatrix}. (29) Thus we have ( id + FF 1 ) ​ ( 𝟎 , 𝒖 , t ) = ( FF f ​ ( 𝒖 ) , 𝒖 , t ) . (\id+\mathrm{FF}_{1})({\bm{0}},{\bm{u}},t)=\bigl(\mathrm{FF}_{f}({\bm{u}}),\,{\bm{u}},\,t\bigr).

For the second layer, choose δ ≤ 1 3 \delta\leq\tfrac{1}{3} and M ≥ 1 M\geq 1 and set 𝑾 1 ( 2 ) = [ 2 ​ 𝑰 n 𝟎 n × L M ​ 𝟏 n 2 ​ 𝑰 n 𝟎 n × L M ​ 𝟏 n 𝑰 n 𝟎 n × L 𝟎 n × 1 − 𝑰 n 𝟎 n × L 𝟎 n × 1 ] , 𝒃 ( 2 ) = [ − M ​ 𝟏 n ( 1 − M ) ​ 𝟏 n 𝟎 n 𝟎 n ] , 𝑾 2 ( 2 ) = [ 𝑰 n − 𝑰 n − 𝑰 n 𝑰 n 0 L × n 𝟎 L × n 𝟎 L × n 𝟎 L × n 0 1 × n 𝟎 1 × n 𝟎 1 × n 𝟎 1 × n ] . {\bm{W}}^{(2)}_{1}=\begin{bmatrix}2{\bm{I}}_{n}&\mathbf{0}_{n\times L}&M{\bm{1}}_{n}\\ 2{\bm{I}}_{n}&\mathbf{0}_{n\times L}&M{\bm{1}}_{n}\\ {\bm{I}}_{n}&\mathbf{0}_{n\times L}&\mathbf{0}_{n\times 1}\\ -{\bm{I}}_{n}&\mathbf{0}_{n\times L}&\mathbf{0}_{n\times 1}\end{bmatrix},\ {\bm{b}}^{(2)}=\begin{bmatrix}-M{\bm{1}}_{n}\\ (1-M){\bm{1}}_{n}\\ {\bm{0}}_{n}\\ {\bm{0}}_{n}\end{bmatrix},\ {\bm{W}}^{(2)}_{2}=\begin{bmatrix}\ {\bm{I}}_{n}&-{\bm{I}}_{n}&-{\bm{I}}_{n}&\ {\bm{I}}_{n}\ \\ \ \mathbf{0}_{\,L\times n}&\mathbf{0}_{\,L\times n}&\mathbf{0}_{\,L\times n}&\mathbf{0}_{\,L\times n}\\ \ \mathbf{0}_{\,1\times n}&\mathbf{0}_{\,1\times n}&\mathbf{0}_{\,1\times n}&\mathbf{0}_{\,1\times n}\end{bmatrix}. (30) and define FF 2 ​ ( 𝒚 ) ≔ 𝑾 2 ( 2 ) ​ ReLU ​ ( 𝑾 1 ( 2 ) ​ 𝒚 + 𝒃 ( 2 ) ) \mathrm{FF}_{2}({\bm{y}})\coloneqq{\bm{W}}^{(2)}_{2}\,\mathrm{ReLU}\big({\bm{W}}^{(2)}_{1}{\bm{y}}+{\bm{b}}^{(2)}\big) . Then it holds that, for 𝒛 ≔ FF f ​ ( 𝒖 ) {\bm{z}}\coloneqq\mathrm{FF}_{f}({\bm{u}}) , FF 2 ​ ( 𝒛 , 𝒖 , t ) \displaystyle\mathrm{FF}_{2}({\bm{z}},{\bm{u}},t) = 𝑾 2 ( 2 ) ​ [ ReLU ⁡ ( 2 ​ 𝒛 + M ​ t ​ 𝟏 n − M ​ 𝟏 n ) ReLU ⁡ ( 2 ​ 𝒛 + M ​ t ​ 𝟏 n + ( 1 − M ) ​ 𝟏 n ) ReLU ⁡ ( 𝒛 ) ReLU ⁡ ( − 𝒛 ) ] \displaystyle={\bm{W}}^{(2)}_{2}\,\begin{bmatrix}\mathrm{ReLU}\bigl(2{\bm{z}}+Mt{\bm{1}}_{n}-M{\bm{1}}_{n}\bigr)\\ \mathrm{ReLU}\bigl(2{\bm{z}}+Mt{\bm{1}}_{n}+(1-M){\bm{1}}_{n}\bigr)\\ \mathrm{ReLU}({\bm{z}})\\ \mathrm{ReLU}(-{\bm{z}})\end{bmatrix} (31) = 𝑾 2 ( 2 ) ​ [ ReLU ⁡ ( 2 ​ 𝒛 + M ⁡ ( t − 1 ) ​ 𝟏 n ) ReLU ⁡ ( 2 ​ 𝒛 + 1 + M ⁡ ( t − 1 ) ​ 𝟏 n ) 𝒛 𝟎 ] \displaystyle={\bm{W}}^{(2)}_{2}\,\begin{bmatrix}\mathrm{ReLU}\bigl(2{\bm{z}}+M(t-1){\bm{1}}_{n}\bigr)\\ \mathrm{ReLU}\bigl(2{\bm{z}}+1+M(t-1){\bm{1}}_{n}\bigr)\\ {\bm{z}}\\ {\bm{0}}\end{bmatrix} (32) = [ ReLU ⁡ ( 𝒛 − δ + M ⁡ ( t − 1 ) ) − ReLU ⁡ ( 𝒛 − 1 + M ⁡ ( t − 1 ) ) − 𝒛 𝟎 0 ] , \displaystyle=\begin{bmatrix}\mathrm{ReLU}\bigl({\bm{z}}-\delta+M(t-1)\bigr)-\mathrm{ReLU}\bigl({\bm{z}}-1+M(t-1)\bigr)-{\bm{z}}\\ \mathbf{0}\\ 0\end{bmatrix}, (33) where it satisfies that ReLU ⁡ ( 𝒛 − δ + M ⁡ ( t − 1 ) ) − ReLU ⁡ ( 𝒛 − δ − 1 + M ⁡ ( t − 1 ) ) = { 𝒆 ⁡ ( f ⁡ ( 𝒙 ) ) if t = 1 , 0 if t = 0 . \mathrm{ReLU}\bigl({\bm{z}}-\delta+M(t-1)\bigr)-\mathrm{ReLU}\bigl({\bm{z}}-\delta-1+M(t-1)\bigr)\;=\;\begin{cases}\;\bm{e}(f({\bm{x}}))&\text{if}\quad t=1,\\ \;{\bm{0}}&\text{if}\quad t=0.\end{cases} (34)

Therefore, the composition satisfies ( id + FF 2 ) ∘ ( id + FF 1 ) ​ ( 𝟎 , 𝒖 , t ) = ( t ⋅ 𝒆 ⁡ ( f ⁡ ( 𝒙 ) ) , 𝒖 , t ) . (\id+\mathrm{FF}_{2})\circ(\id+\mathrm{FF}_{1})\bigl({\bm{0}},{\bm{u}},t\bigr)=(t\cdot\bm{e}(f({\bm{x}})),\ {\bm{u}},\ t). ∎

### B.4 Proof for Theorem 3.5

###### Proof.

Let G n = ( V n , E n ) G_{n}=(V_{n},E_{n}) be a computation graph, where ℱ = { f 1 , f 2 , … , f | ℱ | } \mathcal{F}=\{f_{1},f_{2},\dots,f_{|\mathcal{F}|}\} . 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}|} indicating the function assigned to v v from the finite set ℱ \mathcal{F} . Let v 1 , v 2 , … , v | V n | v_{1},v_{2},\dots,v_{|V_{n}|} denote a fixed topological ordering of V n V_{n} , with inputs appearing first and outputs last. For each function f i ∈ ℱ f_{i}\in\mathcal{F} , let C f i ( n ) ≔ max { | pred ( v ) | : v ∈ V n , 𝒆 ( v ) = 𝒆 i } . C_{f_{i}}(n)\coloneqq\max\{\,|\mathrm{pred}(v)|:v\in V_{n},\;\bm{e}(v)={\bm{e}}_{i}\,\}. and define C sum ​ ( n ) ≔ ∑ f ∈ ℱ C f ​ ( n ) , C max ​ ( n ) ≔ max f ∈ ℱ ⁡ C f ​ ( n ) . C_{\mathrm{sum}}(n)\coloneqq\sum_{f\in\mathcal{F}}C_{f}(n),\ C_{\mathrm{max}}(n)\coloneqq\max_{f\in\mathcal{F}}C_{f}(n).

Let the precision be s ⁡ ( n ) = C ⋅ ⌈ log 2 ⁡ n ⌉ s(n)=C\cdot\lceil\log_{2}n\rceil where C ∈ ℕ C\in\mathbb{N} is a sufficiently large integer such that 2 s ⁡ ( n ) ≥ n C 2^{s(n)}\;\geq\;n^{C} exceeds the maximum polynomial step bound under consideration. We denote by pred ⁡ ( v i ) ∈ 𝔽 s ⁡ ( n ) C max ​ ( n ) \mathrm{pred}(v_{i})\in\mathbb{F}_{s(n)}^{C_{\max}(n)} the vector of predecessor indices of node v i v_{i} ; that is, if v i v_{i} has d ≤ C max ​ ( n ) d\leq C_{\max}(n) incoming edges from nodes v j 1 , … , v j d v_{j_{1}},\dots,v_{j_{d}} , then pred ⁡ ( v i ) = ( j 1 , … , j d , 𝟎 ) \mathrm{pred}(v_{i})=(j_{1},\dots,j_{d},{\bm{0}}) , where zeros are used for padding so that the length is exactly C max ​ ( n ) C_{\max}(n) .

Let the vocabulary be 𝒱 = Σ {\mathcal{V}}=\Sigma . At decoding step k k , the model has access to the concatenated sequence ( x 1 , x 2 , … , x n , y 1 , y 2 , … , y k ) ∈ Σ n + k , (x_{1},x_{2},\ldots,x_{n},\,y_{1},y_{2},\ldots,y_{k})\in\Sigma^{n+k}, (35) where x = ( x 1 , … , x n ) ∈ Σ n x=(x_{1},\ldots,x_{n})\in\Sigma^{n} denotes the input, and y i y_{i} is the token generated at the i i -th CoT 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 every intermediate output satisfies y i = v n + i ​ ( x ) . y_{i}=v_{n+i}(x). Under this assumption, we prove by induction that the model generates the next token correctly, i.e., y k + 1 = v n + k + 1 ​ ( x ) . y_{k+1}=v_{n+k+1}(x).

##### Embedding

The embedding at position i ∈ [ n + k ] i\in[n+k] , denoted by 𝒉 i ( 0 ) ∈ 𝔽 s ⁡ ( n ) m {\bm{h}}^{(0)}_{i}\in\mathbb{F}_{s(n)}^{m} , where m ≔ | Σ | + | ℱ | + ( 1 + C max ​ ( n ) ) ​ s ​ ( n ) + | Σ | ​ C sum ​ ( n ) m\coloneq|\Sigma|+|\mathcal{F}|+(1+C_{\mathrm{max}}(n))s(n)+|\Sigma|C_{\mathrm{sum}}(n) is defined as 𝒉 i ( 0 ) = ( 𝒆 ⁡ ( v i ​ ( x ) ) , 𝒆 ⁡ ( v i + 1 ) , 𝗌𝖻𝗂𝗇 s ⁡ ( n ) ​ ( i ) , 𝐬𝐛𝐢𝐧𝐩𝐫𝐞𝐝 s ⁡ ( n ) ​ ( v i + 1 ) , 0 | Σ | ​ C sum ​ ( n ) ) , {\bm{h}}^{(0)}_{i}=\left(\bm{e}(v_{i}(x)),\,\bm{e}(v_{i+1}),\,\mathsf{sbin}_{s(n)}(i),\,\bm{\mathrm{sbinpred}}_{s(n)}(v_{i+1}),\,{\bm{0}}_{|\Sigma|C_{\mathrm{sum}}(n)}\right), (36) where 𝒆 : Σ → { 0 , 1 } | Σ | \bm{e}\colon\Sigma\to\{0,1\}^{|\Sigma|} denote the one-hot encoding of the symbol and, 𝐬𝐛𝐢𝐧𝐩𝐫𝐞𝐝 s ⁡ ( n ) ​ ( v ) ∈ 𝔽 s ⁡ ( n ) C max ​ ( n ) ⋅ s ​ ( n ) \bm{\mathrm{sbinpred}}_{s(n)}(v)\in\mathbb{F}^{C_{\mathrm{max}}(n)\cdot s(n)}_{s(n)} encodes the binary representations of the predecessor indices: 𝐬𝐛𝐢𝐧𝐩𝐫𝐞𝐝 s ⁡ ( n ) ​ ( v i ) ≔ ( 𝗌𝖻𝗂𝗇 s ⁡ ( n ) ​ ( pred ​ ( v i ) 0 ) , … , 𝗌𝖻𝗂𝗇 s ⁡ ( n ) ​ ( pred ​ ( v i ) C max ​ ( n ) ) ) . \bm{\mathrm{sbinpred}}_{s(n)}(v_{i})\coloneq\left(\mathsf{sbin}_{s(n)}(\mathrm{pred}(v_{i})_{0}),\,\ldots,\,\mathsf{sbin}_{s(n)}(\mathrm{pred}(v_{i})_{C_{\mathrm{max}}(n)})\right). (37)

This embedding is constructed, for z ∈ Σ z\in\Sigma , as 𝐖𝐄 ⁡ ( z ) = ( 𝒆 ⁡ ( z ) , 0 ) , 𝐏𝐄 ⁡ ( i ) = ( 𝟎 , 𝒆 ⁡ ( v i + 1 ) , 𝗌𝖻𝗂𝗇 s ⁡ ( n ) ​ ( i ) , 𝐬𝐛𝐢𝐧𝐩𝐫𝐞𝐝 s ⁡ ( n ) ​ ( v i + 1 ) , 0 ) . \mathbf{WE}(z)=\left(\bm{e}(z),\,{\bm{0}}\right),\quad\mathbf{PE}(i)=\left({\bm{0}},\,\bm{e}(v_{i+1}),\,\mathsf{sbin}_{s(n)}(i),\,\bm{\mathrm{sbinpred}}_{s(n)}(v_{i+1}),\,\mathbf{0}\right). (38)

##### Attention layer

The first attention layer consists of C max ⁡ ( n ) C_{\mathrm{max}(n)} heads. The h h -th head is configured to attend to the position corresponding to the h h -th predecessor. Specifically, for each position i i and head h ∈ [ C max ⁡ ( n ) ] h\in[C_{\mathrm{max}(n)}] , the attention vectors are defined as: 𝒒 i , h \displaystyle{\bm{q}}_{i,h} = 𝗌𝖻𝗂𝗇 s ⁡ ( n ) ​ ( pred ​ ( v i + 1 ) h ) ⌢ ​ 1 s ⁡ ( n ) , \displaystyle={\mathsf{sbin}_{s(n)}(\mathrm{pred}(v_{i+1})_{h})}^{\frown}{1_{s(n)}}, (39) 𝒌 i , h \displaystyle{\bm{k}}_{i,h} = 2 s ⁡ ( n ) + 1 ⋅ 𝗌𝖻𝗂𝗇 s ⁡ ( n ) ​ ( i ) ⌢ ​ ( − 1 s ⁡ ( n ) ) , \displaystyle=2^{s(n)+1}\cdot{\mathsf{sbin}_{s(n)}(i)}^{\frown}{(-1_{s(n)})}, (40) 𝐯 i , h \displaystyle\mathbf{v}_{i,h} = 𝒆 ​ ( v i ​ ( x ) ) , \displaystyle=\bm{e}(v_{i}(x)), (41) where vectors of different lengths are zero-padded to match the dimension. By Lemma B.5 , each attention head of the last position i = n + k i=n+k retrieves the predecessor’s value of v n + k v_{n+k} 𝒂 n + k , h = 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) h ​ ( x ) ) {\bm{a}}_{n+k,h}=\bm{e}\bigl(v_{\mathrm{pred}(v_{n+k+1})_{h}}(x)\bigr) (42) With an appropriate output projection 𝑶 {\bm{O}} such that 𝑶 ⁡ ( 𝒂 i , 1 , … , 𝒂 i , H ) = ( 𝟎 , 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) 0 ​ ( x ) ) , … , 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) C max ⁡ ( n ) ​ ( x ) ) , 0 ) , {\bm{O}}({\bm{a}}_{i,1},\ldots,{\bm{a}}_{i,H})=({\bm{0}},\,\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{0}}(x)),\,\ldots,\,\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{C_{\mathrm{max}(n)}}}(x)),\,{\bm{0}}), (43) the hidden state at position n + k n+k after the attention layer is given by 𝒉 n + k ( 0.5 ) = ( 𝒆 ⁡ ( v n + k ​ ( x ) ) , 𝒆 ⁡ ( v n + k + 1 ) , 𝗌𝖻𝗂𝗇 s ⁡ ( n ) ​ ( n + k ) , 𝐬𝐛𝐢𝐧𝐩𝐫𝐞𝐝 s ⁡ ( n ) ​ ( v n + k + 1 ) CLOSE , \displaystyle{\bm{h}}^{(0.5)}_{n+k}=\Big(\bm{e}(v_{n+k}(x)),\,\bm{e}(v_{n+k+1}),\,\mathsf{sbin}_{s(n)}(n+k),\,\bm{\mathrm{sbinpred}}_{s(n)}(v_{n+k+1}),\, (44) OPEN 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) 0 ​ ( x ) ) , … , 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) C max ⁡ ( n ) ​ ( x ) ) ⏟ updated , 𝟎 C sum ​ ( n ) ​ | Σ | ) . \displaystyle\underbrace{\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{0}}(x)),\,\ldots,\,\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{C_{\mathrm{max}(n)}}}(x))}_{\text{updated}}\,,{\bm{0}}_{C_{\mathrm{sum}}(n)|\Sigma|}\Big). (45) The second and third attention layers are disabled (i.e., all attention weights are set to zero).

##### Feed-forward layer

By Lemma B.7 , a single feed-forward layer can approximate multiple functions by partitioning the input into blocks. The first feed-forward layer then places the arguments, gathered by attention, into the correct positions. By Lemma B.6 , where the vector 𝒆 i {\bm{e}}_{i} therein corresponds to 𝟏 | Σ | {\bm{1}}_{|\Sigma|} here, the hidden state at the last position, denoted by 𝒉 n + k ( 1 ) {\bm{h}}^{(1)}_{n+k} , becomes ( ( 𝒉 n + k ( 0.5 ) ) 1 : r , ∥ j = 1 | ℱ | ( 𝒆 ( v pred ​ ( v n + k + 1 ) 1 ( x ) ) ⋅ 1 , … , 𝒆 ( v pred ​ ( v n + k + 1 ) C j ​ ( n ) ( x ) ) ⋅ 1 ) , \displaystyle\Big(({\bm{h}}^{(0.5)}_{n+k})_{1:r},\big\|_{j=1}^{|\mathcal{F}|}\,\big(\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{1}}(x))\cdot 1,\ldots,\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{C_{j}(n)}}(x))\cdot 1\Big), (46) where r = | Σ | + | ℱ | + ( 1 + C max ) ​ ( n ) ​ s ​ ( n ) r=|\Sigma|+|\mathcal{F}|+(1+C_{\mathrm{max}})(n)s(n) .

By Assumption 3.4 and Lemmas B.8 and B.7 , there exist feed-forward layers FF 2 , FF 3 : 𝔽 s ⁡ ( n ) m → 𝔽 s ⁡ ( n ) m \mathrm{FF}_{2},\mathrm{FF}_{3}:\mathbb{F}_{s(n)}^{m}\to\mathbb{F}_{s(n)}^{m} such that, for every input tuple 𝒙 j ≔ ( x 1 ( j ) , … , x C j ​ ( n ) ( j ) ) ∈ Σ C j ​ ( n ) {\bm{x}}_{j}\coloneq(x^{(j)}_{1},\ldots,x^{(j)}_{C_{j}(n)})\in\Sigma^{C_{j}(n)} for j ∈ [ | ℱ | ] j\in[|\mathcal{F}|] and every 𝒕 ∈ { 0 , 1 } | ℱ | {\bm{t}}\in\{0,1\}^{|\mathcal{F}|} , the composition ℱ ​ ℱ ≔ ( id + FF 3 ) ∘ ( id + FF 2 ) \mathcal{FF}\coloneq(\id+\mathrm{FF}_{3})\circ(\id+\mathrm{FF}_{2}) satisfies ℱ ℱ ( ∗ , 𝒕 , ∗ , ∥ j = 1 | ℱ | ( 𝒆 ( x 1 ( j ) ) , … , 𝒆 ( x C j ​ ( n ) ( j ) ) ) ) = ( ∑ j | ℱ | 𝒕 j ⋅ 𝒆 ( f j ( 𝒙 j ) ) , ∗ ) , \mathcal{FF}\left(*,\,{\bm{t}},\,*,\,\big\|_{j=1}^{|\mathcal{F}|}\,(\bm{e}(x^{(j)}_{1}),\ldots,\bm{e}(x^{(j)}_{C_{j}(n)}))\right)\;=\;\left(\sum^{|\mathcal{F}|}_{j}{\bm{t}}_{j}\cdot\bm{e}(f_{j}({\bm{x}}_{j})),*\right), (47) zwhere ∗ * denotes an unspecified vector. Since the second and third attention layers are disabled, after the third layer, applying the second and third feed-forward layers FF 2 , FF 3 \mathrm{FF}_{2},\mathrm{FF}_{3} , the hidden state becomes 𝒉 n + k ( 3 ) \displaystyle{\bm{h}}^{(3)}_{n+k} = ℱ ​ ℱ ​ ( 𝒉 n + k ( 1 ) ) \displaystyle=\mathcal{FF}\left({\bm{h}}^{(1)}_{n+k}\right) (48) = ℱ ℱ ( ∗ , 𝒆 ( v n + k + 1 ) , ∗ , ∥ j = 1 | ℱ | ( 𝒆 ( v pred ​ ( v n + k + 1 ) 1 ( x ) ) , … , 𝒆 ( v pred ​ ( v n + k + 1 ) C j ​ ( n ) ( x ) ) ) ) \displaystyle=\mathcal{FF}\left(*,\,\bm{e}(v_{n+k+1}),\,*,\,\bigg\|_{j=1}^{|\mathcal{F}|}\,\big(\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{1}}(x)),\ldots,\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{C_{j}(n)}}(x))\big)\right) (49) = ( ∑ j = 1 | ℱ | 𝒆 ​ ( v n + k + 1 ) j ⋅ 𝒆 ⁡ ( f j ​ ( 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) 1 ​ ( x ) ) , … , 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) C j ​ ( n ) ​ ( x ) ) ) ) , ∗ ) \displaystyle=\Biggl(\sum_{j=1}^{|\mathcal{F}|}\bm{e}(v_{n+k+1})_{j}\cdot\bm{e}\Big(f_{j}\big(\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{1}}(x)),\ldots,\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{C_{j}(n)}}(x))\big)\Big),\;*\Biggr) (50) = ( 𝒆 ⁡ ( f l ​ ( 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) 1 ​ ( x ) ) , … , 𝒆 ⁡ ( v pred ​ ( v n + k + 1 ) C l ​ ( n ) ​ ( x ) ) ) ) , ∗ ) , where 𝒆 ⁡ ( v n + k + 1 ) = 𝒆 l \displaystyle=\Biggl(\bm{e}\Big(f_{l}\big(\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{1}}(x)),\ldots,\bm{e}(v_{\mathrm{pred}(v_{n+k+1})_{C_{l}(n)}}(x))\big)\Big),\;*\Biggr),\,\text{where}\quad\bm{e}(v_{n+k+1})={\bm{e}}_{l} (51) = ( 𝒆 ⁡ ( v n + k + 1 ​ ( x ) ) , ∗ ) . \displaystyle=\Big(\bm{e}\big(v_{n+k+1}(x)\big),\;*\Big). (52)

##### Output layer

The final output is given by 𝒉 n + k = 𝐎𝐔𝐓 ⁡ ( 𝒉 n + k ( 3 ) ) = [ 𝑰 | Σ | 𝟎 ] ​ 𝒉 n + k ( 3 ) = 𝒆 ⁡ ( v n + k + 1 ​ ( x ) ) . {\bm{h}}_{n+k}=\mathbf{OUT}({\bm{h}}^{(3)}_{n+k})=\begin{bmatrix}{\bm{I}}_{|\Sigma|}&\mathbf{0}\end{bmatrix}{\bm{h}}^{(3)}_{n+k}=\bm{e}\!\left(v_{n+k+1}(x)\right). (53) The decoding function then outputs the symbol corresponding to the maximum score, y n + k + 1 = Dec ⁡ ( 𝒉 n + k ) = arg ⁡ max j ∈ [ | Σ | ] ⁡ 𝒆 ⁡ ( v n + k + 1 ​ ( x ) ) = v n + k + 1 ​ ( x ) . y_{n+k+1}=\mathrm{Dec}({\bm{h}}_{n+k})=\arg\max_{j\in[|\Sigma|]}\bm{e}(v_{n+k+1}(x))=v_{n+k+1}(x). (54) By induction on k k , the model computes the values at all nodes in topological order. The parameter size of the model is determined by the requirements of the feedforward layers, O ⁡ ( ff ​ _ ​ param ​ ( G n ) ) . O(\mathrm{ff\_param}(G_{n}))\,. While the dimensions and heads of the attention layers depend on C max ​ ( n ) C_{\max}(n) , which is precisely what is already required for the feedforward layers to approximate the target functions. ∎

### B.5 Proof of Theorem 3.6 for Looped Transformer

###### Proof.