Physics of Language Models: Part 1, Learning Hierarchical Language Structures β part 2/7
Cleaned-text code points [15939, 31863). Bibliography omitted here; source.md includes it.
A symbol x i x_{i} in a sample ( x , π , π° ) βΌ L β‘ ( π’ ) (x,\mathfrak{p},\mathfrak{s})\sim L(\mathcal{G}) is the NT boundary / NT end at level β β [ L β 1 ] \ell\in[L-1] if π β β ( i ) β π β β ( i + 1 ) \mathfrak{p}_{\ell}(i)\not=\mathfrak{p}_{\ell}(i+1) or i = len β ( x ) i=\textbf{len}(x) . We denote π β β ( i ) = def π x i β is the NT boundary at level β β \mathfrak{b}_{\ell}(i)\stackrel{{\scriptstyle\mathrm{\scriptscriptstyle def}}}{{=}}\mathds{1}_{x_{i}\text{ is the NT boundary at level }\ell} as the NT-end boundary indicator function. The deepest NT-end of i i isβ see also Figure 2 β π β― ( i ) = min β β { 2 , 3 , β¦ , L β 1 } { π β ( i ) = 1 } or β₯ if set is empty . \textstyle\mathfrak{b}^{\sharp}(i)=\min_{\ell\in\{2,3,\dots,L-1\}}\{\mathfrak{b}_{\ell}(i)=1\}\hskip 10.00002pt\text{or $\bot$ if set is empty }\kern 5.0pt.
The πΌπΏππ₯ \mathsf{cfg3} synthetic CFG family. We focus on seven synthetic CFGs of depth L = 7 L=7 detailed in Section A.1 . The hard datasets πΌπΏππ₯π» , πΌπΏππ₯π , πΌπΏππ₯π , πΌπΏππ₯π , πΌπΏππ₯πΏ \mathsf{cfg3b},\mathsf{cfg3i},\mathsf{cfg3h},\mathsf{cfg3g},\mathsf{cfg3f} have sizes ( 1 , 3 , 3 , 3 , 3 , 3 , 3 ) (1,3,3,3,3,3,3) and increasing difficulties πΌπΏππ₯π» < πΌπΏππ₯π < πΌπΏππ₯π < πΌπΏππ₯π < πΌπΏππ₯πΏ \mathsf{cfg3b}<\mathsf{cfg3i}<\mathsf{cfg3h}<\mathsf{cfg3g}<\mathsf{cfg3f} . The easy datasets πΌπΏππ₯πΎπ£ \mathsf{cfg3e1} and πΌπΏππ₯πΎπ€ \mathsf{cfg3e2} have sizes ( 1 , 3 , 9 , 27 , 81 , 27 , 9 ) (1,3,9,27,81,27,9) and ( 1 , 3 , 9 , 27 , 27 , 9 , 4 ) (1,3,9,27,27,9,4) respectively. The sequences generated by these CFGs are up to 3 6 = 729 3^{6}=729 in length. Typically, the learning difficulty of CFGs inversely scales with the number of NT/T symbols, assuming other factors remain constant, because having more NT/T symbols makes the language less ambiguous and more easily parsed using greedy (see Figure 4 and we discuss more in Appendix G ). We thus primarily focus on πΌπΏππ₯π» , πΌπΏππ₯π , πΌπΏππ₯π , πΌπΏππ₯π , πΌπΏππ₯πΏ \mathsf{cfg3b},\mathsf{cfg3i},\mathsf{cfg3h},\mathsf{cfg3g},\mathsf{cfg3f} .
### 2.2 Why Such CFGs
We use CFG as a proxy to study rich, recursive structure reasoning in languagesβfrom logics and grammars to formats and patterns. Those structures are diverse yet strict (e.g., in a CFG describing chapter numbers, Chapter 3.1 can be only followed by Chapter 3.1.1, Chapter 4 or Chapter 3.2, not others). The CFGs we consider are non-trivial, with over 2 270 > 10 80 2^{270}>10^{80} strings in πΌπΏππ₯πΏ \mathsf{cfg3f} among a total of over 3 300 > 10 140 3^{300}>10^{140} possible strings of length 300 or more (see our entropy estimation later in Figure 4 ). The probability of a random string belonging to this language is nearly zero, and a random completion of a valid prefix is unlikely to satisfy the CFG. In particular, Figure 31 in the appendix shows that πΌπΏππ₯πΏ \mathsf{cfg3f} cannot be learned by transformers (much) smaller than GPT2-small. In contrast, the English CFG (e.g., derived from Penn TreeBank) can be learned to good accuracy using tiny GPT2 models with βΌ \sim 100k parameters β so it is too easy for our interpretability purpose.
To obtain clean interpretability result and facilitate clearer analysis of learned representations across processing levels, we selected a CFG family with aβcanonical representationβ (e.g., layered CFG). This layered structure, while simplified, allows for direct probing of per-level NT symbol encodings and attention patterns at distinct hierarchical depths, aiding our interpretability goals. This controlled design allows us to demonstrate a strong correlation between the CFG representation and the hidden states in the learned transformer. We also create additional CFG families to examine βnot-so-canonicalβ CFG trees, with results deferred to Appendix G (see an example in Figure 3 ). We do not claim our results encompass all CFGs; our chosen CFGs are already challenging for a transformer to learn and can lead to clean hierarchical interpretability results.
## 3 Results 1-3: Transformer Can Learn Such CFGs
Before we analyze how transformers perform structure reasoning on such CFGs, we have to first verify that they at least can learn such CFGs. In this section, we generate a large corpus { x ( i ) } i β [ N ] \{x^{(i)}\}_{i\in[N]} from a synthetic CFG language L β‘ ( π’ ) L(\mathcal{G}) in Section 2.1 , and pretrain a (decoder-only) transformer model F F on this corpus, treating each terminal symbol as a separate token, using an auto-regressive task (see Appendix A.3 for details). We then evaluate how well the model learns such L β‘ ( π’ ) L(\mathcal{G}) .
Models. We denote the GPT2 small architecture (12-layer, 12-head, 768-dimensions) as πΆπΏπ \mathtt{GPT} [ 28 ] and implemented its two modern variants. We denote πΆπΏπ \mathtt{GPT} with relative positional attention [ 14 ] as πΆπΏπ ππΎπ
\mathtt{GPT}_{\mathsf{rel}} , and πΆπΏπ \mathtt{GPT} with rotary attention [ 32 , 10 ] as πΆπΏπ πππ \mathtt{GPT}_{\mathsf{rot}} . For purposes in later sections, we introduce two weaker variants. πΆπΏπ πππ \mathtt{GPT}_{\mathsf{pos}} replaces the attention matrix with a matrix based solely on tokensβ relative positions, while πΆπΏπ πππ \mathtt{GPT}_{\mathsf{uni}} uses a constant, uniform average of past tokens from various window lengths as the attention matrix. Detailed explanations of these variants are in Section A.2 .
We quickly summarize our findings and then elaborate them in details. Result 1-3 ( Figure 4 ) . The GPT models (except the original absolute embedding variant) can effectively learn our synthetic CFGs. Given any prefix, they can generate completion strings β’ that can perfectly adhere to the CFG rules most of the time, (accuracy) β’ that are sufficiently diverse in the CFG language, and (diversity) β’ that closely follow the probabilistic distribution of the CFG language. (probability) Moreover, one had better use rotary or relative attentions; the original πΆπΏπ \mathtt{GPT} (with absolute positional embedding) performs even worse than πΆπΏπ πππ \mathtt{GPT}_{\mathsf{uni}} (with uniform attention).
Result 1: Completion accuracy. We evaluate F F by letting it generate completions for prefixes x : c = ( x 1 , x 2 , β― , x c ) x_{:c}=(x_{1},x_{2},\cdots,x_{c}) from strings x x freshly sampled from L β‘ ( π’ ) L(\mathcal{G}) . The generation accuracy is measured as ππ« x βΌ L β‘ ( G ) + randomness of F [ ( x : c , F ( x : c ) ) β L ( π’ ) ] \operatornamewithlimits{\mathbf{Pr}}_{x\sim L(G)\,+\,\textrm{randomness of $F$}}[(x_{:c},F(x_{:c}))\in L(\mathcal{G})] . We use multinomial sampling without beam search for generation. 5 5 5 The last softmax layer converts the model outputs into a probability distribution over (next) symbols. We follow this distribution to generate the next symbol, reflecting the unaltered distribution learned by the transformer. This is the source of the βrandomness of F F β and is often referred to as using βtemperature Ο = 1 \tau=1 .β
Figure 4 (left) shows the generation accuracies for cuts c = 0 c=0 and c = 50 c=50 . The c = 0 c=0 result tests the modelβs ability to generate a sentence in the CFG, while c = 50 c=50 tests that to complete a sentence. 6 6 6 πΌπΏππ₯ \mathsf{cfg3} family is large enough to ensure a negligible chance of a freshly sampled prefix of length 50 being seen during pretraining. The results show that the pretrained GPT models can often generate strings that perfectly adhere to the CFG rules for the πΌπΏππ₯ \mathsf{cfg3} data family.
Result 2: Generation diversity. Could it be possible that the pretrained GPT models only memorized a small subset of strings from the CFG? We evaluate this by measuring the diversity of its generated strings. High diversity suggests a better understanding of the CFG rules.
We consider two methods to estimate diversity. One is to estimate the distributionβs entropy, which provides a rough estimate of (the log 2 \log_{2} of) the support size, see the middle of Figure 4 . The other is to use birthday paradox to theoretically lower bound the support size [ 7 ] . This allows us to make precise claims, such as in the πΌπΏππ₯πΏ \mathsf{cfg3f} dataset, there are at least 4 Γ 10 8 4\times 10^{8} distinct sentential forms derivable from a symbol at levels 1 to 5 or levels 2 to 6; not to say from the root to level 7. Details are in Appendix B . Our general conclusion is that the pre-trained model does not rely on simply memorizing a small set of patterns to achieve high completion accuracy.
Result 3: Distribution comparison. To fully learn a CFG, it is crucial to also learn the probabilistic distribution. One naive approach is to compare the marginal distributions p β‘ ( a , i ) p(a,i) , for the probability of symbol a β ππ β a\in\mathbf{NT}_{\ell} appearing at position i i . We observe a strong alignment between the generation probabilities and the ground-truth, included in Appendix B.2 .
Another approach is to use the standard KL-divergence formula to compare the next-token prediction probability (as predicted by the transformer model) and the ground-truth. Let p β p^{*} denote the distribution over strings in the true CFG and p p that from the transformer model. Let S = { x ( i ) } i β [ M ] S=\big\{x^{(i)}\big\}_{i\in[M]} be samples from the true CFG distribution. Then, the KL-divergence can be estimated as follows: 7 7 7 Similar formula was also used in [ 12 ] . 1 | S | β x β S 1 len β ( x ) + 1 β i β [ len β ( x ) + 1 ] β t β π βͺ { πΎππ } ππ« p β [ t β£ x 1 , β¦ , x i β 1 ] log ππ« p β [ t β£ x 1 , β¦ , x i β 1 ] ππ« p [ t β£ x 1 , β¦ , x i β 1 ] \textstyle\frac{1}{|S|}\sum_{x\in S}\frac{1}{\textbf{len}(x)+1}\sum_{i\in[\textbf{len}(x)+1]}\sum_{t\in\mathbf{T}\cup\{\mathsf{eos}\}}\operatornamewithlimits{\mathbf{Pr}}_{p^{*}}[t\mid x_{1},\dots,x_{i-1}]\log\frac{\operatornamewithlimits{\mathbf{Pr}}_{p^{*}}[t\mid x_{1},\dots,x_{i-1}]}{\operatornamewithlimits{\mathbf{Pr}}_{p}[t\mid x_{1},\dots,x_{i-1}]} (Above, ππ« p [ t β£ x 1 , β¦ , x i β 1 ] \operatornamewithlimits{\mathbf{Pr}}_{p}[t\mid x_{1},\dots,x_{i-1}] is the next-token distribution predicted by the model, and ππ« p β [ t β£ x 1 , β¦ , x i β 1 ] \operatornamewithlimits{\mathbf{Pr}}_{p^{*}}[t\mid x_{1},\dots,x_{i-1}] is that from the ground-truth. 8 8 8 There are many dynamic programming methods to compute ππ« p β [ t β£ x 1 , β¦ , x i β 1 ] \operatornamewithlimits{\mathbf{Pr}}_{p^{*}}[t\mid x_{1},\dots,x_{i-1}] exactly; which one to use is irrelevant. ) In Figure 4 (right) we compute such KL-divergence using M = 20000 M=20000 samples.
Connection to DP. Result 1-3 (e.g., learning the CFGβs next-token distribution) is merely a small step towards showing that the model employs a DP-like approach. Dynamic programming (e.g., the inside-outside algorithm [ 9 ] ) can compute next-token distributions of CFGs, and such algorithms can be implemented using nonlinear neural networks like transformers, achieving a global minimum in the auto-regressive training objective. 9 9 9 This has been carefully explored for masked language modeling case in Zhao et al. [40] . However, the mere existence of a dynamic-programming transformer to obtain the training objectiveβs global minimum is not satisfactory. Does employing an AdamW stochastic optimizer for 100k iterations on the training objective yield such an algorithm? The remainder of this paper will delve deeper to address this question.
Other Applications of Results 1β3. While not the focus of this paper, our constructed CFGs also serve as a quick testbed for comparing architecture designs. For instance, the strong performance of uniform attention aligns with the effectiveness of ALiBi [ 27 ] and H-Alibi [ 16 ] , and has motivated our follow-up work on designing Transformer architectures that explicitly leverage short-window uniform attention [ 1 ] . Additional robustness experiments for uniform attentionβacross data complexities and model sizesβare included in Appendix H .
## 4 Results 4-5: How Do Transformers Learn CFGs?
In this section, we delve into the learned representation of the transformer to understand how it encodes CFGs. We employ various measurements to probe the representation and gain insights.
Recall classical way to solve CFGs. Given CFG π’ \mathcal{G} , the classical way to reason about if a sequence x x satisfies L β‘ ( π’ ) L(\mathcal{G}) is to use dynamic programming (DP) [ 29 , 31 ] . One possible implementation of DP involves using the function π£π― β‘ ( i , j , a ) \mathsf{DP}(i,j,a) , which determines whether or not x i + 1 , x i + 1 β β¦ , x j x_{i+1},x_{i+1}\dots,x_{j} can be generated from symbol a a following the CFG rules. From this DP representation, a DP recurrent formula can be easily derived. 10 10 10 For example, one can compute π£π― β‘ ( i , j , a ) = 1 \mathsf{DP}(i,j,a)=1 if and only if there exists i = i 1 < i 2 < β― < i k = j i=i_{1}<i_{2}<\cdots<i_{k}=j such that π£π― β‘ ( i r , i r + 1 , b r ) = 1 \mathsf{DP}(i_{r},i_{r+1},b_{r})=1 for all r β [ k β 1 ] r\in[k-1] and a β b 1 , b 2 , β¦ , b k a\to b_{1},b_{2},\dots,b_{k} is a rule of the CFG. Implementing this naively would result in a O β‘ ( len 4 ) O(\textbf{len}^{4}) algorithm for CFGs with a maximum rule length of 3 3 . However, it can be implemented more efficiently with O β‘ ( len 3 ) O(\textbf{len}^{3}) time by introducing auxiliary nodes (e.g., via binarization). In the context of this paper, any sequence x βΌ L β‘ ( π’ ) x\sim L(\mathcal{G}) that satisfies the CFG must satisfy the following conditions: π β β ( i ) = 1 , π β β ( j ) = 1 , β k β ( i , j ) , π β β ( k ) = 0 β and β π° β β ( j ) = a βΉ π£π― β‘ ( i , j , a ) = 1 \mathfrak{b}_{\ell}(i)=1,\mathfrak{b}_{\ell}(j)=1,\forall k\in(i,j),\mathfrak{b}_{\ell}(k)=0\text{ and }\mathfrak{s}_{\ell}(j)=a\implies\mathsf{DP}(i,j,a)=1 (4.1) (recall the NT-boundary π β \mathfrak{b}_{\ell} and the NT-ancestor π° β \mathfrak{s}_{\ell} notions from Section 2.1 ). Note that ( 4.1 ) is not an βif and only ifβ condition because there may be a subproblem π£π― β‘ ( i , j , a ) = 1 \mathsf{DP}(i,j,a)=1 that does not lie on the final CFG parsing tree but is still locally parsable by some valid CFG subtree. However, ( 4.1 ) provides a βbackboneβ of subproblems, where verifying all π£π― β‘ ( i , j , a ) = 1 \mathsf{DP}(i,j,a)=1 values in this backbone certifies that the sentence x x is a valid string from L β‘ ( π’ ) L(\mathcal{G}) . It is worth mentioning that there are exponentially many implementations of the same DP algorithm 11 11 11 Each inner loop of the dynamic programming can proceed in any arbitrary order, not limited to k = i . . j k=i..j or k = j . . i k=j..i , and the algorithm can prune and break early. This gives a safe estimate of at least ( n ! ) Ξ© β‘ ( n 2 ) (n!)^{\Omega(n^{2})} possible implementations. Furthermore, there are at least 2 Ξ© β‘ ( n ) 2^{\Omega(n)} ways to perform binarization, meaning to break length-3 rules to length-2 ones. This is just to detect if a given string of length n n belongs to the CFG. and not all ( i , j , a ) (i,j,a) tuples need to be computed in π£π― β‘ ( i , j , a ) \mathsf{DP}(i,j,a) . Only those in the βbackboneβ are necessary.
Connecting to transformer. In this section, we investigate whether pre-trained transformer F F also implicitly encodes the NT ancestor and boundary information, which forms the basis for its structure reasoning capabilities. If so, it suggests the model contains sufficient information to support all the π£π― β‘ ( i , j , a ) \mathsf{DP}(i,j,a) values in the backbone. This is a significant finding, considering that transformer F F is trained solely on the auto-regressive task without any exposure to NT information. If the model encodes NT ancestor and boundary information after pretraining (as demonstrated in Results 4-5), this means it internally possesses the structural knowledge necessary not only for generation but also to certify the grammatical correctness of sentences according to the CFG. That is, its internal states effectively represent the parse tree.
### 4.1 Result 4: Transformerβs Last Layer Encodes NT Ancestors/Boundaries