Transformer Working Memory Enables Regular Language Reasoning And Natural Language Length Extrapolation (2023)

Abstract: Unlike recurrent models, conventional wisdom has it that Transformers cannot perfectly model regular languages. Inspired by the notion of working memory, we propose a new Transformer variant named RegularGPT. With its novel combination of Weight-Sharing, Adaptive-Depth, and Sliding-Dilated-Attention, RegularGPT constructs working memory along the depth dimension, thereby enabling efficient and successful modeling of regular languages such as PARITY. We further test RegularGPT on the task of natural language length extrapolation and surprisingly find that it rediscovers the local windowed attention effect deemed necessary in prior work for length extrapolation.

## introduction it is long believed that working memory (wm), a term coined in 1960s to liken human minds to computers, plays an important role in humans' reasoning ability and the guidance of decision-making behavior @xcite @xcite @xcite @xcite . while no single definition encompasses all applications of wm @xcite , the following one should be shared by all of the theories of interest: working memory is a system of components that holds a limited amount of information temporarily in a heightened state of availability for use in ongoing processing. @xcite wm is instantiated in the two major driving forces of sequence modeling: recurrent neural networks'(rnn) @xcite @xcite short term memory modulated by their recurrent nature and gate design @xcite @xcite , and transformers' @xcite salient tokens heightened by self-attention. in reality, self-attention often attends broadly @xcite , violating the limited amount of information notion of wm. our hypothesis is that such violation is to blame for transformers' failure on algorithmic reasoning of regular languages @xcite such as parity, a seemingly simple task that checks if the number of 1s in a bit string is even. surprisingly, a transformer can only count the number of 1s correctly when the sequence length is held fixed at training sequence length t tr , and it fails miserably when the testing sequence length extrapolates to t @xmath0. in contrast, an rnn can extrapolate perfectly. the goal of this work is therefore to improve transformers' wm by limiting the amount of accessible information at a time. existing attempts that use a combination of scratchpad and recency biases @xcite @xcite is not optimal as it completely foregoes the parallelization property of a transformer, making it as computationally inefficient as an rnn. this begs the question: does there exist a more efficient transformer working memory design? the answer is affirmative thanks to the proposed reg-ulargpt, which boils down to the three design choices: weight-sharing, adaptive-depth, and sliding-dilated-attention; each of them has been proposed previously but it is the unique combination that sparks the successful and efficient learning of regular languages. we will further demonstrate its: 1) similar recursive parallel structure as linear rnn @xcite , resulting in log t tr or log t ex layers, and 2) generalizability by showing strong performance on the task of transformer natural language length extrapolation @xcite @xcite . in this work, we use [n ] to denote the list of non-negative integers [0, . . . , n -1]. the transformer model used in this work is always causal. it takes in an input sequence of @xmath1, passes them through a fixed amount of l transformer layers, and finally computes the distribution over the vocabulary v via the prediction head w o . ## proposed architecture of regulargpt we present our modifications to the vanilla transformer below. only the related operations will be expanded, and we follow all the other details of gpt2 @xcite . ## connection to prior work sliding-dilated-attention this special attention pattern dates back to pre-transformer era such as wavenet @xcite with dilated convolution. it can also be viewed as a special form of longformer attention pattern with systematic dilation @xcite . 2 limiting the 2 the original longformer also adopts dilated attention on a few heads at higher layers but without the systematic pattern range of attention in lower layers of a transformer is also corroborated in rae and razavi (2020b), where they find such design does not deteriorate the performance. adaptive-depth and weight-sharing al-bert @xcite and universal transformer @xcite share the parameters across layers. the weight sharing design makes them compatible with the idea of adaptive computation time @xcite and dynamic halting @xcite , which allocate different computational budget depending on the complexity of tasks @xcite . however, they lack the special sliding-dilated-attention design that is necessary for ruling out naive solutions. @xcite for @xmath2: ## natural language experiments given that regulargpt has been battle-tested on the main experiment of regular languages, we now shift gear to benchmark its performance in the natural language scenario. given a model trained on sequences of length t tr , we test it on much longer sequences of length t ex t tr during inference, and the goal is to observe similar perplexities. to optimize efficiency, we employ a random selection process to extract 1,000 chunks, each with t ex tokens from the testing set. subsequently, we calculate the average perplexity of the last tokens within these chunks to ensure each of them has t ex -1 tokens as the context, thereby avoiding the issue of early token curse @xcite . we compare our model against the existing methods that are known to demonstrate the ability of length extrapolation including t5 @xcite , alibi @xcite , and ker-ple @xcite . 4 to counteract the loss of expressive power due to weight sharing, we thicken each layer of regulargpt to k as detailed in §3. in it is worth noting that 128/12 performs relatively well and is close to previous methods designed specifically for the task of natural language extrapolation. we will analyze its inner workings in depth in figure 4 and §7, in which we find that regu-largpt learns the similar local receptive field as prior work, which is likely the key to its successful natural language extrapolation performance. ## conclusion this paper introduces regulargpt, a novel variant of the transformer architecture inspired by the notion of working memory that can effectively model regular languages with high efficiency. theoretical explanations and accompanying clustering visualizations are presented to illustrate how regulargpt captures the essence of regular languages. moreover, regulargpt is evaluated on the task of natural language length extrapolation, revealing its intriguing rediscovery of the local windowed attention effect previously observed in related research. notably, regulargpt establishes profound connections with various existing architectures, thereby laying the groundwork for the development of future transformer models that facilitate efficient algorithmic reasoning and length extrapolation.