TRAMS: Training-free Memory Selection for Long-range Language Modeling
Authors: Haofei Yu
Cunxiang Wang
Yue Zhang
Wei Bi
Venue: EMNLP 2023 Findings (2023)
Keywords: Language Model; Inference Strategy; Long-context Modeling

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.

TRAMS: Training-free Memory Selection for Long-range Language Modeling
Haofei Yu Cumxiang Wang*, Yue Zhang*, Wei Bi
Language Technologies Institute, Carnegie Mellon University, USA School of Engineering, Westlake University, China Tencent AI Lab, China haofeiyēcs. cmu edu.edu, {wangcunxiang, zhangyue victoriabi@tencents com
cdu
Abstract
Key
The Transformer architecture is crucial for nu- merous Al models, but still faces challenges in long-range language modeling. Though several specific transformer architectures have been designed to tackle issues long-range de- pendencies, existing methods like Transformer- XI. are plagued by a high percentage of inef- fective memories. In this study, we present plug-and-play strategy, known TRAining- free Memory Selection (TRAMS). that selects to- kens participating in attention calculation based on one simple metric. This strategy allows us to keep tokens that are likely have high atten- tion score with the current queries and ignore the other ones. We have tested our approach on the word-level benchmark (WikiText-109) and the character-level benchmark (enwik8), and the results indicate an improvement with- out having additional training or adding addi- tional parameters.
Query
Memory Selection
selott best QK' Oracle
select best based on TRAMS TRAMS
Figure I: Two memory selection methods: For oracle, selects memories with the highest attention scores after computing QKT. TRAMS, selects important key/value pairs that are independent of queries based our self-defined metric before computing QKT.
the attention module. In general, the time com- plexity for calculation is (NV2d) if a transformer model with d dimensions is set up with input consisting of N tokens.
1 Introduction
Transformer-basel models (Kenton and Toutanova, 2019; Liu et al., 2019; Raffel et al., 2020; Lan et al., 2019; Brown et al., 2020) have achieved remark- able performance over the past few years. The key component of these model architectures is the atten- tion mechanism (Vaswani et al., 2017). However, the original attention design struggles to efficiently handle long sequences, which becomes particularly problematic in scenarios such as document-level translation (Werlen et al., 2018; Kim et al., 2019) and large-scale text generation (Zhou et al., 2023), its time and space computation costs increase quadratically with the sequence length (Tay al., 2022). The primary factor for this elevated com- putational complexity can be traced back to the multiplication between queries and keys used in
To tackle this computation bottleneck, numerous efforts have been made. The first line 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 expres- sion of attention have been explored widely (Wang et al., 2020; Peng et al., 2022b,a; Choromanski et al., 2021; Zheng et al., 2022b,a). The second line of work is to keep the calculation expression the same and use an external structure like hash function (Kitaev et al., 2019; Daras et al., 2020), clustering (Roy et al., 2021; Vyas et al., 2020) and memory selector (Pietruszka al., 2022: Dai al., 2019; Bertsch et al., 2023; Sukhbaatar et al., 2021, 2019; Child et al., 2019) find the suitable sub- set of queries and keys in the long sequence for attention calculation.
*ork done during internship Teneent Al Lab.
Co-first Author. The correponding author.

---

Our work falls into the second category, in which we propose a training-free free memory selection mechanism to select suitable tokens for attention computation. Specifically, we focus on pushing Transformer-XL (Dai et al., 2019) architecture to a better position by selecting higher-quality tokens inside its memory. Based on our initial investiga- tion, we construct memory subset by selecting 50% of the memories with the largest attention values and maintaining the same performance. indicates that a large portion of information in memory not fully utilized. This motivates us to explore better methods to optimize memory usage. Illustrated in Figure 1, we propose TRAining- free Memory Selection method (TRAMS) that can be directly plugged into memory-based long-range language models and reduces the time complex- ity of computing attention matrix. Through ex- periments on two language modeling benchmark datasets, namely word-level Wiki Text-1000 (Mer- ity et al., 2016) and character-level enwik8 (Ma- honey, 2011), we achieve an improvement in the model's performance, as demonstrated by 0.19 perplexity (ppl) drop in WikiText-103 and 0.017 reduction in bits-per-character (bpc) in enwik8.
2.2 Attention Reformulation
In a memory-augmented language model, the stan- dard attention mechanism between input hidden states and memory hidden states can be written as:
QK
Attn(h, m) softmax
)
(2)
where Q = hWQ is the product of target token hidden states h and query projection matrix Wo: K mW is the product of memory token hidden states m and key projection matrix WK: V mWy is also the product of memory token hidden states and value projection matrix Wv. Different from the well-known attention score calculation, we can compute this attention formula in a different order but maintain the same result:
QK (m
-W)(WWWWm)) =m)(mW
(3)
Thus, we define Q'= h as the reformulated query for this attention expression and K' mW, K Wq as the reformulated keys for attention. With this reformulation, we transfer all attention-related para- metric information onto reformulated key vectors.
To our knowledge, we are the first to design a training-free memory selection method based on Transformer-XL. architecture.
2.3 Transformer Hidden Space
Since h the input of the current transformer layer and also the output of the previous transformer layer, it is the result of the last layer's Layernorm operation. We can define the coordinate- av- erage of h as μ and the coordinate standard deviation of h a. Expressions be written
2 Method
2.1
Problem Definition
We use h Nxd to represent the input hidden states for the attention module, e Nxd to rep- resent the output hidden states for the attention module, Mxd to represent the memory hid- den states used in the attention calculation. We use Wo. WK. K. Wv to represent the trainable projection matrix in the attention module. We define d for the dimension of the model, M for the memory size, and N for the input size. The attention calculation process can be formally written as o Attn(h. h. m.).
1 (hh 11 (4) / Σ ha x O, a= V dd
Since the mean value for the hidden states h is around zero, we can confirm the hidden states vec- tors are approximately orthogonal to the vector and the 1.2 norm of hidden states is around Vd.
With the above annotations, the problem of mem- ory selection can be defined as choosing subset of hidden states memory m from the memory m that brings the minimum difference to the transformer layer output but with a smaller memory size.
With this approximation, we can expand our re- formulated attention score as:
=KK== (m))
= =1|''II-|["| cos (Q',K")
m arg min Attn (h, m) - Attn(,, m)] (1) mom
/Vd-l|K' cos (Q', K*)
(5)
where ||Q']] stands the L2 norm for Q' and [K'|| stands for the L.2 norm for K" Based on Fig 2,
SSurce code for this paper is available at http:///gtth..com//aeei//j///

---

Key Norm Prob Density
0.6
Query Norm Prob Densit
Spearman Correlation (%)
100
0.5
0.4
2
60
0.3
0.2
40
0.1
20
cos (Q', XO)
KT
0.0
10
30
50
100
Figure 2: Norm distribution of reformulated Q' and K'. The red distribution represents the query norm. The blue distribution represents the key norm.
Highest Attention Score (%)
Figure 3: Spearman Correlation Score on different rank- ing metrics with the groundtruth one.
we see that reformulated query norm ||Q'll has a much sharper distribution compared with key norm ||K|]], indicating reformulated query norm can be approximated by a constant factor.
This metric allows identify tokens that can produce very high attention scores when paired with the appropriate query (owing to a high value of || [*|l) and very low scores when paired with an unsuitable query. This is due to the near orthogonal- ity to the query space, as indicated by small angle with i, which is orthogonal to the query space.
2.4 Training-free Memory Selection (TRAMS)
Our target for memory selection is to recover the complete attention score with as few memory to- kens as possible. This problem is equivalent to finding the subset of memory tokens that have the highest attention scores with queries. We propose a heuristic method to perform token-level selec- tion based on a neenoryindeependent metric this section.
Experiments
We introduce the compared methods and report the main results and analysis on different attention variants for inference in this section. Datasets de- tails 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 Appendix B.
There are two crucial components for calculating the attention score after approximating ||Q'|l with constant factor: the norm of the reformulated keys |||]||| and the angles between the reformu- lated keys and queries arccos (Q', K"), which is proved in Khandelwal et al. (2019). Commonly, we believe that arccos (Q', K") is the more impor- tant factor in general. Yet, if we use the ranking of attention score value for all query and key pairs as ground-truie ranking, based on Fig 3, we empiri- cally discovered that rankings based on key norms and rankings based on angles produce close Spear- man correlation scores when only taking the high- est 1 1 attention scores into account. Therefore, it indicates that we can rank our memory tokens based on || solely to gain a relatively good per- formance when we desire top 1 1% attention scores with queries in our memories instead of all. Additionally, we discovered that relying solely on large norm isn't sufficient as constraint. Specifically, keys that are nearer to i tend to yield a higher attention score. To address this, we in- troduce a combined metric: s = cos (K', 1) [|*]]
3.1 Compared Methods
Transformer+RPE. (Vaswani et al., 2017): the vanilla transformer baseline with relative position embedding that is the same as Transformer-XI. Therefore, the only difference between this model and Transformer-Xe -XL is the additional memories. More information related to relative position em- bedding can be seen Appendix C.
Transformer-XL. (Dai et al., 2019): a specific- designed architecture for long-range language mod- eling. It includes relative position embedding and recurrent memories per layer. Memory slots are filled with hidden states from previous time steps.
3.2 Experimental Settings
We compare our methods with the Transformer- XL (Dai al., 2019) under the same size of mem- ory (m 200) for attention calculation. For the input token length n for both models, we keep the same as in (Dai et al., 2019) (n 64). Additionally,

---

