Original: N OT A LL M EMORIES ARE C REATED E QUAL : L EARNING TO E XPIRE
Revised: N OT A LL M EMORIES ARE C REATED E QUAL : L EARNING TO E XPIRE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Attention mechanisms have shown promising results in sequence modeling tasks that require long-term memory.
Revised: Attention mechanisms have shown promising results in sequence modeling tasks that require long-term memory.

Original: Recent work has investigated mechanisms to reduce the computational cost of preserving and storing the memories (Rae et al., 2020).
Revised: Recent work has investigated mechanisms to reduce the computational cost of preserving and storing the memories (Rae et al., 2020).

Original: However, not all content in the past is equally important to remember
Revised: However, not all content in the past is equally important to remember

Original: We propose Expire-Span a method that learns to retain the most important information and expire the irrelevant information
Revised: We propose Expire-Span a method that learns to retain the most important information and expire the irrelevant information

Original: This enables Transformers to scale to attend to tens of thousands of previous timesteps efﬁciently as not all hidden states from previous timesteps are preserved
Revised: This enables Transformers to scale to attend to tens of thousands of previous timesteps efﬁciently as not all hidden states from previous timesteps are preserved

Original: We demonstrate that Expire
Revised: We demonstrate that Expire

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Transformer architectures (Vaswani et al., 2017) have demonstrated strong performance across a variety of tasks (Devlin et al., 2019; Roller et al., 2020;
Revised: Transformer architectures (Vaswani et al., 2017) have demonstrated strong performance across a variety of tasks (Devlin et al., 2019; Roller et al., 2020;

Original: Brown et al., 2020), including those that require learning long term relationships (Zhang et al., 2018; Fan et al., 2019a; Izacard & Grave, 2020).
Revised: Brown et al., 2020), including those that require learning long term relationships (Zhang et al., 2018; Fan et al., 2019a; Izacard & Grave, 2020).

Original: Recent work has focused on scaling attention mechanisms efﬁciently to longer memory sizes, enabling large improvements on long context tasks (Dai et al., 2019; Sukhbaatar et al., 2019a).
Revised: Recent work has focused on scaling attention mechanisms efﬁciently to longer memory sizes, enabling large improvements on long context tasks (Dai et al., 2019; Sukhbaatar et al., 2019a).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Other work (Lample et al., 2019) approaches this by considering how to efﬁciently search large memories.
Revised: Other work (Lample et al., 2019) approaches this by considering how to efﬁciently search large memories.

Original: We will focus on learning what to forget, and thereby reducing the computational burden of the model easing the challenges of the search problem.
Revised: We will focus on learning what to forget, and thereby reducing the computational burden of the model easing the challenges of the search problem.

Original: We propose E XPIRE -S PAN , a straightforward extension to attention mechanisms, that learns when to expire unneeded memories.
Revised: We propose E XPIRE -S PAN , a straightforward extension to attention mechanisms, that learns when to expire unneeded memories.

Original: By expiring memories that are no longer useful, E XPIRE -S PAN enables scaling to memories tens of thousands of timesteps long.
Revised: By expiring memories that are no longer useful, E XPIRE -S PAN enables scaling to memories tens of thousands of timesteps long.

Original: This learnable mechanism allows the model to adjust the span size as needed, selecting which information is critical to retain and forgetting the rest.
Revised: This learnable mechanism allows the model to adjust the span size as needed, selecting which information is critical to retain and forgetting the rest.

Original: 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.
Revised: 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.

Original: After the E XPIRE -S PAN runs out, the memory will be forgotten, but in a gradually differentiable way to retain end-to-end training with backpropagation.
Revised: After the E XPIRE -S PAN runs out, the memory will be forgotten, but in a gradually differentiable way to retain end-to-end training with backpropagation.

Original: This process is done independently for each layer, allowing different layers to specialize in different time-scales.
Revised: This process is done independently for each layer, allowing different layers to specialize in different time-scales.

Original: We demonstrate that E XPIRE -S PAN can distinguish between critical and irrelevant information on several illustrative tasks in natural language processing and reinforcement learning.
Revised: We demonstrate that E XPIRE -S PAN can distinguish between critical and irrelevant information on several illustrative tasks in natural language processing and reinforcement learning.

Original: We then show time
Revised: We then show time

Original: Expire-span eheheheheh
Revised: Expire-span eheheheheh

Original: Attention q t query o t output h t on long-context language modeling benchmarks and a frame-by-frame colliding objects task that
Revised: Attention q t query o t output h t on long-context language modeling benchmarks and a frame-by-frame colliding objects task that

Original: E XPIRE -S PAN can scale to memories in the tens of thousands — by expiring irrelevant information, capacity is freed to have even larger memory.
Revised: E XPIRE -S PAN can scale to memories in the tens of thousands — by expiring irrelevant information, capacity is freed to have even larger memory.

Original: Finally, we analyze the information retained and expired by E XPIRE -S PAN models, to understand the importance of long context memory.
Revised: Finally, we analyze the information retained and expired by E XPIRE -S PAN models, to understand the importance of long context memory.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: The
Revised: The

Original: Transformer (Vaswani et al., 2017) took full advantage of this approach.
Revised: Transformer (Vaswani et al., 2017) took full advantage of this approach.

Original: 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).
Revised: 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).

Original: However, extending the memory span is computationally expensive due to the quadratic time and space complexity of self-attention.
Revised: However, extending the memory span is computationally expensive due to the quadratic time and space complexity of self-attention.

Original: Various work has focused on reducing this complexity and increasing memory capacity.
Revised: Various work has focused on reducing this complexity and increasing memory capacity.

Original: 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.
Revised: 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.

Original: Other work sparsiﬁes attention by computing fewer tokens (Fan et al., 2019c), often by using ﬁxed attention masks (Child et al., 2019) or sliding windows and dilation (Beltagy et al., 2020).
Revised: Other work sparsiﬁes attention by computing fewer tokens (Fan et al., 2019c), often by using ﬁxed attention masks (Child et al., 2019) or sliding windows and dilation (Beltagy et al., 2020).

Original: The BP Transformer (Ye et al., 2019) structures tokens as a binary tree, so some tokens have coarse attention.
Revised: The BP Transformer (Ye et al., 2019) structures tokens as a binary tree, so some tokens have coarse attention.

Original: These works focus on learning what to attend to, but searching larger and larger memories is very difﬁcult.
Revised: These works focus on learning what to attend to, but searching larger and larger memories is very difﬁcult.

Original: In contrast, we focus on learning to expire what is irrelevant.
Revised: In contrast, we focus on learning to expire what is irrelevant.

Original: Compressive Transformer (Rae et al., 2020) reduces the number of memories by replacing every few memories with a single compressed one.
Revised: Compressive Transformer (Rae et al., 2020) reduces the number of memories by replacing every few memories with a single compressed one.

Original: A disadvantage of this is that all memories have the same compression ratio, so relevant memories are equally compressed.
Revised: A disadvantage of this is that all memories have the same compression ratio, so relevant memories are equally compressed.

Original: Another line of work investigates linear-time attention mechanisms.
Revised: Another line of work investigates linear-time attention mechanisms.

Original: Wu et al.
Revised: Wu et al.

Original: (2018) replace selfattention with convolutions that run in linear time, but the scalability to long context tasks remains limited.
Revised: (2018) replace selfattention with convolutions that run in linear time, but the scalability to long context tasks remains limited.

Original: Wang et al.
Revised: Wang et al.

Original: (2020) propose linear time attention by decomposing attention into multiple smaller attentions, that recombine to form a low-rank factorization of the original attention.
Revised: (2020) propose linear time attention by decomposing attention into multiple smaller attentions, that recombine to form a low-rank factorization of the original attention.

