NOT ALL MEMORIES ARE CREATED EQUAL: LEARNING TO EXPIRE

Abstract

Attention mechanisms have shown promising results in sequence modeling tasks that require long-term memory. Recent work has investigated mechanisms to reduce the computational cost of preserving and storing the memories (Rae et al., 2020). However, not all content in the past is equally important to remember. We propose Expire-Span, a method that learns to retain the most important information and expire the irrelevant information. This enables Transformers to scale to attend to tens of thousands of previous timesteps efficiently, as not all hidden states from previous timesteps are preserved. We demonstrate that Expire-Span can help models identify and retain critical information and show it can achieve state of the art results on long-context language modeling, reinforcement learning, and algorithmic tasks. Finally, we show that Expire-Span can scale to memories that are tens of thousands in size, which is helpful on incredibly long context tasks such as character-level PG-19 and a frame-by-frame moving objects task.

1. INTRODUCTION

Transformer architectures (Vaswani et al., 2017) have demonstrated strong performance across a variety of tasks (Devlin et al., 2019; Roller et al., 2020; Brown et al., 2020) , including those that require learning long term relationships (Zhang et al., 2018; Fan et al., 2019a; Izacard & Grave, 2020) . Recent work has focused on scaling attention mechanisms efficiently to longer memory sizes, enabling large improvements on long context tasks (Dai et al., 2019; Sukhbaatar et al., 2019a) . However, a critical component of human memory is not just the ability to remember, but also forgetting irrelevant information to focus on the salient, relevant bits. Most studies of long-term memory in humans indicate that not everything is remembered (Murre & Dros, 2015) -instead, only vivid, remarkable memories are retained from the far past (Wixted, 2004) . Standard Transformer architectures lack the ability to search over extremely large memories, as the self-attention mechanism is computationally intensive and the storage cost of preserving the large memory grows quickly. Recent work (Child et al., 2019; Rae et al., 2020) has proposed learning how to extend to greater context through sparse mechanisms or through compression, to more compactly represent the past. However, there exists a fundamental problem with large memories beyond strict computational concerns: as the amount of information stored increases, deciding which information is relevant becomes more challenging. Other work (Lample et al., 2019) approaches this by considering how to efficiently search large memories. We will focus on learning what to forget, and thereby reducing the computational burden of the model easing the challenges of the search problem. We propose EXPIRE-SPAN, a straightforward extension to attention mechanisms, that learns when to expire unneeded memories. By expiring memories that are no longer useful, EXPIRE-SPAN enables scaling to memories tens of thousands of timesteps long. This learnable mechanism allows the model to adjust the span size as needed, selecting which information is critical to retain and forgetting the rest. More concretely, we augment the self-attention with a simple predictor that outputs an expiration value for each hidden state that determines how long a memory should be retained and accessible to the model. After the EXPIRE-SPAN runs out, the memory will be forgotten, but in a gradually differentiable way to retain end-to-end training with backpropagation. This process is done independently for each layer, allowing different layers to specialize in different time-scales. We demonstrate that EXPIRE-SPAN can distinguish between critical and irrelevant information on several illustrative tasks in natural language processing and reinforcement learning. We then show

