Long-Range Language Modeling with Selective Cache (2023)

Abstract: The computational cost of transformer-based language models grows quadratically with the sequence length. In this paper, we introduce the selective cache, which stores the selected key-value pairs from the previous context. By selecting important key-value pairs the model makes better use of the cache so that in limited cache size, a longer context history can be stored. We design three kinds of selection methods. The first is based on human language processing. The key-value pairs are selected if they correspond to tokens that are fixated longer, as recorded in eye-tracking-while-reading experiments. We also incorporate the cognitively-inspired selection process into the language model as a trainable process, resulting in two additional methods with improved performance. The selection task is converted into a pruning task so they can be trained with differentiable masks. We demonstrate that the proposed selective cache improves the language modeling performance across different datasets. With the same number of stored key-value pairs (cache size), our selective cache outperforms XL cache and compressive cache by considerable margins.

## introduction transformers @xcite have been used as the backbone architecture of various models and achieve state-of-the-art results in a wide range of tasks @xcite @xcite . compared to other architectures such as recurrent neural networks, its attention mechanism directly connects long-distance elements in the input sequence. thus, it better captures long-range dependencies. however, it also constitutes the bottleneck in transformers. the time and space complexity of transformer attention is o(n 2 ). hence, the computational cost becomes prohibitively expensive or even makes the task infeasible when processing long sequences. extensive research is performed to address this problem, introducing a number of x-former models @xcite @xcite . while many proposed models successfully mitigate the problem to some extent, there are still open limitations such as the storage cost of preserving the whole sequence in memory and the inability to be directly inserted into existing pretrained models. in this paper, we propose a novel cache that stores past key-value pairs, namely the selective cache, for auto-regressive language modeling (lm). this method aims to make better use of the cache introduced in transformer-xl @xcite by selecting more informative key-value pairs. different from prior works such as @xcite , it distills the previous context by extracting or selecting, instead of averaging. in other words, some key-value pairs are kept in the cache as it is, while others are simply discarded. it is utilized as easily as the cache in transformer-xl, that is, the stored keys and values are simply prepended to the current keys and values. we explore three selection methods for the cache: (1) using eye fixation duration as importance scores, (2) using neural networks which learn to select, and (3) replace. the first one uses eye-movement data. we simply consider the total reading time of each word. long-fixated words are important for language comprehension @xcite and are usually difficult to infer from the context @xcite . furthermore, fixation duration correlates with the gradient-based saliency in transformer language models @xcite . thus, selecting key-value pairs associated with long-fixated tokens is a cognitively plausible modeling decision. the second method is to train a selector neural network. by learn-ing to prune key-value pairs, the selector learns to select. it is based on a neural network pruning method @xcite , which uses differentiable stochastic masks and l 0 penalty. the selector learns to mask out unimportant key-value pairs so it can be used to select them. finally, the third method uses the same technique but applies the mask and its opposite mask simultaneously, in order to train a replacer network that directly learns the trade-off. the motivation is to simulate the actual situation where adding new items to the cache results in discarding old ones. in addition, we introduce a novel method for training the selector and replacer network. by sampling a previous snippet of context to train the network and stopping gradient during cache updating, we avoid backpropagation through time (bptt), while the network can still learn to model long-term dependency. experimental results show that using these methods improves the model's performance, especially the latter two. they bring about considerable performance gain against xl cache and compressive cache (we use these two terms to refer to the methods used in @xcite with the same total cache size. moreover, thanks to the enhanced interpretability of selective cache, we can directly investigate which tokens are stored. we find that the trained selector and replacer tend to keep named entities and rare tokens, which is a common tendency found in many prior works @xcite @xcite . we also show that the selective cache can preserve information from a very distant context, or even from an infinite distance, depending on the characteristics of the processed document. ## related work transformers have difficulty processing long sequences. there are a lot of works trying to address this limitation. linformer @xcite and performer @xcite introduce new self-attention mechanisms, which can approximate the original full attention with linear space and time complexity. many works utilize sparsity to achieve efficient long-range attention, including strided attention @xcite , global attention by a few tokens @xcite @xcite and random attention to a limited number of tokens @xcite . moreover, suitable sparsity patterns can be learned for each attention head @xcite . some other works @xcite use clustering techniques to partition elements in the sequence and perform intra-cluster attention. besides, hierarchical structure @xcite is incorporated into the attention mechanism to reduce computational complexity. long documents can also be split into segments and processed in a recurrent manner. block-recurrent transformer @xcite use a set of hidden representations to store past information and use cross-attention to interact with them, while memory recurrent transformer @xcite utilizes dedicated memory tokens to do so. transformer-xl @xcite achieves this segment-level recurrence by the cache mechanism, which stores the past hidden states and uses them as an extended context when the model processes the next input sequence. compressive transformer @xcite extends the transformer-xl with a secondary cache which compresses the old hidden states. on the other hand, memorizing transformer @xcite expands the cache to an enormous size, and uses an approximate k-nearestneighbor (knn) search to retrieve key-value pairs from the cache efficiently. ∞-former @xcite utilizes a continuous-space attention mechanism and represents the input sequence as a continuous signal, so that the long-term memory can be represented with lower precision. expire-span @xcite computes a span for each hidden state that determines how long it should stay in memory. in this paper, we make use of eye-tracking data, which has been incorporated into many nlp frameworks @xcite . eye-tracking data provides advantages to nlp models in terms of both their performance and interpretability. research has demonstrated that incorporating eyetracking features can enhance prediction accuracy in tasks like named entity recognition @xcite , part-ofspeech tagging @xcite , sentiment analysis @xcite and general nlp benchmark tasks @xcite . in the meantime, eye-tracking data is utilized to explore the correlation between human behavior and neural attention @xcite @xcite . we also use a pruning technique based on stochastic masks @xcite , which is mainly used in neural network pruning @xcite and interpretation @xcite . different from these approaches, we aim to improve performance by using it to select important parts of previous context. even though the technique is used for different purposes in different scenarios, it actually does the same thing, i.e., learns what is less important. ## model in this paper, we use decoder-only transformers @xcite to perform auto-regressive lm tasks. long documents are split into segments of 512 tokens. the segments are not shuffled and fed into the model sequentially. in other words, the language model processes the document step by step, one segment at a time, as is done in @xcite . at each step, transformer-xl @xcite ) caches and fixes (stops gradient) the hidden state sequence computed for the current segment. in the next step, it is reused to extend the context. although in @xcite , the hidden states are cached, we follow the practice in @xcite and save the key-value pairs into the cache for the purpose of efficiency. when doing attention, the keys and values are prepended to the current keys and values. when the xl cache size c xl is greater than the input segment length c inp , the xl cache is a first-in-first-out (fifo) queue. when using cached representations, it is necessary to use relative position embeddings. we use t5 relative position embeddings @xcite , which adds different biases to different relative offsets when doing attention. when switching to a new document, the cache may contain some content from the old document, we apply document masks to solve this problem. concretely, for the cached key-value pairs, we keep track of their document ids. each token can only attend to other tokens with the same document id. these are the common settings used for all our experiments involving the cache, including baselines and proposed models. ## experiments we evaluate the models described above on three datasets using two different model sizes. the experiments are all language modeling and the performance is evaluated by perplexity per token. besides, we also do various qualitative and quantitative analyses on the models. from these analyses, we provide strong evidence of the effectiveness of the proposed models. ## conclusion in this work, we propose the selective cache that stores selected key-value pairs from previous context for language models, which significantly increases the length of the context that a language model can attend to. moreover, we propose three different ways to select important key-value pairs, namely using fixation duration, the selector network and the replacer network. the experimental results from different datasets and model sizes demonstrate the effectiveness of the proposed approaches. without increasing total cache size, the selective cache outperforms xl cache by considerable margins. moreover, further analyses reveal important characteristics of the selective cache, such as the neural network-based selection tends to select named entities and rare tokens, as shown by the increased interpretability of our models.