SlowGuess/NAACL_2022_OCR huggingface dataset:
We follow the standard architecture of the Transformer-XL (Dai et al., 2019) that has different configurations for different tasks. Specifically, on Wikitext-103, we use a 16-layer Transformer with 10 attention heads and head dimension 41 equipped with adaptive embeddings (Baevski and Auli, 2019). We control the target sequence length to be 150 and the memory length 150 for all models following the setting of Dai et al. (2019). For the Compressive Transformer and  $\infty$ -former, we additionally use an external memory of size 150 following the setting of Martins et al. (2021). On the text8 and enwik8 datasets, we use a 12-layer Transformer with 8 heads and head dimension 64. The length of the target sequence and the recurrence memory are both set to 512. In the main results we use the identical evaluation setting to the training phase on all datasets and do not use a longer memory. We use the Pytorch framework (Paszke et al., 2019) and Apex for mixed-precision training. In practice, we found that calculating the exponentials (\$3.2) may lead to numerical overflow in mixed-precision mode, so we compute the logarithm of the exponential sum using logsumexp and logaddexp operator. Further details of the dataset and the hyperparameter settings are described in the Appendix C.
We show the results of word-level language modeling benchmark Wikitext-103 in Table 1. We first observe that all the models extended with memories significantly outperforms Transformer+RPE. Under the same memory length, LaMemo outperforms Transformer-XL with a clear margin, which demonstrates the effectiveness of learning dynamic memory representations over static ones. When
Table 1: Word-level language modeling results on Wikitext-103. We report ppl (perplexity) on dev and test set. We also report the number of parameters, memory size, external memory size, and the number of FLOPS (floating-point operations) for computing one step prediction on average.
Model	dev bpc	test bpc
Dataset: text8
Transformer+RPE	1.232	1.303
Transformer-XL (Dai et al., 2019)	1.172	1.239
LaMemo	1.128	1.196
Dataset: enwik8
Transformer+RPE	1.253	1.240
Transformer-XL (Dai et al., 2019)	1.150	1.128
LaMemo	1.129	1.107
Table 2: Character-level language modeling results on text8 and enwik8. We report bpc (bits-per-character) on the dev and test set.
compared to the compressive memory and the unbounded memory that take longer contexts into account, LaMemo still achieves lower perplexity. This indicates that the look-ahead memory allows the language model to exploit the recent contexts to gain performance, while simply increasing the context length yields marginal improvement. This is in accordance with previous findings of how language models utilize contexts (Khandelwal et al., 2018; Sun et al., 2021). In terms of the parameters, LaMemo has the same number of parameters as the Transformer-XL while other baselines use additional parameters in CNN to compress or smooth the hidden states. Lastly, we show the number of FLOPS necessary for computing one step prediction.  $\infty$ -former has the highest number of FLOPS for resampling enough points from the continuous signal to update the memory using smoothing techniques. LaMemo also incurs additional computations to re-contextualize the memory under the current context. Note that although the Compressive Transformer has lower number of FLOPS than LaMemo, it has an external memory that consumes more GPU memory.
We also present the results of character-level language modeling on text8 and enwik8 datasets in Table 2. We observe similar trends as the results on
5The external memory consists of 150 compressed vectors for Compressive Transformer, and 150 radial basis functions for  $\infty$ -former respectively.
Configuration	Encoding	dev ppl	test ppl
Full	Ours	22.98	23.77
w/o mem interp	Ours	23.67	24.90
w/o look-ahead	Ours	23.42	24.56
Full	Dai et al. (2019)	FAIL	FAIL
Table 3: Ablation study on Wikitext-103. We investigate three model configurations and two encoding schemes.
the word-level benchmark, where LaMemo outperforms Transformer-XL by 0.04 on text8 and 0.02 on enwik8 with the same context length. Additionally, we observe that all models exhibit overfitting on text8, which might be caused by the extremely small vocabulary size of the dataset.
We conduct ablation studies on Wikitext-103 to examine the effects of the proposed techniques, i.e., look-ahead attention, memory interpolation, and disentangled relative positional encodings.
We use the same model architecture and the same target and memory length as the main results. We first study three configurations, including (1) using the Full model setting, (2) ablating the memory interpolation module (w/o mem interp), i.e., set the memorizing coeffcient  $\alpha_{\tau - 1} = 0$ , and (3) ablating the look-ahead attention (w/o look-ahead), i.e., only use the causal context representations  $C_{\tau - 1}^{\rightarrow}$  in each layer. As shown in the First three rows in Table 3, both the memory interpolation and the look-ahead attention are indispensable for achieving the best performance. Additionally, we found that cancelling out memory interpolation leads to a worse performance, which indicates that the distant past still provides additional information beyond the current context.
The second study targets at studying different encoding schemes. We substitute our encodings with the RPE of Transformer-XL Dai et al. (2019) and run multiple experiments with 3 different random seeds, but all the models fail to converge. We plot the training curves using two encodings in Figure 8 in Appendix B, where we observe that our disentangled RPE is more stable during training and achieves lower perplexity.
5 Extrapolating to Longer Contexts
In this section, we extrapolate the models to longer contexts during inference to study the effect of dynamic contextualization to the distant past.
Figure 5: Test perplexity of LaMemo and Transformer-XL when extrapolating to longer contexts during inference, where  $m$  is the ratio of the memory length to the target length.
Figure 6: The memorizing coefficient  $\alpha$  of different layers in a 16-layer model with a same memory and target length of 150. Smaller index means older memory.
We fix the length of the target sequence to 64 and extrapolate the trained models to longer memory length  $64 \times m$  during inference, where  $m = 1, \dots, 10$ . We compare the perplexity of LaMemo and Transformer-XL trained on Wikitext-103 when augmented by a memory with different length. As shown in Figure 5, LaMemo consistently achieves lower perplexity than Transformer-XL when extrapolating to longer contexts, while the performance of both models saturate when  $m$  is over 7. Additionally, we observe that the gap of perplexity between the two models increases when taking longer contexts into account. This demonstrates the effectiveness of dynamically refreshing the distant memory representations under the current context.
6 Attention Analysis
In this section, we analyze the attention distribution of LaMemo to validate the effectiveness of utilizing bi-directional contexts with look-ahead attention.
We first visualize the memorizing coefficient  $\alpha$  which stands for the portion of the past activations in the current memory representations. As show in Figure 6, we plot  $\alpha$  in different layers as a function of the memory index averaged on 100 text segments. We observe that in lower layers the memory mainly attends to the past ( $\alpha \approx 1.0$ ). We conjecture that long-term bi-directionality is not necessary for low-level representations such as lexical features. In higher layers, the memory substantially utilizes the future contents to refresh the high-level representations, especially for the old memory state with a small memory index.
Next, we visualize the attention weight distribution on the context tokens when predicting each target token in Figure 1. For every token, we take the maximal attention weight in each interval of 5 tokens on its left and scale to a context length of 100. The result indicates that LaMemo learns better memory representations by attending to the right-side tokens, which increases the memory utilization when predicting the target token.
We present the generated texts of LaMemo and Transformer-XL trained on Wikitext-103 in Appendix D. Both models maintain a memory size of 512, and we seed them with the same context randomly sampled from the test set and generate 256 tokens using top-p sampling (Holtzman et al., 2020) with  $p = 0.95$ .
The Transformer (Vaswani et al., 2017), with its pair-wise modeling ability of the input, becomes prevailing for sequence modeling, especially long sequence processing tasks, such as long text generation (Tan et al., 2021; Ji and Huang, 2021), long document QA (Beltagy et al., 2020; Ainslie et al., 2020), language modeling (Dai et al., 2019; Rae et al., 2020), video processing (Wu et al., 2019), and etc. Specifically, language modeling (Merit et al., 2017) which requires processing documents with thousands of tokens has become a natural
testbed for benchmarking this long-term processing ability. However, due to the quadratic time and space complexity of self-attention, scaling to inputs with thousands of tokens is computationally prohibitive.
One line of work investigated the linear-time attention mechanism to mitigate the scability issue of Transformer. Linformer (Wang et al., 2020) projects the inputs to lower dimension in length and approximates the full attention with a low-rank factorization. Linear Transformer (Katharopoulos et al., 2020) regards the self-attention as a kernel function and uses a linear dot-product as a substitute. Choromanski et al. (2021) and Peng et al. (2021) proposed to approximate the softmax more precisely with the expectation of the dot-product of random features. Although achieving substantial improvements on benchmarks designated for long inputs (Tay et al., 2021). These methods, however, focus on approximating the full attention with low-rank factorizations or kernel functions, which compromise the expressiveness and robustness of the original softmax attention, are reported to be inferior to the simple local attentions on real world language processing tasks (Xiong et al., 2021).
Our work falls in another line, which augments the Transformer with a parametrized memory to store critical history information. Memory-augmented networks (Graves et al., 2014; Weston et al., 2015; Sukhbaatar et al., 2015) have been studied in the context of recurrent neural networks for a long time, but are mostly restricted to small and synthetic datasets. With the rapid development of Transformer, various works start to adapt memories to this architecture.
Dai et al. (2019) first extended Transformer with a recurrence memory that caches hidden states computed in previous steps for the target tokens to attend to. Rae et al. (2020) further extended the context with an external memory that stores compressed hidden states at the temporal level. Martins et al. (2021) used continuous space attention to attend over the old history and updated the memory with recent hidden states to enable unbounded memory capacity. Wu et al. (2021) proposed to use the encoder-decoder architecture to encode the memory states with previous text segments and pass this memory to future time steps. Instead of using a fixed-size attention span for different layers, Sukhbaatar et al. (2019) and Correia et al. (2019) proposed to learn dynamic attention spans for dif
Due to the space limit, we only sample 8 layers from all the 16 layers.
ferent attention heads, which greatly reduced the computations. These works focused on enabling the Transformer to access contents in long distance, but did not consider to learn better memory representations by refreshing the old memory under the current context. Our work is orthogonal to learning adaptive attention spans and can be combined with this technique to reduce the complexity.
We present LaMemo, a memory mechanism that allows the memory states to incrementally attend to the right-side tokens and interpolates with the old memory states on the left side, which enables the memory to interact with bi-directional contexts with a complexity linear in memory length. Experiments on three language modeling datasets demonstrate the superiority of LaMemo over baselines with various types of memory mechanisms. We also found that LaMemo increases the utilization of older memory states when predicting the target tokens, and yields a higher performance boost when extrapolating to longer memory length, which indicates the effectiveness of recontextualizing the memory under the current context.
This work was supported by the National Science Foundation for Distinguished Young Scholars (with No. 62125604) and the NSFC projects (Key project with No. 61936010 and regular project with No. 61876096). This work was also supported by the Guoqiang Institute of Tsinghua University, with Grant No. 2019GQG1 and 2020GQG0005. This work was also sponsored by Tsinghua-Toyota Joint Research Fund.
Joshua Ainslie, Santiago Ontañón, Chris Alberti, Va-clav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. 2020. ETC: encoding long and structured inputs in transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pages 268-284. Association for Computational Linguistics.
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. Layer normalization. CoRR, abs/1607.06450.
Alexei Baevski and Michael Auli. 2019. Adaptive input representations for neural language modeling. In 7th
International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015. Neural machine translation by jointly learning to align and translate. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. CoRR, abs/2004.05150.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. 2021. Rethinking attention with performers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Gonçalo M. Correia, Vlad Niculae, and André F. T. Martins. 2019. Adaptively sparse transformers. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019, pages 2174-2184. Association for Computational Linguistics.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G. Carbonell, Quoc Viet Le, and Ruslan Salakhutdinov. 2019. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pages 2978-2988. Association for Computational Linguistics.
Alex Graves, Greg Wayne, and Ivo Danihelka. 2014. Neural tuning machines. CoRR, abs/1410.5401.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 770-778. IEEE Computer Society.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. 2020. The curious case of neural text degeneration. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Curtis Hawthorne, Andrew M. Dai, Matthew D. Hoffman, and Douglas Eck. 2018. An improved relative self-attention mechanism for transformer with application to music generation. CoRR, abs/1809.04281.
Haozhe Ji and Minlie Huang. 2021. Discodvt: Generating long text with discourse-aware discrete variational transformer. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November; 2021, pages 4208-4224. Association for Computational Linguistics.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. 2020. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 5156-5165. PMLR.
Urvashi Khandelwal, He He, Peng Qi, and Dan Jurafsky. 2018. Sharp nearby, fuzzy far away: How neural language models use context. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July 15-20, 2018, Volume 1: Long Papers, pages 284-294. Association for Computational Linguistics.
Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
Matt Mahoney. 2011. Large text compression benchmark.
Pedro Henrique Martins, Zita Marinho, and Andre F. T. Martins. 2021.  $\infty$ -former: Infinite memory transformer. CoRR, abs/2109.00301.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net.
Tomás Mikolov and Geoffrey Zweig. 2012. Context dependent recurrent neural network language model. In 2012 IEEE Spoken Language Technology Workshop (SLT), Miami, FL, USA, December 2-5, 2012, pages 234-239. IEEE.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 8024-8035.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. 2021. Random feature attention. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. 2020. Compressive transformers for long-range sequence modelling. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. 2018. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT, New Orleans, Louisiana, USA, June 1-6, 2018, Volume 2 (Short Papers), pages 464-468. Association for Computational Linguistics.
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. 2019. Adaptive attention span in transformers. In Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28-August 2, 2019, Volume 1: Long Papers, pages 331-335. Association for Computational Linguistics.
Sainbayar Sukhbaatar, Arthur Szlam, Jason Weston, and Rob Fergus. 2015. End-to-end memory networks. In Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pages 2440-2448.
Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke, and Mohit Iyyer. 2021. Do long-range language models actually use long-range context? In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021, pages 807-822. Association for Computational Linguistics.
Bowen Tan, Zichao Yang, Maruan Al-Shedivat, Eric P. Xing, and Zhiting Hu. 2021. Progressive generation of long text with pretrained language models. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021, Online, June 6-11, 2021, pages 4313-4324. Association for Computational Linguistics.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. 2021. Long range arena: A benchmark for efficient transformers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. 2020. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04768.
Jason Weston, Sumit Chopra, and Antoine Bordes. 2015. Memory networks. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
Chao-Yuan Wu, Christoph Feichtenhofer, Haoqi Fan, Kaiming He, Philipp Krahenbuhl, and Ross B. Girshick. 2019. Long-term feature banks for detailed video understanding. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pages 284-293. Computer Vision Foundation / IEEE.
Qingyang Wu, Zhenzhong Lan, Jing Gu, and Zhou Yu. 2021. Memformer: The memory-augmented transformer.
Wenhan Xiong, Barlas Oguz, Anchit Gupta, Xilun Chen, Diana Liskovich, Omer Levy, Wen-tau Yih, and Yashar Mehdad. 2021. Simple local attentions remain competitive for long-context tasks. CoRR, abs/2112.07210.
A Derivation of Memory Interpolation
We derive Eq. 10 into the form of standard self-attention in the following:
$$
C _ {\tau - 1} ^ {\leftrightarrow} = \alpha_ {\tau - 1} \mathrm {s g} (C _ {\tau - 1} ^ {\rightarrow}) + (1 - \alpha_ {\tau - 1}) C _ {\tau - 1} ^ {\leftarrow}.
$$
We consider the  $i$ -th row of  $C_{\tau -1}^{\leftrightarrow}$ , denoted as  $c_{i}^{\leftrightarrow}$ . We omit the stop-grad operation  $\mathrm{sg}(\cdot)$  and substitute  $\alpha$  with the result from Eq. 11:
$$
\begin{array}{l} \boldsymbol {c} _ {i} ^ {\leftrightarrow} = \alpha_ {i} \boldsymbol {c} _ {i} ^ {\rightarrow} + (1 - \alpha_ {i}) \boldsymbol {c} _ {i} ^ {\leftarrow} \\ = \frac {s _ {i} ^ {\rightarrow}}{s _ {i} ^ {\rightarrow} + s _ {i} ^ {\leftarrow}} \pmb {c} _ {i} ^ {\rightarrow} + \frac {s _ {i} ^ {\leftarrow}}{s _ {i} ^ {\rightarrow} + s _ {i} ^ {\leftarrow}} \pmb {c} _ {i} ^ {\leftarrow}, \\ \end{array}
$$
where  $s_i^{\rightarrow}$ $s_i^{\leftarrow}$  is the denominator of the softmax when computing  $\pmb{c}_i^{\rightarrow}$ $c_{i}^{\leftarrow}$  respectively:
$$
\begin{array}{l} s _ {i} ^ {\rightarrow} = \sum_ {j \leq i} \exp \left(\frac {\boldsymbol {q} _ {i} ^ {\prime \top} \boldsymbol {k} _ {j} ^ {\prime}}{\sqrt {d}}\right) = \sum_ {j \leq i} \sin (\boldsymbol {q} _ {i} ^ {\prime}, \boldsymbol {k} _ {j} ^ {\prime}), \\ s _ {i} ^ {\leftarrow} = \sum_ {j > i} \exp \left(\frac {\boldsymbol {q} _ {i} ^ {\top} \boldsymbol {k} _ {j}}{\sqrt {d}}\right) = \sum_ {j > i} \sin (\boldsymbol {q} _ {i}, \boldsymbol {k} _ {j}), \\ \end{array}
$$