Not All Memories are Created Equal: Learning to Expire

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 on long-context language modeling benchmarks and a frame-by-frame colliding objects task that EXPIRE-SPAN can scale to memories in the tens of thousands — by expiring irrelevant information , capacity is freed to have even larger memory . Finally , we analyze the information retained and expired by EXPIRE-SPAN models , to understand the importance of long context memory . 2 RELATED WORK . Memory is crucial for many tasks and has been studied in recurrent networks ( Elman , 1990 ; Hochreiter & Schmidhuber , 1997 ; Mikolov et al. , 2010 ) for a long time . The development of memory augmented networks ( Graves et al. , 2014 ; Sukhbaatar et al. , 2015b ) made it possible to store large quantities of information and selectively access them using attention ( Bahdanau et al. , 2015 ) . The Transformer ( Vaswani et al. , 2017 ) took full advantage of this approach . Processing long sequences with Transformers is an active area with applications in language understanding ( Brown et al. , 2020 ) , reinforcement learning ( Parisotto et al. , 2019 ) , video processing ( Wu et al. , 2019 ) , and protein folding ( Rives et al. , 2019 ) . However , extending the memory span is computationally expensive due to the quadratic time and space complexity of self-attention . Various work has focused on reducing this complexity and increasing memory capacity . Dynamic attention spans , such as Adaptive-Span ( Sukhbaatar et al. , 2019a ) and Adaptively Sparse Transformer ( Correia et al. , 2019 ) , focus on learning which heads can have shorter spans of attention , but can only extend to spans of a few thousand . Other work sparsifies attention by computing fewer tokens ( Fan et al. , 2019c ) , often by using fixed attention masks ( Child et al. , 2019 ) or sliding windows and dilation ( Beltagy et al. , 2020 ) . The BP Transformer ( Ye et al. , 2019 ) structures tokens as a binary tree , so some tokens have coarse attention . These works focus on learning what to attend to , but searching larger and larger memories is very difficult . In contrast , we focus on learning to expire what is irrelevant . Compressive Transformer ( Rae et al. , 2020 ) reduces the number of memories by replacing every few memories with a single compressed one . A disadvantage of this is that all memories have the same compression ratio , so relevant memories are equally compressed . Another line of work investigates linear-time attention mechanisms . Wu et al . ( 2018 ) replace selfattention with convolutions that run in linear time , but the scalability to long context tasks remains limited . Wang et al . ( 2020 ) propose linear time attention by decomposing attention into multiple smaller attentions , that recombine to form a low-rank factorization of the original attention . Those methods , however , focus on making attention more efficient without actually reducing the number of memories . Further , as our goal is to reduce the number of memories that feed to self-attention by learning to expire , EXPIRE-SPAN can be easily combined with these efficiency improvements . For a review of further recent Transformer variants , see Tay et al . ( 2020 ) . 3 BACKGROUND . Transformer architectures have been widely used as decoder-only auto-regressive models for sequential tasks . Each Transformer decoder is made of a stack of identical layers , composed of a multi-head self-attention sublayer followed by a feedforward sublayer . The output of each timestep is the hidden state hlt at layer l , which is then projected to key k , value v , and query q vectors : qlt =W l qh l t , k l t =W l kh l t , v l t =W l vh l t. where W represents the weight . Going forward , we focus on a single layer and omit the layer index l for brevity . Information from previous timesteps is accessed through attention a to create output o : at , i = Softmaxi∈Ct ( q > t ki ) , ot =Wo ∑ i∈Ct at , ivi . The set Ct indicates which memories can be accessed at time t , which is the focus on this work . The space and time complexity of self-attention is linearly correlated to |Ct| , making it an important metric of efficiency . For the rest of the paper , we will refer to |Ct| as the memory size . Including all previous timesteps in self-attention by settingCt = { 1 , . . . , t−1 } results in a quadratic complexityO ( T 2 ) to compute the full attention over a sequence of length T . Fixed-spans ( Dai et al. , 2019 ) take a more scalable approach such that Ct = { t−L , . . . , t− 1 } so the attention is restricted to previous L steps . The total complexity in this case is O ( TL ) , where L is the attention span . Adaptive-Span ( Sukhbaatar et al. , 2019a ) further improves upon this by learning an optimal span L per attention head from data , which results in small L values for many heads . Compression approaches ( Rae et al. , 2020 ) reduce memory size by compressing multiple timesteps into a single memory , with complexity O ( TL/c ) , where c is the compression rate . However , in all these approaches , all memories are treated equally without regards to their importance to the task . In this work , we focus on distinguishing between relevant and irrelevant memories by learning to expire unneeded information — by expiring , the remaining attention on relevant information can scale beyond existing long context memory approaches . 4 EXPIRE-SPAN . We detail the EXPIRE-SPAN mechanism and how to integrate it into Transformer architectures to focus attention on relevant information and expire the rest . We describe how to scale EXPIRE-SPAN and practically train with drastically longer memory spans . 4.1 METHOD . EXPIRE-SPAN , depicted in Figure 1 , allows models to selectively forget memories that are no longer relevant . We describe it in the context of a single Transformer layer and omit the layer index l for brevity . Our goal is to reduce the size of Ct defined in Section 3 for more efficiency without performance degradation . For each memory hi , we will compute a scalar EXPIRE-SPAN ei ∈ [ 0 , L ] : ei = Lσ ( w > hi + b ) . Here w , b represent trainable parameters , σ is the sigmoid function , and L is the maximum span . This expire-span ei determines how long hi should be kept and included in Ct. At time t , the remaining span of hi is rt , i = ei− ( t−i ) . When rt , i becomes negative , it indicates the memory hi is expired and can be removed from Ct . This can be implemented by updating attention weights with a binary masking function mt , i : a′t , i = mt , iat , i∑ j mt , jat , j , ot = ∑ i a′t , ivi where mt , i = { 1 if rt , i > 0 0 otherwise , However , with such a discrete masking , the expire-span ei will not receive any gradient for training . Therefore , we use a masking function that smoothly transitions from 0 to 1 as shown in Figure 2 : mt , i = max ( 0 , min ( 1 , 1 + rt , i/R ) ) , where R is a hyper-parameter that determines the length of a ramp that is bounded between 0 to 1 . This function has non-zero gradient for values in [ −R , 0 ] to train ei , but also can take a value of 0 , which is necessary for expiring memories . Thus Ct = { i | mt , i > 0 } . Sincemt , i is a monotonically decreasing function of t , once a memory is expired , it can be permanently deleted . x m ( x ) 1 −R Figure 2 : Soft Mask Function Our goal is to reduce average memory size , which is directly related with the average EXPIRE-SPAN . 1 T ∑ t |Ct| = 1 T ∑ t ∑ i < t 1mt , i > 0 = 1 T ∑ i ( R+ ∑ t > i 1rt , i > 0 ) = 1 T ∑ i ( R+ ∑ t > i 1ei > t−i ) = R− 1 + 1 T ∑ i beic Thus , we add an auxiliary term to the loss function to penalize the L1-norm of EXPIRE-SPANS : Ltotal = Ltask + α ∑ i ei/T , where α > 0 is a hyperparameter . This loss term will decrease the span of memories that contribute less to the main task , resulting in a model with a small memory that focuses on the most relevant information . Note that w and b are the only new parameters , and are negligible in size compared to the total number of parameters in standard neural models .

Summary: The paper proposes a method for overcoming the long-term memory bottleneck of transformers. The idea is to assign a value (expire-span) to each formed memory, which indicates how long the memory should be stored and be available for the transformer to access it. The authors demonstrate the performance of their approach on a set of