Grokking of Hierarchical Structure in Vanilla Transformers
Year: 2023

Abstract:
For humans, language production and comprehension are sensitive to the hierarchical structure of sentences. In natural language processing, past work has questioned how effectively neural sequence models like transformers capture this hierarchical structure when generalizing to structurally novel inputs. We show that transformer language models can learn to generalize hierarchically after training for extremely long periods—far beyond the point when indomain accuracy has saturated. We call this phenomenon structural grokking. On multiple datasets, structural grokking exhibits inverted U-shaped scaling in model depth: intermediatedepth models generalize better than both very deep and very shallow transformers. When analyzing the relationship between model-internal properties and grokking, we find that optimal depth for grokking can be identified using the tree-structuredness metric of Murty et al. (2023). Overall, our work provides strong evidence that, with extended training, vanilla transformers discover and use hierarchical structure.

heading: 1 Introduction
text: Although human language is produced as a linear sequence, it is hierarchically organized. Smaller units compose to form larger constituents. The ability to infer this hierarchical structure underlies our ability to produce and understand new sentences (Chomsky, 1965; Crain and Nakayama, 1987). In this paper, we investigate whether standard neural transformer models (Vaswani et al., 2017) can also generalize hierarchically when trained on language processing tasks (Fig 1). Our main finding is that hierarchical generalization in transformers does occur, but very slowly: performance on structurally novel sentences increases gradually, long after performance on sentences from the training distribution has plateaued. We term this phenomenon structural grokking, by analogy to existing findings on simple classification tasks (Power et al., 2022).
On two datasets, we show that structural grokking exhibits inverted U-shaped scaling behavior as a function of model depth: hierarchical generalization improves, then declines, as we train deeper models. Prior work suggests that a number of model-internal properties might track the emergence of hierarchical structure in transformers, including weight norms (Merrill et al., 2021; Liu et al., 2022; Power et al., 2022), attention sparsity (Merrill et al., 2021), and functional treestructuredness (Murty et al., 2023). We find that functional tree-structuredness is uniquely able to predict structural grokking—while weight norms and attention sparsity increase monotonically in model depth, tree-structuredness is highest for models of the optimal depth for structural grokking.
Our results challenge findings from prior work (Mueller et al., 2022; Petty and Frank, 2021) claiming that ordinary transformers completely fail on the tests of hierarchical generalization that we study. We attribute these failures to early stopping based on in-domain validation performance, which signif-
439
icantly underestimates hierarchical generalization due to structural grokking. On the datasets where this prior work reports generalization accuracies below 20%, simply by training for longer, mean accuracy across random seeds reaches 80%, and several seeds achieve near-perfect generalization performance. Past findings are also partially explained by U-shaped scaling: this work uses models that are too shallow (Mueller et al., 2022; Petty and Frank, 2021) or too deep (Mueller et al., 2022). Our results align with past findings on the role of extended training in other language processing problems (Csordás et al., 2021; Hoffmann et al., 2022).

heading: 2 Background
text: Transformers Given a sequence of tokens w≤i = w1, w2, . . . , wi, where each token is drawn from a fixed vocabulary V , an L-layer transformer language model (LM) fLθ outputs a distribution over the next token wi+1 ∈ V , fLθ (w≤i) ∈ R|V |. A key part of the architecture is a sequence of L self-attention layers, where layer p computes contextual vectors of token k as a non-linear parametric function of a convex combination of contextual vectors of tokens w≤k from the previous layer, where coefficients apk ∈ Rk are known as the attention distribution. The LM weights are learned by maximizing the log probability of the correct continuation wk+1, given prefix w≤k.
Hierarchical structure in transformers While unsupervised pre-training of transformers has led to state-of-the-art transfer learning results across NLP, the architecture itself has been claimed to lack human-like inductive biases toward hierarchical structure (Tran et al., 2018; Hahn, 2020; Petty and Frank, 2021; Mueller et al., 2022). We revisit these claims in this work.
To understand whether a given model has a bias for acquiring hierarchical structure, we follow McCoy et al. (2020) and evaluate generalization in models trained on ambiguous tasks in which training data is consistent with both a “hierarchical rule” as well as a “non-hierarchical rule” (Fig 1). To test if the hierarchical rule has been acquired, we test generalization on a separate out-of-distribution test set, constructed such that only learners that have acquired the hierarchical rule are successful.
Grokking Power et al. (2022) identify the phenomenon of grokking on small algorithmic datasets
where they find that test performance improves long after training performance has saturated. We hypothesize a similar structural grokking, where the model groks hierarchical structure long after in-domain validation performance has saturated, and consequently, hierarchical generalization can continue to improve with extended training.