Original: Those methods, however, focus on making attention more efﬁcient without actually reducing the number of memories.
Revised: Those methods, however, focus on making attention more efﬁcient without actually reducing the number of memories.

Original: Further, as our goal is to reduce the number of memories that feed to self-attention by learning to expire, E XPIRE -S PAN can be easily combined with these efﬁciency improvements.
Revised: Further, as our goal is to reduce the number of memories that feed to self-attention by learning to expire, E XPIRE -S PAN can be easily combined with these efﬁciency improvements.

Original: For a review of further recent Transformer variants, see Tay et al.
Revised: For a review of further recent Transformer variants, see Tay et al.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Transformer architectures have been widely used as decoder-only auto-regressive models for sequential tasks.
Revised: Transformer architectures have been widely used as decoder-only auto-regressive models for sequential tasks.

Original: Each Transformer decoder is made of a stack of identical layers, composed of a multi-head self-attention sublayer followed by a feedforward sublayer.
Revised: Each Transformer decoder is made of a stack of identical layers, composed of a multi-head self-attention sublayer followed by a feedforward sublayer.

Original: The output of each timestep is the hidden state h lt at layer l , which is then projected to key k , value v , and query q vectors:
Revised: The output of each timestep is the hidden state h lt at layer l , which is then projected to key k , value v , and query q vectors:

Original: [Equation]
Revised: [Equation]

Original: where W represents the weight.
Revised: where W represents the weight.

Original: Going forward, we focus on a single layer and omit the layer index l for brevity.
Revised: Going forward, we focus on a single layer and omit the layer index l for brevity.

Original: Information from previous timesteps is accessed through attention a to create output o :
Revised: Information from previous timesteps is accessed through attention a to create output o :

Original: [Equation]
Revised: [Equation]

Original: i
Revised: i

Original: [Equation]

Original: Softmax

Original: [Equation]
Revised: [Equation]

Original: The set C t indicates which memories can be accessed at time t , which is the focus on this work.
Revised: The set C t indicates which memories can be accessed at time t , which is the focus on this work.

Original: The space and time complexity of self-attention is linearly correlated to | C t | , making it an important metric of efﬁciency.
Revised: The space and time complexity of self-attention is linearly correlated to | C t | , making it an important metric of efﬁciency.

Original: For the rest of the paper, we will refer to | C t | as the memory size .
Revised: For the rest of the paper, we will refer to | C t | as the memory size .

Original: Including all previous timesteps in self-attention by setting C t = { 1 , . . .
Revised: Including all previous timesteps in self-attention by setting C t = { 1 , . . .

Original: , t − 1 } results in a quadratic complexity
Revised: , t − 1 } results in a quadratic complexity

Original: O
Revised: O

Original: ( T 2 ) to compute the full attention over a sequence of length T .
Revised: ( T 2 ) to compute the full attention over a sequence of length T .

Original: Fixed-spans (Dai et al., 2019) take a more scalable approach such that C t = { t − L, . .
Revised: Fixed-spans (Dai et al., 2019) take a more scalable approach such that C t = { t − L, . .

Original: , t − 1 } so the attention is restricted to previous L steps.
Revised: , t − 1 } so the attention is restricted to previous L steps.

Original: The total complexity in this case is O ( TL ) , where L is the attention span.
Revised: The total complexity in this case is O ( TL ) , where L is the attention span.

Original: Adaptive-Span (Sukhbaatar et al., 2019a) further improves upon this by learning an optimal span
Revised: Adaptive-Span (Sukhbaatar et al., 2019a) further improves upon this by learning an optimal span

Original: L per attention head from data, which results in small L values for many heads.
Revised: L per attention head from data, which results in small L values for many heads.

Original: 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.
Revised: 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.

Original: However, in all these approaches, all memories are treated equally without regards to their importance to the task.
Revised: However, in all these approaches, all memories are treated equally without regards to their importance to the task.

Original: 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.
Revised: 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.

Original:  E XPIRE -S PAN
Revised:  E XPIRE -S PAN

Original: We detail the E XPIRE -S PAN mechanism and how to integrate it into Transformer architectures to focus attention on relevant information and expire the rest.
Revised: We detail the E XPIRE -S PAN mechanism and how to integrate it into Transformer architectures to focus attention on relevant information and expire the rest.

Original: We describe how to scale E XPIRE -S PAN and practically train with drastically longer memory spans.
Revised: We describe how to scale E XPIRE -S PAN and practically train with drastically longer memory spans.

Original: 4. M ETHOD
Revised: 4. M ETHOD

Original: E XPIRE -S PAN , depicted in Figure 1, allows models to selectively forget memories that are no longer relevant.
Revised: E XPIRE -S PAN , depicted in Figure 1, allows models to selectively forget memories that are no longer relevant.

Original: We describe it in the context of a single Transformer layer and omit the layer index l for brevity.
Revised: We describe it in the context of a single Transformer layer and omit the layer index l for brevity.

Original: Our goal is to reduce the size of C t deﬁned in Section 3 for more efﬁciency without performance degradation.
Revised: Our goal is to reduce the size of C t deﬁned in Section 3 for more efﬁciency without performance degradation.

Original: For each memory h i , we will compute a scalar E XPIRE -S
Revised: For each memory h i , we will compute a scalar E XPIRE -S

Original: PAN e i ∈ [0 , L ] :
Revised: PAN e i ∈ [0 , L ] :

Original: [Equation]
Revised: [Equation]

Original: Here w , b represent trainable parameters, σ is the sigmoid function, and L is the maximum span.
Revised: Here w , b represent trainable parameters, σ is the sigmoid function, and L is the maximum span.

Original: This expire-span e i determines how long h i should be kept and included in C t .
Revised: This expire-span e i determines how long h i should be kept and included in C t .

Original: At time t , the remaining span of h i is r t,i =
Revised: At time t , the remaining span of h i is r t,i =

Original: e i − ( t − i ) .
Revised: e i − ( t − i ) .

Original: When r t,i becomes negative, it indicates the memory h
Revised: When r t,i becomes negative, it indicates the memory h

Original: i is expired and can be removed from C t .
Revised: i is expired and can be removed from C t .

Original: This can be implemented by updating attention weights with a binary masking function m t,i :
Revised: This can be implemented by updating attention weights with a binary masking function m t,i :

Original: [Equation]
Revised: [Equation]

Original: where m t,i = (cid:26) 1 if r t,i >0 otherwise ,
Revised: where m t,i = (cid:26) 1 if r t,i >0 otherwise ,

Original: However, with such a discrete masking, the expire-span e i will not receive any gradient for training.
Revised: However, with such a discrete masking, the expire-span e i will not receive any gradient for training.

Original: Therefore, we use a masking function that smoothly transitions from 0 to 1 as shown in Figure 2:
Revised: Therefore, we use a masking function that smoothly transitions from 0 to 1 as shown in Figure 2:

Original: [Equation]
Revised: [Equation]

Original: where R is a hyper-parameter that determines the length of a ramp that is bounded between 0 to 1.
Revised: where R is a hyper-parameter that determines the length of a ramp that is bounded between 0 to 1.

Original: This function has non-zero gradient for values in [ − R, 0] to train e i , but also can take a value of 0, which is necessary for expiring memories.
Revised: This function has non-zero gradient for values in [ − R, 0] to train e i , but also can take a value of 0, which is necessary for expiring memories.

Original: Thus C t = { i | m t,i > 0 } .
Revised: Thus C t = { i | m t,i > 0 } .

