L a M emo: Language Modeling with Look-Ahead Memory (2022)

Abstract: Although Transformers with fully connected self-attentions are powerful to model long-term dependencies, they are struggling to scale to long texts with thousands of words in language modeling. One of the solutions is to equip the model with a recurrence memory. However, existing approaches directly reuse hidden states from the previous segment that encodes contexts in a uni-directional way. As a result, this prohibits the memory to dynamically interact with the current context that provides up-to-date information for token prediction. To remedy this issue, we propose Look-Ahead Memory (LaMemo) that enhances the recurrence memory by incrementally attending to the right-side tokens and interpolating with the old memory states to maintain long-term information in the history. LaMemo embraces bi-directional attention and segment recurrence with an additional computation overhead only linearly proportional to the memory length. Experiments on widely used language modeling benchmarks demonstrate its superiority over the baselines equipped with different types of memory mechanisms.

## introduction language modeling is an important task that tests the ability of modeling long-term dependencies by predicting the current token based on the previous context @xcite . recently, transformer-based language models achieved remarkable performance by enabling direct interaction between long-distance word pairs. however, as the computation overhead grows with the length of the input sequence, transformers can only process a fixed length segment at a time. to allow long-term information flow across individual segments, existing approaches augment the model with a recurrence memory that stores hidden states computed in previous time steps @xcite and their compressions @xcite for the target tokens to attend to. one limitation of this approach is that the recurrence memory is only aware of older contexts since they are previously computed to predict the next word from left to right. as a result, distant memory states become outdated and less activated by the current context, as illustrated in figure 1 . when humans read or write a document, they maintain a memory that records important information from the past and often refresh them under the current context to keep it up-to-date. in this paper, we propose look-ahead memory (lamemo) where memory states "look ahead" to future time steps by attending to the token representations on their right side to provide up-to-date contextualization. 3 to maintain information from the long-term history, we propose memory interpolation to take both past and future tokens into consideration, which mimics the bi-directional attention. note that, directly applying bi-directional attention to update the memory representations brings an additional complexity of o(m 2 ) (m is the mem-ory length). this is expensive when the memory is very long. lamemo incrementally attends to the right and accumulate the weighted attention sum from previous segments to simulate the full attention in only o(m × n ) complexity (n is the target sequence length), which does not increase the attention complexity of transformer-xl, namely o(n 2 + m × n ). we provide an illustration of this mechanism in figure 3 . another technique proved to be effective in language modeling is the relative positional encoding @xcite @xcite , which biases the pair-wise attention score purely based on the relative distance of the two tokens. however its ability to generalize to the attention of the future tokens remains unknown, since both the distance and the direction need to be taken into consideration. in preliminary experiments, we observed the unstability of directly applying the relative positional encoding of @xcite to this setting. we propose a simple yet effective modification based on @xcite that disentangles the bias of the relative distance and the attention direction which facilitates the training of lamemo. we give both theoretical and empirical analysis to the unstability issue and demonstrate the effectiveness of the proposed disentangled relative positional encoding method. to sum up, our contributions are as follows: ## method in this section, we describe our method in detail with our motivation to learn better representations for the memory. ## experiments we evaluate lamemo on both word-level and character-level language modeling tasks and compare with existing transformer baselines augmented with different types of memory. ## extrapolating to longer contexts in this section, we extrapolate the models to longer contexts during inference to study the effect of dynamic contextualization to the distant past. we fix the length of the target sequence to 64 and extrapolate the trained models to longer memory length 64 × m during inference, where @xmath0, • • • , 10. we compare the perplexity of lamemo and transformer-xl trained on wikitext-103 when augmented by a memory with different length. as shown in figure 5 , lamemo consistently achieves lower perplexity than transformer-xl when extraploating to longer contexts, while the performance of both models saturate when m is over 7. additionally, we observe that the gap of perplexity between the two models increases when taking longer contexts into account. this demonstrates the effectiveness of dynamically refreshing the distant memory representations under the current context. ## attention analysis in this section, we analyze the attention distribution of lamemo to validate the effectiveness of utilizing bi-directional contexts with look-ahead attention. we first visualize the memorizing coefficient α which stands for the portion of the past activations in the current memory representations. as show in figure 6 , we plot α in different layers as a function of the memory index averaged on 100 text segments. 6 we observe that in lower layers the memory mainly attends to the past (α ≈ 1.0). we conjecture that long-term bi-directionality is not necessary for low-level representations such as lexical features. in higher layers, the memory substantially utilizes the future contents to refresh the high-level representations, especially for the old memory state with a small memory index. next, we visualize the attention weight distribution on the context tokens when predicting each target token in figure 1 . for every token, we take the maximal attention weight in each interval of 5 tokens on its left and scale to a context length of 100. the result indicates that lamemo learns better memory represetations by attending to the right-side tokens, which increases the memory utilization when predicting the target token. ## case study we present the generated texts of lamemo and transformer-xl trained on wikitext-103 in appendix d. both models maintain a memory size of 512, and we seed them with the same context randomly sampled from the test set and generate 256 tokens using top-p sampling @xcite with @xmath1.95. ## related work the transformer @xcite , with its pair-wise modeling ability of the input, becomes prevailing for sequence modeling, especially long sequence processing tasks, such as long text generation @xcite , long document qa @xcite , language modeling @xcite , video processing @xcite , and etc. specifically, language modeling @xcite which requires processing documents with thousands of tokens has become a natural testbed for benchmarking this long-term processing ability. however, due to the quadratic time and space complexity of self-attention, scaling to inputs with thousands of tokens is computationally prohibitive. one line of work investigated the linear-time attention mechanism to mitigate the scability issue of transformer. linformer @xcite projects the inputs to lower dimension in length and approximates the full attention with a low-rank factorization. linear transformer @xcite regards the self-attention as a kernel function and uses a linear dot-product as a substitute. @xcite and @xcite proposed to approximate the softmax more precisely with the expectation of the dot-product of random features. although achieving substantial improvements on benchmarks designated for long inputs @xcite . these methods, however, focus on approximating the full attention with low-rank factorizations or kernel functions, which compromise the expressiveness and robustness of the original softmax attention, are reported to be inferior to the simple local attentions on real world language processing tasks @xcite . our work falls in another line, which augments the transformer with a parametrized memory to store critical history information. memoryaugmented networks @xcite @xcite have been studied in the context of recurrent neural networks for a long time, but are mostly restricted to small and synthetic datasets. with the rapid development of transformer, various works start to adapt memories to this architecture. @xcite first extended transformer with a recurrence memory that caches hidden states computed in previous steps for the target tokens to attend to. @xcite further extended the context with an external memory that stores compressed hidden states at the temporal level. @xcite used continuous space attention to attend over the old history and updated the memory with recent hidden states to enable unbounded memory capacity. @xcite proposed to use the encoder-decoder architecture to encode the memory states with previous text segments and pass this memory to future time steps. instead of using a fixed-size attention span for different layers, @xcite and @xcite proposed to learn dynamic attention spans for dif-ferent attention heads, which greatly reduced the computations. these works focused on enabling the transformer to access contents in long distance, but did not consider to learn better memory representations by refreshing the old memory under the current context. our work is orthogonal to learning adaptive attention spans and can be combined with this technique to reduce the complexity. ## conclusion we present lamemo, a memory mechanism that allows the memory states to incrementally attend to the right-side tokens and interpolates with the old memory states on the left side, which enables the memory to interact with bi-directional contexts with a complexity linear in memory length. experiments on three language modeling datasets demonstrate the superiority of lamemo over baselines with various types of memory mechanisms. we also found that lamemo increases the utilization of older memory states when predicting the target tokens, and yields a higher performance boost when extrapolating to longer memory length, which indicates the effectiveness of recontextualizing the memory under the current context.