Physics of Language Models: Part 1, Learning Hierarchical Language Structures — part 3/7

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

Let l l be the last layer of the transformer (other layers are studied in Appendix C.2 ). Given an input string x x , we denote the hidden state of the transformer at layer l l and position i i as E i ​ ( x ) ∈ ℝ d E_{i}(x)\in\mathbb{R}^{d} . We first investigate whether a linear function can predict ( 𝔟 1 ​ ( i ) , … , 𝔟 L ​ ( i ) ) i ∈ [ len ​ ( x ) ] \big(\mathfrak{b}_{1}(i),\dots,\mathfrak{b}_{L}(i)\big)_{i\in[\textbf{len}(x)]} and ( 𝔰 1 ​ ( i ) , … , 𝔰 L ​ ( i ) ) i ∈ [ len ​ ( x ) ] \big(\mathfrak{s}_{1}(i),\dots,\mathfrak{s}_{L}(i)\big)_{i\in[\textbf{len}(x)]} using the full ( E i ​ ( x ) ) i ∈ [ len ​ ( x ) ] \big(E_{i}(x)\big)_{i\in[\textbf{len}(x)]} . If so, it implies that the last-layer hidden states encode the CFG’s structural information up to a linear transformation .

Multi-head linear probing (full). Due to the high dimensionality of this linear function (e.g., len ​ ( x ) = 300 \textbf{len}(x)=300 and d = 768 d=768 yield 300 × 768 300\times 768 dimensions) and variable string lengths , we propose a multi-head linear function for efficient learning. We consider a set of linear functions f r : ℝ d → ℝ | 𝐍𝐓 | f_{r}\colon\mathbb{R}^{d}\to\mathbb{R}^{|\mathbf{NT}|} , where r ∈ [ H ] r\in[H] and H H is the number of “heads”. To predict any 𝔰 ℓ ​ ( i ) \mathfrak{s}_{\ell}(i) , we apply: G i ​ ( x ) = ∑ r ∈ [ H ] , k ∈ [ len ​ ( x ) ] w r , i → k ⋅ f r ​ ( E k ​ ( x ) ) ∈ ℝ | 𝐍𝐓 | \displaystyle\textstyle G_{i}(x)=\sum_{r\in[H],k\in[\textbf{len}(x)]}w_{r,i\to k}\cdot f_{r}(E_{k}(x))\in\mathbb{R}^{|\mathbf{NT}|} (4.2) where w r , i → k = def exp ⁡ ( ⟨ P i , r , P k , r ⟩ ) ∑ k ′ ∈ [ len ​ ( x ) ] exp ⁡ ( ⟨ P i , r , P k ′ , r ⟩ ) w_{r,i\to k}\stackrel{{\scriptstyle\mathrm{\scriptscriptstyle def}}}{{=}}\frac{\exp(\langle P_{i,r},P_{k,r}\rangle)}{\sum_{k^{\prime}\in[\textbf{len}(x)]}\exp(\langle P_{i,r},P_{k^{\prime},r}\rangle)} for trainable parameters P i , r ∈ ℝ d ′ P_{i,r}\in\mathbb{R}^{d^{\prime}} . G i G_{i} can be seen as a “multi-head attention” over linear functions. We train G i ​ ( x ) ∈ ℝ | 𝐍𝐓 | G_{i}(x)\in\mathbb{R}^{|\mathbf{NT}|} using the cross-entropy loss to predict ( 𝔰 ℓ ​ ( i ) ) ℓ ∈ [ L ] \big(\mathfrak{s}_{\ell}(i)\big)_{\ell\in[L]} . Despite having multiple heads, G i ​ ( x ) G_{i}(x) is still a linear function over ( E k ​ ( x ) ) k ∈ [ len ​ ( x ) ] (E_{k}(x))_{k\in[\textbf{len}(x)]} as the linear weights w r , i → k w_{r,i\to k} depend only on positions i i and k k , not on x x . Similarly, we train G i ′ ​ ( x ) ∈ ℝ L G^{\prime}_{i}(x)\in\mathbb{R}^{L} using the logistic loss to predict the binary values ( 𝔟 ℓ ​ ( i ) ) ℓ ∈ [ L ] \big(\mathfrak{b}_{\ell}(i)\big)_{\ell\in[L]} . Details are in Section A.4 .

