Do Transformers Need Deep Long-Range Memory? (2020)

Abstract: Deep attention models have advanced the modelling of sequential data across many domains. For language modelling in particular, the Transformer-XL — a Transformer augmented with a long-range memory of past activations — has been shown to be state-of-the-art across a variety of well-studied benchmarks. The Transformer-XL incorporates a long-range memory at every layer of the network, which renders its state to be thousands of times larger than RNN predecessors. However it is unclear whether this is necessary. We perform a set of interventions to show that comparable performance can be obtained with 6X fewer long range memories and better performance can be obtained by limiting the range of attention in lower layers of the network.

## introduction when we read a book, we maintain representations of the characters and events in the text that help us understand the story. we do this with a selective memorisation process; most of the finer details of the text are quickly forgotten and we retain a relatively compact representation of the book's details. early models of natural language used recurrent neural networks (rnns) such as the long short-term memory @xcite which emulated this selective memory approach by modelling the past in a compact state vector. the model learns to store relevant information within its state implicitly in order to optimise the task loss. the lstm has reigned as a state-of-the-art language model for over two decades since its inception in the '90s @xcite and is arguably the most ubiquitous neural sequence model. unlike human memory systems, however, the lstm struggles to reason over long-range contexts when reading text. this has been observed in multiple contexts. in the carefully curated lambada benchmark @xcite which tests language model predictions on sections of book text that have long term structure as decided by human raters, lstms completely fail. namely lstms guess the correct word 0% of the time, where humans are considered to be above 70% accuracy. for regular language modelling, @xcite observed that an lstm augmented with attention would rarely attend beyond seven preceding words of context. samples from lstms language models quickly devolve into generic text devoid of an overall theme. this has lead many to wonder whether there is any non-negligible long-range signal in the task of language modelling. recently we have seen that deep attention models can draw long-range signal from text, even when the objective is as simple as next-word prediction. with the advent of the transformer @xcite , significant gains in language modelling performance can be obtained by extending the models' attention to thousands of words. the transformer-xl @xcite , a transformer variant specialised for long-range sequence modelling via the introduction of a cache of past activations, obtained state-of-the-art results in the four major lm benchmarks -ptb @xcite , lm1b @xcite , enwik8 (hutter, 2012), and wikitext @xcite . in the case of the latter two, @xcite showed the model effectively used over one thousand words of context, and the resulting samples reflect a thematic consistency spanning paragraphs. when transformers are paired with long contexts and a large amount of data, e.g. gpt-2 @xcite and megatron @xcite , the resulting samples are remarkable in their long-range consistency and stylistic realism. however transformers abandon the compact and selective representation of the past. they store a hidden activation at every time-step (up to a given attention range) and every layer within the network. this can consume orders of magnitude more space than prior rnn hidden states, or the original text. e.g. a typical state-of-the-art lstm language model state size may range from 4kb @xcite to model wikipedia articles to 64kb @xcite to model newsand is never greater than 1mb. whereas a current state-of-the-art 18-layer transformer-xl state size for wikipedia articles is 112mb. the state is so large because a separate memory (e.g. 1600 vectors of size @xmath0. if this were found to be unnecessary then we can reduce the state's memory considerably. in this paper we investigate a simple question: can we use short-range attention for the majority of layers in the transformer and recover the same performance? the hypothesis is that this should be possible, because many steps of reasoning will only involve short-range correlations, i.e. to piece characters together to form words or phrases. we find indeed it is possible. we recover comparable performance for long-range language modelling by using a small fraction (1/6th) of long-range memories to the baseline transformerxl. crucially, we find it matters where long-range memories are placed in the network. placing them in the lower layers of the network is ineffective; placing them in the latter layers or interleaved across the network works much better. we show that such a model trains with 2x less time and memory, due to the reduction in expensive attention operations. ## background the transformer is a deep neural network for processing sequences @xcite , it processes a window of n consecutive inputs x t-n , . . . , x t in parallel. at each layer it reasons over time using multi-head attention which we will briefly describe. for a given layer l, let h @xmath1, and @xmath2. let k be the number of attention heads, then q i , k i , v @xmath3, keys, and values per attention head. these are defined to be q @xmath4, k @xmath5, and v @xmath6. the attention head output is defined to be, where σ(•) is defined to be the softmax operator. attention is the linear combination of each attention head, @xmath7. the attention operation consumes o(n) compute per step and thus o(n 2 ) for the window of inputs at each layer. the transformer-xl (txl) proposes concatenating the past activations from the same window @xmath8. this results in an attention cost of o(n(n + m)) which can be significantly cheaper than processing all n + m inputs in parallel, which would require o((n + m) 2 ). the txl's memory can be considered to be a state, alike to an rnn. however it requires a considerable space: l × m × d. ## experiments we investigate whether the transformer-xl can perform comparably with fewer long-range memory (lrm) layers on the two prominent longrange language modelling benchmarks, enwik8 and wikitext-103. ## results we plot the enwik8 learning curves for a subset of layer variants in figure 2 . the worst-performing, is the variant with a single long-term memory at the lowest layer (black curve). however perhaps more surprisingly, we see a model with 12 lrms at the lower layers of the network is actually worse than a model with a single lrm on the final layer we followed-up by running an additional arrangement of only placing lrms in the middle layers and found this to be worse than interleaved or final (1.01bpc for 4 long-range memories) which shows there is significant benefit to having some long-range memories in the higher layers. crucially, we are able to match (and slightly exceed) the full model's test performance with 12 lrms, and even a model with 4 lrms is very close (0.9846 w/ 24 vs 0.9916 w/ 4 interleaved). it is worth noting that our txl baseline actually outperforms the published version on enwik8: 0.985 bpc (ours) vs 0.993 @xcite , which provides credence to the quality of the experimental setup. we also inspect word-level language modelling on wikitext-103, using the same 18-layer trans-formerxl parameters @xcite . we obtain a baseline test perplexity of 18.3 (matching the published value), and obtain 18.4 and 18.6 for interleaved and last-layer spacing respectively when using l/6 (i.e. 3) lrms. we also try placing 3 lrms on the first three layers and obtain 20.1 perplexity. we remark that (i) long-range memory is important for a significant improvement in performance, (ii) it is better to not place lrms in the shallow layers, and (iii) it is not necessary to have as many long-range memories as model-layers for comparable modelling performance. ## related work there have been several recent works exploring deep sequence models with a small attention window per layer. @xcite proposed the dynamic convolution, where the model directly produces a set of weights over a sequence in memory and then combines them with a convolution. the attention window is thus restricted to the convolution kernel size -a couple of words. @xcite show comparable performance to the transformer at sentence-level machine translation. however they do not investigate longer-context applications. @xcite propose shortening the range of attention for transformers by compressing the distant past. they find the first layers of the model are the most compressible, and obtain state-of-theart in several long-range language model benchmarks (wikitext-103 and enwik8). however they do not consider restricting the range of attention for a subset of layers to save compute and space. @xcite propose an adaptive attention scheme for the transformerxl where the model can learn to modulate the size of its attention window per attention head. they observe the neural network converges to using smaller attention spans for lower layers in the network, which adds additional evidence to the finding that long-range memories are not useful in these lower layers. be-cause @xcite place the range of attention in the optimisation problem it is very flexible. in this study we promote interpretability by making a set of direct interventions to the memory size across layers. this does result in less generality, as we explicitly create two types of attention ranges, where adaptive attention can select many. however ultimately the two approaches of generality and interpretability complement one another. @xcite show that one can train a transformer by having all layers attend to a single memory that is the linear combination of all layers' memories. thus at training all layers' memories are maintained, but at evaluation or generation time there can be a single memory. this gives evidence that we do not need to store many separate representations for long-range memory to perform well at test time, but the approach does require storing them during training -and incurs significant slowdown to the model. ## discussion we explore a set of interventions to the transformer-xl's architecture that are very simple to implement, i.e. a few lines of code, but shed light on the fundamental workings of the model when modelling long sequences of text. in our set of interventions, we only modify the flow of information within the network, versus the number of trainable parameters. thus we do not have confounding factors of varying network capacity. our finding is that we do not need long-range memories at every layer of the network. comparable performance can be obtained with a fraction (1/6th) of long-range memories if they are spaced equally across the network, or in the latter layers. we hypothesise this is because modelling long-range correlations is best done when representations are first formed from short-range correlations. we also find a real performance drop using a single long-range memory, proving long-range dependency is not superfluous to the task. this study has implications for practitioners interested in speeding up deep transformer-xl models. there have been a number of long-range transformer variants published in the past year @xcite @xcite which aim to extend the range of attention via sparsity or compression. however these models maintain the use of uniform memory capacity for each layer. here we show that long-range attention does not need to be scaled for every layer, and thus these architectures can be further sped-up with this observation. this study also has implications for researchers using a single long-range memory, which has typically been the approach in traditional rnn + attention systems. for example, the differentiable neural computer @xcite and recent memory-augmented agents for reinforcement learning, which utilise a distinct working memory with a single long-range episodic memory @xcite . perhaps performance could be improved by adding additional layers of episodic memories. the practice of storing deep long-range memories is not scalable if we wish for neural networks to have the kinds of large-horizon reasoning that humans possess. we believe the solution of maintaining a small number of long-range memories is a step towards tractable lifelong memory.