Original: Since m t,i is a monotonically decreasing function of t , once a memory is expired, it can be permanently deleted.
Revised: Since m t,i is a monotonically decreasing function of t , once a memory is expired, it can be permanently deleted.

Original: [Figure]
Revised: [Figure]

Original: Our goal is to reduce average memory size, which is directly related with the average E XPIRE -S PAN .
Revised: Our goal is to reduce average memory size, which is directly related with the average E XPIRE -S PAN .

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: Thus, we add an auxiliary term to the loss function to penalize the L1-norm of E XPIRE -S PANS :
Revised: Thus, we add an auxiliary term to the loss function to penalize the L1-norm of E XPIRE -S PANS :

Original: L total =
Revised: L total =

Original: L task + α (cid:80) i e i /T , where α > 0 is a hyperparameter.
Revised: L task + α (cid:80) i e i /T , where α > 0 is a hyperparameter.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 4. A DDING E XPIRE -S PAN TO T RANSFORMERS
Revised: 4. A DDING E XPIRE -S PAN TO T RANSFORMERS

Original: We describe how E XPIRE -S PAN can be utilized within Transformer self-attention layers to decrease the memory size and focus the memory on the most salient, relevant information, expiring the rest.
Revised: We describe how E XPIRE -S PAN can be utilized within Transformer self-attention layers to decrease the memory size and focus the memory on the most salient, relevant information, expiring the rest.

Original: We also discuss practical training concerns, such as efﬁciency and regularization.
Revised: We also discuss practical training concerns, such as efﬁciency and regularization.

Original: Modiﬁcations to Multi-Head Attention Self-attention consists of multiple heads that have different keys, values, and queries.
Revised: Modiﬁcations to Multi-Head Attention Self-attention consists of multiple heads that have different keys, values, and queries.

Original: However, they all share one underlying memory, so a memory cannot be removed if it is used by any of the heads.
Revised: However, they all share one underlying memory, so a memory cannot be removed if it is used by any of the heads.

Original: Thus, we compute an E XPIRE -S PAN at each layer that is shared amongst the heads.
Revised: Thus, we compute an E XPIRE -S PAN at each layer that is shared amongst the heads.

Original: Block Parallel We use the caching mechanism (Dai et al., 2019), where a block of timesteps
Revised: Block Parallel We use the caching mechanism (Dai et al., 2019), where a block of timesteps

Original: B =
Revised: B =

Original: [ t, . . .
Revised: [ t, . . .

Original: , t + b − 1] is processed in parallel for efﬁciency — once a block is computed, its hidden states [ h t , . .
Revised: , t + b − 1] is processed in parallel for efﬁciency — once a block is computed, its hidden states [ h t , . .

Original: , h t + b − 1 ] are cached so that future blocks can attend to them.
Revised: , h t + b − 1 ] are cached so that future blocks can attend to them.

Original: This means a memory can only be deleted if it is not used by any of the queries in B .
Revised: This means a memory can only be deleted if it is not used by any of the queries in B .

Original: Concretely, h i will be deleted when m t,i = 0 where t is the ﬁrst token of B .
Revised: Concretely, h i will be deleted when m t,i = 0 where t is the ﬁrst token of B .

Original: However, this is not a concern for very long-term memories where L (cid:29) B .
Revised: However, this is not a concern for very long-term memories where L (cid:29) B .

Original: Position Embedding Relative position embeddings (Shaw et al., 2018) make it possible to condition on the ordering of inputs by modifying the attention to a t,i = Softmax ( q (cid:62)
Revised: Position Embedding Relative position embeddings (Shaw et al., 2018) make it possible to condition on the ordering of inputs by modifying the attention to a t,i = Softmax ( q (cid:62)

Original: t k i + q (cid:62) t p t − i ) .
Revised: t k i + q (cid:62) t p t − i ) .

Original: However, because this second term is computed for the whole block in parallel for efﬁciency, it can become expensive for a large L even when the average memory size | C t | is small.
Revised: However, because this second term is computed for the whole block in parallel for efﬁciency, it can become expensive for a large L even when the average memory size | C t | is small.

Original: Our solution is to remove position embeddings from older memories i < t − b (where b is the block size), which empirically does not affect performance.
Revised: Our solution is to remove position embeddings from older memories i < t − b (where b is the block size), which empirically does not affect performance.

Original: The computational complexity of the position embeddings is then O ( b ) , thus allowing us to increase the maximum span L .
Revised: The computational complexity of the position embeddings is then O ( b ) , thus allowing us to increase the maximum span L .

Original: Loss Computation The L1-norm loss for E XPIRE -S PAN must be computed for every memory h i .
Revised: Loss Computation The L1-norm loss for E XPIRE -S PAN must be computed for every memory h i .

Original: A simple way is to compute it on tokens within the current block B .
Revised: A simple way is to compute it on tokens within the current block B .

Original: This empirically results in poor performance — a possible explanation is that the time between positive and negative gradients on e i may become too distant.
Revised: This empirically results in poor performance — a possible explanation is that the time between positive and negative gradients on e i may become too distant.

Original: Negative gradients that increase e i only come from the main loss L task through the masking function m t,i , which has non-zero gradients only when memory h
Revised: Negative gradients that increase e i only come from the main loss L task through the masking function m t,i , which has non-zero gradients only when memory h

Original: i is about to expire with 0 < m t,i < 1 for t ∈ B .
Revised: i is about to expire with 0 < m t,i < 1 for t ∈ B .

Original: For a large L (cid:29) B , h i may have been computed many blocks before and since then the model weights would have changed.
Revised: For a large L (cid:29) B , h i may have been computed many blocks before and since then the model weights would have changed.

Original: In contrast, the positive gradients that decrease e i are computed on the current block i ∈ B .
Revised: In contrast, the positive gradients that decrease e i are computed on the current block i ∈ B .

Original: To remove this discrepancy, we compute the auxiliary loss on e i at the same time as negative gradients when 0 < m t,i < 1 for t ∈ B .
Revised: To remove this discrepancy, we compute the auxiliary loss on e i at the same time as negative gradients when 0 < m t,i < 1 for t ∈ B .

Original: Regularization A potential challenge in exceptionally long memory is greater capacity to overﬁt.
Revised: Regularization A potential challenge in exceptionally long memory is greater capacity to overﬁt.

Original: As E XPIRE -S PAN can scale to memories in the tens of thousands, it can overﬁt to learning speciﬁc span sizes on the training set that do not generalize.
Revised: As E XPIRE -S PAN can scale to memories in the tens of thousands, it can overﬁt to learning speciﬁc span sizes on the training set that do not generalize.

Original: As a form of regularization, we propose to randomly shorten the memory during training.
Revised: As a form of regularization, we propose to randomly shorten the memory during training.

Original: For each batch, we sample l ∼ U (0 , L ) and set a t,i = 0 for all t − i > l only during training.
Revised: For each batch, we sample l ∼ U (0 , L ) and set a t,i = 0 for all t − i > l only during training.

Original: This way, the model cannot assume the memory will always contain speciﬁc information, as the memory is randomly shortened.
Revised: This way, the model cannot assume the memory will always contain speciﬁc information, as the memory is randomly shortened.

Original: This can be seen as a form of structured dropout (Fan et al., 2019b) applied to the memory size.
Revised: This can be seen as a form of structured dropout (Fan et al., 2019b) applied to the memory size.

Original: [Figure]
Revised: [Figure]

Original: Training with Small Initial Spans
Revised: Training with Small Initial Spans

Original: E XPIRE -S PAN scales to long attention spans as it quickly learns to expire irrelevant content.
Revised: E XPIRE -S PAN scales to long attention spans as it quickly learns to expire irrelevant content.

