Grokking of Hierarchical Structure in Vanilla Transformers (2023)

Abstract: For humans, language production and comprehension is 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 in-domain accuracy has saturated. We call this phenomenon structural grokking. On multiple datasets, structural grokking exhibits inverted U-shaped scaling in model depth: intermediate-depth 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 CITATION. Overall, our work provides strong evidence that, with extended training, vanilla transformers discover and use hierarchical structure.

## 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 @xcite . in this paper, we investigate whether standard neural transformer models @xcite 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 @xcite . her vultures don't comfort the dogs that do wait. don't her vultures comfort the dogs that do wait? your xylophone who doesn't eat does swim. does your xylophone who doesn't eat swim? output the most frequent bracket at this index. ## experiments datasets since our goal is to understand hierarchical generalization in transformers, we use two datasets from @xcite and additionally evaluate on a simple bracket-tracking task. for dyck, models are trained to predict next tokens in strings drawn from dyck 20,10 , the language of well-nested brackets with 20 types and max nesting depth of 10. we evaluate generalization to structurally unobserved strings in dyck 20,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 @xcite . ## 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 @xcite . 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.