heading: 3 Experiments
text: Datasets Since our goal is to understand hierarchical generalization in transformers, we use two datasets from (McCoy et al., 2020) and additionally evaluate on a simple bracket-tracking task. For Dyck, models are trained to predict next tokens in strings drawn from Dyck20,10, the language of well-nested brackets with 20 types and max nesting depth of 10. We evaluate generalization to structurally unobserved strings in Dyck20,10 (see Fig 1 for examples and Appendix A for details). For the McCoy et al. (2020) datasets, in Question-Formation, models must convert English sentences into questions and, in Tense-Inflection, models must map from sentences and tense markers to appropriately re-inflected sentences. We evaluate generalization on the out-of-distribution test set from McCoy et al. (2020).
Model We train transformer LMs with {2, 4, 6, 8, 10} layers (see Appendix B for more details). For each depth, we train models with 10 random seeds for 300k (400k for Dyck) steps. Given the input sentence (or prefix in the case of Dyck) we decode greedily from the model at test time. For Dyck, we report the accuracy of generating the correct closing bracket type by ranking among closing brackets, given an input prefix from the language. As done in prior work (McCoy et al., 2020; Petty and Frank, 2021; Mueller et al., 2022), for QuestionFormation, we report first word accuracy of the decoded question, and for Tense-Inflection, we report the fraction of test inputs for which the target verb is correctly inflected.

heading: 3.1 Main Results
text: Transformers exhibit structural grokking We first present results obtained with the best model depth on all datasets in Fig 2. We find clear evidence of structural grokking: Across datasets, generalization improves many training steps after indistribution accuracy has saturated, sometimes approaching perfect accuracy.
Tense-Inflection
(a) Inverted U-shaped scaling (b) Dynamics of model internal properties
Question-Formation
Av er
ag e
su cc
es s
Av er
ag e
su cc
es s
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
tre e
st ru
ct ur
ed ne ss tre e st ru ct ur ed ne ss w ei gh t n or m w ei gh t n or m at te nt io n sp ar si ty at te nt io n sp ar si ty
thousand steps thousand steps thousand steps
10 layer > 6 layer > 2 layer 6 layer > 10 layer > 2 layer10 layer > 6 layer > 2 layer
6 layer > 4 layer > 2 layer 6 layer > 4 layer > 2 layer 4 layer > 6 layer > 2 layer
Figure 3: (a) Inverted U-shaped laws for grokking: On Question-Formation (top) and Tense-Inflection (bottom), we find that both very small and very deep models either fail to exhibit structural grokking or do so infrequently, compared to an in-between optimal model depth. (b) While weight norms and attention sparsity increase for all models and do not differentiate between different sizes, tree-structuredness is highest for the optimal model depth.
Early stopping considered harmful Next, we compare generalization accuracy obtained by early stopping on in-domain validation accuracy (as done in Petty and Frank (2021); Mueller et al. (2022)) to longer training runs (Fig 2). Early stopping leads to vastly underestimating generalization. For instance, average generalization goes up from <40%, <50% to <90%, <80% on Question-Formation and TenseInflection, respectively.
Inverted U-shaped scaling On QuestionFormation and Tense-Inflection, we train models of increasing depths from 2 to 10 layers. For each depth, we report the fraction of seeds (out of 10) where generalization accuracy eventually crosses 80%, in Fig 3a. We find an inverted
U-shaped scaling behavior—very shallow and very deep models are unsuccessful, while most seeds generalize in models of intermediate depth. This may also explain why prior work that either used very shallow models (1–3-layer transformers in Petty and Frank (2021); Mueller et al. (2022)) or very deep models (12-layer transformers in Mueller et al. (2022)) failed to generalize well.

heading: 4 Analysis
text: Given that structural grokking occurs only in a subset of model architectures, can we identify when it has happened (or predict when it will occur)? Several model-internal properties have been claimed to relate to either grokking or emergent hierarchical
structure in transformers.
Weight Norms Recent work (Power et al., 2022; Liu et al., 2022) identifies the L2 norm of parameter weights as an important quantity for grokking. For instance, Power et al. (2022) find weight decay to improve grokking speed and Liu et al. (2022) identify a “goldilocks zone” in weight norm space where grokking occurs. More generally, norm growth over the course of training has been studied as a key factor in neural network generalization (Soudry et al., 2018).
Attention Sparsity Merrill et al. (2021) prove that norm growth in transformers leads to attention saturation, an important property for emergent linguistic structure (Merrill et al., 2022). As a proxy for attention sparsity of fLθ , we compute the negative mean entropy of all distributions {apk}.
Tree-structuredness McCoy et al. (2020) show that tree-structured encoders such as Tai et al. (2015) show near perfect hierarchical generalization. While transformers are relatively unconstrained, recent evidence suggests that, when trained on language data, they implictly implement (approximately) tree-structured computations. In particular, the tree projection method of Murty et al. (2023) precisely characterizes the extent to which a transformer’s internal computation on an input can be approximated with a tree-structured neural encoding, providing a tree-structuredness score (tscore) for any transformer, and a binary tree that best approximates its computation on an input string (see Appendix C for details). To evaluate whether these trees correspond to human notions of syntax, we additionally compare recovered trees to gold-standard ones (tparseval, Black et al., 1991).