Original: However, at the beginning of training, the long span can use large quantities of GPU memory.
Revised: However, at the beginning of training, the long span can use large quantities of GPU memory.

Original: To circumvent this, we initialize the bias term with a negative value.
Revised: To circumvent this, we initialize the bias term with a negative value.

Original: This prevents large memory usage at the beginning of training, after which the model quickly learns to expire and the memory usage is no longer problematic.
Revised: This prevents large memory usage at the beginning of training, after which the model quickly learns to expire and the memory usage is no longer problematic.

Original:  E XPERIMENTS AND R ESULTS
Revised:  E XPERIMENTS AND R ESULTS

Original: We show that E XPIRE -S PAN can focus on salient information on various constructed tasks that necessitate expiration.
Revised: We show that E XPIRE -S PAN can focus on salient information on various constructed tasks that necessitate expiration.

Original: Then, we highlight the scalability of E XPIRE -S PAN when operating on extremely large memories on different tasks.
Revised: Then, we highlight the scalability of E XPIRE -S PAN when operating on extremely large memories on different tasks.

Original: Additional experiments and details are in the Appendix.
Revised: Additional experiments and details are in the Appendix.

Revised: 5. B ASELINES

Revised: We compare our method against several baselines from Section 3 that takes different approaches to limit the memory size.

Revised: Transformer-XL (Dai et al., 2019) corresponds to the ﬁxed-spans approach where simply the last L memories are kept.

Revised: Our Transformer-XL implementation also serves as a base model for all the other baselines to guarantee that the only difference among them is how memories are handled.

Original:  5. I MPORTANCE OF E XPIRATION : I LLUSTRATIVE T ASKS
Revised: The other baselines are Adaptive-Span (Sukhbaatar et al., 2019a) and Compressive Transformer (Rae et al., 2020). For Compressive Transformer, we implemented the mean-pooling version, which is shown to give a good performance despite its simplicity. 5. I MPORTANCE OF E XPIRATION : I LLUSTRATIVE T ASKS

Original: Walking down a Corridor To illustrate a case where proper expiration of unnecessary memories is critical, we begin with an RL gridworld task.
Revised: Walking down a Corridor To illustrate a case where proper expiration of unnecessary memories is critical, we begin with an RL gridworld task.

Original: In the Corridor task, depicted in Figure 3 (left), the agent is placed at one end of a very long corridor, next to an object that is either red or blue.
Revised: In the Corridor task, depicted in Figure 3 (left), the agent is placed at one end of a very long corridor, next to an object that is either red or blue.

Original: The agent must walk down the corridor and go to the door that corresponds to the color of the object that it saw at the beginning to receive +1 reward.
Revised: The agent must walk down the corridor and go to the door that corresponds to the color of the object that it saw at the beginning to receive +1 reward.

Original: The requirement on the memory is low: the agent should remember the object color so it can walk through the correct door.
Revised: The requirement on the memory is low: the agent should remember the object color so it can walk through the correct door.

Original:  E XPIRE -S PAN models can take advantage of this fact and keep the memory size small regardless of the corridor length, which can vary between 3 and 200.
Revised: [Figure] E XPIRE -S PAN models can take advantage of this fact and keep the memory size small regardless of the corridor length, which can vary between 3 and 200.

Original: This is conﬁrmed in the results shown in Figure 4 (left) where the E XPIRE -S PAN models achieve high performance on this task with very small memories.
Revised: This is conﬁrmed in the results shown in Figure 4 (left) where the E XPIRE -S PAN models achieve high performance on this task with very small memories.

Original: Without the ability to forget, the standard Transformers require large memory for storing all navigation steps that grow with the corridor length.
Revised: Without the ability to forget, the Transformer-XL models require large memory for storing all navigation steps that grow with the corridor length.

Original: Multi-Room Portal Task
Revised: Multi-Room Portal Task

Original: Next, we analyze the performance of E XPIRE -S PAN on a reinforcement learning memorization task.
Revised: Next, we analyze the performance of E XPIRE -S PAN on a reinforcement learning memorization task.

Original: Agents in a gridworld are tasked with navigating through multiple rooms separated by different doors, depicted in Figure 5 (left).
Revised: Agents in a gridworld are tasked with navigating through multiple rooms separated by different doors, depicted in Figure 5 (left).

Original: Each room has two exit doors with different colors — one door portals to the adjacent room, while the other portals the agent back to the start.
Revised: Each room has two exit doors with different colors — one door portals to the adjacent room, while the other portals the agent back to the start.

Original: However, which door works in which room is randomized for each episode.
Revised: However, which door works in which room is randomized for each episode.

Original: Thus, the only way to visit more rooms is by trial-and-error, where agents need to remember the sequence of correct doors to successfully navigate to the end.
Revised: Thus, the only way to visit more rooms is by trial-and-error, where agents need to remember the sequence of correct doors to successfully navigate to the end.

Original: We show results in Figure 5 (right).
Revised: We show results in Figure 5 (right).

Original: Baseline models need longer memory to perform better and visit more rooms — as each room requires many navigation steps, this requires more memory.
Revised: Transformer-XL models need longer memory to perform better and visit more rooms — as each room requires many navigation steps, this requires more memory.

Original: However, those navigation steps are irrelevant because the agent only needs to memorize which doors are visited.
Revised: However, those navigation steps are irrelevant because the agent only needs to memorize which doors are visited.

Original: Since E XPIRE -S PAN models can discard irrelevant memories and focus its memory on memorizing the exact sequence of doors, they achieve strong performance with much smaller memory compared to the standard Transformer.
Revised: Since E XPIRE -S PAN models can discard irrelevant memories and focus its memory on memorizing the exact sequence of doors, they achieve strong performance with much smaller memory compared to the Transformer-XL baseline.

Original: Receiving and Executing Instructions To illustrate a more difﬁcult task where a model must learn to recognize relevant memories and expire the rest, we use a dialogue-based story generation task from the LIGHT (Urbanek et al., 2019) text world game environment.
Revised: Receiving and Executing Instructions To illustrate a more difﬁcult task where a model must learn to recognize relevant memories and expire the rest, we use a dialogue-based story generation task from the LIGHT (Urbanek et al., 2019) text world game environment.

Original: The model visits various locations and receives instructions of the form can you tell the [butler] that the [town ofﬁcial] wants to see them? .
Revised: The model visits various locations and receives instructions of the form can you tell the [butler] that the [town ofﬁcial] wants to see them? .

Original: When the model is in a location where the butler is present, they must execute the instruction by generating You tell the butler “town ofﬁcial wants to see you!” .
Revised: When the model is in a location where the butler is present, they must execute the instruction by generating You tell the butler “town ofﬁcial wants to see you!” .

Original: Between receiving and executing, thousands of words of distractor text exist.
Revised: Between receiving and executing, thousands of words of distractor text exist.

Original: The model must learn what is relevant to retain and expire the distractors.
Revised: The model must learn what is relevant to retain and expire the distractors.

Original: Note multiple instructions can be in queue for execution.
Revised: Note multiple instructions can be in queue for execution.

Original: We experiment with a dataset where the average distance between receiving and executing instructions is around 950 distractor sequences.
Revised: We experiment with a dataset where the average distance between receiving and executing instructions is around 950 distractor sequences.

Original: Models are trained as language models, but evaluated only on their success in executing the instruction.
Revised: Models are trained as language models, but evaluated only on their success in executing the instruction.

Original: Task details and model architecture are provided in the Appendix.
Revised: Task details and model architecture are provided in the Appendix.

