LaMemo: Language Modeling with Look-Ahead Memory

Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5747 - 5762 July 10-15, 2022 ©2022 Association for Computational Linguistics LaMemo: Language Modeling with Look-Ahead Memory Haozhe Ji , Rongsheng Zhang , Zhenyu Yang , Zhipeng Hu , Minlie Huang ∗ The CoAI group, DCST, Institute for Artificial Intelligence, State Key Lab of Intelligent Technology and Systems, Beijing National Research Center for Information Science and Technology, Tsinghua University, Beijing 100084, China Fuxi AI Lab, NetEase Inc., China, Guangdong OPPO Mobile Telecommunications Corp., Ltd. jhz20@mails.tsinghua.edu.cn, aihuang@tsinghua.edu.cn {zhangrongsheng, zphu}@corp.netease.com, yangzhenyu@oppo.com Abstract Although Transformers with fully connected self-attentions are powerful to model long-term dependencies, they are struggling to scale to long texts with thousands of words in language modeling. One of the solutions is to equip the model with a recurrence memory. However, ex- isting approaches directly reuse hidden states from the previous segment that encodes con- texts in a uni-directional way. As a result, this prohibits the memory to dynamically interact with the current context that provides up-to- date information for token prediction. To rem- edy this issue, we propose Look-Ahead Memory (LaMemo)1 that enhances the recurrence mem- ory by incrementally attending to the right-side tokens, and interpolating with the old mem- ory states to maintain long-term information in the history. LaMemo embraces bi-directional attention and segment recurrence with an addi- tional computation overhead only linearly pro- portional to the memory length. Experiments on widely used language modeling benchmarks demonstrate its superiority over the baselines equipped with different types of memory.2 1 Introduction Language modeling is an important task that tests the ability of modeling long-term dependencies by predicting the current token based on the previous context (Mikolov and Zweig, 2012; Merity et al., 2017). Recently, Transformer-based language mod- els achieved remarkable performance by enabling direct interaction between long-distance word pairs. However, as the computation overhead grows with the length of the input sequence, Transformers can only process a fixed length segment at a time. To allow long-term information flow across individual segments, existing approaches augment the model ∗ Corresponding author 1We are also inspired by the French word “La Mémoire”, meaning “the memory”. 2Source code available at https://github.com/ thu-coai/LaMemo. -100 -90 -80 -70 -60 -50 -40 -30 -20 -10 Context 10−2 10−1 Attention LaMemo Transformer-XL Figure 1: Attention weights on the context (in log-scale) in the final layer of Transformer-XL and LaMemo av- eraged on 15K tokens. Transformer-XL quickly loses attention to older contexts, while LaMemo maintains awareness to the history with the grow of the context length. with a recurrence memory that stores hidden states computed in previous time steps (Dai et al., 2019) and their compressions (Rae et al., 2020; Martins et al., 2021) for the target tokens to attend to. One limitation of this approach is that the recur- rence memory is only aware of older contexts since they are previously computed to predict the next word from left to right. As a result, distant memory states …

Content truncated to 500 words