Grokking of Hierarchical Structure in Vanilla Transformers

Introduction

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 tree-structuredness (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 significantly 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).

Background

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.

Experiments

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 Question-Formation, 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.

1 Main Results

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 in-distribution accuracy has saturated, sometimes approaching perfect accuracy.

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 Tense-Inflection, respectively.

Inverted U-shaped scaling

On Question-Formation 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.

Analysis

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.

Recent work (Power et al., 2022; Liu et al., 2022) identifies the $L_{2}$ 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 $f_{\theta}^{L}$, we compute the negative mean entropy of all distributions $\{\mathbf{a}^{p}_{k}\}$.

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 ($t_{\text{score}}$) 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 ($t_{\text{parseval}}$, Black et al., 1991).

1 Results

We characterize the dynamics of weight norms (normalized by number of layers to compare different model depths), attention sparsity, and tree-structuredness, 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.

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), $t_{\text{score}}{}$ 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 $t_{\text{parseval}}{}$ 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.

Conclusion

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.

Reproducibility

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

Acknowledgements

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.

References

Appendix A Dataset Details

All statistics are in Table 1. For Question-Formation and Tense-Inflection, we use splits as given in McCoy et al. (2020) with no additional preprocessing. We give details of Dyck below.

Appendix B Model Details

We use a transformer language model with the following hyperparameters:

Tied input and output matrices as done in Press and Wolf (2017)

Next, we use the following hyperpameters for optimization:

AdamW ($\beta_{1}$: 0.9, $\beta_{2}$: 0.999, $\epsilon$: 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 $L_{2}$ norm of 10.

Appendix C Functional Tree-Structuredness

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:

where $g_{\phi}$ is the class of tree structured encoders that processes sentence $S$ according to bottom-up trees $T(S)$, and $\mathcal{L}$ is a distance function between vector outputs of $f$ and $g_{\phi}$ on spans from the binary tree $T$. TP minimizes Equation 1 approximately, and recovers an approximate $\widehat{T}_{\text{proj}}$. The tree score over a dataset $\mathcal{D}$ is defined as

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 $\widehat{T}_{\text{proj}}$ against ground truth gold syntax trees, $T_{g}$, when available, Murty et al. (2023) define

Appendix D Training Loss Curves

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.