Original: We illustrate in Figure 4 (right)  that E XPIRE -S PAN is much more successful at solving this task than the baseline as it can focus on the speciﬁc instruction lines.
Revised: We illustrate in Figure 4 (right) and Table 4 that E XPIRE -S PAN is much more successful at solving this task than Transformer-XL, Adaptive-Span, and Compressive Transformer as it can focus on the speciﬁc instruction lines.

Revised: For Adaptive-Span and Compressive Transformer, we experiment with various memory sizes to tune the maximum span.

Original: 5. S CALABILITY OF E XPIRE -S PAN
Revised: 5. S CALABILITY OF E XPIRE -S PAN

Original: We analyze the scalability of E XPIRE -S PAN .
Revised: We analyze the scalability of E XPIRE -S PAN .

Original: On a copy task, we train models with spans up to 128K timesteps.
Revised: On a copy task, we train models with spans up to 128K timesteps.

Original: Then, we show the utility of E XPIRE -S PAN on two character level language modeling tasks — enwiki8 and PG-19, and colliding objects task that is processed frame by frame.
Revised: Then, we show the utility of E XPIRE -S PAN on two character-level language modeling tasks — enwiki8 and PG-19, and colliding objects task that is processed frame by frame.

Original: Extremely Long Copy To illustrate the scalability of E XPIRE -S PAN , we construct a copy task where the model sees a sequence of A very far in the past.
Revised: Extremely Long Copy To illustrate the scalability of E XPIRE -S PAN , we construct a copy task where the model sees a sequence of A very far in the past.

Original: The rest of the characters are B . 
Revised: The rest of the characters are B . The

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Original: The model must copy the correct quantity of A .
Revised: [Figure] [Table] [Figure] model must copy the correct quantity of A .

Original: We design the task such that long span (up to 128K) can be required, as the A tokens are very far into the past.
Revised: We design the task such that long span (up to 128K) can be required, as the A tokens are very far into the past.

Original: In Table 1, we show that only by scaling the attention span to 128K is it possible to achieve improved performance.
Revised: In Table 1, we show that only by scaling the attention span to 128K is it possible to achieve improved performance.

Original: We compare to astandard
Revised: We compare to a

Original: Transformer with 2048 attention limit and a E XPIRE -S PAN model with smaller span.
Revised: Transformer-XL baseline with 2048 attention limit and a E XPIRE -S PAN model with smaller span.

Original: enwiki8 We test on the enwiki8 character-level language modeling benchmark (Mahoney, 2011).
Revised: enwiki8 We test on the enwiki8 character-level language modeling benchmark (Mahoney, 2011).

Original: We set the max span L for E XPIRE -S PAN to 16K.
Revised: We set the maximum span L for E XPIRE -S PAN to 16K.

Original: We compare E XPIRE -S PAN to existing work in Table 2.
Revised: We compare E XPIRE -S PAN to existing work in Table 2.

Original: E XPIRE -S PAN achieves strong performance even with the smallest model size by far. This [Figure] [Figure]
Revised: E XPIRE -S PAN outperforms similarly sized baselines by a large margin, and matches the performances of much larger models.

Original: Model Params [Figure] Trans
Revised: This indicates that models can learn to expire relevant information and encode long context effectively, even on competitive language modeling benchmarks.

Original: [Table] 277M 0. [Table] 95M 0.
Revised: Next, we compare the performance of E XPIRE -S PAN with Adaptive-Span and Transformer-XL baselines when we vary the average span size (see Figure 7).

Original: [Table] 181M [Table] All
Revised: Models with E XPIRE -S PAN achieve stronger results — when comparing at any given memory size, E XPIRE -S PAN has better performance than both baselines.

Original: [Table] Attention [Table] 114M 0.
Revised: Further, the performance of models with E XPIRE -S PAN does not vary widely even if the memory size is drastically reduced 1 .

Original:  Compressive Trans [Table] 277M 0. Routing Trans.
Revised: Finally, we compare the performance of E XPIRE -S PAN with the Compressive Transformer baseline.

Original: [Table]

Original: [Table]

Original: 77M

Original: [Figure] [Figure] 38M [Figure]
Revised: We vary the memory size of Compressive Transformer and compress with rate 4.

Original: [Figure]

Original: indicates thatmodels can learn to expire relevant information and encode long context effectively, even on competitive language modeling benchmarks. Next, we compare the performance of E XPIRE -S PAN with Adaptive-Span and the standard Transformer when we vary the average span size (see Figure 7). Models with E XPIRE -S PAN  achieve stronger results — when comparing at any given memory size, E XPIRE -S PAN has better performance than both baselines. Further, the performance of models with E XPIRE -S PAN does not vary widely even if the memory size is drastically reduced.
Revised: We ﬁnd that E XPIRE -S PAN models achieve much better performance, as shown in Table 4 with less GPU memory and faster training time per batch.

Original: Character Level PG-

Revised: Character Level PG-

Original: We use the PG-19 (Rae et al., 2020) language modeling benchmark and convert it to model characters.
Revised: We use the PG-19 (Rae et al., 2020) language modeling benchmark and convert it to model characters.

Original: We use all characters in the training set, creating a vocabulary size of 3506 characters.
Revised: We use all characters in the training set, creating a vocabulary size of 3506 characters.

Original: We compare the performance of Transformer with an attention over the last 1K and 2K timesteps with an E XPIRE -S PAN model with maximum spans at 8K, 16K, and 32K. As shown in Figure 6, E XPIRE -S PAN models train stably with very long span, and E XPIRE -S PAN is substantially better than the Transformer baseline. 
Revised: We compare the performance of Transformer-XL with an attention over the last 1K and 2K timesteps with an E XPIRE -S PAN model with maximum spans at 8K, 16K, and 32K. As shown in Figure 6, E XPIRE -S PAN models train stably with very long span, and E XPIRE -S PAN is substantially better than the Transformer baseline. The performance improves with larger spans, [Table]

Original: The performance improves with larger spans, though extending to 32K does not increase performance beyond 16K. 
Revised: [Table] though extending to 32K does not increase performance beyond 16K. When comparing to Adaptive Span in Table [Table]

Revised: Adaptive-Span has the same performance but requires slightly more memory while bit faster.

Original: Colliding Objects Another setting where learning which long context may be important is in video understanding.
Revised: Colliding Objects Another setting where learning which long context may be important is in video understanding.

Original: Despite video data being memory intensive, salient events might be localized [Table] [Figure]
Revised: Despite video data being memory intensive, salient events might be localized in space and time.

Original: In (a) , one layer of the model has two focus areas, “Egypt” and “Alexander”. In (b) , we replace “Egypt” with “somewhere” and the focus disappears. In (c) , we insert “SpongeBob SquarePants” and the model attends to the full entity.
Revised: We test our model on a task where two objects move around and collide

Original: [Figure]
Revised: [Table]

Original: in space and time. We test our model on a task where two objects move around and collide, and the goal is to reason about the location of speciﬁed-color collisions.
Revised: and the goal is to reason about the location of speciﬁed-color collisions.

Original: Objects have a color that can randomly change.
Revised: Objects have a color that can randomly change.

Original: We divide a grid into four quadrants and the model is asked to recall the quadrants of the last collision of a speciﬁed-color pair.
Revised: We divide a grid into four quadrants and the model is asked to recall the quadrants of the last collision of a speciﬁed-color pair.

Original: Because the collisions are rare, and collisions of speciﬁc colors are even rarer, the model must process a large quantity of frames.
Revised: Because the collisions are rare, and collisions of speciﬁc colors are even rarer, the model must process a large quantity of frames.

Original: We display the task in Figure 8 and results in Table 3.
Revised: We display the task in Figure 8 and results in Table 3.