heading: 4.1 Results
text: We characterize the dynamics of weight norms (normalized by number of layers to compare different model depths), attention sparsity, and treestructuredness, by computing these quantities every 3k gradient updates for Question-Formation and Tense-Inflection. For data-dependent properties such as attention sparsity and tree-structuredness, we sample 10k examples from the training data. We plot these quantities for the smallest model, the largest model for which at least one run shows successful grokking, and for the optimal model depth, in Fig 3b.
Tense-Inflection (a) Inverted U-shaped scaling (b) Dynamics of model internal properties Question-Formation Av er ag e su cc es s Av er ag e su cc es s 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 tre e st ru ct ur ed ne ss tre e st ru ct ur ed ne ss w ei gh t n or m w ei gh t n or m at te nt io n sp ar si ty at te nt io n sp ar si ty thousand steps thousand steps thousand steps 10 layer > 6 layer > 2 layer 6 layer > 10 layer > 2 layer10 layer > 6 layer > 2 layer 6 layer > 4 layer > 2 layer 6 layer > 4 layer > 2 layer 4 layer > 6 layer > 2 layer
Optimal models are most tree-structured Weight norms and attention sparsity grow for all model settings in both datasets. However, these properties by themselves are unable to predict that both shallow and deep models fail—shallow models learn the sparsest solutions as well as solutions with largest weight norms, but never generalize hierarchically. As noted by Murty et al. (2023), tscore improves over time for all models, indicating increased tree-structuredness over time. For both datasets, the “optimal” model learns the most tree-structured solution compared to both deep and shallow models. Liu et al. (2022) note that, on algorithmic tasks, grokking “coincides with the emergence of structure in embeddings”. Similarly, for language tasks, we find that structural grokking coincides with the emergence of tree structured internal computations.
Transformers are surprisingly effective at structure induction From the dynamics of tparseval in Fig 4, we note that all models, regardless of whether they generalize or not, learn structures that are close to ground truth syntax, sometimes outperforming a right-branching baseline. McCoy et al. (2020) note that tree-structured encoders only generalize when structured according to correct parse trees. Here, we find that all transformers learn correct tree structures, but only the ones that are the most tree-structured generalize best.

heading: 5 Conclusion
text: This work shows that transformers are capable of exhibiting structure-sensitive “hierarchical generalization” via a grokking mechanism. Their overall learning behavior gradually shifts from memorization (high in-domain accuracy, poor out-of-domain accuracy) to generalization (high in-domain and
out-of-domain accuracy). While we show such behavior on relatively small datasets with small models, we believe these results may have broader implications, as training for longer has been shown to help even for web-scale language modeling (Hoffmann et al., 2022) and on compositional generalization tasks (Csordás et al., 2021). Structural grokking happens most often at “medium-sized” model depths, and both very shallow and very deep models fail to exhibit it. While properties previously connected with linguistic generalization in transformers such as weight norms and attention sparsity do not differentiate good architectures from bad ones, functional tree-structuredness of the transformer can well predict the optimal model depth. While there are clear limitations to the transformer architecture (such as the inability to implement unbounded recursion), our results show that it may have stronger inductive biases than previously believed: With sufficient training, transformers can represent hierarchical sentence structure and use this structure to generalize correctly.

heading: 6 Reproducibility
text: All code and data for these experiments is available at https://github.com/MurtyShikhar/ structural-grokking.git.

heading: 7 Acknowledgements
text: SM was funded by a gift from Apple Inc. CM is a fellow in the CIFAR Learning in Machines and Brains program. We thank John Hewitt, Belinda Li, Rishi Bommasani and members of the Stanford NLP group for feedback on the paper.
Limitations
Our work has the following limitations. First, we only evaluate generalization on datasets based on English language. Second, we show structural grokking on three datasets, and while we believe this to be a general phenomenon, we leave investigating similar behavior on other datasets for future work. Next, we also do not study the effect of training data size on structural grokking, and do not investigate whether transformers learn to grok hierarchical structure in low data regimes. Finally, all datasets here are based on context-free grammars, either similar to or taken directly from prior work, and we believe constructing similar generalization benchmarks on real language data is a good avenue for future work.