Model
#Total Mem
#Selected Mem
PPL
21.2
MikiText-103
Porplexity
24.1
Transforme++R1
29.14
Transfomner-XL
200
24.17
24
TRAMS
600
200
23.98
Transformer XL
23.9
enwik8
TRVS
300
500
Transforer+RPP
1,240
Transforner-XL.
200
1.215
Memory Pool Size M
TRAMS
600
200
1.198
Figure 4: Ablation study on memory pool size M when we fix m=200 and n=64.
Table 1: Model performance the word-level WikiText- 103 and the character-level enwik8 datasets,
24.6
Transforma-Xi XL.
TRANS
Perplexity
24.4
the memory selection process is performed on a memory pool with the size of M. Our model and the Transformer-XL share the model parameters but have different inference strategies.
24.2
24
23.8
23.6
500
600
3.3 Main Results
Selected Memory Size
The main results of WikiText-103 and enwik8 datasets are shown in Table 1. Without additional training or additional parameters, we gain 0.19 im- provement in perplexity and 0.017 improvement for bit-per-character with our TRAMS mechanism. We implement p-test by inferencing multiple model checkpoints and prove that our results significant (p 0.05).
Figure 5: Ablation study on selected memory size when we fix M=600 and n=64.
24.4
Tranformer-Xi
24.3
TRAMS
Perplexity
24.2
24.1
24
23.9
4 Discussions
23.8
16
32
64
128
Input Length
Is TRAMS vulnerable to the selection of hyperpa- rameters? There are three hyper-parameters in TRAMS: the memory pool size M that TRAMS is able to select from; the selected memory size that is used in the forward process; and the input to- ken size n that is involved in both backward and forward process.
Figure 6: Ablation study on target length when fix M=400 and m=200.
What is the inference cost compared to Transformer-XL. Since there is no training part in our model, we focus discussing the inference cost. Compared with Transformer-XL, our model requires storing a larger memory pool do mem- ory selection. Therefore, the memory cost of our method would be larger. When 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. Table shows the GPU memory cost and wall-clock time for the Transformer-XX baseline and our model. Our model requires slightly more GPU memory usage and around 50% additional inference time for memory selection.
From the ablation study M, Figure 4 sug- gests an optimal range between 300 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 wit- nesses substantial drop in perplexity when the memory size selected is about 25%. Selecting larger portion does not yield further improvement. This is consistent with Figure 3, where TRAMS ex- cels by concentrating on the top 10% of results. Lastly, the study on n, Figure shows that as the target token length decreases, the efficacy of memory selection improves.

---

Model
Peak GPU Mem (MB)
Wall-clock Time
Memory Sequence Segment
Transforme-XI. TRAMS
33.27
49.55
...Sioon Stephens which was performed in 2001 at the Royal Court Theatre. He had guest role the television series Judge John Deed in 2082. In 2004 Boulter landed role "Craig" in the episode "Teddy's Story" of the television series The Long Firm; he starred alongside actors Mark Strong and
Table Results on GPU peak memory usage and wall- clock inference time on 10x
24.2
Perplexity
Derek Jacobi. He was cast in the 2085 theatre productions of the Philip Ridley play Mercury Fur, which was performed at the Drum Theatre in Plymouth and the <unk> Chocolate Factory in London. was directed by John Tiffany
24.1
w/ menoyy selection on layer
wlo momory selection
24
and starred alongside Ben Whishaw Shane Zaza, Harry Kent, Fraser Ayres, Sophie Stanton, and Dominic Hall. <eos> 2006, Boulter starred alongside Whishaw in the play Citizenship wr written by Mark Ravenhill
13
16
Layer Number
Figure 7: Ablation Study on Layer- Layer-wise Importance on 10x-
Target Sequence Segment
How does TRAMS benefit from memory selection? Memory selection helps the model pick tokens with higher attention scores with the queries, thus in- creasing the average memory utilization. Quantita- tively, our method improves the average attention probability by 24.25% for the same size of memory compared with Transformer-Xi
He appeared in the television series Judge John Deed in 2002
Table 3: Case Study for memory selection from WikiText-1033, text indicates that this word mem- ory sequence is selected and used in the forward pass. text indicates that this word in the target sequence benefits from the memory.
Does each layer hold the same importance? Based Figure
show the ablation study when applying memory selection on each layer while remaining other layers the same. There is an observable drop when we apply the memory selection on the deeper layers starting from Layer 13 while we do not observe a clear influence when applying memory selection on shallow layers.
memory selector. Our experiments indicate that this method offers a simple yet effective means of identifying valuable memory tokens. Explor- ing optimal memory selection strategies for large language models is a promising avenue for future research. Additionally, integrating trainable pa- rameters into these models as memory selectors presents another exciting direction for future work.
5
Case Study
Limitations
To have an understanding of what kind context should be selected, we provide one example case to understand specifically what kind of tokens in the memory would be selected. Based on Table 3, we can see that most of the selected memory tokens how-frequency words. Those low-free wwwffreuenccy words like "John" in the memory would be beneficial for the prediction of "John" in the target sequence.
Our study has a couple of main limitations. First, we are currently focusing on the Transformer-XL. architecture, but there are many other models with different sizes we haven't 't tried. It indicates that our findings could be limited to typical transformer architecture. Second, our method has many hy- perparameters including M, m, and Adjust- ing them can greatly change how well our model works. careful calibration thus necessary, and one must tread cautiously to strike a balance and achieve the desired performance, which could be time-consuming and computationally expensive.
6 Conclusion
In this work, formulate the problem of mem- ory selection in transformer architecture and refor- mulate the attention calculation process to obtain our self-defined queries and keys. After that, we propose a inceeeeedeee metric that utilizes memory hidden states to implement a training-free
Ethics Statement
There are no recognized potential risks.

---

References
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Man- dar Joshi, Danqi Chen, Omer Levy, Mike Lewis. Luke Zettlemoyer, and Veselin Stoyanov. 2019. Roberta: A robustly optimized bert pretraining ap- proach. arXiv preprint 1rXiv:1907.1100
Amanda Bertsch. Uni Alon, Graham Neubig, and Matthew R Gormley. 2023. Unlimiformer: Long-
range transformers with unlimited length input. arXiv preprint arX.1:2305.01625.
Matt Mahoney, 2011. Large text compression bench- mark.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systemss 33:1877-1901.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016. Pointer sentinel mixture mod- els. In International Conference on Learning Repre- sentations.
Rewon Child, Scott Gray, Alec Radford, and Πya Sutskever. 2019. Generating long se-
Hao Peng, Jungo Kasai, Nikolaos Pappas, Dani Yogatama, Zhaofeng Wu, Lingpeng Kong, Roy Schwartz, and Noah A Smith. 2022a. Abc: Attention with bounded-memory control. In Proceedings of the 60th Annual Meeting of the Association for Compu- tational Linguistics (Volume 1: Long Papers), pages 7469-7483.
quences with sparse transformers. arXiv preprint arXiv:1904. 10509.
Krzysztof Choromanski, Haoxian Chen, Han Lin, Yuanzhe Ma. Arijit Sehanobish, Deepali Jain, Michael Ryoo. Jake Varley, Andy Zeng, Valerii Likhosherstov, et al. 2021. Hybrid random features. arXiv preprint arXiv: 10.04677
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. 20226. Random feature attention. In International Confer- ence on Learning Representations.
Zihang Dai, Zhilin Yang. Yiming Yang, Jaime G Car- bonell, Quoe Le, and Ruslan Salakhutdinov. 2019. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978-2988. 2988.
Michat Pietruszka, Lukasz Borchmann, and Lukasz Garncarek, 2022. Sparsifying transformer models with trainable representation pooling. In Proceed- ings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pa- pers), pages 8616-8633.
Giannis Daras, Nikita Kitaev. Augustus Odena, and Alexandros G Dimakis. 2020. Smyrf-efficient atten- tion using asymmetric clustering. Advances in Neu- ral Information Processing Systems. 33:6476 6489.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text trans- former. The Journal of Machine Learning Research, 21(1):5485-55 55511
Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. 2019. Bert: Pre-training of deep bidirec- tional transformers for language understanding. In Proceedings of NAACI-HILTT, pages 4171-41866
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2019. Generalization through memorization: Nearest neighbor language models. In International Conference Learning Representations.
Aurko Roy. Mohammad Saffar, Ashish Vaswani, and David Grangier. 2021. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53- 68.
Yunsu Kim, Due Thanh Tran, and Hermann Ney. 2019, When and why is document-level context useful in neural machine translation? In Proceedings of the Fourth Workshop on Discourse in Machine Transla- tion (DiscoMT 2019), pages 24-34.
Sainbayar Sukhbaatar, Édouard Grave, Piotr Bo- janowski, and Armand Joulin. 2019. Adaptive at- tention span in transformers. In Proceedings of the 57th Annual Meeting of the Association for Compu- tational Linguistics, pages 331-335.
Diederik P Kingma and Jimmy Ba. 2014. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980
Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur Szlam, Jason Weston, and Angela Fan. 2021. Not all memories are created equal: Learning to forget by expiring. In International Conference Machine Learning, pages 9902-9912. PMLR
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2019. Reformer: The efficient transformer. In Inter- national Conference on Learning Representations.
Yi Tay, Mostafa Dehghami, Dara Bahri, and Donald Met- zler. 2022. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. 2019. Albert: A lite bert for self-supervised learning of language representations. In International Confer- ence on Learning Representations.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all

---