Original: The task requires many frames, so long context is very beneﬁcial — we see that as the E XPIRE -S PAN maximal span increases, performance steadily rises.
Revised: The task requires many frames, so long context is very beneﬁcial — we see that as the E XPIRE -S PAN maximal span increases, performance steadily rises.

Original: We train the 64K with the random drop regularization method described in Section 4.2.
Revised: We train the 64K with the random drop regularization method described in Section 4.2.

Original: When extending to extremely long span, 64k in size, we ﬁnd the strongest performance, which matches the size of the largest attention limit reported to date (Kitaev et al., 2019).
Revised: When extending to extremely long span, 64k in size, we ﬁnd the strongest performance, which matches the size of the largest attention limit reported to date (Kitaev et al., 2019).

Revised: As a baseline, we train an Adaptive-Span model with maximum spans limited to 16k.

Original:  A NALYSIS AND D ISCUSSION
Revised: While this model perform similar to the 16k E XPIRE -S PAN model, it uses 44% more memory (17.2gb compared to 11.9gb) and takes 2 . 8 × longer to train on a batch (365ms compared to 130ms). A 32k Adaptive-Span model run out of memory in the same setting where we trained our 64k E XPIRES PAN model, which shows the efﬁciency of our approach.

Revised: We also included a Compressive Transformer baseline with a maximum span of 8.5k (512 normal memories and 2k memories compressed by 4 × ).

Revised: This model performs worse than our model, and increasing its span does not help.

Revised: It uses the same GPU memory as our 16k model, but runs 2 .

Revised: 5 × slower, despite having a smaller span.

Revised: 5. E FFICIENCY OF E XPIRE -S PAN

Revised: We describe the efﬁciency of E XPIRE -S PAN compared to various baselines, including AdaptiveSpan, Compressive Transformer, and Transformer-XL.

Revised: We quantify with two metrics: (1) Peak

Revised: [Figure]

Revised: In (a) , one layer of the model has two focus areas, “Egypt” and “Alexander”.

Revised: In (b) , we replace “Egypt” with “somewhere” and the focus disappears.

Revised: In (c) , we insert “SpongeBob

Revised: SquarePants” and the model attends to the full entity.

Revised: [Figure]

Revised: GPU memory usage and (2) training time per batch (comparing ﬁxed batch size).

Revised: Results are shown in Table 4 and indicate that across all tasks, E XPIRE -S PAN uses substantially less GPU memory and processes each batch more quickly than all baselines, while achieving better performance in most of the tasks.

Revised: Transformer-XL as a baseline cannot adapt to the data at all, so it becomes slow and inefﬁcient quite quickly.

Revised: Adaptive-Span can adapt to the data and adjust its memory size, but this memory size is ﬁxed after training and does not have the dynamic adjustment of Expire-Span (where memory depends on context).

Revised: Finally, Compressive Transformer does compress the past memories, but it compresses them all the same amount.

Revised: In contrast, E XPIRE -S PAN can forget irrelevant content, which both improves performance by focusing on salient information, but also reduces the load on the GPU and allows for faster processing per batch.

Revised: Overall, across all tasks, we ﬁnd that E XPIRES PAN achieves strong performance with GPU memory savings and low time per batch.

Revised:  A NALYSIS AND D ISCUSSION

Original: We analyze the information retained and expired by E XPIRE -S PAN models.
Revised: We analyze the information retained and expired by E XPIRE -S PAN models.

Original: See the appendix for more analysis and ablations on the importance of regularization.
Revised: See the appendix for more analysis and ablations on the importance of regularization.

Original: Retaining Salient Information We analyze what is retained by an E XPIRE -S PAN model on enwiki8.
Revised: Retaining Salient Information We analyze what is retained by an E XPIRE -S PAN model on enwiki8.

Original: In Figure 9 (a), we show that the model retains information about named entities Egypt and
Revised: In Figure 9 (a), we show that the model retains information about named entities Egypt and

Original: Alexander the Great by attending to them (darker color).
Revised: Alexander the Great by attending to them (darker color).

Original: Next, we analyze how attention changes when we artiﬁcially edit the past text.
Revised: Next, we analyze how attention changes when we artiﬁcially edit the past text.

Original: In Figure 9 (b), we replace Egypt with somewhere , and this generic word is expired.
Revised: In Figure 9 (b), we replace Egypt with somewhere , and this generic word is expired.

Original: In Figure 9 (c), we edit Egypt to Spongebob SquarePants , which is a rare named entity, and the model attends strongly to this rather than expiring.
Revised: In Figure 9 (c), we edit Egypt to Spongebob SquarePants , which is a rare named entity, and the model attends strongly to this rather than expiring.

Original: In addition to named entities, E XPIRE -S PAN also focuses on spaces, newlines, and section titles, all of which can retain information about words, sentences or sections.
Revised: In addition to named entities, E XPIRE -S PAN also focuses on spaces, newlines, and section titles, all of which can retain information about words, sentences or sections.

Original: These can vary between different layers, indicating that E XPIRE -S PAN models use the memory at each layer differently.
Revised: These can vary between different layers, indicating that E XPIRE -S PAN models use the memory at each layer differently.

Original: See the Appendix for details.
Revised: See the Appendix for details.

Original: Importance of Long Term Memory
Revised: Importance of Long Term Memory

Original: Lastly, we analyze which predictions beneﬁt the most from memory capacity.
Revised: Lastly, we analyze which predictions beneﬁt the most from memory capacity.

Original: We take an E XPIRE -S PAN model trained on enwiki8 and artiﬁcially limit the span size at inference.
Revised: We take an E XPIRE -S PAN model trained on enwiki8 and artiﬁcially limit the span size at inference.

Original: We compare which predictions decreased in accuracy.
Revised: We compare which predictions decreased in accuracy.

Original: In Figure 10, we limit the span to 1024 tokens, and see that models have a higher loss when predicting the named entity Guinea coast compared to the 16k span.
Revised: In Figure 10, we limit the span to 1024 tokens, and see that models have a higher loss when predicting the named entity Guinea coast compared to the 16k span.

Original: Guinea coast was mentioned 3584 tokens earlier, which indicates that long attention is often necessary to predict words mentioned in far away context.
Revised: Guinea coast was mentioned 3584 tokens earlier, which indicates that long attention is often necessary to predict words mentioned in far away context.

Original: In general, we found that rare tokens and structural information about documents, such as section headings or document titles, required longer attention span.
Revised: In general, we found that rare tokens and structural information about documents, such as section headings or document titles, required longer attention span.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We present E XPIRE -S PAN , an operation that can be added to any attention mechanism to enable models to learn what to forget.
Revised: We present E XPIRE -S PAN , an operation that can be added to any attention mechanism to enable models to learn what to forget.

Original: By expiring irrelevant information, models can scale attention to tens of thousands of past memories.
Revised: By expiring irrelevant information, models can scale attention to tens of thousands of past memories.

Original: We highlight the strong performance of E XPIRE -S PAN in language modeling, reinforcement learning, object collision, and algorithmic tasks, and use it to attend to over tens of thousands of past memories.
Revised: We highlight the strong performance of E XPIRE -S PAN in language modeling, reinforcement learning, object collision, and algorithmic tasks, and use it to attend to over tens of thousands of past memories.

Original: The scalability of E XPIRE -S PAN has strong potential for allowing models to be applied to more challenging, human-like tasks that would require expiration.
Revised: The scalability of E XPIRE -S PAN has strong potential for allowing models to be applied to more challenging, human-like tasks that would require expiration.