Using such multi-head linear probing, we discover that:

Result 4 ( Figure 5 ) . Pre-training allows GPT models to almost perfectly encode the NT ancestor 𝔰 ℓ ​ ( i ) \mathfrak{s}_{\ell}(i) and NT boundary 𝔟 ℓ ​ ( i ) \mathfrak{b}_{\ell}(i) information in the last transformer layer’s hidden states ( E k ​ ( x ) ) k ∈ [ len ​ ( x ) ] (E_{k}(x))_{k\in[\textbf{len}(x)]} , up to a linear transformation. (See Figure 5 for comparison against randomly-initialized 𝙶𝙿𝚃 𝗋𝖺𝗇𝖽 \mathtt{GPT}_{\mathsf{rand}} or encoder model 𝚍𝚎𝙱𝙴𝚁𝚃𝚊 \mathtt{deBERTa} , which fail to recover deep NT structure.) 11 11 footnotetext: 𝚍𝚎𝙱𝙴𝚁𝚃𝚊 \mathtt{deBERTa} is a modern variant of BERT, equipped with relative attentions. It is expected that encoder models may not learn deep NT information, because in a masked-language modeling (MLM) task, the model only needs to figure out the missing token from its surrounding, say, 20 tokens. This can be done by pattern matching, as opposed to global planning like dynamic programming. But, do we need this full layer for linear probing? We explore next.

### 4.2 Result 5: NT Ancestors are Encoded At NT Boundaries

In Result 4, we used the full hidden layer, ( E i ​ ( x ) ) i ∈ [ len ​ ( x ) ] \big(E_{i}(x)\big)_{i\in[\textbf{len}(x)]} , to predict ( 𝔰 ℓ ​ ( i ) ) ℓ ∈ [ L ] \big(\mathfrak{s}_{\ell}(i)\big)_{\ell\in[L]} for each position i i . This is essential since it’s information-theoretically impossible to extract all of i i ’s NT ancestors by only reading E i ​ ( x ) E_{i}(x) or even all hidden states to its left , especially if x i x_{i} is the start of a string or a subtree in the CFG. But, how about those ones information-theoretically possible? In particular, how about predicting 𝔰 ℓ ​ ( i ) \mathfrak{s}_{\ell}(i) at locations i i with 𝔟 ℓ ​ ( i ) = 1 \mathfrak{b}_{\ell}(i)=1 — i.e., at the end of the CFG subtrees.

Multi-head linear probing (diagonal). We consider a neighborhood of position i i in the hidden states, say E i ± 1 ​ ( x ) E_{i\pm 1}(x) , and use that for linear probing. In symbols, we replace w r , i → k w_{r,i\to k} in ( 4.2 ) with zeros for | i − k | > 1 |i-k|>1 (tridiagonal masking), or with zeros for i ≠ k i\neq k (diagonal masking). G i ​ ( x ) = ∑ r ∈ [ H ] , k ∈ [ len ​ ( x ) ] , | i − k | ≤ δ w r , i → k ⋅ f r ​ ( E k ​ ( x ) ) ∈ ℝ | 𝐍𝐓 | where δ = 0 or 1 \displaystyle\textstyle G_{i}(x)=\sum_{r\in[H],k\in[\textbf{len}(x)],|i-k|\leq\delta}w_{r,i\to k}\cdot f_{r}(E_{k}(x))\in\mathbb{R}^{|\mathbf{NT}|}\hskip 20.00003pt\text{where $\delta=0$ or $1$ } (4.3)

