A Transformer with Stack Attention (2024)

Abstract: Natural languages are believed to be (mildly) context-sensitive. Despite underpinning remarkably capable large language models, transformers are unable to model many context-free language tasks. In an attempt to address this limitation in the modeling power of transformer-based language models, we propose augmenting them with a differentiable, stack-based attention mechanism. Our stack-basedattention mechanism can be incorporated into any transformer-based language model and adds a level of interpretability to the model. We show that the addition of our stack-based attention mechanism enables the transformer to model some, but not all, deterministic context-freelanguages.

## introduction language models (lms) based on the transformer architecture @xcite have shown great empirical success at a wide range of nlp tasks @xcite @xcite . however, recent theoretical @xcite and empirical @xcite @xcite research suggests that language models based on transformers show difficulty in learning basic algorithmic patterns. a prime example is the dyck-n language, i.e., the language of balanced parentheses of depth ď n. when n ą 1, it has been argued that transformers are theoretically @xcite and empirically @xcite unable to learn a dyck-n language. additionally, @xcite report that transformer-based lms fail to learn four deterministic context-free (dcf) tasks. the authors of this work contend that the resolution of this insufficiency is paramount if human-level language understanding is to be achieved by computers. indeed, @xcite famously argues that human language has many context-free traits; see also @xcite . @xcite goes further and argues that snippets of swiss german are even higher on the chomksy hierarchy. the scientific question treated in this paper is whether there exists a minimal modification to the transformer architecture that does allow it to learn a larger swathe of the formal languages most closely associated with human language. specifically, in this paper, we augment the transformer architecture with a novel stack attention mechanism that enables it to learn certain cf languages empirically. our stack attention mechanism simulates a stack by maintaining a probability distribution over which of the subsequently observed tokens is at the top element of the stack. in turn, this probability distribution serves as an attention mechanism. compared to @xcite , which also applies stack augmentation to the transformer, our stack attention is more space efficient and allows for easier interpretation through visualizing the attention weights. we incorporate our innovation into the transformer by adding a stack attention sub-layer to each layer, rather than completely replacing the standard attention. augmenting models in a modular way like this allows for direct integration with pre-trained transformer-based lms. we evaluate our stack-augmented transformer through comparison with a standard transformer on four dcf tasks taken from @xcite . we find that the stack-augmented transformer performs substantially better than the standard transformer on two of the four dcf tasks. nevertheless, in contrast to dusell and chiang (2024), who claim their architecture can recognize the entire class of cf languages, we find transformers with our stack attention still struggle on two dcf tasks that involve modular arithmetic. ## preliminaries in this section, we provide the necessary technical background for our exposition. we first review the self-attention mechanism and then introduce the transformer architecture. ## a transformer with stack attention recently @xcite showed that transformers are not able to learn several non-regular dcf languages. inspired by the fact that pushdown automata @xcite , automata that employ a single stack, can model cf languages (evey, 1963), we introduce a novel stack attention mechanism that emulates the functionality of a stack and integrate it into the transformer architecture, aiming to enable it to learn some cf languages. ## deterministic cf tasks we now discuss several tasks that are encodable by deterministic context-free grammars. ## language modeling we consider masked language modeling using roberta @xcite and autoregressive language modeling using gpt-2 @xcite . following the experimental setup proposed by previous authors @xcite , we experiment on the penn treebank (ptb), licensed through the ldc @xcite , and wikitext-2 @xcite . we consider models both trained from scratch and fine-tuned from pre-trained weights. the pre-trained models and datasets are obtained from huggingface @xcite . see appendix b for more details about setup and hyperparameters. the results in tab. 3 are mixed. our major finding is that transformers benefit from the stack attention when training data is scarce, but the benefits gradually diminish as the size of training data grows. more concretely, when the models are trained from scratch, the addition of our stack attention mechanism does result in a noticeable benefit ## discussion from the results described in §4.3 and §5, we observe two trends: we interpret these trends as support for the idea that stack attention improves the representational capacity of a transformer language model and, additionally, confers an inductive bias to the transformer architecture that allows it to better learn certain context-free tasks more efficiently. the larger representational capacity explains why the performance on certain tasks, i.e., rs and sm, improves drastically with the addition of stack attention and the better inductive bias explains why transformer language models with stack attention perform better with less training data on the english modeling task. however, the fact that a vanilla transformer language model performs on par with stack attention when modeling larger english language datasets suggests that a good inductive bias is not needed for larger data sets. this suggests that, in contrast to the viewpoint of traditional linguistic theory @xcite , models that are higher on the chomsky hierarchy are not necessary for developing a good statistical language model. we believe this claim is consistent with the literature, in which few successful large language models are endowed with a syntactic bias. however, there are many smaller syntax-infused language models @xcite that do work well on smaller data, as ours does. ## conclusion we propose a novel implementation of a differentiable stack and show that a transformer augmented with such stacks can solve certain deterministic context-free tasks that are beyond the capability of standard transformers. however, unlike a stack rnn, the stack transformer cannot model the entire class of deterministic context-free languages. the model can be fully parallelized, leading to great improvement in time efficiency. it would be interesting to explore this possibility in the future.