Original: A A PPENDIX A. A DDITIONAL E XPERIMENTAL R ESULTS A.1. W IKITEXT -103 L ANGUAGE M ODELING
Revised: A A PPENDIX A. A DDITIONAL E XPERIMENTAL R ESULTS A.1. W IKITEXT -103 L ANGUAGE M ODELING

Original: The Wikitext-103 word-level language modeling benchmark (Merity et al., 2016) consists of a collection of Wikipedia articles and a ﬁxed vocabulary size of 270K. We set the max attention span for E XPIRE -S PAN to 8K.
Revised: The Wikitext-103 word-level language modeling benchmark (Merity et al., 2016) consists of a collection of Wikipedia articles and a ﬁxed vocabulary size of 270K. We set the max attention span for E XPIRE -S PAN to 8K.

Original: We compare E XPIRE -S PAN to existing work in Table 4 and show that even fairly small models trained with E XPIRE -S PAN achieve competitive results.
Revised: We compare E XPIRE -S PAN to existing work in Table 5 and show that even fairly small models trained with E XPIRE -S PAN achieve competitive results.

Original: Next, we analyze the performance of E XPIRE -S PAN on Wikitext-103 as the memory size increases.
Revised: Next, we analyze the performance of E XPIRE -S PAN on Wikitext-103 as the memory size increases.

Original: We compare to a standard Transformer model in Figure 11 — even with far smaller memory, E XPIRE -S PAN performs much better.
Revised: We compare to a standard Transformer model in Figure 11 — even with far smaller memory, E XPIRE -S PAN performs much better.

Original: A.1. E XPIRE - SPAN ANALYSIS ON ENWIK
Revised: A.1. E XPIRE - SPAN ANALYSIS ON ENWIK

Original: In Figure 12, we analyze multiple layers of the model and show that different layers identify different types of information.
Revised: In Figure 12, we analyze multiple layers of the model and show that different layers identify different types of information.

Original: Several layers retain summarizing information about sentences or sections by attending strongly to spaces, new lines, and section titles.
Revised: Several layers retain summarizing information about sentences or sections by attending strongly to spaces, new lines, and section titles.

Original: A.1. I MPORTANCE OF S TRUCTURED D ROPOUT FOR R EGULARIZATION
Revised: A.1. I MPORTANCE OF S TRUCTURED D ROPOUT FOR R EGULARIZATION

Original: We analyze the importance of structured dropout to regularize the large memory capacity provided by E XPIRE -S PAN .
Revised: We analyze the importance of structured dropout to regularize the large memory capacity provided by E XPIRE -S PAN .

Original: In an experiment on enwiki8, Figure 13 shows that loss on a portion of validation data was incredibly large.
Revised: In an experiment on enwiki8, Figure 13 shows that loss on a portion of validation data was incredibly large.

Original: This part corresponds to a 66K token long table.
Revised: This part corresponds to a 66K token long table.

Original: We hypothesize that the model likely never encountered such a table during training.
Revised: We hypothesize that the model likely never encountered such a table during training.

Original: During validation, this caused all nontable tokens to expire.
Revised: During validation, this caused all nontable tokens to expire.

Original: Without regularizing the model memory size during training, the model can easily overﬁt.
Revised: Without regularizing the model memory size during training, the model can easily overﬁt.

Original: [Table]
Revised: [Table]

Original: ., 2020) 277M
Revised: ., 2020) 277M

Original: [Figure]
Revised: [Figure]

Original: Routing Trans.
Revised: Routing Trans.

Original: (Roy et al., 2020)
Revised: (Roy et al., 2020)

Original: [Figure]
Revised: [Figure]

Original: E XPIRE -S PAN
Revised: E XPIRE -S PAN

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Expire
Revised: Expire

Original: A.1. C OLLIDING O BJECTS E ASY
Revised: A.1. C OLLIDING O BJECTS E ASY

Original: We experiment with an easier version of the Colliding Objects task where objects do not have colors.
Revised: We experiment with an easier version of the Colliding Objects task where objects do not have colors.

Original: The model has to predict either the last collision, or a mapping of the last 3 collisions.
Revised: The model has to predict either the last collision, or a mapping of the last 3 collisions.

Original: In contrast to the harder task, there are no color switches and any collision prediction is valid.
Revised: In contrast to the harder task, there are no color switches and any collision prediction is valid.

Original: [Table]
Revised: [Table]

Original: A. A DDITIONAL I MPLEMENTATION D ETAILS A.2. LIGHT T ASKS
Revised: A. A DDITIONAL I MPLEMENTATION D ETAILS A.2. LIGHT T ASKS

Original: Instruction Following Task
Revised: Instruction Following Task

Original: We train 6-layer models with a hidden size of 512 and 8 attention heads.
Revised: We train 6-layer models with a hidden size of 512 and 8 attention heads.

Original: To train, we use the Adam optimizer with a learning rate of 7e-4 and 8000 warmup updates.
Revised: To train, we use the Adam optimizer with a learning rate of 7e-4 and 8000 warmup updates.

Original: We set the Expire-Span ramp to 64 and the Expire-Span loss to 2e-6.
Revised: We set the Expire-Span ramp to 64 and the Expire-Span loss to 2e-6.

Original: A.2. R EINFORCEMENT L EARNING
Revised: A.2. R EINFORCEMENT L EARNING

Original: We used MazeBase (Sukhbaatar et al., 2015a) to construct tasks in grid world.
Revised: We used MazeBase (Sukhbaatar et al., 2015a) to construct tasks in grid world.

Original: Agents can observe its surrounding 3 × 3 area and move in the four cardinal directions.
Revised: Agents can observe its surrounding 3 × 3 area and move in the four cardinal directions.

Original: Every objects and their properties are described by words such as “agent”,“block”, “blue”, etc.
Revised: Every objects and their properties are described by words such as “agent”,“block”, “blue”, etc.

Original: Thus, the input to the model is a binary tensor of size 3 × 3 × vocabulary-size.
Revised: Thus, the input to the model is a binary tensor of size 3 × 3 × vocabulary-size.

Original: We train 2-layer Transformers with 64 hidden units using actor-critic algorithm.
Revised: We train 2-layer Transformers with 64 hidden units using actor-critic algorithm.

Original: We used a BPTT length of 100, and an entropy cost of 0.0005.
Revised: We used a BPTT length of 100, and an entropy cost of 0.0005.

Original: Corridor Task
Revised: Corridor Task

Original: The corridor length is sampled from U (3 , 200) .
Revised: The corridor length is sampled from U (3 , 200) .

Original: All models are trained for 100M steps.
Revised: All models are trained for 100M steps.

Original: We used RMSProp optimizer with a learning rate of 0.0001 and a batch size of 64.
Revised: We used RMSProp optimizer with a learning rate of 0.0001 and a batch size of 64.

Original: For the expire-span models, we set the maximum span L to 200, the loss coefﬁcient α to 5e-6 and the ramp length R to 16.
Revised: For the expire-span models, we set the maximum span L to 200, the loss coefﬁcient α to 5e-6 and the ramp length R to 16.

Original: Multi-Room Portal
Revised: Multi-Room Portal

Original: In this task, there are 50 rooms sequentially connected together.
Revised: In this task, there are 50 rooms sequentially connected together.

Original: Each room is 5 × 5 in size, and have two doors with different colors.
Revised: Each room is 5 × 5 in size, and have two doors with different colors.

Original: If agent go to the correct door, it will be teleported to the next room, but if it is the wrong door, the agent will be teleported back to the ﬁrst room and have to start over.
Revised: If agent go to the correct door, it will be teleported to the next room, but if it is the wrong door, the agent will be teleported back to the ﬁrst room and have to start over.