Result 5 ( Figure 6 + 7 ) . For GPT models, the information of position i i ’s NT ancestor/boundary is locally encoded around position i ± 1 i\pm 1 when i i is on the NT boundary. This is because: • At NT boundaries (i.e., 𝔟 ℓ ​ ( x ) = 1 \mathfrak{b}_{\ell}(x)=1 ), we discover that diagonal or tridiagonal multi-head linear probing ( 4.3 ) is adequate for accurately predicting the NT ancestors 𝔰 ℓ ​ ( x ) \mathfrak{s}_{\ell}(x) (see Figure 7 ). • Such masking is also sufficient for accurately predicting NT boundaries 𝔟 ℓ ​ ( i ) \mathfrak{b}_{\ell}(i) (deferred to Figure 19 in Appendix C.1 ). In contrast, encoder models like 𝚍𝚎𝙱𝙴𝚁𝚃𝚊 \mathtt{deBERTa} do not store deep NT information at the NT boundaries.

Related work. Linear probing at least traces back to Hewitt and Manning [15] , who examines the correlation between BERT’s hidden states and the parse tree distance metric (similar to NT-distance in our language). Subsequent studies [ 30 , 40 , 20 , 18 , 34 , 36 , 8 ] also explored probing techniques to suggest that BERT-like transformers can approximate CFGs from natural languages .

Our approach differs not only in the multi-head probing formula that we proposed; also that we use synthetic data to demonstrate that linear probing can almost perfectly recover NT ancestors and boundaries, even for complex and ambiguous CFG strings exceeding hundreds of tokens (c.f. English CFG has an average length of 28, see Appendix G ). We focus on training generative decoder-only models; an encoder-based model like BERT [ 17 ] or its modern variant 𝚍𝚎𝙱𝙴𝚁𝚃𝚊 \mathtt{deBERTa} [ 14 ] may not learn deep (i.e., close to the CFG root) NT information very well, as shown in Result 4 - 5 .

Our results, along with Section 5 next, shall provide evidence that generative language models like GPT-2 employ a DP-like approach to generate CFGs, while encoder-based models trained via MLM struggle to learn more complex/deeper CFGs.

## 5 Results 6-9: How Do Transformers Learn NTs?

We now delve into the attention patterns, which reveal the model’s reasoning mechanisms. We demonstrate that these patterns mirror the CFG’s syntactic structure and rules, with the transformer employing different attention heads to reason with NTs at different CFG levels.

### 5.1 Result 6: Position-Based Attention

We first note that the transformer’s attention weights are primarily influenced by the tokens’ relative distance. This holds true even when trained on the CFG data with absolute positional embedding. This implies that the transformer learns the CFG’s regularity and periodicity through positional information, which it then uses for generation.

Formally, let A l , h , j → i ​ ( x ) A_{l,h,j\to i}(x) for j ≥ i j\geq i represent the attention weight for positions j → i j\to i at layer l l and head h h of the transformer, on input sequence x x . For each layer l l , head h h , and distance p ≥ 0 p\geq 0 , we compute the average of the partial sum ∑ 1 ≤ i ′ ≤ i A l , h , j → i ′ ​ ( x ) \sum_{1\leq i^{\prime}\leq i}A_{l,h,j\to i^{\prime}}(x) over all data x x and pairs i , j i,j with j − i = p j-i=p . We plot this cumulative sum for l , h , p l,h,p in Figure 8 . We observe a strong correlation between the attention pattern and the relative distance p = j − i p=j-i . The attention pattern is also multi-scale , with some attention heads focusing on shorter distances and others on longer ones.

Motivated by this, we explore whether using position-based attention is sufficient to learn CFGs. In Figure 4 , we find that 𝙶𝙿𝚃 𝗉𝗈𝗌 \mathtt{GPT}_{\mathsf{pos}} (or even 𝙶𝙿𝚃 𝗎𝗇𝗂 \mathtt{GPT}_{\mathsf{uni}} ) performs well, surpassing the vanilla 𝙶𝙿𝚃 \mathtt{GPT} , but not reaching the full potential of 𝙶𝙿𝚃 𝗋𝖾𝗅 \mathtt{GPT}_{\mathsf{rel}} . This supports the superior practical performance of relative-position based transformer variants (such as 𝙶𝙿𝚃 𝗋𝖾𝗅 , 𝙶𝙿𝚃 𝗋𝗈𝗍 \mathtt{GPT}_{\mathsf{rel}},\mathtt{GPT}_{\mathsf{rot}} , 𝚍𝚎𝙱𝙴𝚁𝚃𝚊 \mathtt{deBERTa} ) over their base models ( 𝙶𝙿𝚃 \mathtt{GPT} or BERT). On this other hand, this also indicates that position-based attention alone is not enough for transformers to learn CFGs.

