TRAMS : Training-free Memory Selection for Long-range Language Modeling (2023)

Abstract: The Transformer architecture is crucial for numerous AI models, but it still faces challenges in long-range language modeling. Though several specific transformer architectures have been designed to tackle issues of long-range dependencies, existing methods like Transformer-XL are plagued by a high percentage of ineffective memories. In this study, we present a plug-and-play strategy, known as TRAining-free Memory Selection (TRAMS), that selects tokens participating in attention calculation based on one simple metric. This strategy allows us to keep tokens that are likely to have a high attention score with the current queries and ignore the other ones. We have tested our approach on the word-level benchmark (WikiText-103) and the character-level benchmark (enwik8), and the results indicate an improvement without having additional training or adding additional parameters.

## introduction transformer-based models @xcite @xcite @xcite have achieved remarkable performance over the past few years. the key component of these model architectures is the attention mechanism @xcite . however, the original attention design struggles to efficiently handle long sequences, which becomes particularly problematic in scenarios such as document-level translation @xcite and large-scale text generation @xcite , as its time and space computation costs increase quadratically with the sequence length @xcite . the primary factor for this elevated computational complexity can be traced back to the multiplication between queries and keys used in figure 1 : two memory selection methods: for oracle, it selects memories with the highest attention scores after computing qk ⊺ . for trams, it selects important key/value pairs that are independent of queries based on our self-defined metric before computing qk ⊺ . the attention module. in general, the time complexity for calculation is o(n 2 d) if a transformer model with d dimensions is set up with an input consisting of n tokens. to tackle this computation bottleneck, numerous efforts have been made. the first line of work is to find a new efficient expression to compute the attention score. despite the advancements made, these methods often compromise performance, thus paving the way for alternative solutions. efficient architectures that provide an approximate expression of attention have been explored widely @xcite @xcite . the second line of work is to keep the calculation expression the same and use an external structure like hash function @xcite , clustering @xcite and memory selector @xcite @xcite @xcite to find the suitable subset of queries and keys in the long sequence for attention calculation. our work falls into the second category, in which we propose a training-free memory selection mechanism to select suitable tokens for attention computation. specifically, we focus on pushing transformer-xl @xcite architecture to a better position by selecting higher-quality tokens inside its memory. based on our initial investigation, we construct a memory subset by selecting 50% of the memories with the largest attention values and maintaining the same performance. it indicates that a large portion of information in memory is not fully utilized. this motivates us to explore better methods to optimize memory usage. illustrated in figure 1 , we propose a trainingfree memory selection method (trams) that can be directly plugged into memory-based long-range language models and reduces the time complexity of computing attention matrix. through experiments on two language modeling benchmark datasets, namely word-level wikitext-103 @xcite and character-level enwik8 (mahoney, 2011), we achieve an improvement in the model's performance, as demonstrated by a 0.19 perplexity (ppl) drop in wikitext-103 and a 0.017 reduction in bits-per-character (bpc) in enwik8. to our knowledge, we are the first to design a training-free memory selection method based on transformer-xl architecture. 1 2 method ## experiments we introduce the compared methods and report the main results and analysis on different attention variants for inference in this section. datasets details for wikitext-103 and enwik8 benchmarks and their evaluation metric details are included in appendix a. the details of the model that we built memory selection on can be seen in appendix b. ## discussions is trams vulnerable to the selection of hyperparameters? there are three hyper-parameters in trams: the memory pool size m that trams is able to select from; the selected memory size m that is used in the forward process; and the input token size n that is involved in both backward and forward process. from the ablation study on m , figure 4 suggests an optimal range between 300 to 400 for the memory pool size. beyond this range, enlarging the memory pool often leads to the selection of irrelevant tokens, deteriorating our performance. regarding m, figure 5 indicates that trams witnesses a substantial drop in perplexity when the memory size selected is about 25%. selecting a larger portion does not yield further improvement. this is consistent with figure 3 , where trams excels by concentrating on the top 10% of results. lastly, in the study on n, figure 6 shows that as the target token length decreases, the efficacy of memory selection improves. 200 300 400 500 600 23.6 23.8 24 24.2 24.4 24.6 selected memory size m perplexity transformer-xl trams figure 5: ablation study on selected memory size m when we fix @xmath0. 16 32 64 128 23.8 23.9 24 24.1 24.2 24.3 24.4 input length n perplexity transformer-xl trams figure 6: ablation study on target length n when we fix @xmath1. what is the inference cost compared to transformer-xl? since there is no training part in our model, we focus on discussing the inference cost. compared with transformer-xl, our model requires storing a larger memory pool to do memory selection. therefore, the memory cost of our method would be larger. when it comes to timing cost, our model has an additional memory token norm computation memory sorting operations, and memory selection operations for each layer. these extra operations require extra inference time. ## case study to have an understanding of what kind of context should be selected, we provide one example case to understand specifically what kind of tokens in the memory would be selected. based on ## conclusion in this work, we formulate the problem of memory selection in transformer architecture and reformulate the attention calculation process to obtain our self-defined queries and keys. after that, we propose a query-independent metric that utilizes memory hidden states to implement a training-free memory sequence segment ...simon stephens , which was performed in 2001 at the royal court theatre. he had a guest role in the television series judge john deed in 2002. in 2004 boulter landed a role as "craig" in the episode "teddy's story" of the television series the long firm; he starred alongside actors mark strong and derek jacobi. he was cast in the 2005 theatre productions of the philip ridley play mercury fur, which was performed at the drum theatre in plymouth and @xmath2. he was directed by john tiffany and starred alongside ben whishaw , shane zaza, harry kent, fraser ayres, sophie stanton, and dominic hall. <@xmath3, boulter starred alongside whishaw in the play citizenship written by mark ravenhill ... target sequence segment he appeared in the television series judge john deed in 2002 ...