Original: Which of the two doors is correct in each room is randomly decided and ﬁxed throughout the episode.
Revised: Which of the two doors is correct in each room is randomly decided and ﬁxed throughout the episode.

Original: This information is not visible to the agent, thus can only be discovered by trial and error within each episode.
Revised: This information is not visible to the agent, thus can only be discovered by trial and error within each episode.

Original: The current room number is visible to the agent.
Revised: The current room number is visible to the agent.

Original: When the agent successfully transitions from the k -th room to the next, it receives a reward of 0 .
Revised: When the agent successfully transitions from the k -th room to the next, it receives a reward of 0 .

Original: 1 k .
Revised: 1 k .

Original: The episode ends if the agent makes two mistakes in the same room, reaches the last room, or when the number of steps reach 1000.
Revised: The episode ends if the agent makes two mistakes in the same room, reaches the last room, or when the number of steps reach 1000.

Original: A reward discount of 0.98 is used.
Revised: A reward discount of 0.98 is used.

Original: All models are trained with
Revised: All models are trained with

Original: Adam optimizer with a learning rate of 0.0005, and a batch size of 1024, with gradients are clipped at 0.1.
Revised: Adam optimizer with a learning rate of 0.0005, and a batch size of 1024, with gradients are clipped at 0.1.

Original: We set L = 100 , R = 16 and α = 1e-6 for the expire-span models.
Revised: We set L = 100 , R = 16 and α = 1e-6 for the expire-span models.

Original: A.2. C OLLISION T ASK
Revised: A.2. C OLLISION T ASK

Original: At the start of the simulation, each particle samples a Gaussian Normal velocity and position uniform inside a 16 × 16 grid.
Revised: At the start of the simulation, each particle samples a Gaussian Normal velocity and position uniform inside a 16 × 16 grid.

Original: At each time step the particles’ position is updated by adding its velocity (unless it would go off the grid, in which case its velocity is re-sampled).
Revised: At each time step the particles’ position is updated by adding its velocity (unless it would go off the grid, in which case its velocity is re-sampled).

Original: There are 5 different colors, and a particle can change its color randomly at each step with 0.05 probability.
Revised: There are 5 different colors, and a particle can change its color randomly at each step with 0.05 probability.

Original: A collision happens when the two particles have the same rasterized locations, but it does not affect the movement.
Revised: A collision happens when the two particles have the same rasterized locations, but it does not affect the movement.

Original: Given a question specifying two colors, the task is to report in which of the four quadrants of the grid the last collision of the speciﬁed-colors occurred.
Revised: Given a question specifying two colors, the task is to report in which of the four quadrants of the grid the last collision of the speciﬁed-colors occurred.

Original: To make the task easier to learn, 40% of the queries will have the matching colors as the last collision.
Revised: To make the task easier to learn, 40% of the queries will have the matching colors as the last collision.

Original: The model is given an input sequence of tokens that has 8 entries per timestep.
Revised: The model is given an input sequence of tokens that has 8 entries per timestep.

Original: The ﬁrst 4 are the rounded and rasterized ( x, y ) locations of the two particles, and next 2 are tokens representing the colors of the particles.
Revised: The ﬁrst 4 are the rounded and rasterized ( x, y ) locations of the two particles, and next 2 are tokens representing the colors of the particles.

Original: The last 2 entries are “question” tokens that specify the colors of the collision.
Revised: The last 2 entries are “question” tokens that specify the colors of the collision.

Original: The output sequence has a token for each quadrant.
Revised: The output sequence has a token for each quadrant.

Original: We generate 50M steps for training, which equals to 400M entries.
Revised: We generate 50M steps for training, which equals to 400M entries.

Original: Easy Version: The particles have no color in this version.
Revised: Easy Version: The particles have no color in this version.

Original: There are two types of questions, in which the task is to report either
Revised: There are two types of questions, in which the task is to report either

Original: in which
Revised: in which

Original: of the
Revised: of the

Original: four
Revised: four

Original: quadrants of the grid
Revised: quadrants of the grid

Original: the last
Revised: the last

Original: collision occurred
Revised: collision occurred

Original: the label mapping of the last 3 collisions.
Revised: the label mapping of the last 3 collisions.

Original: A.2. L ANGUAGE M ODELING D ETAILS
Revised: A.2. L ANGUAGE M ODELING D ETAILS

Original: enwiki8 We train 12-layer models with a hidden size of 512 and 8 attention heads.
Revised: enwiki8 We train 12-layer models with a hidden size of 512 and 8 attention heads.

Original: To train, we use Adam optimizer with a learning rate of 7e-4, a batch size of 512, a block size of 512 andwarmup updates.
Revised: To train, we use Adam optimizer with a learning rate of 7e-4, a batch size of 512, a block size of 512 andwarmup updates.

Original: All models are trained for 100k updates.
Revised: All models are trained for 100k updates.

Original: The model in Table 2 is further ﬁne-tuned for another 10k updates with a 10x smaller LR.
Revised: The model in Table 2 is further ﬁne-tuned for another 10k updates with a 10x smaller LR.

Original: The baseline models used for comparison are the same size model following the same training protocol.
Revised: The baseline models used for comparison are the same size model following the same training protocol.

Original: PG-19 Besides the attention span, all model parameters and training parameters were held constant.
Revised: PG-19 Besides the attention span, all model parameters and training parameters were held constant.

Original: Each model had 12 layers, a hidden size of 512, a feed forward hidden size of 2048, 8 attention heads, and processed 512 characters at a time.
Revised: Each model had 12 layers, a hidden size of 512, a feed forward hidden size of 2048, 8 attention heads, and processed 512 characters at a time.

Original: We initialized the weights using a uniform distribution as described by Glorot & Bengio (2010), used dropout and attention dropout of 0.2, clipped the gradients at 0.3, warmed up the learning rate linearly for 8000 steps, and used cosine annealing to decay the learning rate after warmup (Loshchilov & Hutter, 2016).
Revised: We initialized the weights using a uniform distribution as described by Glorot & Bengio (2010), used dropout and attention dropout of 0.2, clipped the gradients at 0.3, warmed up the learning rate linearly for 8000 steps, and used cosine annealing to decay the learning rate after warmup (Loshchilov & Hutter, 2016).

Original: For the E XPIRE -S PAN models, we used a ramp of 128 and an expiration loss coefﬁcient of 1e-6.
Revised: For the E XPIRE -S PAN models, we used a ramp of 128 and an expiration loss coefﬁcient of 1e-6.

Original: Wikitext-103 All models have 8 layers and 1024 hidden units (4096 in ReLU layers).
Revised: Wikitext-103 All models have 8 layers and 1024 hidden units (4096 in ReLU layers).

Original: In addition to the dropout of 0.3 applied to attention and ReLU activation, outputs from the embedding layer and the last layer had a dropout of 0.2.
Revised: In addition to the dropout of 0.3 applied to attention and ReLU activation, outputs from the embedding layer and the last layer had a dropout of 0.2.

Original: We used the adaptive input Baevski & Auli (2019) and the adaptive softmax Grave et al.
Revised: We used the adaptive input Baevski & Auli (2019) and the adaptive softmax Grave et al.

Original: (2017) for reducing the number of parameters within word embeddings.
Revised: (2017) for reducing the number of parameters within word embeddings.

Original: The models are trained for 300k updates with a block size of 256, and gradients are clipped at 0.1.
Revised: The models are trained for 300k updates with a block size of 256, and gradients are clipped at 0.1.

Original: The other hyperparameters are the same as the enwik8 experiments.
Revised: The other hyperparameters are the same as the enwik8 experiments.