### 5.2 Result 7-9: Boundary-Based Attention

Next, our idea is to remove the position-bias from the attention to examine the remainder. We discover that the transformer also learns a strong boundary-based attention pattern, where tokens on the NT-end boundaries typically attend to the “most adjacent” NT-end boundaries , see Figure 2 for an illustration. This pattern enables the transformer to effectively learn the hierarchical and recursive structure of the CFG, and generate output tokens based on the NT symbols and rules.

Formally, let A l , h , j → i ​ ( x ) A_{l,h,j\to i}(x) for j ≥ i j\geq i denote the attention weight for positions j → i j\to i at layer l l and head h h of the transformer, on input sequence x x . Given a sample pool { x ( n ) } n ∈ [ N ] ∈ L ⁡ ( 𝒢 ) \{x^{(n)}\}_{n\in[N]}\in L(\mathcal{G}) , we compute for each layer l l , head h h , 12 12 12 Throughout this paper, we use ⟦ ⋅ ⟧ \llbracket\cdot\rrbracket to denote multi-sets that allow multiplicity, such as ⟦ 1 , 2 , 2 , 3 ⟧ \llbracket 1,2,2,3\rrbracket . This allows us to conveniently talk about its set average. A ¯ l , h , p = A v e r a g e ⟦ A l , h , j → i ( x ( n ) ) ∣ n ∈ N , 1 ≤ i ≤ j ≤ len ( x ( n ) ) s.t. j − i = p ⟧ , \overline{A}_{l,h,p}=Average\llbracket A_{l,h,j\to i}(x^{(n)})\mid n\in{N},1\leq i\leq j\leq\textbf{len}(x^{(n)})\text{ s.t. }j-i=p\rrbracket\kern 5.0pt, which represents the average attention between any token pairs of distance p p over the sample pool. To remove position-bias, we focus on B l , h , j → i ​ ( x ) = def A l , h , j → i ​ ( x ) − A ¯ l , h , j − i B_{l,h,j\to i}(x)\stackrel{{\scriptstyle\mathrm{\scriptscriptstyle def}}}{{=}}A_{l,h,j\to i}(x)-\overline{A}_{l,h,j-i} in this subsection. Our observation can be broken down into three steps.

Result 7 ( Figure 9(a) ) . B l , h , j → i ​ ( x ) B_{l,h,j\to i}(x) exhibits a strong bias towards tokens i i at NT ends . This can be seen in Figure 9(a) , where we present the average value of B l , h , j → i ​ ( x ) B_{l,h,j\to i}(x) over data x x and pairs i , j i,j where i + δ i+\delta is the deepest NT-end at level ℓ \ell (symbolically, 𝔟 ♯ ​ ( i + δ ) = ℓ \mathfrak{b}^{\sharp}(i+\delta)=\ell ). The attention weights are highest when δ = 0 \delta=0 and decrease rapidly for surrounding tokens.

Result 8 ( Figure 9(b) ) . B l , h , j → i ​ ( x ) B_{l,h,j\to i}(x) favors pairs i , j i,j both at NT ends at some level ℓ \ell . This can be seen in Figure 9(b) , where we show the average value of B l , h , j → i ​ ( x ) B_{l,h,j\to i}(x) over data x x and pairs i , j i,j where 𝔟 ℓ ​ ( i + δ 1 ) = 𝔟 ℓ ​ ( j + δ 2 ) = 1 \mathfrak{b}_{\ell}(i+\delta_{1})=\mathfrak{b}_{\ell}(j+\delta_{2})=1 for δ 1 , δ 2 ∈ { − 1 , 0 , 1 } \delta_{1},\delta_{2}\in\{-1,0,1\} . It is maximized when δ 1 = δ 2 = 0 \delta_{1}=\delta_{2}=0 .

