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: To help Transformer learn long sequence efficiently, the paper performs attention on selective timesteps that have high expire-span scores. For each timestep, the expire-span score is computed by mapping the corresponding hidden feature to a number, which is learnt during training. Soft masking is applied to make the learning differentiable. An additional loss is introduced to reduce the average span, making the attention sparse. The proposed attention is integrated into each layer of Transformer and tested on several synthetic tasks and two language modelling datasets, yielding promising results.