heading: A Dataset Details
text: All statistics are in Table 1. For QuestionFormation and Tense-Inflection, we use splits as given in McCoy et al. (2020) with no additional preprocessing. We give details of Dyck below.
Dyck Details We construct our Dyck dataset by sampling 200k strings from Dyck20,10, the language of well-nested brackets with 20 different bracket types and nesting depth atmost 10. For each string, we define its structure as a binary vector of 0s and 1s. For instance the structure of “(([]))” is “11110000”. To construct a generalization set, we sample strings with unobserved structures i.e. strings whose 0-1 structure does not match the structure of any of the training strings. Since the objective at test time is to measure closing bracket accuracy, we only rank model probability among all closing brackets, and only evaluate on prefixes where the opening bracket is atleast 10 tokens away from its corresponding closing bracket.

heading: B Model Details
text: We use a transformer language model with the following hyperparameters:
• Number of attention heads = 4
• Hidden dimensionality = 512
• Tied input and output matrices as done in Press and Wolf (2017)
Next, we use the following hyperpameters for optimization:
• AdamW (β1: 0.9, β2: 0.999, ϵ: 1e-7), with learning rates in {1e-4, 5e-5, 1e-5}, noting that 1e-4 works best for all experiments. We use a linear warmup scheduler warming up from 0 to the final learning rate over 10k gradient steps.
• We clip gradients to have a max L2 norm of 10.
• We use a batch size of 8.

heading: C Functional Tree-Structuredness
text: Tree Projections (TP; Murty et al. (2023)) measure how well computations performed by a given transformer f can be approximated with tree-structured encoders. To do this, TP solves the following optimization problem:
ϕproj, Tproj ≜ argmin ϕ,T L(f, gϕ, T ), (1)
where gϕ is the class of tree structured encoders that processes sentence S according to bottom-up trees T (S), and L is a distance function between vector outputs of f and gϕ on spans from the binary tree T . TP minimizes Equation 1 approximately, and recovers an approximate T̂proj. The tree score over a dataset D is defined as tscore ≜ ∑
S∈D ETSCI(S, T )− SCI(S, T̂proj(S)) |D| ,
(2)
where SCI (span contextual invariance) is the distance between contextual and context-free vector representations of all spans p in T (for more details, see Murty et al. (2023)). In particular, SCI score for a sentence S structured according to T (S) is
SCI(S, T ) ≜ ∑
s∈T d(vSp , ṽp) (3)
for some suitably chosen distance function d (here, cosine similarity). To measure the bracketing F1 score (PARSEVAL; Black et al. (1991)) of the induced tree projection of the transformer T̂proj against ground truth gold syntax trees, Tg, when available, Murty et al. (2023) define
tparseval ≜ PARSEVAL(T̂proj, Tg,D). (4)

heading: D Training Loss Curves
text: We explore the hypothesis that syntactic grokking is simply a result of the training loss continuing to decrease, even after in-domain validation performance has saturated in Fig 5. We note that training losses generally saturate before in-domain validation performance saturates (also noted in Power et al. (2022)). Next, we also find that all models, regardless of whether they grok or not, eventually get to comparable training losses. We conclude that the inverted U-shaped trend is not an artifact of poorly optimized models.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
Section-6
7 A2. Did you discuss any potential risks of your work? Left blank.
3 A3. Do the abstract and introduction summarize the paper’s main claims? Yes (in Section-1 and Abstract)
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts? We trained several transformer models, the checkpoints of which will be available upon de-anonymization
3 B1. Did you cite the creators of artifacts you used? We cite all prior work whose datasets we use in Sections-2 and 3
B2. Did you discuss the license or terms for use and / or distribution of any artifacts? Not applicable. Left blank.
B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Not applicable. Left blank.
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Not applicable. Left blank.
B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? Not applicable. Left blank.
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Left blank.
C 3 Did you run computational experiments? Left blank.
3 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? Left blank.
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found hyperparameter values? Left blank.
3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Left blank.
C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Not applicable. Left blank.
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? Not applicable. Left blank.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? Not applicable. Left blank.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? Not applicable. Left blank.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? Not applicable. Left blank.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? Not applicable. Left blank.

heading: Future Work
text: Second, we show structural grokking on three datasets, and while we believe this to be a general phenomenon, we leave investigating similar behavior on other datasets for future work. Next, we also do not study the effect of training data size on structural grokking, and do not investigate whether transformers learn to grok hierarchical structure in low data regimes. Finally, all datasets here are based on context-free grammars, either similar to or taken directly from prior work, and we believe constructing similar generalization benchmarks on real language data is a good avenue for future work