Result 9 ( Figure 9(c) ) . B l , h , j → i ​ ( x ) B_{l,h,j\to i}(x) favors “adjacent” NT-end token pairs i , j i,j . Above, we define “adjacency” as follows. We introduce B l , h , ℓ ′ → ℓ , r end → end B^{\text{end }\to\text{end }}_{l,h,\ell^{\prime}\to\ell,r} to represent the average value of B l , h , j → i ​ ( x ) B_{l,h,j\to i}(x) over samples x x and token pairs i , j i,j that are at the deepest NT-ends on levels ℓ , ℓ ′ \ell,\ell^{\prime} respectively (symbolically, 𝔟 ♯ ​ ( i ) = ℓ ∧ 𝔟 ♯ ​ ( j ) = ℓ ′ \mathfrak{b}^{\sharp}(i)=\ell\wedge\mathfrak{b}^{\sharp}(j)=\ell^{\prime} ), and are at a distance r r based on the ancestor indices at level ℓ \ell (symbolically, 𝔭 ℓ ​ ( j ) − 𝔭 ℓ ​ ( i ) = r \mathfrak{p}_{\ell}(j)-\mathfrak{p}_{\ell}(i)=r ). We observe that B l , h , ℓ ′ → ℓ , r end → end B^{\text{end }\to\text{end }}_{l,h,\ell^{\prime}\to\ell,r} decreases as r r increases, and is highest when r = 0 r=0 (or r = 1 r=1 for pairs ℓ ′ → ℓ \ell^{\prime}\to\ell without an r = 0 r=0 entry). 13 13 13 For any token pair j → i j\to i with ℓ = 𝔟 ♯ ​ ( i ) ≥ 𝔟 ♯ ​ ( j ) = ℓ ′ \ell=\mathfrak{b}^{\sharp}(i)\geq\mathfrak{b}^{\sharp}(j)=\ell^{\prime} — meaning i i is at an NT-end closer to the root than j j — it satisfies 𝔭 ℓ ​ ( j ) − 𝔭 ℓ ​ ( i ) ≥ 1 \mathfrak{p}_{\ell}(j)-\mathfrak{p}_{\ell}(i)\geq 1 so their distance r r is strictly positive.

In conclusion, tokens corresponding to NT-ends at level ℓ ′ \ell^{\prime} statistically have higher attention weights to their most adjacent NT-ends at every level ℓ \ell , even after removing position-bias . 14 14 14 Without removing position-bias, such a statement might be meaningless as the position-bias may favor “adjacent” anything, including NT-end pairs.

### 5.3 Connection to Dynamic Programming (DP)

Dynamic programming involves storage of intermediate results and a recurrent formula to combine them. While identifying a specific DP implementation within transformers is infeasible due to numerous possibilities ( Footnote 11 ), we can probe for crucial commonalities. Section 4 demonstrated that transformers encode the DP’s storage backbone—all necessary 𝖣𝖯 ⁡ ( i , j , a ) \mathsf{DP}(i,j,a) values on the correct CFG parse tree—independent of any specific DP implementation. Regarding the recurrent formula (e.g., 𝖣𝖯 ⁡ ( k , j , a ) \mathsf{DP}(k,j,a) derived from 𝖣𝖯 ⁡ ( k , i , b ) ∧ 𝖣𝖯 ⁡ ( i , j , c ) \mathsf{DP}(k,i,b)\wedge\mathsf{DP}(i,j,c) for rule a ↦ b , c a\mapsto b,c ), 𝖣𝖯 ⁡ ( k , i , b ) \mathsf{DP}(k,i,b) is stored near i i , while 𝖣𝖯 ⁡ ( k , j , a ) \mathsf{DP}(k,j,a) and 𝖣𝖯 ⁡ ( i , j , c ) \mathsf{DP}(i,j,c) are near j j ( Result 5 ). This necessitates a memory read from i i at j j ( j → i j\to i ). Indeed, for adjacent NT-ends i , j i,j at the same level, GPT models exhibit such j → i j\to i attention ( Result 8 ), suggesting an information flow consistent with DP. See Figure 10 (top).