you need. Advances in neural information processing systems, 30.
Ba, 2014) optimize with a batch size=60, learning rate=2.5e-4, target length=150, mem- ory 1ength=150, and a cosine scheduler without warmup steps.
Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. 2020. Fast transformers with clustered tention. Advances in Neural Information Processing Systems, 33:21665-21 21674.
When it comes to a different dataset, we use different Transformer-Xi architecture. For WikiText-103, we use 16-layer transformer ar- chitecture with 10 heads, 410 hid dim, 0.1 dropout ratio, 0.0 attention dropout ratio, 2100 inner dim, and adaptive softmax mechanism. For enwik8, we propose 12-layer transformer architecture with 8 heads, 512 hid dim, 0.1 dropout ratio, 0.0 attention dropout ratio, and 2048 inner dim. Both models are trained for 350K steps.
Sinong Wang, Belinda Li, Madian Khabsa. Han Fang, and Hao Ma. 2020. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 0400:00000...88
Lesly Miculicich Werlen, Dhananjay Ram, Nikolaos Pappas, and James Henderson. 2018. Document- level neural machine translation with hierarchical attention networks. In Proceedings of the 2018 Con- ference on Empirical Methods in Natural Language Processing, pages 2947-2954.
A batch size=10 and target length=150 are fixed for all inference experiments avoid unfair com- parison. All experiments including training and inference are conducted using 4 2080Ti GPUs. It takes 280 GPU hours train the enwik8 model checkpoint. takes GPU hours to train the WikiText-103 model checkpoint.
Lin Zheng. Chong Wang, and Lingpeng Kong. 2022a. Linear complexity randomized self-attention mech- anism. In International Conference on Machine Learning, pages
PMLR.
Lin Zheng, Jianbo Yuan, Chong Wang, and Lingpeng Kong. 2022b. Efficient attention via control vari- ates. In The Eleventh International Conference on Learning Representations.
C Relative Position Embedding
Wangchunshu Zhou, Yuchen Eleanor Jiang, Peng Cui, Tianaa Wang, Zhenxin Xiao, Yifan Hou, Ryan Cot- terell, and Mrinmaya Sachan. 2023. Recurrentgpt: Interactive generation of (arbitrarily) long text. arXiv preprint 23055 1330.
Concerning positional encodings, maintain the same results with -raLsformer-XL. The positional encodings include learnable parameters of R u, and v. Typically, Rı- is derived from a Icarnable r network included in the model. The advantage of using this design when computing the attention score is that it avoids temporal confusion caused by indexing the same position and considers the relative distance between two tokens. The formula for attention score calculation with relative position embedding can be written as:
A Dataset and Evaluation Metrics
Wiki (Merity et al., 2016) is a commonly used word-level language modeling benchmark. It has an average length of 3.6 thousand tokens per article and includes 28 thousand Wikipedia articles. This word-level datuset has vocabulary size of around 260K. We use the same data pre-processing setting in Dai et al. (2019) for this dataset. We use perplexityy ty as our metric.
A d
X W W X, X, W W W X, + W "R,,,
Ri
(6)
Enwik8 (Mahoney, 2011) is a character- lan- guage modeling benchmark. This dataset contains 100M unprocessed Wikipedia characters. The train set, dev set, and test set include 80M, 10M, and 10M characters separately. enwik8 has pre- processing stage and is directly used. bpc (bit per character) is defined as an evaluation metric and we report results on both the dev set and test set.
Moreover, after doing ablation studies on rela- tive position embedding, we found that R con- tributes the most to the result and u, only has small influence the final performance. The existence of Ri- leads the exponentially de- cayed attention probability distribution related a memory position. As a result, we base our mem- ory selection on the A which includes positional information instead of the pure X W W X,. To be noticed, all concepts related to qK are all equipped with position embedding instead sim- ple dot product.
B Training Configurations
Since we do inference experiments based on trained model, we separately train two Transformer- XL models for WikiText-103 and enwik8. For the training stage, we use Adam (Kingma and


Reviews:
[{"review_id": "Bx1hqXOz7P", "reviewer": "Reviewer_ncPF", "summary": "", "questions": "- $\\textbf{Question A}$: As mentioned in \"Reasons to reject\", what is the PPL of TRAMS when the memory size=150? Same question to enwiki8, what is the BPC of TRAMS when the memory size=512?\n- $\\textbf{Question B}$: Following the above question, is TRAMS sensitive to \"Mem Size\"? Are there any reasons for setting \"Mem Size\" to 200?\n- $\\textbf{Question C}$:  In L68-L70, \"only 25% of the top memories are essential for maintaining the performance\". As 25% comes from WikiText-103, does the proportion of redundant memory dataset specific?\n- $\\textbf{Question D}$(optional): In Section 2.4, how the number of selected memory tokens is decided after ranking? Is it a fixed number? Just curious why `topk_num` is set to 50 in L356 (https://anonymous.4open.science/r/TRAMS-768A/utils/modeling_transfo_xl.py#L356).", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for providing valuable suggestions and encouraging us to provide more details related to the experiments.\n\n# Rebuttal Point1\n\n**_To respond to the first reason to reject:_**\n\n**_“Table 1 needs more experimental results to convince readers that the proposed method TRAMS improves Transformer-XL. For example, in the WikiText-103 dataset, the PPL of 12L Transformer-XL 24.17 differs from the reported PPL 24.0 in the Transformer-XL (Dai et al., 2019) paper. Please refer to \"Ours - Transformer-XL Standard\" in Table 1 of Dai's work. The dissimilarity comes from different memory sizes, 200 in this work and 150 in the original Transformer-XL (see run_wt103_base.sh: https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/run_wt103_base.sh). Experiments on the proposed method TRAMS with the same memory size as Transformer-XL by datasets are required.”_**\n\nThank you for bringing up the important issue of fair baseline comparison in our research. When preparing for the final version of our paper, we would put the target token length and memory token length number as part of our experimental setting section to avoid confusing people on that. We would add this comparison in the main experimental table to show that our paper is working on a successfully reproduced baseline Transformer-XL checkpoint.\n\n**Token Length Roles:** The discrepancy in Perplexity (PPL) values in the Wikitext-103 dataset can be attributed primarily to variations in two settings: target token length and memory token length. The target token length relates to the number of tokens that participate actively in the backward process of the model. Memory token length, on the other hand, means tokens that are omitted from this backward process and have stop-gradient operations applied. These differences form the basis for varying PPL values.\n\n**Reproduced Baselines:** To validate the implementation of our Transformer-XL baseline model, originally developed by Dai et al., we undertook a two-pronged approach. Firstly, we meticulously aligned our model's configurations with the settings outlined in the original Transformer-XL paper. Secondly, we evaluated our model's performance against the official Transformer-XL baseline reproduced in LaMemo.\n\n**Reproduced Baselines Details:** The original Transformer-XL model was trained with 150 target tokens and 150 memory tokens. During inference, however, the model deploys 64 target tokens and 640 memory tokens—a shift from the 150+150 memory + target token lengths used during training. This setting is documented in the [Transformer-XL repository](https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/run_wt103_base.sh). Conversely, the LaMemo model by Ji et al. maintains a consistent setting for both training and inference, employing 150 target tokens and 150 memory tokens as mentioned in the section 4.3 in their original paper. Given these divergent configurations, we felt it necessary to evaluate our model against both these established baselines.\n\n**Experimental Results:** We display the results of our experiments with the replicated Transformer-XL model in Table1 and Table2. Our findings indicate that our re-implemented model performs at par with, or even exceeds, the baselines established in both the Dai et al. and Ji et al. papers. These results are consistent across the enwik8 and Wikitext-103 datasets. As a result, we can assert with confidence that the improvements observed in our version of the model are indeed valid.\n\nTable1: Performance comparison between different token length model settings on Wikitext-103 dataset\n\n| Model Setting                    | Target Token Length | Memory Token Length | PPL(↓)    |\n| -------------------------------- | ------------------- | ------------------- | --------- |\n| Transformer-XL (Dai et al. 2019) | 64                  | 640                 | 24.0      |\n| Transformer-XL (our reproduced)  | 64                  | 640                 | **23.66** |\n| Transformer-XL (Ji et al. 2022)  | 150                 | 150                 | 24.56     |\n| Transformer-XL (our baseline)    | 150                 | 150                 | **24.26** |\n| Transformer-XL (our baseline)    | 64                  | 200                 | **24.17** |\n\nTable2: Performance comparison between different token length model settings on enwik-8 dataset\n\n| Model Setting                    | Target Token Length | Memory Token Length | BPC(↓)    |\n| -------------------------------- | ------------------- | ------------------- | --------- |\n| Transformer-XL (Dai et al. 2019) | 80                  | 2100                | **1.06**  |\n| Transformer-XL (our reproduced)  | 80                  | 2100                | 1.094     |\n| Transformer-XL (Ji et al. 2022)  | 512                 | 512                 | 1.239     |\n| Transformer-XL (our reproduced)  | 512                 | 512                 | **1.122** |\n| Transformer-XL (our reproduced)  | 64                  | 200                 | 1.215     |\n\n[1] Zihang Dai et. al. Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context [[https://arxiv.org/abs/1901.02860](https://arxiv.org/abs/1901.02860)]\n\n[2] Haozhe Ji et. al. LaMemo: Language Modeling with Look-Ahead Memory [[https://arxiv.org/abs/2204.07341](https://arxiv.org/abs/2204.07341)] \n\n# Rebuttal Point2\n\n**_To answer question A:_**\n\n**_“Question A: As mentioned in \"Reasons to reject\", what is the PPL of TRAMS when the memory size=150? Same question to enwiki8, what is the BPC of TRAMS when the memory size=512?”_**\n\nThat is a very good question. I think this question is crucial. \n\n**Memory Pool Size Influence:** The performance of TRAMS is not solely dependent on a fixed memory size but also hinges on the pool size from which memory tokens are selected. For instance, with a memory size of 150, TRAMS utilizes 150 memory tokens for inference. However, the source of these tokens—whether they are drawn from a pool of 300 or more—can affect performance. Given these fluctuations, we find it essential to discuss TRAMS' performance in relation to different sizes of memory pools. Therefore, below we show two tables discussing the performance of our method in different memory pool sizes.\n\n**Wikitext-103 Results:** Table 3 varies the size of the memory pool, which has a significant impact on perplexity. A memory pool size of 350 appears to be the most effective. When the memory pool size is too small, perplexity doesn't decrease substantially. This is intuitive, as a smaller pool offers fewer candidate tokens for selection. On the other hand, when the memory pool size is too large, perplexity also fails to decrease appreciably. This is likely because a larger pool may include too many irrelevant tokens that are far removed from the current context, leading the model to make incorrect selections more often. \n\n**Enwik8 Results:** Table 4 compares different ratios of selected memory size to actual used memory size. We observe that TRAMS consistently outperforms the Transformer-XL baseline when the ratio is set at 512, regardless of the target token size. However, as the memory pool size increases—reaching as large as 1536—we notice a decline in performance compared to the Transformer-XL baseline. Interestingly, a \"U-shaped\" curve emerges when we analyze the effects of memory pool size on the enwik8 dataset as well. This indicates that on enwik8 dataset, it is important to find a suitable memory pool size for memory selection. One noticeable thing is that the performance improvement of TRMAS with 512 memory size is not as good as the number mentioned in our paper. We think the reason is because enwik8, as a character-level LM task, has a limited vocab, and it is difficult for TRAMS to find typically useful tokens when having 512 available memories.\n\n\n\nTable3: Performance comparison with different memory pool sizes for memory selection in TRAMS for wikitext-103 dataset\n\n| Settings       | Memory Pool Size | Selected Memory Size/Actual Used Memory Size | Target Token Size | PPL(↓)    |\n| -------------- | ---------------- | -------------------------------------------- | ----------------- | --------- |\n| Transformer-XL | —                | 150                                          | 150               | 24.26     |\n| TRAMS          | 200              | 150                                          | 150               | 24.14     |\n|                | 250              | 150                                          | 150               | 24.08     |\n|                | 300              | 150                                          | 150               | 24.07     |\n|                | 350              | 150                                          | 150               | **24.06** |\n|                | 400              | 150                                          | 150               | 24.09     |\n|                | 450              | 150                                          | 150               | 24.11     |\n|                | 500              | 150                                          | 150               | 24.13     |\n|                | 550              | 150                                          | 150               | 24.16     |\n|                | 600              | 150                                          | 150               | 24.19     |\n\nTable4: Performance comparison with different memory pool sizes for memory selection in TRAMS for enwik8 dataset\n\n| Settings       | Memory Pool Size | Selected Memory Size/Actual Used Memory Size | Target Token Size | BPC(↓)    |\n| -------------- | ---------------- | -------------------------------------------- | ----------------- | --------- |\n| Transformer-XL | —                | 512                                          | 512               | 1.122     |\n| TRAMS          | 1024             | 512                                          | 512               | **1.117** |\n|                | 1536             | 512                                          | 512               | 1.123     |\n| Transformer-XL | —                | 512                                          | 80                | 1.137     |\n| TRAMS          | 1024             | 512                                          | 80                | **1.132** |\n|                | 1536             | 512                                          | 80                | 1.144     |\n\n# Rebuttal Point3\n\n**_To answer question B (part1):_**\n\n**_“Question B: Following the above question, is TRAMS sensitive to \"Mem Size\"?”_**\n\nThe simple answer is **not** sensitive to “memory size”. Below is a detailed illustration.\n\n**“Mem Size” Definition:** The “Meme size” mentioned in our paper means “selected memory size / actual used memory size” in the TRAMS. This term refers to the subset of tokens selected from a larger pool of nearest memory tokens for use during the inference process. In contrast, in the Transformer-XL model, \"selected memory size/ actual used memory size\" pertains to the nearest memory tokens themselves.\n\n**Sensitivity Test:** In response to the previous question, we explored the sensitivity of performance to changes in memory pool size. Now, we shift our focus to understanding how varying the selected memory size impacts performance, while keeping the memory pool size constant. From the findings presented in Table 5, as referred to in answer to question B, both TRAMS and Transformer-XL exhibit similar trends as the memory size changes. Specifically, when the memory size is on the smaller side—ranging from 150 to 200—both models experience noticeable fluctuations in performance. In contrast, when the memory size is larger—within the 300 to 400 range—performance changes are relatively negligible.\n\n**TRAMS vs Transformer-XL:** Upon examining the performance curves for TRAMS and Transformer-XL in Table 5, it becomes clear that TRAMS experiences a steeper decline in perplexity (PPL) as memory size increases. For instance, as the memory size rises from 150 to 200, PPL falls by 0.32 for Transformer-XL, as opposed to a more significant drop of 0.52 for TRAMS. Conversely, when the memory size escalates from 350 to 400, Transformer-XL sees only a modest 0.06 decline in PPL, while TRAMS remains stable, showing no change.\n\nWe think sensitivity is definitely an important part of our method. We would add one subsection at the discussion section to answer the question \"Is our method sensitive to memory size?\".\n\n# Rebuttal Point4\n\n**_To answer question B (part2):_**\n\n**_“Are there any reasons for setting \"Mem Size\" to 200?”_**\n\nThank you for pointing out the need for clarification of setting \"Mem Size\" to 200.\n\n**Reason for 200:** The choice of using 200 memory tokens as a hyper-parameter is primarily heuristic. We noticed a significant performance drop when pairing a target token length of 64 with a memory token length of 200. As a result, we decided to retain these hyper-parameters for our experiments.\n\nTo provide further insights into this heuristic selection, we conducted an ablation study with additional experimental results:\n\n1. **Impact of Memory Pool Size (Table 5):** Initially, we experimented with varying the memory pool sizes while keeping the selected memory length constant.\n2. **Impact of Selected Memory Size (Table6):** Secondly, we varied the selected memory size while maintaining a fixed selection ratio.\n3. **Impact of Absolute Memory Size (Table7):** Next, we experimented with the same selection ratio but altered the absolute memory length to assess its impact.\n4. **Importance of Target Token Lengths (Table 8):** Finally, we examined the importance of different target token lengths by varying them while keeping the memory pool size and selected memory size constant.\n\n\n\nTable5: Ablation Study when having fixed memory pool size and different selected memory size\n\n| Target Token Length | Memory Pool Size | Selected Memory Size/Actual Used Memory Size | Selection Ratio | Transformer-XL PPL(↓) | TRAMS PPL(↓) | TRAMS PPL Drops compared with Transformer-XL (↑) |\n| ------------------- | ---------------- | -------------------------------------------- | --------------- | --------------------- | ------------ | ------------------------------------------------ |\n| 64                  | 600              | 150                                          | 25.0%           | 24.49                 | 24.49        | 0                                                |\n| 64                  | 600              | 175                                          | 29.17%          | 24.31                 | 24.19        | **0.12**                                         |\n| 64                  | 600              | 200                                          | 33.3%           | 24.17                 | 24.05        | **0.12**                                         |\n| 64                  | 600              | 225                                          | 37.5%           | 24.05                 | 23.97        | 0.08                                             |\n| 64                  | 600              | 250                                          | 41.7%           | 23.95                 | 23.91        | 0.04                                             |\n| 64                  | 600              | 300                                          | 50.0%           | 23.83                 | 23.82        | 0.01                                             |\n| 64                  | 600              | 350                                          | 58.3%           | 23.75                 | 23.74        | 0.01                                             |\n| 64                  | 600              | 400                                          | 66.7%           | 23.69                 | 23.70        | -0.01                                            |\n| 64                  | 600              | 450                                          | 75%             | 23.67                 | 23.67        | 0                                                |\n| 64                  | 600              | 500                                          | 83.3%           | 23.66                 | 23.64        | 0.02                                             |\n| 64                  | 600              | 550                                          | 91.7%           | **23.65**             | **23.63**    | 0.02                                             |\n| 64                  | 600              | 600                                          | 100.0%          | **23.65**             | 23.65        | 0                                                |\n\nTable6: Ablation Study when having fixed selected memory size and different memory pool size\n\n| Target Token Length | Memory Pool Size | Selected Memory Size/Actual Used Memory Size | Selection Ratio | Transformer-XL PPL(↓) | TRAMS PPL(↓) | TRAMS PPL Drops compared with Transformer-XL (↑) |\n| ------------------- | ---------------- | -------------------------------------------- | --------------- | --------------------- | ------------ | ------------------------------------------------ |\n| 64                  | 200              | 200                                          | 100%            | 24.17                 | 24.17        | 0                                                |\n| 64                  | 250              | 200                                          | 80.0%           | 24.17                 | 24.04        | 0.13                                             |\n| 64                  | 300              | 200                                          | 66.7%           | 24.17                 | 23.99        | 0.18                                             |\n| 64                  | 350              | 200                                          | 57.14%          | 24.17                 | **23.98**    | **0.19**                                         |\n| 64                  | 400              | 200                                          | 50%             | 24.17                 | **23.98**    | **0.19**                                         |\n| 64                  | 450              | 200                                          | 44.4%           | 24.17                 | 24.00        | 0.17                                             |\n| 64                  | 500              | 200                                          | 40%             | 24.17                 | 24.01        | 0.16                                             |\n| 64                  | 550              | 200                                          | 36.4%           | 24.17                 | 24.03        | 0.14                                             |\n| 64                  | 600              | 200                                          | 33.3%           | 24.17                 | 24.05        | 0.12                                             |\n\nTable7: Ablation Study when having fixed selection ratio but different memory pool sizes and selected memory size\n\n| Target Token Length | Memory Pool Size | Selected Memory Size/Actual Used Memory Size | Selection Ratio | Transformer-XL PPL(↓) | TRAMS PPL(↓) | TRAMS PPL Drops compared with Transformer-XL (↑) |\n| ------------------- | ---------------- | -------------------------------------------- | --------------- | --------------------- | ------------ | ------------------------------------------------ |\n| 64                  | 400              | 200                                          | 50%             | 24.17                 | 23.98        | **0.19**                                         |\n| 64                  | 800              | 400                                          | 50%             | 23.69                 | **23.83**    | -0.14                                            |\n| 64                  | 1200             | 600                                          | 50%             | **23.65**             | 23.99        | -0.34                                            |\n| 64                  | 1600             | 800                                          | 50%             | 23.72                 | 24.24        | -0.52                                            |\n\n\n\nTable8: Ablation Study when having different target lengths but with fixed memory pool size and selected memory size\n\n| Target Token Length | Memory Pool Size | Selected Memory Size/Actual Used Memory Size | Selection Ratio | Transformer-XL PPL(↓) | TRAMS PPL(↓) | TRAMS PPL Drops compared with Transformer-XL (↑) |\n| ------------------- | ---------------- | -------------------------------------------- | --------------- | --------------------- | ------------ | ------------------------------------------------ |\n| 16                  | 400              | 200                                          | 50%             | 24.30                 | 24.06        | **0.24**                                         |\n| 32                  | 400              | 200                                          | 50%             | 24.25                 | 24.03        | 0.22                                             |\n| 64                  | 400              | 200                                          | 50%             | 24.17                 | 23.98        | 0.19                                             |\n| 128                 | 400              | 200                                          | 50%             | **24.04**             | **23.90**    | 0.14                                             |\n\n\n\n**Table 5 Insights:**\n\nThe benefits of selecting more tokens to the current TRAMS model diminish as more tokens are included. The top-k selected tokens are most impactful during the inference process when k is relatively small. This finding aligns with the results in Figure 3 of the paper.\n\n**Table 6 Insights:**\n\nWhen using a larger memory pool, there seems to be an optimal level of selection ratio, around 50%, that yields the best performance. To be noticed, the selection ratio is selecting a shared key/value vector for all queries. The selection ratio mentioned by Figure 1 in the paper is selecting unique key/value vector sets for each unique query. With too few memory tokens, it does not provide many choices. With too many, the tokens become outdated and less relevant to the current context. This results in a \"U-shaped\" curve in terms of perplexity reduction.\n\n**Table 7 Insights:**\n\nWhen the selection ratio is held constant, only the top-k tokens (when k is small) contribute significantly to the model's performance. This suggests that the attention distribution in Transformer-XL is rather sharp than flat.\n\n**Table 8 Insights:**\n\nAs the target token length decreases, the token selection process becomes more effective. This can be attributed to the fact that longer target token lengths amplify the influence of position embedding, which can negatively impact the token selection process and make the model make mistakes when choosing the correct tokens.\n\nThese ablation studies offer nuanced insights into various aspects of token selection and memory pool dynamics in TRAMS and Transformer-XL models. We will include these discussions in the revised paper upon acceptance. When having an extra page for the final version, we would draw graphs for those tables to allow reader easy to understand the sweet spot phenomenon and tendency when trying to select hyper-parameters.\n\n# Rebuttal Point5\n\n**_To answer question C:_**\n\n**_“Question C: In L68-L70, \"only 25% of the top memories are essential for maintaining the performance\". As 25% comes from WikiText-103, does the proportion of redundant memory dataset specific?”_**\n\nWe'd like to clarify this issue and I think your suggestion is very useful. \n\n**Clarification on 25% Top Memory Selection Ratio:** Firstly, we have to clarify an ambiguous point in our paper regarding the selection of top memories for each query token. Assume we have 3 query tokens \\{ $Q_1$, $Q_2$, $Q_3$ \\} and 4 key/value tokens \\{ $KV_1$, $KV_2$, $KV_3$, $KV_4$ \\}. We pick the best key/value tokens for each query token—such as  $KV_2$  and  $KV_3$  for  $Q_1$ , $KV_3$ and  $KV_4$ for $Q_2$, and $KV_1$  and $KV_2$  for $Q_3$. We conclude that only 25% of top memories are actually useful. This conclusion is straightforward as it merely requires the computation of the attention map [query_len $\\times$ key\\_len] matrix and the selection of top-$n$ elements in each row.\n\n**TRAMS Memory Selection:** Our proposed TRAMS method simplifies this process. For the same 3 query tokens and 4 key/value tokens, TRAMS selects \\( $KV_2$, $KV_3$ \\), and all queries \\( $Q_1$, $Q_2$, $Q_3$ \\) attend to these two key-value vectors. Instead of computing a full attention map, we compute the norms of all key candidates using our reformulated attention mechanism. While we may not achieve identical performance when using only 25% of the memory, we observe improved performance compared to the Transformer-XL baseline with equivalent memory usage. This shows that our memory selection mechanism positively impacts model performance.\n\n**Memory Selection Comparison:** There are two ways to do memory selection.\n\nThe first one is selecting tokens with highest attention scores (need to compute attention scores). The second one is TRAMS (not need to compute attention scores)\n\n- **Selecting tokens with highest attention scores:**\n\n  - Compute $ QK^T $ matrix\n\n  - Choose a unique key set index for each query token\n\n- **TRAMS:**\n  - Compute the norm of $K'$  using our reformulated attention mechanism\n\n  - Choose one key set index for all query tokens\n\n**Experimental Setup:** We conducted our experiments on the enwik8 dataset to evaluate the performance of computing attention scores in contrast to the TRAMS.\n\n\n\nTable9: Ablation Study on memory redundancy in enwik8 dataset\n\n| target length | overall available memory length | attention score-based memory length | Top Attention Ratio | TRAMS BPC(↓) |\n| ------------- | ------------------------------- | ----------------------------------- | ------------------- | ------------ |\n| 80            | 600                             | 50                                  | 8.30%               | 1.144        |\n| 80            | 600                             | 100                                 | 16.67%              | 1.135        |\n| 80            | 600                             | 200                                 | 33.33%              | 1.128        |\n| 80            | 600                             | 300                                 | 50.00%              | 1.125        |\n| 80            | 600                             | 400                                 | 66.67%              | 1.124        |\n| 80            | 600                             | 600                                 | 100%                | 1.123        |\n\n**Memory Redundancy Ratio Comparison:** Based on the data in Table 9, the optimal proportion of redundant memory size varies depending on the dataset. For the enwik8 dataset, for instance, the ideal level of redundancy is from 16% to 50%. This variation can be attributed to differing levels of **token diversity** across datasets. Take the wikitext-103 dataset, which is used for word-level language modeling tasks, for example. This dataset has a vocabulary size of 267,735, encompassing all the words that appear within it. On the other hand, the enwik8 dataset is used for character-level language modeling and has a much smaller vocabulary size of just 205, including all characters in the dataset. Therefore, the disparity in token diversity between the datasets affects the efficacy of our method, leading to different optimal memory selection ratios for each.\n\nWe will include these discussions in the revised paper upon acceptance. When having an extra page for the final version, we would add a graph of enwik8 dataset memory redundancy in the introduction section to create a stronger motivation.\n\n# Rebuttal Point6\n\n**_To answer question D:_**\n\n**_“Question D(optional): In Section 2.4, how the number of selected memory tokens is decided after ranking? Is it a fixed number? Just curious why topk_num is set to 50 in L356 ([https://anonymous.4open.science/r/TRAMS-768A/utils/modeling_transfo_xl.py#L356](https://anonymous.4open.science/r/TRAMS-768A/utils/modeling_transfo_xl.py#L356)).”_**\n\nThank you for raising the question about implementation details. \n\n\n\n**Ablation Study for Memory Selection:** The number of memory tokens we chose for our experiments was guided by empirical testing. Our underlying intuition suggests that both extremely small and large memory selection ratios are unlikely to provide significant gains over the original Transformer-XL baseline. For a systematic analysis detailing the impact of various hyper-parameters, readers can refer to Tables 5, 6, 7, and 8.\n\n\n\n**Code Implementation:** When it comes to the code implementation, the approach to memory use during inference is bifurcated into two segments. To illustrate, let's assume a Transformer-XL model with a target token length of 150 and an available memory pool size of 600. From this pool, we might opt to select 200 tokens. This two-step process is outlined as follows:\n\n1. Initial Selection: Firstly, we directly select the nearest 150 tokens from the memory pool, aligning with the target length of 150.\n2. Subsequent Selection: Then, we pick an additional 50 tokens from the remaining 450-memory-token pool. This is why the code includes a parameter topk_num=50.\n\n\n\n**Influence of Positional Embeddings:** The rationale for this particular method of memory token selection stems from the Transformer-XL's architecture, specifically its use of relative positional embeddings. To visualize, let's represent the relative position embedding for the nth token as POS_n. The pattern for these embeddings would look like:\n\n| M3   | M2   | M1   | T1   | T2   | T3   | T4   |\n| ---- | ---- | ---- | ---- | ---- | ---- | ---- |\n| POS4 | POS3 | POS2 | POS1 | x    | x    | x    |\n| POS5 | POS4 | POS3 | POS2 | POS1 | x    | x    |\n| POS6 | POS5 | POS4 | POS3 | POS2 | POS1 | x    |\n| POS7 | POS6 | POS5 | POS4 | POS3 | POS2 | POS1 |\n\nIn this arrangement, the position embeddings for the first few \"target length\" tokens in the memory (like M1, M2, and M3) carry significant positional information. Neglecting to include these tokens would be imprudent given their importance in the model's calculations.\n\nWe will include these discussions in the revised paper upon acceptance. When having an extra page for the final version, we would add a positional embedding influence heatmap graph at the appendix. Moreover, we would add a subsection in the discussion section and discuss how many nearest token should we choose.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper provides constructive information which gives readers a clear view of how the proposed method TRAMS is formed. \n2. The proposed memory-selection algorithm is simple to plug into existing transformers-based models, which is novel and effective in improving memory efficiency.", "reasons_to_reject": "1. Table 1 needs more experimental results to convince readers that the proposed method TRAMS improves Transformer-XL. For example, in the WikiText-103 dataset, the PPL of 12L Transformer-XL 24.17 differs from the reported PPL 24.0 in the Transformer-XL (Dai et al., 2019) paper. Please refer to \"Ours - Transformer-XL Standard\" in Table 1 of Dai's work. The dissimilarity comes from different memory sizes, 200 in this work and 150 in the original Transformer-XL (see run_wt103_base.sh: https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/run_wt103_base.sh). Experiments on the proposed method TRAMS with the same memory size as Transformer-XL by datasets are required.", "questions_to_authors": "- $\\textbf{Question A}$: As mentioned in \"Reasons to reject\", what is the PPL of TRAMS when the memory size=150? Same question to enwiki8, what is the BPC of TRAMS when the memory size=512?\n- $\\textbf{Question B}$: Following the above question, is TRAMS sensitive to \"Mem Size\"? Are there any reasons for setting \"Mem Size\" to 200?\n- $\\textbf{Question C}$:  In L68-L70, \"only 25% of the top memories are essential for maintaining the performance\". As 25% comes from WikiText-103, does the proportion of redundant memory dataset specific?\n- $\\textbf{Question D}$(optional): In Section 2.4, how the number of selected memory tokens is decided after ranking? Is it a fixed number? Just curious why `topk_num` is set to 50 in L356 (https://anonymous.4open.science/r/TRAMS-768A/utils/modeling_transfo_xl.py#L356).", "paper_topic_and_main_contributions": "The authors proposed a training-free memory selector which targets improving the memory efficiency of Transformer-XL (Dai et al., 2019)'s. The contribution of this work includes introducing a ranking-based method, selecting top-k memory tokens via $\\text{cos} \\langle K', \\vec{\\mathbb{1}}\\rangle ||K'||$, where $||K'||$ is the reformulated key derived in equation 4. According to the results in Table 1, applying the proposed memory selector to the Transformer-XL reduces the PPL of the Standard Transformer-XL by 0.19 on the WikiText-103 dataset and the BPC of a 12L Transformer-XL by 0.017 on the enwik8 dataset respectively, with the memory size set to 200.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OWtX8q0xzz", "reviewer": "Reviewer_XG1k", "summary": "", "questions": "See Reasons To Reject.  I will reconsider my score in the rebuttal.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal Point1\n\n**_To respond to the first reason to reject:_**\n\n**_“Incomplete Content: Given that only a portion of the paper is available, it's challenging to understand the full context, especially the theoretical foundations and the broader implications of the results. This makes it hard to gauge the complete impact of the paper.”_**\n\nRegarding the **full context** of our paper, we are primarily driven by the observed memory redundancy in the Transformer-XL model. Our main contribution is the introduction of a novel memory selection method tailored for models that utilize transformer attention mechanisms and external working memory during the inference phase. These motivations and contributions are clearly articulated in the introduction, detailed in the methodology section, and their effectiveness is substantiated in the experimental section.\n\nOn the topic of **theoretical foundations**, we have framed the Transformer memory selection challenge as an optimal set selection optimization task. Our proposed TRAMS method is designed as an approximation algorithm to address this optimization challenge. While we acknowledge that our approximation might not come with a theoretical performance guarantee, its efficacy is evident from our empirical results. We'd like to emphasize that our submission is to a conference that focuses on \"empirical\" methods in NLP. As such, we believe it would be disproportionate to dismiss our work solely based on the absence of theoretical guarantees.\n\nLastly, concerning the **broader implications** of our findings, our research introduces a training-free, plug-and-play technique. This is significant, especially when considering that previous studies have seldom treated memory selection as a standalone function rooted in key/value vectors. Our methodology provides a fresh perspective on the representation spaces molded by attention mechanisms. We are confident that our research will pave the way for further innovative studies in this area.\n\n# Rebuttal Point2\n\n**_To respond to the second reason to reject:_**\n\n**_“Lack of Detailed Discussion: While the experimental setup is comprehensive, there seems to be a lack of in-depth discussion on the results. Such a discussion would provide insights into why certain observations were made and the implications of the results.”_**\n\nFirstly, we would like to clarify the depth of our current discussions. Our paper encompasses a thorough analysis, including topics such as _\"how the model benefits from memory selection\", \"the comparative inference cost with Transformer-XL\", and an exploration into whether \"each layer holds equal importance\"_. We are somewhat perplexed about the specific areas of discussion deemed lacking. We kindly request more explicit feedback on the specific discussions or analyses that the reviewer found absent. We are more than willing to undertake additional experiments or discussions to address these gaps. For instance, Reviewer ncPF expressed concerns about the sensitivity of our configuration setting. To address this, we have incorporated relevant results. We direct the reviewer to our response to question B in Reviewer ncPF's section, where we delve into a detailed ablation study concerning hyper-parameter selection.\n\n# Rebuttal Point3\n\n**_To respond to the third reason to reject:_**\n\n**_“Potential Over-reliance on Experimental Results: While empirical evaluations are important, it would be beneficial if the paper also provided theoretical justifications or proofs for the proposed method. This would strengthen the paper's claims and make the results more robust.”_**\n\nFrankly speaking, our approach is rooted in intuitive reasoning, focusing on entity-based tokens at the semantic space's extremities. Theoretically establishing a gap boundary is challenging for our context. We've conducted rigorous empirical evaluations, aligning with the conference's focus on empirical methods in NLP. While theoretical proofs are valuable, our empirical validations offer significant insights in the dynamic domain of NLP. We believe our empirical approach provides robust contributions to the field.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. **Comprehensive Experimental Setup**: The paper provides a detailed experimental setup which allows readers to understand the conditions under which the proposed method was tested. This ensures reproducibility and gives credibility to the results presented.\n2. **Comparison with State-of-the-art**: The authors have taken the effort to compare their method with existing techniques. This is crucial for the research community to understand the relative advantages of the new method.\n3. **Clarity**: From the sections provided, it appears that the authors have made an effort to be clear in their presentation, making it accessible to readers who might be new to the topic.", "reasons_to_reject": "1. **Incomplete Content**: Given that only a portion of the paper is available, it's challenging to understand the full context, especially the theoretical foundations and the broader implications of the results. This makes it hard to gauge the complete impact of the paper.\n2. **Lack of Detailed Discussion**: While the experimental setup is comprehensive, there seems to be a lack of in-depth discussion on the results. Such a discussion would provide insights into why certain observations were made and the implications of the results.\n3. **Potential Over-reliance on Experimental Results**: While empirical evaluations are important, it would be beneficial if the paper also provided theoretical justifications or proofs for the proposed method. This would strengthen the paper's claims and make the results more robust.", "questions_to_authors": "See Reasons To Reject.  I will reconsider my score in the rebuttal.", "paper_topic_and_main_contributions": "1. Introduction of a novel approach or algorithm for memory selection without the need for training.\n2. Extensive experimental setup to evaluate the performance of the proposed method.\n3. A comparison of the proposed method with existing state-of-the-art techniques, showing its effectiveness.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hY5l87Raiz", "reviewer": "Reviewer_Rdm5", "summary": "", "questions": "A. How do you see this work fitting it with LLMs of today? Do you see a way to integrate any of the conclusions of your work to more modern architectures such as LLaMA, Falcon, etc?\n\nB. Do you think that it is possible to reduce the additional GPU memory and compute time required by your heuristic? As it stands, it seems that the benefit from your method in performance is very little in comparison to the cost.\n\nC. Why only 200 memory tokens were examined?\n\nD. Is it possible with your method to get better results using less memory tokens?", "limitations": "", "rating": 3, "confidence": 2, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for providing valuable suggestions and encouraging us to provide more details related to the experiments.\n\n# Rebuttal Point1\n\n**_To respond to the first reason to reject:_**\n\n**_“Experiments show a very narrow image of the benefits of the proposed method. A single architecture was used with memory size of 200 tokens.”_**\n\nThank you for highlighting this concern. We'd like to address the mentioned issues in the following aspects:\n\n**A Single Architecture**: We believe Transformer-XL is a well-known architecture in the long-context language modeling field, comparing all related works. Thus, choosing Transformer-XL as the baseline can better show the efficiency of our method. Given the page limit of a short paper, we want to show the most important results in our current experiment section. TRAMS could potentially be layered atop more advanced versions of Memory Transformers, such as LaMemo and Infinity-Former. If our work is accepted with one more page, we are willing to add such results.\n\n**Memory Size Sensitivity**: As for the memory size, we have added some sensitivity analysis. Please kindly refer to our answer (Part 2) to question B for Reviewer ncPF (**Rebuttal Point4 for Reviewer ncPF**) to check the comprehensive analysis.\n\n**Case Study**: To give a broader image of our benefits beyond the quantitative experimental details, we have also undertaken qualitative experiments in the form of case studies. Specifically, we wish to shed light on the types of memory tokens that are predominantly chosen during the memory selection process. **Two examples randomly sampled from the wikitext-103 dataset are shown at the end of our rebuttal.** Based on our case studies, we observed that named entities such as \"John Deed\" and \"Du Fu\" are more likely to be selected over high-frequency words like \"You,\" \"he,\" and \"me.\" This observation underlines how our method enhances performance. Specifically, low-frequency entities within the target tokens benefit from memory tokens that are also low-frequency entities. This suggests that our memory selection mechanism not only improves model performance but also adds an interpretative layer to the model's decisions.This observation underlines how our method enhances performance. Specifically, low-frequency entities within the target tokens benefit from memory tokens that are also low-frequency entities. This suggests that our memory selection mechanism not only improves model performance but also adds an interpretative layer to the model's decisions.\n\n[1] Haozhe Ji et. al: LaMemo: Language Modeling with Look-Ahead Memory [[https://arxiv.org/abs/2204.07341](https://arxiv.org/abs/2204.07341)] \n\n[2] Zihang Dai et al: Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context [[https://arxiv.org/abs/1901.02860](https://arxiv.org/abs/1901.02860)] \n\n[3] Pedro Henrique Martins et al: ∞-former: Infinite Memory Transformer [[https://arxiv.org/abs/2109.00301](https://arxiv.org/abs/2109.00301)] \n\n# Rebuttal Point2\n\n**_To respond to the second reason to reject (part1):_**\n\n**_“(Unable to validate baseline results since) authors used an unorthodox Tranformer-XL architecture that wasn't mentioned in the original paper”_**\n\nThank you for bringing up the concern regarding the validity of our baseline reproduction.\n\n**Comparison of Architecture:** Firstly, it's important to note that the reformulated transformer attention proposed in our paper maintains the same parameters and architectural elements as the standard Transformer-XL model. The distinction lies solely in the computation order.\n\n**Computation Order Change:** In the standard Transformer-XL, hidden states for each layer are initially mapped as Key and Query. Following this, the model calculates the dot product between the Key and Query. However, in our reformulated attention mechanism, we project the hidden states based on  $W_k$  and  $W_q$ , allowing us to calculate  $hW_k^TW_q$  before proceeding to compute the dot product with the original hidden states $h$.\n\n**Purpose of the Conceptual Change:** This reordering in computation enables us to develop our unique scoring method based solely on $hW_k^TW_q$, rather than requiring both $hW_q^T$  and $hW_k^T$. In essence, this adjustment is a conceptual modification rather than an architectural one. Through this conceptual change, we can naturally introduce a scoring function that relies on the norm computation of a single vector $hW_k^TW_q$.\n\n# Rebuttal Point3\n\n**_To respond to the second reason to reject (part2):_**\n\n**_“Unable to validate baseline results (since authors used an unorthodox Tranformer-XL architecture that wasn't mentioned in the original paper)”_**\n\nThanks for pointing out this important baseline validation problem.\n\nTo respond to the baseline validation, this question is accidentally overlapped with reject reason one mentioned by reviewer ncPF (**Rebuttal Point1 for Reviewer ncPF**). We provide thorough experimental results (Table 1 and Table 2 mentioned in response to the first reason to reject in review1) for fair comparison. This serves to validate the performance of our baseline model.\n\n# Rebuttal Point4\n\n**_To respond to the third reason to reject:_**\n\n**_“Insufficient comparison to alternative methods mentioned in the introduction and related work.”_**\n\nThank you for highlighting the need for a more exhaustive comparison with alternative methods in the introduction and related works. We organize alternative methods we found besides those mentioned in the paper during the rebuttal period. We would have a more comprehensive comparison for the introduction and related works in our final version.\n\n**Method Classification:** Recent developments in this field have been abundant, prompting us to categorize related works into distinct approaches. The first general strategy aims to create more efficient Transformer architectures capable of accommodating longer context inputs. The second employs additional components, such as token or memory corpuses, to enhance the model's ability to capture long-distance dependencies. We present some alternative methods discovered during the rebuttal period for further context.\n\n(Tay et al.) provides a clear roadmap for different Transformer architecture variants. In its roadmap, Transformer-XL is classified as one type of recurrence-augmented Transformer architecture.\n\n**Extending Context Window:** Tay et al. offer a comprehensive roadmap delineating various Transformer architecture variants. Within this framework, Transformer-XL is categorized as a recurrence-augmented type. Numerous works like Funnel Transformer (Dai et al.), Poolingformer (Zhang et al.), Longformer (Beltagy et al.), Big Bird (Zaheer et al.), and ETC (Ainslie et al.) propose methods centered around memory selection or downsampling to create more adaptive Transformer architectures. These can be fine-tuned for tasks requiring longer context windows. Additionally, innovative projects like Random Feature Attention (Peng et al.), Linear Randomized Attention (Zheng et al.), and EVA (Zheng et al.) focus on optimizing attention modules within the Transformer framework. Recent advancements such as Focused Transformer (Tworkowski et al.) and Position Interpolation (Chen et al.) leverage well pre-trained large language models to extend context windows. These initiatives enable the models to assimilate more information from extended contextual windows.\n\n**Utilizing External Module:** Another avenue to enhance long-distance dependency recognition involves external memory modules. KNN-LM (Khandelwal et al.) and Memorizing Transformer (Wu et al.) utilize k-nearest-neighbor algorithms to bolster the capabilities of the original Transformer model. Unlimiformer (Bertsch et al.) employs a modified version of the Memorizing Transformer to extend the model's context capacity indefinitely with the aid of external memory storage.\n\n**State Space Research:** A related but distinct research trajectory that deserves mention focuses on the use of structured state spaces explicitly designed for long-context modeling (Gu et al.).\n\n\n\n**Alternative Methods:**\n\n[1] Yi Tay et al. Efficient Transformer: A Survey [[https://arxiv.org/abs/2009.06732](https://arxiv.org/abs/2009.06732)] \n\n[2] Zihang Dai et al. Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing [[https://arxiv.org/pdf/2006.03236.pdf](https://arxiv.org/pdf/2006.03236.pdf)]\n\n[3] Hang Zhang et al. Poolingformer: Long Document Modeling with Pooling Attention [[https://arxiv.org/pdf/2105.04371.pdf](https://arxiv.org/pdf/2105.04371.pdf)] \n\n[4] Iz Beltagy et al. Longformer: The Long-Document Transformer [[https://arxiv.org/pdf/2004.05150.pdf](https://arxiv.org/pdf/2004.05150.pdf)] \n\n[5] Manzil Zaheer et al. Big Bird: Transformers for Longer Sequences [[https://arxiv.org/abs/2007.14062](https://arxiv.org/abs/2007.14062)] \n\n[6] Joshua Ainslie et al. ETC: Encoding Long and Structured Inputs in Transformers [[https://arxiv.org/pdf/2004.08483.pdf](https://arxiv.org/pdf/2004.08483.pdf)] \n\n[7] Szymon Tworkowski et al. Focused Transformer: Contrastive Training for Context Scaling [[https://arxiv.org/abs/2307.03170](https://arxiv.org/abs/2307.03170)] \n\n[8] Shouyuan Chen et al. Extending Context Window of Large Language Models via Positional Interpolation [[https://arxiv.org/abs/2306.15595](https://arxiv.org/abs/2306.15595)] \n\n[9] Hao Peng et al. Random Feature Attention [[https://arxiv.org/abs/2103.02143](https://arxiv.org/abs/2103.02143)] \n\n[10] Lin Zheng et al. Efficient Attention via Control Variates [[https://arxiv.org/abs/2302.04542](https://arxiv.org/abs/2302.04542)] \n\n[11] Lin Zheng et al. Linear Complexity Randomized Self-attention Mechanism [[https://arxiv.org/abs/2204.04667](https://arxiv.org/abs/2204.04667)] \n\n[12] Albert Gu et al. Efficiently Modeling Long Sequences with Structured State Spaces [[https://arxiv.org/pdf/2111.00396.pdf](https://arxiv.org/pdf/2111.00396.pdf)] \n\n[13] Urvashi Khandelwal et al. Generalization through Memorization: Nearest Neighbor Language Models [[https://arxiv.org/abs/1911.00172](https://arxiv.org/abs/1911.00172)] \n\n[14] Amanda Bertsch et al. Unlimiformer: Long-Range Transformers with Unlimited Length Input [[https://arxiv.org/abs/2305.01625](https://arxiv.org/abs/2305.01625)] \n\n[15] Yuhuai Wu et al. Memorizing Transformers [[https://arxiv.org/abs/2203.08913](https://arxiv.org/abs/2203.08913)] \n\n# Rebuttal Point5\n\n**_To answer Question A:_**\n\n**_A. How do you see this work fitting it with LLMs of today? Do you see a way to integrate any of the conclusions of your work to more modern architectures such as LLaMA, Falcon, etc?_**\n\nThank you for your insightful comments on the integration of TRAMS and LLMs.\n\nOur model can still be useful in the era of LLMs for the following reasons:\n\n**Easy Implementation with No Extra Training Required:** A common challenge in adapting new methods to LLMs is that they often necessitate additional training designs. Our approach, however, eliminates this need. Specifically designed to focus on memory selection, our TRAMS method can be directly plugged into current LLMs for either pre-training, supervised finetuning, or only the inference stage.\n\n**Memory Selection Remains Crucial:** Current LLMs need to allow long context as inputs and long text as outputs. Both features need to involve potentially large attention matrix computations. Therefore, it is important to investigate memory selection techniques such as our work,  to allow useful information to be selected from the long context/generated output.\n\n**Improved Representational Capabilities of LLMs**: The increased potency of LLMs suggests that they offer a superior representational space compared to smaller models like the Transformer-XL. With a more structured semantic landscape, we have a better chance of performing selections based on individual representation vectors without having to rely on query representation.\n\nIt can be a promising direction for our future work.\n\n# Rebuttal Point6\n\n**_To answer Question B:_**\n\n**_“Question B. Do you think that it is possible to reduce the additional GPU memory and compute time required by your heuristic? As it stands, it seems that the benefit from your method in performance is very little in comparison to the cost.”_**\n\nThanks for pointing out the issue of additional cost.\n\n**Space and Time Complexity:** To explain Table 2 in the original paper in more detail, the main difference between TRAMS and Transformer-XL is that , for space complexity, TRAMS requires a larger memory pool compared with Transformer-XL for each layer before the attention mechanism. For time complexity, TRAMS requires additional computation of the norm of each reformulated key vector when compared with Transformer-XL. Therefore, TRAMS brings an extra space complexity of  $\\mathcal{O}((M-m) \\times L)$, where  $M$ represents overall memory pool size for TRAMS,  $m$ represents memory length for Transformer-XL and $L$ stands for the number of layers. TRAMS brings an extra time complexity of $\\mathcal{O}(M \\times L)$, where $M$ represents memory pool size and $L$ stands for the number of layers.\n\nHowever, we think it is possible to reduce the additional GPU memory and compute time required by my heuristic. Therefore, we propose two methods that can reduce the GPU memory and compute time:\n\n**Selective Layer Memory Selection:** As indicated in Figure 4 of our paper, the last few layers exert a greater influence on performance than the preceding layers. One strategy to improve is to perform memory selection only on the last few layers. This change reduces the time complexity from $\\mathcal{O}(m \\times L)$ to $\\mathcal{O}(m \\times l)$, where $l$ denotes the selected layers and $L$ represents all the transformer layers.\n\n\n\nTable 10: GPU and Wall-Clock Time for Selective Layer Memory Management\n\n| Model Setting                                 | GPU Peak Memory (MB) | Wall-clock Time (s) | Perplexity |\n| --------------------------------------------- | -------------------- | ------------------- | ---------- |\n| Transformer-XL(mentioned in the paper)        | 3529                 | 33.27               | 24.17      |\n| TRAMS (on last five layers)                   | 3590                 | 43.95               | 24.00      |\n| TRMS (on all layers) (mentioned in the paper) | 3719                 | 49.55               | 23.98      |\n\n**Dynamic Online Memory Selection:** Another method to further minimize GPU memory usage involves dynamic memory management. We can select the top-$k$ tokens from the target length tokens and maintain a dynamic list of these top-$k$ tokens. It reduces the time complexity from $\\mathcal{O}(m)$ to $\\mathcal{O}(n)$, where $m$ stands for memory length and $n$ represents the target length. We leave the implementation and discussion of this online memory selection method as part of our future work due to limited rebuttal time.\n\nWe will include this discussion in the final version.\n\n# Rebuttal Point7\n\n**_To answer Question C:_**\n\n**_Question C. Why only 200 memory tokens were examined?_**\n\n Please kindly refer to our answer to question B in reviewer ncPF (**Rebuttal Point4 for Reviewer ncPF**).\n\n# Rebuttal Point8\n\n**_To answer Question D:_**\n\n**_Question D. Is it possible with your method to get better results using less memory tokens?_**\n\nThanks for pointing out a valuable direction for future work.\n\n**Conclusion:** It is not easy to accomplish it with the current TRAMS method. First, the Transformer-XL architecture is sensitive to a small memory (like around 100). However, a better scoring function for memory selection may help, which is promising for further exploration. We provide a detailed analysis below:\n\n\n\nTable11: Discussion when continuously decreasing the memory size\n\n| Target Token Length | Memory Pool Size | Selected Memory Size/Actual Used Memory Size | Transformer-XL PPL(↓) | TRAMS PPL(↓) | TRAMS PPL drops compared with Transformer-XL (↑) |\n| ------------------- | ---------------- | -------------------------------------------- | --------------------- | ------------ | ------------------------------------------------ |\n| 64                  | 600              | 200                                          | 24.17                 | 24.05        | 0.12                                             |\n| 64                  | 600              | 160                                          | 24.41                 | 24.27        | 0.24                                             |\n| 64                  | 600              | 100                                          | 25.08                 | 24.95        | 0.13                                             |\n\n**Transformer-XL Sensitivity**: From Table 11, we find that when the memory size changes from 200 to 160, the TRAMS perplexity change increases by 0.22 points. The TRAMS method can use memory selection to improve performance, but we cannot select too few tokens. The reason is that when the memory size is very small, like 100, the nearest tokens are almost the most useful tokens for generation. The benefits of extracting one token from the context history are smaller than those of remaining with the nearest token.\n\n**Future Work**: We believe that there should be a way to develop a better scoring function for memory selection such that we can select fewer memory tokens and get better performance at the same time. For example, we can directly optimize the set selection problem, which is to identify a subset of memory tokens that maximizes the remaining attention value. However, it is not trivial to find a solution for this optimization problem and we leave it as our future work.\n\n# Rebuttal Point9\n\n**_To respond to missing reference:_**\n\n**_“Missing References:_**\n\n**_I believe that the following works are also related to this paper and would like to hear the authors opinion of how they relate to them:_**\n\n**_A. Generating Long Sequences with Sparse Transformers_**\n\n**_B. Fast Transformer Decoding: One Write-Head is All You Need”_**\n\nThanks for providing the missing references, we will include them in the revised paper. Furthermore, we will provide a more comprehensive comparison with the two papers in the revised paper.\n\nWhen it comes to the relationship between our paper and theirs, both papers propose new Transformer architectures for longer and faster generation tasks. Both of them are super interesting and related to our paper’s topic. \n\nThe “**Sparse Transformer**” paper also focuses on memory selection via masking and targets at providing a more efficient transformer architecture. One major difference is that Sparse Transformer is trying to give different masks for each different query. Moreover, the selected memory directly participates in the forward-backward training process. Therefore, it is a sparse architecture instead of a “memory selection” process that selects side information in an additional part of the original model.\n\nThe “**Fast Transformer Decoding**” paper proposes a variant of attention called multi-query attention that shares a single set of keys and values. Similar to ours, this paper tries to propose a new attention architecture. The major difference between ours and theirs is that our motivation is the memory redundancy phenomenon discovered in Transformer-XL architecture. Their motivation is that limited memory-bandwidth would cause repeated loading of key/value tensors and slow inference. \n\n\n\n# Rebuttal Appendix\n\nWe listed two case study examples that are mentioned in the **Rebuttal for current Reviewer** here.\n\n**Case1:**\n\n**Sequence segments that in the memory pool (top-k memory tokens are emphasized here):**\n\n_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 the &lt;unk> Chocolate Factory in London . He was directed by **John** Tiffany and starred alongside Ben **Whishaw** , Shane Zaza , Harry Kent , Fraser Ayres , Sophie Stanton and Dominic Hall . &lt;eos> In 2006 , Boulter starred alongside **Whishaw** in the play Citizenship written by Mark **Ravenhill**. He appeared on a 2006 episode of the television series , Doctors , followed by a role in the 2007 theatre production of How to Curse directed by Josie Rourke . How to Curse was performed at Bush Theatre in the London Borough of Hammersmith and Fulham . Boulter starred in two films in 2008 , Daylight Robbery by filmmaker **Paris** **&lt;unk>**, and Donkey Punch directed by Olly Blackburn . In May 2008 , Boulter made a guest appearance on a two @-@ part episode arc of the television series Waking the **Dead** , followed by an appearance on the television series Survivors in November 2008 . He had a recurring role in ten episodes of the television series Casualty in 2010 , as \" Kieron Fletcher \" . Boulter starred in the 2011 film Mercenaries directed by Paris &lt;unk> . &lt;eos> &lt;eos> = = Career = = &lt;eos> &lt;eos> &lt;eos> = = = 2000 – 2005 = = = &lt;eos> &lt;eos> In 2000 Boulter had a guest @-@ starring role on the television series The Bill ; he portrayed \" Scott Parry \" in the episode , \" In Safe Hands \" . Boulter starred as \" Scott \" in the play Herons written by Simon Stephens , which was performed in 2001 at the Royal Court Theatre . A review of Boulter 's performance in The Independent on Sunday described him as \" horribly menacing \" in the role , and he received critical reviews in The Herald , and Evening Standard . He appeared_\n\n**Sequence segments that we are handling (target tokens that can be potentially benefited from the selected memories are highlighted):**\n\n_in the television series Judge **John Deed** in 2002 as \" **&lt;unk>** Armitage \" in the episode \" Political &lt;unk> \" , and had a role as a different character \" Toby Steele \" on The Bill . &lt;eos> He had a recurring role in 2003 on two episodes of The Bill , as character \" Connor Price \" . In 2004 **Boulter** landed_\n\n**Case2:**\n\n**Sequence segments that in the memory pool (top-k memory tokens are emphasized here):**\n\n_criticized old commentaries **during** _the _Yuan Dynasty were too unfathomable . Matsuo **Bashō** , the greatest haiku poet , was also strongly influenced by Du Fu ;**in** Oku no **Hosomichi** , his masterpiece , he cites the first two lines of A Spring **View** ( &lt;unk> **)** before a haiku as its introduction and also many of his other haiku have similar wording and themes . It is said that when he died in Osaka **during** a long travel , a copy of Du Fu 's poetry was found with him as one of a few precious items which he was able to carry around . **&lt;eos> &lt;eos>** = = Translation = = &lt;eos> &lt;eos> A variety of styles have been used in efforts to translate **Du Fu** 's work into **English** . As Burton Watson remarks in The Selected Poems of Du **Fu** , \" There are many different ways to approach the problems involved in translating Du Fu , which is why we need as many different translations as possible \" ( p. &lt;unk> ) . The translators have **had** to contend with bringing out the formal constraints of the original without sounding laboured to a Western ear ( particularly when translating regulated **verse** , or &lt;unk> ) , and accommodating the complex allusions contained particularly in the later works ( Hawkes writes that \" his poems do not as a rule come through very well in translation \" — p. ix ) . One extreme on each issue is represented by Kenneth Rexroth 's One Hundred Poems From the Chinese . His are free translations , which seek to conceal the &lt;unk> through &lt;unk> and expansion and contraction of the content ; his responses to the allusions are firstly to omit most of these poems from his selection , and secondly to \" translate out \" the references in those works which he does select . &lt;eos> Other translators have placed much greater weight on trying to convey a sense of the poetic forms used by Du Fu . Vikram Seth in Three Chinese Poets uses English @-@ style rhyme schemes , whereas Keith Holyoak in Facing the Moon approximates the Chinese rhyme scheme ; both use end @-@ stopped lines and preserve some degree of parallelism . In The Selected Poems of Du Fu , Burton Watson follows the &lt;unk> quite strictly , persuading the_\n\n**Sequence segments that we are handling (target tokens that can be potentially benefited from the selected memories are highlighted):**\n\n_western reader to adapt to the poems rather than vice versa . Similarly , he deals with the allusion of the later works by combining literal translation with extensive annotation . &lt;eos> In 2015 , Stephen Owen published translations , with facing Chinese texts , of the complete poetry of **Du Fu** in six volumes , with extensive scholarly apparatus , which emphasized &lt;unk>_", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Paper is well written and clear, experiments are well presented and proper statistics tests were used to strengthen the claims of the paper\n2. The paper show that with proper heuristics you can improve an already pretrained model by making better use of its context memory. Also, it shows that most tokens that are saved in memory are redundant and can be discarded saving precious GPU memory and memory transactions.", "reasons_to_reject": "1. Experiments show a very narrow image of the benefits of the proposed method. A single architecture was used with memory size of 200 tokens.\n2. Unable to validate baseline results since authors used an unorthodox Tranformer-XL architecture that wasn't mentioned in the original paper\n3. Insufficient comparison to alternative methods mentioned in the introduction and related work.", "questions_to_authors": "A. How do you see this work fitting it with LLMs of today? Do you see a way to integrate any of the conclusions of your work to more modern architectures such as LLaMA, Falcon, etc?\n\nB. Do you think that it is possible to reduce the additional GPU memory and compute time required by your heuristic? As it stands, it seems that the benefit from your method in performance is very little in comparison to the cost.\n\nC. Why only 200 memory tokens were examined?\n\nD. Is it possible with your method to get better results using less memory tokens?", "paper_topic_and_main_contributions": "The paper offers a new heuristic, based on theory and empirical that supports, for calculating which tokens should be saved in memory when using a Transformer-XL model to improve model performance without requiring any form of retraining or calibration of the model. The authors show that their new heuristic improves the model performance on 2 key datasets, WikiText-103 and enwik8.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "Table 1: The first row model field (Transformer+RPE (Vaswani et al., 2017)) is out of bounds and make the table unclear for reading.\n\nCode link is not working"}, "extra_text": {"Missing_References": "I believe that the following works are also related to this paper and would like to hear the authors opinion of how they relate to them:\n\nA. [Generating Long Sequences with Sparse Transformers](https://arxiv.org/abs/1904.10509)\n\nB. [Fast Transformer Decoding: One Write-Head is All You Need](https://arxiv.org/abs/1911.02150)"}}]

Author rebuttal:

Meta-review: This paper proposes a training-free memory selector to improve memory efficiency in Transformer-XL models, which includes a ranking-based method to select the top-k memory tokens. The proposed method, TRAMS, is demonstrated to be effective in improving memory efficiency through experiments. 

Strengths:
* The paper is well-structured and clearly explains the proposed method.
* TRAMS is a novel and effective memory-selection algorithm that is easy to integrate with existing transformer-based models.
* The experimental setup is comprehensive and the authors compare their method with state-of-the-art techniques.

Weaknesses:
* There is a lack of detailed discussion on the results and their implications.
* The paper would benefit from more experimental results and theoretical justifications to support the claims.
* Incomplete content makes it difficult to understand the full context and impact of the paper.

Overall speaking, this paper gives some insightful notes on the method and has clear goodness. However, details and many contents need to be enhanced.
Decision: Accept-Findings