Pushdown Layers: Encoding Recursive Structure in Transformer Language Models (2023)

Abstract: Recursion is a prominent feature of human language, and fundamentally challenging for self-attention due to the lack of an explicit recursive-state tracking mechanism. Consequently, Transformer language models poorly capture long-tail recursive structure and exhibit sample-inefficient syntactic generalization. This work introduces Pushdown Layers, a new self-attention layer that models recursive state via a stack tape that tracks estimated depths of every token in an incremental parse of the observed prefix. Transformer LMs with Pushdown Layers are syntactic language models that autoregressively and synchronously update this stack tape as they predict new tokens, in turn using the stack tape to softly modulate attention over tokens—for instance, learning to “skip” over closed constituents. When trained on a corpus of strings annotated with silver constituency parses, Transformers equipped with Pushdown Layers achieve dramatically better and 3-5x more sample-efficient syntactic generalization, while maintaining similar perplexities. Pushdown Layers are a drop-in replacement for standard self-attention. We illustrate this by finetuning GPT2-medium with Pushdown Layers on an automatically parsed WikiText-103, leading to improvements on several GLUE text classification tasks.

## introduction an important property of human language and thought is recursion, which allows us to compose and reason about complex objects in terms of simpler constituents @xcite . while extensively studied in natural language syntax and semantics, recursion is also a key component of several other aspects of intelligent behaviors including mathematical reasoning, programming, and goaldirected planning. most recursion-capable systems model recursive processes via a stack memory, which is updated as new computation is performed. for instance, a programming language may implement recursion by maintaining a run-time stack of caller-callee frames, storing intermediate outputs in the stack, and updating the stack as new function calls are made. similarly, a shift-reduce parser implements recursion through a stack of intermedi-ate constituents, shifting tokens onto the stack as they are observed, and occasionally reducing stack elements into constituents as they are completed. in contrast, the self-attention mechanism underlying modern neural sequence models has no explicit mechanism to maintain a stack memory as it generates strings, and instead relies on hidden representations to implicitly but imperfectly encode such information @xcite . while this encoding can model bounded recursive structure in formal languages @xcite , it is unclear if it is sufficient for robust syntactic generalization, especially under data-constrained settings. in this work, we show that an explicit stack memory mechanism can improve syntactic generalization in transformer language models (lms). we introduce pushdown layers foot_0 , a drop-in replacement for standard self-attention that augments transformer lms with stack memory. this memory is modeled using a stack tape that stores estimated depths of every token in an incremental parse of the observed prefix. the stack tape is updated autoregressively: as new tokens are predicted, transformers with pushdown layers (pushdown transformers) synchronously make probabilistic attachment decisions to either "shift", thus assigning the newly predicted token a depth of 0, or "reduce" with one of the constituents in the prefix so far, updating token depths accordingly (see fig. 1 ). this stack tape is used to additively and softly modulate the attention of the transformer over tokens-for instance, pushdown layers may guide the lm to only attend to head words of constituents, or skip over reduced constituents by decreasing attention. pushdown transformer lms are syntactic language models that learn joint probabilities of sequences and parses in terms of individual word predictions and structure-building operations, and can be trained on any text corpus annotated with constituency parses. but unlike other syntactic language models with structural supervision @xcite @xcite , pushdown layers do not change the output space of the underlying sequence model, and impose no constraints on attention mechanisms-the manner in which pushdown layers use syntactic structure for representation building is learnt purely via gradient descent. ## background multi-head self-attention. transformer language models @xcite are a class of neural sequence models that use multi-head self-attention to obtain contextualized representations of tokens in a sequence, which are then used to predict the next token. in particular, let @xmath0, x 2 , . . . , x n } be an input sequence. let h l @xmath1. then, the hidden representation of the i th token is updated as ## pushdown layers transformer lms with pushdown layers are syntactic language models that generate strings while simultaneously building a parse tree over these strings from left to right. this parse tree is built incrementally by tracking the recursive state of every token, which is synchronously updated along with word-level predictions. this recursive state is represented via our stack tape as tree-depths of every prefix token, and updates are realized with a stack. the contents of the stack tape are used to softly modulate attention over prefix tokens via additive offsets to attention logits (fig. 2 ). ## analysis for all analyses, we use the 16 layer pushdown-lm trained on bllip-lg from section 4.2. parsing. since pushdown-lm is a syntactic language model, we obtain parses via beam search (beam @xmath2, y) under our figure 5 : for the three subject-verb agreement tasks from @xcite , we compute average attention over the distractor noun when the verb is being predicted, for both the base-lm and pushdown-lm (ours). across all variants, we find that our model consistently pulls attention away from distractor nouns. model. however, since this parse is (a) unlabeled and (b) binarized, we perform an unlabeled f1 evaluation (using evalb; collins, 1997) over binarized ground-truth parses from the ptb test set. we also remove instances consisting of unknown words for our model, since our model is trained without any unk tokens, giving us 2335 out of 2416 sentences. we compare our model against @xcite , the parser that was used to annotate training data for pushdown-lm. we also present unlabeled f1 on the auto-parsed bllip-lg test set. from results in ## other related work while recursive structure is fundamental to natural language, modeling such structure is difficult for self-attention. hahn (2020) considers dyck, the simplest formal language with recursive structure, proving that hard attention cannot recognize dyck and soft attention cannot recognize dyck with low cross-entropy. in practice, we find that even simpler languages like parity are challenging for encoder-only transformers @xcite . on the other hand, transformers with decoders have been shown to be turing-complete @xcite , but these constructions rely on the impractical assumption of running the decoder for an unbounded number of steps. in practice, we find that transformer lms struggle with generalization beyond regular languages and tend to learn shortcuts instead @xcite . given these limitations, there is significant interest in inductive biases that encourage recursive structure in transformers. one line of work considers constraining self-attention patterns according to syntactic parses @xcite @xcite , among others). a second line of work adds structure to language modeling by learning joint probabilistic modeling of structure and strings @xcite 2015; choe and charniak, 2016; dyer et al., 2016, among others). both of these ideas are combined in recent work of qian et al. (2021); sartran et al. @xcite , that proposes joint string, parse transformer language models with constrained attention patterns. while pushdown layers are also in this modeling tradition, we do so without operating on long transition actions, and enforce structural constraints via gradient based learning. a separate line of work proposes neural networks augmented with structured memory like stacks @xcite @xcite or random access memories @xcite . such augmented neural networks are vastly better at algorithmic generalization and learning recursive structure @xcite . our work is the first that designs a structured memory (the stack-tape) for transformers, that is updated just like stacks in a shift/reduce manner, but unlike prior work, the specific design of pushdown layers makes training parallelizable. finally, there have been several efforts to add syntactic inductive biases into sequence models (typically rnns) that can acquire and use parse structures in an unsupervised manner @xcite @xcite . we leave unsupervised training of pushdown transformers for future work. ## conclusion we propose pushdown layers, a new kind of selfattention that augments transformer language models with a stack based memory. pushdown layers enable auto-regressive transformers to softly bias attention towards a recursive syntactic computation, through an updatable stack-tape that stores token depths in an incremental parse. when trained on synthetic and natural languages, we find that transformer lms with pushdown layers achieve better generalization to deep recursive structure, as well as better and more sample-efficient syntactic generalization. when pre-trained lms are finetuned with pushdown layers, we obtain improvements on some glue tasks.