Original: N EURAL A TTENTION M EMORY
Revised: N EURAL A TTENTION M EMORY

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: Scaled dot-product attention has become the essence of state-of-the-art deep neural networks for variety of machine learning tasks.
Revised: Scaled dot-product attention has become the essence of state-of-the-art deep neural networks for various machine learning tasks.

Revised: Though its ubiquitous accomplishments, it is inefficient for long sequence tasks and problematic for tasks requiring memory states such as compositional generalization.

Original: We propose a novel perspective of the attention mechanism by reinventing it as a memory architecture for neural networks, namely neural attention memory (NAM).
Revised: We propose a novel perspective of the attention mechanism by reinventing it as a memory architecture for neural networks, namely Neural Attention Memory (NAM).

Original: Following the same query-keyvalue structure, NAM writes a memory matrix via  sum of outer products of value and unit key vectors. Reading the memory can be done by multiplying the matrix with a unit query vector.
Revised: NAM follows the same query-key-value structure by constructing a memory matrix while reducing its computational complexity from quadratic to linear to the sequence length. NAM writes a memory matrix via the sum of outer products of value and unit key vectors, and reads it by multiplying the matrix with a unit query vector.

Original: It is possible to implement a mechanism theoretically equivalent to  attention with NAM primitives, namely normalized outer-product attention.
Revised: Indeed, we show that our normalized outer-product attention mechanism is mathematically equivalent to the conventional attention mechanism. Then, we evaluate a NAM-based

Original: One benefit of NAM is that its computational complexity is linear to the sequence length while the scaled dot-product attention’s complexity is quadratic to it.

Original: We evaluate a NAM-based Transformer on long-range arena tasks and show its efficiency and efficacy.
Revised: Transformer on long-range arena tasks and demonstrate its efficiency and efficacy.

Original: Also, we design two NAM-based memory-augmented neural networks namelyLong short-term attention memory (LSAM) and NAM
Revised: Finally, we propose two NAM-based memory-augmented neural networks, namely

Original:  Turing machine (NAM-TM).
Revised: Long Short-Term Attention Memory (LSAM) and NAM Turing Machine (NAMTM), and test their compositional generalization capability using four different tasks.

Original: LSAM replaces LSTM’s long-term cell state with NAM memory matrix and NAM-TM implements a Turing tape data structure using NAM read/write primitives.
Revised: LSAM replaces LSTM’s long-term cell state with NAM memory matrix and NAM-TM implements a Turing tape data structure using NAM read/write primitives.

Original: We test their capability using four different types of compositional generalization tasks.

Original: The experimental results show that both LSAM and NAM-TM outperform traditional Transformer and LSTM, as well as 
Revised: The experimental results show that the proposed models outperform traditional Transformer and LSTM, as well as DNC.

Original: DNC, a complicated memory-augmented neural network model.

Original: Moreover, NAM opens up possibilities in a variety of machine learning research problems including hierarchical data modeling, efficient edge inference, and few-shot learning.
Revised: NAM opens up possibilities in diverse machine learning research problems, including hierarchical data modeling, efficient edge inference, and few-shot learning.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Scaled dot-product attention (Vaswani et al., 2017) has become a core mechanism of state-of-the-art deep learning models for variety of machine learning tasks, including natural language processing (Devlin et al., 2018), multi-modal task (Li et al., 2019), and graph data processing (Hamilton et al., 2017).
Revised: Scaled dot-product attention (Vaswani et al., 2017) has become a core mechanism of state-of-the-art deep learning models for variety of machine learning tasks, including natural language processing (Devlin et al., 2018), multi-modal task (Li et al., 2019), and graph data processing (Hamilton et al., 2017).

Original: Specifically, the Transformers using the self-attention method based on it has replaced recurrent neural network (RNN)based methods by outperforming them.
Revised: Specifically, the Transformers using the self-attention method have replaced recurrent neural networks (RNN) by outperforming them in most of the tasks.

Original: Despite its success, there exist limitations to the mechanism.
Revised: Despite its success, there exist limitations to the mechanism.

Original: First, it needs the information of the entire sequence to compute one attention so that its computational complexity becomes quadratic to the length of the sequence.
Revised: First, it needs the information of the entire sequence to compute one attention so that its computational complexity becomes quadratic to the length of the sequence.

Original: Hence, it is inefficient for long sequence tasks (Tay et al., 2020) or edge inference environments (Tambe et al., 2020).
Revised: Hence, it is inefficient for long sequence tasks (Tay et al., 2020) or edge inference environments (Tambe et al., 2020).

Original: Also, its stateless design enables efficient parallelism but makes it impossible to solve tasks that require memory states.
Revised: Also, its stateless design enables efficient parallelism but makes it impossible to solve tasks that require memory states.

Original: Hence, Transformers fail to generalize the rules that require inductive bias (Dehghani et al., 2018) or compositional generalization (Lake & Baroni, 2018).
Revised: Hence, Transformers fail to generalize the rules that require inductive bias (Dehghani et al., 2018) or compositional generalization (Lake & Baroni, 2018).

Original: There have been studies designing neural networks with external memory to solve algorithmic tasks where Transformers fail.
Revised: There have been studies designing neural networks with external memory to solve algorithmic tasks where Transformers fail.

Original: These memory-augmented neural networks (MANN) design differentiable read/write functions that can be trained by back-propagation.
Revised: These memory-augmented neural networks (MANN) design differentiable read/write functions that can be trained by backpropagation.

Original: Some of them implement basic data structures like stack (Joulin & Mikolov, 2015) and queue (Grefenstette et al., 2015) while some implement complex memory structures using attention mechanisms (Graves et al., 2014; 2016).
Revised: Some of them implement basic data structures like stack (Joulin & Mikolov, 2015) and queue (Grefenstette et al., 2015) while some implement complex memory structures using attention mechanisms (Graves et al., 2014; 2016).

Original: They outperform generic neural networks in synthetic algorithmic tasks but are considered impractical due to their complexities and inefficiencies.
Revised: They outperform generic neural networks in synthetic algorithmic tasks but are considered impractical due to their complexities and inefficiencies.

Original: In this work, we re-invent the attention mechanism as a memory architecture for neural networks, namely neural attention memory (NAM).
Revised: In this work, we re-invent the attention mechanism as a memory architecture for neural networks, namely neural attention memory (NAM).

Original: NAM’s design objective is to build simple, efficient, yet powerful external memory which also incorporates the attention mechanism.
Revised: NAM’s design objective is to build simple, efficient, yet powerful external memory which also incorporates the attention mechanism.

Original: Following the same query-key-value structure of attention, NAM stores key-value pairs to a memory matrix via additively writing their outer-products.
Revised: Following the same query-key-value structure of attention, NAM stores key-value pairs to a memory matrix via additively writing their outer-products.

Original: Reading the memory matrix is simply done by multiplying the matrix with a unit query vector.
Revised: Reading the memory matrix is simply done by multiplying the matrix with a unit query vector.

Original: We provide mathematical formulation for the read/write primitives, and make theoretical analyses showing that these read and write primitives can replace attention.
Revised: We provide mathematical formulation for the read/write primitives, and make theoretical analyses showing that these read and write primitives can replace attention.

Original: One big benefit of NAM is that it can perform attention in a more efficient way.
Revised: One big benefit of NAM is that it can perform attention in a more efficient way.

Original: NAM-based attention, namely normalized outer-product attention, has linear computational complexity to the sequence length, which is a goal for efficient long-range Transformers (Choromanski et al., 2020; Beltagy et al., 2020; Kitaev et al., 2020).
Revised: NAM-based attention, namely normalized outer-product attention, has linear computational complexity to the sequence length, which is a goal for efficient long-range Transformers (Choromanski et al., 2020; Beltagy et al., 2020; Kitaev et al., 2020).

Original: Mathematically, it is very similar to linear attention (Katharopoulos et al., 2020) in that it computes the multiplication of keys and values first to reduce the complexity.
Revised: Mathematically, it is very similar to linear attention (Katharopoulos et al., 2020) in that it computes the multiplication of keys and values first to reduce the complexity.

Original: We evaluate NAM-based efficient Transformer in long-range arena (Tay et al., 2020) tasks.
Revised: We evaluate NAM-based efficient Transformer in long-range arena (Tay et al., 2020) tasks.

Original: Its efficacy is on par with the base Transformer and Linear Transformer, implying that NAM can be an efficient alternative to  scaled dot-product attention.
Revised: Its efficacy is on par with the base Transformer and Linear Transformer, implying that NAM can be an efficient alternative to the scaled dot-product attention.

Original: Using NAM read/write primitives, we design two memory-augmented neural networks (MANN), namely Long short-term attention memory (LSAM) and NAM Turing machine (NAM-TM).
Revised: Using NAM read/write primitives, we design two memory-augmented neural networks (MANN), namely Long Short-term Attention Memory (LSAM) and NAM Turing Machine (NAM-TM).

Original: LSAM is a generic RNN architecture that replaces LSTM’s long-term cell state with a memory matrix.
Revised: LSAM is a generic RNN architecture that replaces LSTM’s long-term cell state with a memory matrix.

Original: Instead of additively writing a vector cell state, LSAM reads and writes the memory matrix using NAM primitives.
Revised: Instead of additively writing a vector cell state, LSAM reads and writes the memory matrix using NAM primitives.

Original: The design combines strengths of attention and RNN while maintaining the same computational complexity as LSTM.
Revised: The design combines strengths of attention and RNN while maintaining the same computational complexity as LSTM.

Original: NAM-TM is a MANN for algorithmic tasks, constructing a Turing tape structure.
Revised: NAM-TM is a MANN for algorithmic tasks, leveraging a Turing tape structure.

Original: A tape has read and write heads accessing the memory with NAM read/write primitives.
Revised: A tape has read and write heads accessing the memory with NAM read/write primitives.

Original: They can move along the tape with four actions: NO-OP, LEFT, RIGHT, and JUMP.
Revised: They can move along the tape with four actions: NO-OP, LEFT, RIGHT, and JUMP.

Original: The actions are implemented as differentiable functions to enable end-to-end training with backpropagation.
Revised: The actions are implemented as differentiable functions to enable end-to-end training with backpropagation.

Original: We compare LSAM and NAM-TM to others in compositional generalization tasks of number sequence prediction (Nam et al., 2019), sequence reduction, and SCAN (Lake & Baroni, 2018).
Revised: We compare LSAM and NAM-TM to others in compositional generalization tasks of number sequence prediction (Nam et al., 2019), sequence reduction, and SCAN (Lake & Baroni, 2018).

Original: Specifically, we test their zero-shot generalization capability in length by training the models withshorter sequences  and validating them with longer sequences unobservable during training.
Revised: Specifically, we test their zero-shot generalization capability in length by training the models with sequences of limited length and validating them with longer sequences unobservable during training.

Original: The evaluation results show that their computational powers are superior to other neural networks, includingUniversal
Revised: The evaluation results show that their computational powers are superior to other baselines, including

Original:  Transformer (Dehghani et al., 2018) and DNC (Graves et al., 2016).
Revised: Universal Transformer (Dehghani et al., 2018) and DNC (Graves et al., 2016).

Original: While the generic LSAM model consistently outperforms the others, NAM-TM shows even better results at algorithmic tasks.
Revised: While the generic LSAM model consistently outperforms the others, NAM-TM shows even better results at algorithmic tasks.

Original: The results indicate that NAM is a powerful method to implement memory in neural networks.
Revised: The results indicate that NAM is a powerful method to implement memory in neural networks.

Original: The efficient, simple, and flexible structure of NAM opens up new possibilities in multiple machine learning research fields.
Revised: The efficient, simple, and flexible structure of NAM opens up new possibilities in multiple machine learning research fields.

Original: One straightforward application is leveraging NAM’s efficiency for edge inference environment.
Revised: One straightforward application is leveraging NAM’s efficiency for edge inference environment.

Original: Another possibility is using NAM for hierarchical data modeling by generalizing NAM with tensor product .
Revised: Another possibility is using NAM for hierarchical data modeling by generalizing NAM with tensor product .

Original: Moreover, memorization of input-output mapping using NAM can be a solution for  few-shot learning.
Revised: Moreover, memorization of input-output mapping using NAM can be a solution for one-shot and few-shot learning.

Original: The main contributions of this work are as follows:
Revised: The main contributions of this work are as follows:

Original: • We re-invent the attention mechanism as a memory architecture for neural networks, namely neural attention memory (NAM).
Revised: • We re-invent the attention mechanism as a memory architecture for neural networks, namely neural attention memory (NAM).

Original: • We present mathematical basis for NAM read/write primitives, and give theoretical proofs that NAM is equivalent to attention in certain conditions.
Revised: • We present mathematical basis for NAM read/write primitives, and give theoretical proofs that NAM is equivalent to attention in certain conditions.

Original: • We show that NAM can construct an efficient Transformer for long-range sequence tasks.
Revised: • We show that NAM can construct an efficient Transformer for long-range sequence tasks.

Original: • We propose two memory-augmented neural network designs of LSAM and NAM-TM and show their capabilities in compositional generalization tasks.
Revised: • We propose two memory-augmented neural network designs of LSAM and NAM-TM and show their capabilities in compositional generalization tasks.

Original:  B ACKGROUND 2. S CALED DOT - PRODUCT ATTENTION
Revised:  B ACKGROUND 2. S CALED DOT - PRODUCT ATTENTION

Original: Attention mechanisms of deep neural networks (Bahdanau et al., 2014; Luong et al., 2015) provide differentiable methods of choosing (attending) one item from a variable-length sequence.
Revised: Attention mechanisms of deep neural networks (Bahdanau et al., 2014; Luong et al., 2015) provide differentiable methods of choosing (attending) one item from a variable-length sequence.

Original: While there are multiple variations of attention mechanism, most of them share the same high-level structure: 1) compute the attention scores of the items, and 2) return the weighted sum of their vector representations using the scores.
Revised: While there are multiple variations of attention mechanism, most of them share the same high-level structure: 1) compute the attention scores of the items, and 2) return the weighted sum of their vector representations using the scores.

Original: Among the variations, scaled dot-product attention (Vaswani et al., 2017) has been the most successful.
Revised: Among the variations, scaled dot-product attention (Vaswani et al., 2017) has been the most successful.

Original: For each token, there are a key vector and a value vector associated to it.
Revised: For each token, there are a key vector and a value vector associated to it.

Original: Given a query vector, the scores are determined by the scaled dot-product of the query and the keys.
Revised: Given a query vector, the scores are determined by the scaled dot-product of the query and the keys.

Original: Then the output is computed by weighted sum of softmaxed scores and the value vectors.
Revised: Then the output is computed by weighted sum of softmaxed scores and the value vectors.

Original: The self-attention mechanism based on the scaled dot-product attention has proven to be very powerful,even replacing the needs of recurrent neural networks (RNN).
Revised: The self-attention mechanism based on the scaled dot-product attention has proven to be very powerful, replacing the needs of RNNs.

Original: Since attention reaches every element of a sequence in an O (1) path, it avoids the vanishing gradient problem (Hochreiter, 1998), enjoys high parallelism, and allows huge models with deeply stacked layers (Devlin et al., 2018; Brown et al., 2020).
Revised: Since attention reaches every element of a sequence in an O (1) path, it avoids the vanishing gradient problem (Hochreiter, 1998), enjoys high parallelism, and allows huge models with deeply stacked layers (Devlin et al., 2018; Brown et al., 2020).

Original: However, its stateless and parallel architecture also bring multiple limitations.
Revised: However, its stateless and parallel architecture also bring multiple limitations.

Original: First, the computational cost quadratically increases with the sequence length, making it inefficient for longrange contexts (Beltagy et al., 2020; Choromanski et al., 2020; Kitaev et al., 2020; Katharopoulos et al., 2020) and edge inference environments (Tambe et al., 2020).
Revised: First, the computational cost quadratically increases with the sequence length, making it inefficient for long-range contexts (Beltagy et al., 2020; Choromanski et al., 2020; Kitaev et al., 2020; Katharopoulos et al., 2020) and edge inference environments (Tambe et al., 2020).

Original: Also, the memory-less architecture lacks inductive bias (Dehghani et al., 2018), making it impossible to generalize inductive algorithmic rules (Kim et al., 2021).
Revised: Also, the memory-less architecture lacks inductive bias (Dehghani et al., 2018), making it impossible to generalize inductive algorithmic rules (Kim et al., 2021).

Original: There aremultitudes of researches to resolve those issues, but a simple and practical solution is yet to be found.
Revised: There are researches to resolve those issues, but a simple and practical solution is yet to be found.

Original: 2. M EMORY - AUGMENTED NEURAL NETWORK
Revised: 2. M EMORY - AUGMENTED NEURAL NETWORK

Original: In theory, RNNs are proven to be as powerful as Turing machines (Hyötyniemi, 1996).
Revised: In theory, RNNs are proven to be as powerful as Turing machines (Hyötyniemi, 1996).

Original: However, they fail to learn algorithmic patterns that require more that pushdown automata in practice (Nam et al., 2019).
Revised: However, they fail to learn algorithmic patterns that require more that pushdown automata in practice (Nam et al., 2019).

Original: Therefore, there have been efforts to augment external memory architecture to neural networks, as known as memory-augmented neural networks (MANN) (Grefenstette et al., 2015; Joulin &
Revised: Therefore, there have been efforts to augment external memory architecture to neural networks, as known as memory-augmented neural networks (MANN) (Grefenstette et al., 2015; Joulin &

Original: Mikolov, 2015; Graves et al., 2014; 2016).
Revised: Mikolov, 2015; Graves et al., 2014; 2016).

Original: The main challenge for MANNs is to design differentiable read and write functions that can be trained via back-propagation.
Revised: The main challenge for MANNs is to design differentiable read and write functions that can be trained via back-propagation.

Original: Some MANNs use attention mechanism for differentiable read/write functions.
Revised: Some MANNs use attention mechanism for differentiable read/write functions.

Original: For instance, neural turing machine (NTM) (Graves et al., 2014) and differentiable neural computer (DNC) (Graves et al., 2016) leverage attention mechanism for implementing content-based addressing.
Revised: For instance, neural turing machine (NTM) (Graves et al., 2014) and differentiable neural computer (DNC) (Graves et al., 2016) leverage attention mechanism for implementing content-based addressing.

Original: However, the addressing mechanisms are often not powerful enough so that they need to be augmented with complex extras, such as a link matrix with O ( N 2 ) cost.
Revised: However, the addressing mechanisms are often not powerful enough so that they need to be augmented with complex extras, such as a link matrix with O ( N 2 ) cost.

Original: Hence, MANNs often become inefficient and complex so that they are considered impractical outside of the algorithmic task domain.
Revised: Hence, MANNs often become inefficient and complex so that they are considered impractical outside of the algorithmic task domain.

Original:  N EURAL A TTENTION M EMORY
Revised:  N EURAL A TTENTION M EMORY

Original: In this section, we define a memory structure reconstructed from the attention mechanism, namely neural attention memory (NAM).
Revised: In this section, we define a memory structure reconstructed from the attention mechanism, namely neural attention memory (NAM).

Original: We mathematically define the read/write primitives for NAM, and give theoretical proof that it can replace attention.
Revised: We mathematically define the read/write primitives for NAM, and give theoretical proof that it can replace attention.

Original: The main idea of NAM is implementing an attention mechanism via matrix-vector multiplication of a memory matrix M ∈ R d v × d k , a unit query vector q ∈ R d k and a read probability 0 ≤ p r ≤ 1 .
Revised: The main idea of NAM is implementing an attention mechanism via matrix-vector multiplication of a memory matrix M ∈ R d v × d k , a unit query vector q ∈ R d k and a read probability 0 ≤ p r ≤ 1 .

Original: Hereby d v and d k are feature dimensions of value and key vectors respectively.
Revised: Hereby d v and d k are feature dimensions of value and key vectors respectively.

Original: Then, the read operation ( RD ) of NAM is defined as follows.
Revised: Then, the read operation ( RD ) of NAM computes the read vector r ∈ R d v as follows.

Original: [Equation]
Revised: [Equation]

Original: While there can be multiple ways to construct the memory M , the simplest method is sum of outer products (or tensor products) of unit key vectors k i ∈ R d k and value vectors v i ∈ R d v as follows.
Revised: While there can be multiple ways to construct the memory M , the simplest method is sum of outer products (or tensor products) of unit key vectors k i ∈ R d k and value vectors v i ∈ R d v as follows.

Original: [Equation]
Revised: [Equation]

Original: Hereby S denotes the number of tokens (or length) in a sequence.
Revised: Hereby S denotes the number of tokens (or length) in a sequence.

Original: Since both query and key vectors are unit vectors , reading the memory matrix can act as an attention mechanism when the keys are orthonormal as the following theorem holds.
Revised: Since both query and key vectors are unit vectors , reading the memory matrix can act as an attention mechanism when the keys are orthonormal as the following theorem holds.

Original: Theorem 1.
Revised: Theorem 1.

Original: If k 1 . . .
Revised: If k 1 . . .

Original: k S are orthonormal unit vectors and i = 1 ...S, RD( M, k i , 1 )
Revised: k S are orthonormal unit vectors and i = 1 ...S, RD( M, k i , 1 )

Original: =
Revised: =

Original: v i
Revised: v i

Original: The proof of the theorem is trivial because v j k ⊤ j k i is v i when j = i and zero otherwise.
Revised: The proof of the theorem is trivial because v j k ⊤ j k i is v i when j = i and zero otherwise.

Original: However, it is hard to guarantee that all key vectors are orthonormal to each other.
Revised: However, it is hard to guarantee that all key vectors are orthonormal to each other.

Original: Hence, a more general write ( WR ) operation can be defined with a write probability p w and an erase probability p e as below.
Revised: Hence, a more general write ( WR ) operation can be defined with a write probability p w and an erase probability p e as below.

Original: [Equation]
Revised: [Equation]

Original: Verbally speaking, this write operation adds the new key-value pair vk ⊤ to the matrix while erasing the existing value Mk .
Revised: Verbally speaking, this write operation adds the new key-value pair vk ⊤ to the matrix while erasing the existing value Mk .

Original: The sum of outer products (cid:80) v i
Revised: The sum of outer products (cid:80) v i

Original: k ⊤ i is a special case of this when p w = 1 and p e = [-
-] for everyk i and v i .
Revised: k ⊤ i is a special case of this when p w = 1 and p e = [-
-]0 for every i .

Original: It provides a softer guarantee as follows, without requiring orthonormality.
Revised: It provides a softer guarantee as follows, without requiring orthonormality.

Original: Theorem 2.
Revised: Theorem 2.

Original: If k is a unit vector and M ′ =
Revised: If k is a unit vector and M ′ =

Original: WR ( M, k, v, 1 , 1) , RD ( M ′ , k, 1) = v .
Revised: WR ( M, k, v, 1 , 1) , RD ( M ′ , k, 1) = v .

Original: Again, the proof is trivial because M ′ k = Mk + vk ⊤ k − Mkk ⊤ k = v .
Revised: Again, the proof is trivial because M ′ k = Mk + vk ⊤ k − Mkk ⊤ k = v .

Original: In other words, after a WR operation, the RD operation with the same key yields the most recently written value.
Revised: In other words, after a WR operation, the RD operation with the same key yields the most recently written value.

Original: These simple and powerful read/write primitives are also efficient.
Revised: These simple and powerful read/write primitives are efficient in many ways.

Original: First of all, the computational complexity of each read/write operation is O
Revised: First of all, the computational complexity of each read/write operation is O

Original: ( d k × d v ) which is identical to complexity of otherneural network layers and does not grow with thesequence length S .
Revised: ( d k × d v ) which is identical to complexity of other layers and does not depend on the length S .

Original: Moreover, the outer productoperation can be efficiently computed with modern hardware as it  high compute density O ( d k × d v ) and low amount of memory accessO
Revised: Moreover, the outer product can be efficiently computed with modern hardware as it has high compute density O ( d k × d v ) and low amount of memory access

Original:  ( d k + d v ) .
Revised: O ( d k + d v ) .

Original: Finally, the space complexity of the memory matrix O ( d k × d v ) also does not rely on the sequence length, which makes it desirable at resource-constrained edge  environment.
Revised: Finally, the space complexity of the memory matrix O ( d k × d v ) also does not rely on S , which makes it desirable at resource-constrained edge inference environment.

Original:  E FFICIENT T RANSFORMER WITH NAM 4. N ORMALIZED OUTER - PRODUCT ATTENTION
Revised:  E FFICIENT T RANSFORMER WITH NAM 4. N ORMALIZED OUTER - PRODUCT ATTENTION

Original: One popular limitation of the scaled dot-product attention is that its computational complexity is quadratic to the length of the sequence.
Revised: One popular limitation of the scaled dot-product attention is that its computational complexity is quadratic to the length of the sequence.

Original: This becomes problematic when handling the very long sequences (>1000 tokens) and deploying the Transformers for edge inference.
Revised: This becomes problematic when handling the very long sequences (>1000 tokens) and deploying the Transformers for edge inference.

Original: One benefit of NAM is that its computation complexity is linear to to the length of the input sequence.
Revised: One benefit of NAM is that its computation complexity is linear to to the length of the input sequence.

Original: Hence, we can use
Revised: Hence, we can use

Original: NAM as an efficient attention mechanism.
Revised: NAM as an efficient attention mechanism.

Original: For scaled dot-product attention, we first compute dot-products (inner product) of the keys and the query to get the similarity scores.
Revised: For scaled dot-product attention, we first compute dot-products (inner product) of the keys and the query to get the similarity scores.

Original: This computation has quadratic complexity to the sequence length because we need to compute the score of every key-query pair.
Revised: This computation has quadratic complexity to the sequence length because we need to compute the score of every key-query pair.

Original: On the other hand, NAM-based  attention  can reduce the complexity by changing the order of computation, namely normalized outer-product attention .
Revised: On the other hand, NAM-based attention, namely normalized outer-product attention , can reduce the complexity by changing the order of computation.

Original: As we do for NAM,  the keys and queriesare normalized to unit vectors to eliminate the softmax function.
Revised: As we do for NAM, we first normalize the keys and queries to unit vectors to avoid using the softmax function.

Revised: Instead, the scores are computed by inner products of the unit keys and queries, automatically normalized to the range of [ − 1 , 1] .

Original: Then wefirst construct a memory matrix by computing sum of outer products of the keys and the values.
Revised: Then we construct a memory matrix by computing sum of outer products of the keys and the values.

Original: Finally the attention result can be computed by reading the memory matrix with the queries.
Revised: Finally the attention result can be computed by reading the memory matrix with the queries.

Original: Given the query vector q ∈ R d k , and key-value sequences K ∈ R S × d k , V ∈ R S × d v of length S , the normalized outer-product attention can be computed as follows.
Revised: Given the query vector q ∈ R d k , and key-value sequences K ∈ R S × d k , V ∈ R S × d v of length S , the normalized outer-product attention can be computed as follows.

Original: [Equation]
Revised: [Equation]

Original: Hereby µ ( . ) is the unit vector normalization function, p w ∈ [0 , 1] S , p r ∈ [0 , 1] are write and read probabilities, and ⊙ is element-wise multiplication.
Revised: Hereby µ ( . ) is the unit vector normalization function, p w ∈ [0 , 1] S , p r ∈ [0 , 1] are write and read probabilities, and ⊙ is element-wise multiplication.

Original: Note that V ⊤ ( p w ⊙ µ ( K )) is one of the special cases of constructing the memory matrix by setting the erase probabilities to zeros WR operations.
Revised: Note that V ⊤ ( p w ⊙ µ ( K )) is one of the special cases of constructing the memory matrix by setting the erase probabilities to zeros WR operations.

Original: It is possible to create a new type of attention by replacing it with a different writing method to construct the memory matrix.
Revised: It is possible to create a new type of attention by replacing it with a different writing method to construct the memory matrix.

Original: 4. L INEAR SELF - ATTENTION WITH NAM
Revised: 4. L INEAR SELF - ATTENTION WITH NAM

Original: Using normalized outer-product attention, we can implement an efficient self-attention.
Revised: Using normalized outer-product attention, we can implement an efficient self-attention.

Original: The simplest way is setting the read/write probabilities to 1.
Revised: The simplest way is setting the read/write probabilities to 1.

Original: Given queries Q ∈ R S × d k , keys K ∈ R S ×
Revised: Given queries Q ∈ R S × d k , keys K ∈ R S ×

Original: d k , and values V ∈ R S
Revised: d k , and values V ∈ R S

Original: × d v of a sequence with S tokens, NAM self-attention is computed as follows.
Revised: × d v of a sequence with S tokens, NAM self-attention is computed as follows.

Original: [Equation]
Revised: [Equation]

Original: Its computational complexity is O
Revised: Its computational complexity is O

Original: ( Sd k d v ) instead of O
Revised: ( Sd k d v ) instead of O

Original: ( S 2 d k + Sd v ) of scaled dot-product self attention.
Revised: ( S 2 d k + Sd v ) of scaled dot-product self attention.

Original: While it is quadratic to the dimensions d k , d v , it is often not a problem in practice.
Revised: While it is quadratic to the dimensions d k , d v , it is often not a problem in practice.

Original: Most of large Transformers use multiple attention heads and keep the per-head dimension d k to small values like 64.
Revised: Most of large Transformers use multiple attention heads and keep the per-head dimension d k to small values like 64.

Original: Hence, we can keep the computational cost of each attention head and increase the capacity of the model by adding more heads.
Revised: Hence, we can keep the computational cost of each attention head and increase the capacity of the model by adding more heads.

Original: Computationally, this is very similar to the self-attention of Linear Transformer (Katharopoulos et al., 2020).
Revised: Mathematically, this is very similar to the self-attention of Linear Transformer (Katharopoulos et al., 2020).

Original: Linear Transformer has proven to be as effective as other Transformer implementations (Tay et al., 2020) and we expect our Transformer with NAM self-attention to show similar level of efficacy.
Revised: Linear Transformer has proven to be as effective as other Transformer implementations (Tay et al., 2020) and we expect our Transformer with NAM self-attention to show similar level of efficacy.

Original: The differences from Linear Transformer to ours are: 1) we use unit vector normalization instead of ELU kernel function, 2) so that we do not need to compute the causal masking factor
Revised: The differences from Linear Transformer to ours are: 1) we use unit vector normalization instead of ELU kernel function, 2) so that we do not need to compute the causal masking factor

Original: Z (Katharopoulos et al., 2020).
Revised: Z (Katharopoulos et al., 2020).

Original: 4. L ONG - RANGE TASK EVALUATION
Revised: 4. L ONG - RANGE TASK EVALUATION

Original: Although normalized outer-product attention is theoretically capable of replacing scaled dot-product attention, there exists a danger of information loss due to the limited capacity of the memory M .
Revised: Although normalized outer-product attention is theoretically capable of replacing scaled dot-product attention, there exists a danger of information loss due to the limited capacity of the memory M .

Original: Meanwhile, such an information loss is not an issue for Transformers because they utilize hidden activations from the entire sequence.
Revised: Meanwhile, such an information loss is not an issue for Transformers because they utilize hidden activations from the entire sequence.

Original: Also, the orthogonality assumption of the keys cannot hold in practice so that it may suffer from noises.
Revised: Also, the orthogonality assumption of the keys cannot hold in practice so that it may suffer from noises.

Original: Hence we need empirical evidences that NAM self-attention is an effective alternative.
Revised: Hence we need empirical evidences that NAM self-attention is an effective alternative.

Original: [Table]
Revised: [Table]

Original: We test our NAM self-attention based Transformer on long-range arena tasks (Tay et al., 2020) that have been used to compare efficient Transformer architectures.
Revised: We test our NAM self-attention based Transformer on long-range arena tasks (Tay et al., 2020) that have been used to compare efficient Transformer architectures.

Original: Since the previous work has already proven that Linear Transformer is as capable as others, we conduct partial comparison of ours to the original Transformer and Linear Transformer.
Revised: Since the previous work has already proven that Linear Transformer is as capable as others, we conduct partial comparison of ours to the original Transformer and Linear Transformer.

Original: Table 1 shows that NAM can be an efficient alternative to Transformers.
Revised: Table 1 shows that NAM can be an efficient alternative to Transformers.

Original: Despite the information loss and missing normalizing factor Z , the task accuracy of NAM is very similar to the others, even surpassing the vanilla Transformer in some cases.
Revised: Despite the information loss and missing normalizing factor Z , the task accuracy of NAM is very similar to the others, even surpassing the original Transformer in some cases.

Original: We see bigger speedups at the text classification task because the models have much smaller per-head dimension than others (16 vs 64).
Revised: We see bigger speedups at the text classification task because the models have much smaller per-head dimension than others (16 vs 64).

Original: Recall that the computational complexities of both linear transformer and NAM transformer are quadratic to the per-head dimensions.
Revised: Recall that the computational complexities of both linear transformer and NAM transformer are quadratic to the per-head dimensions.

Original: Image classification result shows that smaller per-head dimension can bring bigger speedups for NAM, but may result in lesser capability.
Revised: Hence, the image classification result shows that smaller per-head dimension can bring bigger speedups for NAM, but may result in lesser capacity.

Original: Overall, the normalized outer-product attention is as powerful as the scaled dot-product attention while enjoying greater computational efficiency.
Revised: Overall, the normalized outer-product attention is as powerful as the scaled dot-product attention while enjoying greater computational efficiency.

Original:  We use the code base of the original benchmark 1 and implemented our NAM encoder by modifying the Linear Transformer implementation.
Revised: Setup We use the code base of the original benchmark 1 and implemented our NAM encoder by modifying the Linear Transformer implementation.

Original: The three models share the same hyperparameters and the only differences come from the attention algorithms.
Revised: The three models share the same hyperparameters and the only differences come from the attention algorithms.

Original: The evaluations are run on the Ubuntu 20. system with RTX 3080.
Revised: The evaluations are run on the Ubuntu 20.04 system with RTX 3080.

Original: The results differ from the original work because we used smaller batch sizes due to the limited VRAM capacity.
Revised: The results differ from the original work because we used smaller batch sizes due to the limited VRAM capacity.

Original: All other experimental setups, including hyperparameters, are identical to the original benchmark (Tay et al., 2020).
Revised: All other experimental setups, including hyperparameters, are identical to the original benchmark (Tay et al., 2020).

Original: Details are available in the source code included as a supplementary material.
Revised: Details are available in the source code included as a supplementary material.

Original:  NAM A UGMENTED N EURAL N ETWORKS
Revised:  NAM A UGMENTED N EURAL N ETWORKS

Original: In this section, we design two types of memory-augmented neural networks (MANN) based on NAM, namely Long Short-term Attention Memory (LSAM) and NAM Turing Machine (NAM-TM).
Revised: In this section, we design two types of memory-augmented neural networks (MANN) based on NAM, namely Long Short-term Attention Memory (LSAM) and NAM Turing Machine (NAM-TM).

Original: LSAM is a generic recurrent neural network architecture derived from LSTM (Hochreiter & Schmidhuber, 1997).
Revised: LSAM is a generic recurrent neural network architecture derived from LSTM (Hochreiter & Schmidhuber, 1997).

Original: In LSAM, the long-term cell state of LSTMis replaced with a NAM memory matrix.
Revised: LSAM replaces the long-term cell state of LSTM with a NAM memory matrix.

Original: NAM 
Revised: NAM Turing

Original: Turing Machine is a MANN design for algorithmic tasks inspired by Neural Turing Machine (Graves et al., 2014).
Revised: Machine is a MANN design for algorithmic tasks inspired by Neural Turing Machine (Graves et al., 2014).

Original: Its read and write heads can move along the tape with four actions: left, right, no-op, and jump.
Revised: Its read and write heads can move along the tape with four actions: left, right, no-op, and jump.

Original: Their implementations are based on NAM read/write primitives.
Revised: Implementations of the heads are based on NAM read/write primitives.

Original: We evaluate the models with tasks of number sequence prediction, sequence reduction, and SCAN.
Revised: We evaluate the models with tasks of number sequence prediction, sequence reduction, and SCAN.

Original: We test their compositional generalization capability by splitting the data based on the sequence lengths.
Revised: We test their compositional generalization capability by splitting the data based on the sequence lengths.

Original: 5. L ONG S HORT - TERM A TTENTION M EMORY
Revised: 5. L ONG S HORT - TERM A TTENTION M EMORY

Original: Long Short-term Memory (LSTM) (Hochreiter & Schmidhuber, 1997) leverages two recurrent state vectors: the short-term hidden state and the long-term cell state.
Revised: Long Short-term Memory (LSTM) (Hochreiter & Schmidhuber, 1997) leverages two recurrent state vectors: the short-term hidden state and the long-term cell state.

Original: To mitigate the problems of vanishing/exploding gradients, the cell state is additively updated using forget and input gates.
Revised: To mitigate the problems of vanishing/exploding gradients, the cell state is additively updated using forget and input gates.

Original: Then the output gate selectively reads the cell state to determine the hidden state.
Revised: Then, the output gate selectively reads the cell state to determine the hidden state.

Original: Long Short-term
Revised: Long Short-term

Original: Attention Memory (LSAM) follows the same principle.
Revised: Attention Memory (LSAM) follows the same principle.

Original: Instead of using the vector cell state, it leverages the memory matrix M t ∈ R d 2 which is also additively updated using the NAM write primitive.
Revised: Instead of using the vector cell state, it leverages the memory matrix M t ∈ R d 2 which is also additively updated using the NAM write primitive.

Original: The hidden state h t ∈ R d is retrieved by reading the memory matrix.
Revised: The hidden state h t ∈ R d is retrieved by reading the memory matrix.

Original: Given the input x t ∈ R d , the update rule M t , h t = LSAM ( x t , M t − 1 , h t − 1 ) is defined as follows.
Revised: Given the input x t ∈ R d , the update rule M t , h t = LSAM ( x t , M t − 1 , h t − 1 ) is defined as follows.

Original: [Equation]
Revised: [Equation]

Original: Hereby : is the concatenate operator, σ ( . ) is the sigmoid function, and W qkv , W rw , b qkv , b rw are trainable weights and biases.
Revised: Hereby : is the concatenate operator, σ ( . ) is the sigmoid function, and W qkv , W rw , b qkv , b rw are trainable weights and biases.

Original: Although the memory matrix M t has much higher capacity than a vector cell state, the computational complexity of LSAM is identical to that of LSTM.
Revised: Although the memory matrix M t has much higher capacity than a vector cell state, the computational complexity of LSAM is identical to that of LSTM.

Original: This is because WR and RD have complexity of O ( d 2 ) which is identical to that of matrix-vector multiplication of the weights and the states.
Revised: This is because WR and RD have complexity of O ( d 2 ) which is identical to that of matrix-vector multiplication of the weights and the states.

Original: Like Transformers, we can design multi-headed LSAM by concatenating per-head states M it and h it .
Revised: Like Transformers, we can design multi-headed LSAM by concatenating per-head states M it and h it .

Original: Bidirectional LSAM is also possible by splitting the multiple heads into two directions.
Revised: Bidirectional LSAM is also possible by splitting the multiple heads into two directions.

Original: The backward heads are updated in  opposite  by the rule of
Revised: The backward heads are updated in the opposite direction by the rule of

Original: M jt , h jt = LSAM ( x t , M jt +1 , h jt +1 ) .
Revised: M jt , h jt = LSAM ( x t , M jt +1 , h jt +1 ) .

Original: The LSAM architecture combines the strengths of recurrent neural networks (RNN) and Transformers.
Revised: The LSAM architecture combines the strengths of recurrent neural networks (RNN) and Transformers.

Original: Since LSAM follows the RNN design so that it enjoys strengths of RNNs such as low computational cost for inference and recurrent inductive bias.
Revised: Since LSAM follows the RNN design so that it enjoys strengths of RNNs such as low computational cost for inference and recurrent inductive bias.

Original: Additionally, it benefits from the strengths of attention because reading and writing the memory matrix natively incorporates the attention mechanism.
Revised: Additionally, it benefits from the strengths of attention because reading and writing the memory matrix natively incorporates the attention mechanism.

Original: 5. NAM T URING M ACHINE
Revised: 5. NAM T URING M ACHINE

Original: Neural Turing Machine (NTM) (Graves et al., 2014) is one of the early neural networks that implement external memory structure with differentiable read and write methods.
Revised: Neural Turing Machine (NTM) (Graves et al., 2014) is one of the early neural networks that implement external memory structure with differentiable read and write methods.

Original: It is a basis of Differentiable
Revised: It is a basis of Differentiable

Original: Neural Computer (DNC) (Graves et al., 2016) which has proven to be effective at solving a variety of algorithmic tasks such as answering synthetic questions and finding shortest paths.
Revised: Neural Computer (DNC) (Graves et al., 2016) which has proven to be effective at solving a variety of algorithmic tasks such as answering synthetic questions and finding shortest paths.

Original: Their external memory matrix is accessed by read and write heads using differentiable attention mechanisms.
Revised: Their external memory matrix is accessed by read and write heads using differentiable attention mechanisms.

Original: Hereby we design NAM Turing Machine (NAM-TM) which adopts the design principles of NTM and DNC.
Revised: Hereby we design NAM Turing Machine (NAM-TM) which adopts the design principles of NTM and DNC.

Original: The main idea of NAM-TM is to treat the tape state T =
Revised: The main idea of NAM-TM is to treat the tape state T =

Original: [ v 1 , v 2 , ...v n , 0 , ... ]
Revised: [ v 1 , v 2 , ...v n , 0 , ... ]

Original: ∈ R L × d as a memory matrix as follows.
Revised: ∈ R L × d as a memory matrix as follows.

Original: [Equation]
Revised: [Equation]

Original: Hereby e i are standard basis vectors < 0 , 0 , ... 1 , ..., 0 > of R L where L is the size of the tape.
Revised: Hereby e i are standard basis vectors < 0 , 0 , ... 1 , ..., 0 > of R L where L is the size of the tape.

Original: This tape state can now be accessed by using NAM read/write primitives.
Revised: This tape state can now be accessed by using NAM read/write primitives.

Original: NAM-TM is a differentiable function that takes the tape state T , read and write heads H r , H w ∈ R L , and the input vector x ∈ R d and produces the read output
Revised: NAM-TM is a differentiable function that takes the tape state T , read and write heads H r , H w ∈ R L , and the input vector x ∈ R d and produces the read output

Original: R ∈ R d along with the updated tape and head states T ′ , H ′ r , H ′ w .
Revised: R ∈ R d along with the updated tape and head states T ′ , H ′ r , H ′ w .

Original: [Equation]
Revised: [Equation]

Original: The read and write heads are positional vectors to attend the memory matrix for reading and writing the states.
Revised: The read and write heads are positional vectors to attend the memory matrix for reading and writing the states.

Original: At each time step,  positions can be updated by four actions: LEFT, RIGHT, NO-OP, and JUMP.
Revised: At each time step, the positions can be updated by four actions: LEFT, RIGHT, NO-OP, and JUMP.

Original: They are controlled by a controller neural network nn _ control ( x ) which emits read and write probabilities p r , p w , action probabilities p right , p left , p noop , p jump for each head and a unit jump query vector q jump ∈ R d .
Revised: They are controlled by a controller neural network nn _ control ( x ) which emits read and write probabilities p r , p w , action probabilities p right , p left , p noop , p jump for each head and a unit jump query vector q jump ∈ R d .

Original: Given the controller outputs and the value v = W v x to write, reading and writing the memory are conducted by NAM primitives as follows.
Revised: Given the controller outputs and the value v = W v x to write, reading and writing the memory are conducted by NAM primitives as follows.

Original: [Equation]
Revised: [Equation]

Original: Then, each head is updated to the next position based on the action probabilities.
Revised: Then, each head is updated to the next position based on the action probabilities.

Original: LEFT and
Revised: LEFT and

Original: RIGHT actions can be performed by the differentiable roll function, which is a linear transformation R L −→ R L mapping e ito e i +1 . [Equation]
Revised: RIGHT actions can be performed by the differentiable roll function, which is a linear transformation R L −→ R L mapping e i

Original: The jump position H JUMP is determined by reading the transpose of a key tape K with the unit query vector q jump .
Revised: The jump position H JUMP is determined by reading the transpose of a key tape K with the unit query vector q jump .

Original: The key tape is written in the same way as the tape T , but it stores the corresponding unit key vector derived using the weight matrix W k ∈ R d × d and the unit vector normalization µ ( . ) .
Revised: The key tape is written in the same way as the tape T , but it stores the corresponding unit key vector derived using the weight matrix W k ∈ R d × d and the unit vector normalization µ ( . ) .

Original: [Equation]
Revised: [Equation]

Original: One can understand the transpose of the key tape K ⊤ =
Revised: One can understand the transpose of the key tape K ⊤ =

Original: (cid:80) i e i
Revised: (cid:80) i e i

Original: k ⊤ i as a jump table.
Revised: k ⊤ i as a jump table.

Original: Reading K ⊤ with a key vector k i returns the corresponding position
Revised: Reading K ⊤ with a key vector k i returns the corresponding position

Original: vector e i if the keys are orthonormal.
Revised: vector e i if the keys are orthonormal.

Original: Finally, the next head position H ′ is updated as a weighted sum of the positions.
Revised: Finally, the next head position H ′ is updated as a weighted sum of the positions.

Original: [Equation]

Original: While all of the computations are technically done with a fixed tape length L , none of the trainable parameters depend on the value of L .
Revised: While all of the computations are technically done with a fixed tape length L , none of the trainable parameters depend on the value of L .

Original: That is, a NAM-TM trained on certain length L can be applied to any tape length L ′ withoutany modification nor re-training.
Revised: That is, a NAM-TM trained on certain length L can be applied to any tape length L ′ without modification nor re-training.

Original: Theoretically, it can be extended to infinite-dimensional Hilbert spaces.
Revised: Theoretically, it can be extended to infinite-dimensional Hilbert spaces.

Original: There are multiple strengths in NAM-TM design compared to the memory structures of NTM and
Revised: There are multiple strengths in NAM-TM design compared to the memory structures of NTM and

Original: DNC.
Revised: DNC.

Original: First, unlike NTM and DNC, NAM-TM’s building blocks are simple and computationally efficient.
Revised: First, unlike NTM and DNC, NAM-TM’s building blocks are simple and computationally efficient.

Original: In the previous sections, the efficiency of NAM read/write primitives are already shown.

Original: Additionally, contemporary deep learning frameworks have efficient low-overhead implementations of the roll transformation.

Original: Second, NAM-TM’s addressing mechanism is based on the query-key-value attention mechanism of NAM, which is more powerful than the content-based attention mechanism used in NTM and DNC.
Revised: Second, NAM-TM’s addressing mechanism is based on the query-key-value attention mechanism of NAM, which is more powerful than the content-based attention mechanism used in NTM and DNC.

Original: Finally, NAM-TM’s design is flexible in that it is easy to add/remove transition rules of the read/write heads.
Revised: Finally, NAM-TM’s design is flexible in that it is easy to add/remove transition rules of the read/write heads.

Original: For example, the JUMP transition rule is optional in that a Turing machine only requires LEFT and RIGHT transitions in theory.
Revised: For example, the JUMP transition rule is optional in that a Turing machine only requires LEFT and RIGHT transitions in theory.

Original: One can also add another transition rule for the head positions if the rule can be implemented with differentiable functions.
Revised: One can also add another transition rule for the head positions if the rule can be defined with differentiable functions.

Original: 5. C OMPOSITIONAL G ENERALIZATION T ASKS
Revised: 5. C OMPOSITIONAL G ENERALIZATION T ASKS

Original: We test the computational powers of LSAM and NAM-TM using three types of algorithmic tasks in compositional generalization setups.
Revised: We test the computational powers of LSAM and NAM-TM using three types of algorithmic tasks in compositional generalization setups.

Original: First, number sequence prediction (Nam et al., 2019) task (NSP) is a suite of synthetic problems to predict the following digits of the numerical sequences.
Revised: First, number sequence prediction (Nam et al., 2019) task (NSP) is a suite of synthetic problems to predict the following digits of the numerical sequences.

Original: It can test the compositional generalization capability by testing/validating the models with the longer decimal numbers that are never observed during the training stage.
Revised: It can test the compositional generalization capability by testing/validating the models with the longer decimal numbers that are never observed during the training stage.

Original: In this setup, many models often suffer from drastic fall of test/validation accuracy, due to lack of inductive bias (Kim et al., 2021).
Revised: In this setup, many models often suffer from drastic fall of test/validation accuracy, due to lack of inductive bias (Kim et al., 2021).

Original: We use two representative sequences from NSP: Fibonacci (Fib) and Palindrome (Palin).
Revised: We use two representative sequences from NSP: Fibonacci (Fib) and Palindrome (Palin).

Original: The two tasks require the generalization of digital addition and sequence reversal rules, respectively.
Revised: The two tasks require the generalization of digital addition and sequence reversal rules, respectively.

Original: Next is a sequence reduction (Reduce) task with a simple rule: given the sequence of digits, the target is the reduced sequence by skippingthe zeros.
Revised: Next is a sequence reduction (Reduce) task with a simple rule: given the sequence of digits, the target is the reduced sequence by skipping zeros.

Original: This is a task that can be easily solved by a proper Turing machine.
Revised: This is a task that can be easily solved by a proper Turing machine.

Original: We also use a similar compositional generalization setup by testing/validating using the longer target sequences that are longer than any target sequences in the training dataset.
Revised: We also use a similar compositional generalization setup by testing/validating using the longer target sequences that are longer than any target sequences in the training dataset.

Original: For NSP and Reduce tasks, we use training datasets with d = 1 . . .
Revised: For NSP and Reduce tasks, we use training datasets with d = 1 . . .

Original: 10 decimal digit sequences in little-endian order.
Revised: 10 decimal digit sequences in little-endian order.

Original: Then we validate/test the models with two validation sets (ID, OD-easy) and a test set (OD-hard) for each task.
Revised: Then we validate/test the models with two validation sets (ID, OD-easy) and a test set (OD-hard) for each task.

Original: An in-distribution (ID) validation set consists of d = 5 . . .
Revised: An in-distribution (ID) validation set consists of d = 5 . . .

Original: 10digit sequences, and an out-of-distribution validation set (OD-easy) consists of d = 11 . . .
Revised: 10digit sequences, and an out-of-distribution validation set (OD-easy) consists of d = 11 . . .

Original: 13 -digit sequences.
Revised: 13 -digit sequences.

Original: The harder out-of-distribution test set (OD-hard) consists of d = 14 . . .
Revised: The harder out-of-distribution test set (OD-hard) consists of d = 14 . . .

Original: 16 -digit sequences, challenging the models with generalization to longer contexts training has
Revised: 16 -digit sequences, challenging the models with generalization to longer contexts.

Revised: A training dataset has 25600 samples, and each validation/test set has 2048 samples.

Original: The last task is SCAN (Lake & Baroni, 2018) task which has familiarized the concept of compositional zero-shot generalization.
Revised: The last task is SCAN (Lake & Baroni, 2018) task which has familiarized the concept of compositional zero-shot generalization.

Original: SCAN consists of simplified natural language input sequences and the corresponding output action sequences.
Revised: SCAN consists of simplified natural language input sequences and the corresponding output action sequences.

Original: Among many data split methods, we use train/test split by length 2 to be consistent with the other tasks.
Revised: Among many data split methods, we use train/test split by length 2 to be consistent with the other tasks.

Original: Since SCAN only has a limited number of sequences, we only have two datasets of train and test.
Revised: Since SCAN only has a limited number of sequences, we only have two datasets of train and test.

Original: [Figure]
Revised: [Figure]

Original: We format the problems as masked sequence completion problems as shown in Figure 1.
Revised: We format the problems as masked sequence completion problems as shown in Figure 1.

Original: An input sequence consists of input tokens followed by masks, and an output sequence consists of target tokens following the masks.
Revised: An input sequence consists of input tokens followed by masks, and an output sequence consists of target tokens following the masks.

Original: Since the input and output sequences are 1:1 matched, we can compare with baseline models without sequence-to-sequence structures like DNC.
Revised: Since the input and output sequences are 1:1 matched, we can compare to baseline models without sequence-to-sequence structures like DNC.

Original: We choose four baseline models to compare: a bidirectional Transformer encoder (TF) (Devlin et al., 2018), an 2-layer LSTM model with attention (Bahdanau et al., 2014), an Universal Transformer (Dehghani et al., 2018), and a differentiable neural computer (DNC).
Revised: We choose four baseline models to compare: a bidirectional Transformer encoder (TF) (Devlin et al., 2018), an 2-layer LSTM model with attention (Bahdanau et al., 2014), an Universal Transformer (Dehghani et al., 2018), and a differentiable neural computer (DNC).

Original: The TF model follows the architecture of BERT medium and the hyperparameters of the other models are adjusted to have similar parameter counts.
Revised: The TF model follows the architecture of BERT medium and the hyperparameters of the other models are adjusted to have similar parameter counts.

Original: Our
Revised: Our

Original: LSAM and NAM-TM networks have two LSAM layers and two NAM-TM layers respectively, whose hyperparameters are also adjusted to have similar parameter counts.
Revised: LSAM and NAM-TM networks have two LSAM layers and two NAM-TM layers respectively, whose hyperparameters are also adjusted to have similar parameter counts.

Original: As an ablation study, we also evaluate NAM-TM design without the JUMP transition (No Jmp).
Revised: As an ablation study, we also evaluate NAM-TM design without the JUMP transition (No Jmp).

Original: All experiments are run with
Revised: All experiments are run with

Original: PyTorch 1.10 on the system with Ubuntu 20.04 and RTX 3080.
Revised: PyTorch 1.10 on the system with Ubuntu 20.04 and RTX 3080.

Original: Each experiment takes less than five hours to run 200 training epochs.
Revised: Each experiment takes less than five hours to run 200 training epochs.

Original: The implementation details and hypereparameters can be found at  source code 
Revised: The implementation details and hypereparameters can be found at the source code in the supplementary materials.

Original:  E VALUATION R ESULT
Revised: 5. E VALUATION R ESULT

Original: [Table]
Revised: [Table]

Original: Table 2 shows the sequence accuracy comparison of the models in compositional generalization tasks.
Revised: Table 2 shows the sequence accuracy comparison of the models in compositional generalization tasks.

Original: We report the sequence accuracies of the models, where a wrong prediction of one token is counted as a failed prediction of the entire sequence.
Revised: We report the sequence accuracies of the models, where a wrong prediction of one token is counted as a failed prediction of the entire sequence.

Original: To avoid cherry-picking, the epochs of the best OD-easy validation accuracies are presented for the evaluation results.
Revised: To avoid cherry-picking, the epochs of the best OD-easy validation accuracies are presented for the evaluation results.

Original: However, we present the results of the best test accuracy for SCAN tasks because there are only two available datasets (Train, Test).
Revised: However, we present the results of the best test accuracy for SCAN tasks because there are only two available datasets (Train, Test).

Original: Although LSAM’s architecture is not specifically designed forthose algorithmic tasks, LSAM performs consistently better than the other baselines.
Revised: Although LSAM’s architecture is not specifically designed for algorithmic tasks, LSAM performs consistently better than the other baselines.

Original: Surprisingly, it performs better than DNC, which is a MANN model designed for such algorithmic problems.
Revised: Surprisingly, it performs better than DNC, which is a MANN model designed for such algorithmic problems.

Original: This is not a matter of model capacity because LSAM has a slightly smaller parameter count.
Revised: This is not a matter of model capacity because LSAM has a slightly smaller parameter count.

Original: The results imply that the computational power of NAM memory architecture is superior to that of DNC’s memory architecture.
Revised: The results imply that the computational power of NAM memory architecture is superior to that of DNC’s memory architecture.

Original: As expected, NAM-TM performs significantly better in algorithmic tasks (Palin, Fib, and Reduce) possibly due to its specialized architecture.
Revised: As expected, NAM-TM performs significantly better in algorithmic tasks (Palin, Fib, and Reduce) possibly due to its specialized architecture.

Original: Especially, it finds easier to solve OD-hard problems, whereas the other models experience steep performance decline.
Revised: Especially, it finds easier to solve OD-hard problems, whereas the other models experience steep performance decline.

Original: A potential explanation is that the action-based positional transitions provide robustness in long-context cases.
Revised: A potential explanation is that the action-based positional transitions provide robustness in long-context cases.

Original: However, NAM-TM is not the best-performing model for the SCAN tasks.
Revised: However, NAM-TM is not the best-performing model for the SCAN task.

Original: NAM-TM remains effective at palindrome and reduction tasks even without the JUMP transition (No Jmp).
Revised: NAM-TM remains effective at palindrome and reduction tasks even without the JUMP transition (No Jmp).

Original: Hence, LEFT and RIGHT transitions seem to be enough for emulating simple Turing machines, but more complex transition rules can augment the computational power of NAM-TM.
Revised: Hence, LEFT and RIGHT transitions seem to be enough for emulating simple Turing machines, but more complex transition rules can augment the computational power of NAM-TM.

Original: This suggests extending NAM-TM to variety of tasks by augmenting specialized transition actions.
Revised: This suggests extending NAM-TM to variety of tasks by augmenting specialized transition actions.

Original:  F URTHER A PPLICATIONS
Revised:  F URTHER A PPLICATIONS

Original: Since NAM is a simple and flexible basic block for memory architecture, there are infinite number of  applications.
Revised: Since NAM offers simple and flexible building blocks for memory architecture, there are infinite number of potential applications.

Original: Hereby we suggest a few possible extensions of NAM.
Revised: Hereby we suggest a few possible extensions of NAM.

Original: While we do not evaluate the ideas in this work, we leave them aspossible future work.
Revised: While we do not evaluate the ideas in this work, we leave them as future work.

Original: Hierarchical data modeling Tensor product, the key operation behind NAM, is not limited to two-dimensional outer product.
Revised: Hierarchical data modeling Tensor product, the key operation behind NAM, is not limited to two-dimensional outer product.

Original: Any dimensional tensors T i can construct a memory tensor M by performing sum of tensor products M = (cid:80) i T
Revised: Any dimensional tensors T i can construct a memory tensor M by performing sum of tensor products M = (cid:80) i T

Original: i ⊗ k
Revised: i ⊗ k

Original: i , with the unit key vectors k i .
Revised: i , with the unit key vectors k i .

Original: For example, a document tensor D can constructed by sum of tensor products of sentence-level keys and sentence matrices S i , each of which is also a sum of outer products of word embeddings and word-level keys.
Revised: For example, a document tensor D can constructed by sum of tensor products of sentence-level keys and sentence matrices S i , each of which is also a sum of outer products of word embeddings and word-level keys.

Original: Nested attention to such a hierarchical memory can be conducted by two inner products with a sentence-level unit query vector q s and a word-level unit query vector q w .
Revised: Nested attention to such a hierarchical memory can be conducted by two inner products with a sentence-level unit query vector q s and a word-level unit query vector q w .

Original: Efficient edge inference While Transformers are very successful in many NLP tasks, deploying such models for edge inference has been a challenging task since Transformer’s computation and memory cost per each time step varies with the sequence length.
Revised: Efficient edge inference While Transformers are very successful in many ML tasks, deploying such models for edge inference has been a challenging task since Transformer’s computation and memory cost per each time step varies with the sequence length.

Original: However, NAM’s cost does not depend the sequence length at all.
Revised: However, NAM’s cost does not depend the sequence length at all.

Original: Also, the outer product operation is more compute intensive than the dot products, making NAM more friendly to high-throughput NPUs.
Revised: Also, the outer product operation is more compute intensive than the dot products, making NAM more friendly to high-throughput accelerators.

Original: Hence, NAM can be an efficient Transformer alternative for edge inference.
Revised: Hence, NAM can be an efficient Transformer alternative for edge inference.

Original: Few-shot learning Given a input-output vector pair ( x, y ) , it is possible to conduct one-shot learning of a weight matrix W by WR ( W, y/
Revised: Few-shot learning Given a input-output vector pair ( x, y ) , it is possible to conduct one-shot learning of a weight matrix W by WR ( W, y/

Original: | x | , x/ | x | , 1 , 1) .
Revised: | x | , x/ | x | , 1 , 1) .

Original: This is because Wx is guaranteed to return y by Theorem 2.
Revised: This is because Wx is guaranteed to return y by Theorem 2.

Original: Therefore, we can implement one-shot or few-shot learning as memorization to NAM.
Revised: Therefore, we can implement one-shot or few-shot learning as memorization to NAM.

Original: This aligns well with the human behavior since we conduct few-shot learning by memorizing, not by repetitive training.
Revised: This aligns well with the human behavior since we conduct few-shot learning by memorizing, not by repetitive training.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We proposed a redesign of attention mechanism to construct a differentiable memory for neural networks, namely neural attention memory (NAM).
Revised: We proposed a redesign of attention mechanism to construct a differentiable memory for neural networks, namely neural attention memory (NAM).

Original: Following the same query-key-value structure of scaled dot-product attention, NAM first writes the memory matrix by adding outer products of key-value pairs.
Revised: Following the same query-key-value structure of scaled dot-product attention, NAM first writes the memory matrix by adding outer products of key-value pairs.

Original: Then we can read it by multiplying the matrix with a unit query vector.
Revised: Then we can read it by multiplying the matrix with a unit query vector.

Original: The first strength of NAM is that its computational complexity does not rely on the sequence length.
Revised: The first strength of NAM is that its computational complexity does not rely on the sequence length.

Original: The long-range arena evaluations showed that NAM based attention, namely normalized outer-product attention, is an efficient and effective alternative for scaled dot-product attention.
Revised: The long-range arena evaluations showed that NAM based attention, namely normalized outer-product attention, is an efficient and effective alternative for scaled dot-product attention.

Original: Next, NAM can be a powerful basis for constructing MANN models.
Revised: Next, NAM can be a powerful basis for constructing MANN models.

Original: We designed two NAM-based MANNs: LSAM for generic sequential tasks and NAM-TM for algorithmic tasks.
Revised: We designed two NAM-based MANNs: LSAM for generic sequential tasks and NAM-TM for algorithmic tasks.

Original: In compositional generalization tasks, both outperformed other baselines such as Universal Transformer and DNC, indicating that NAM is a more powerful mechanism for implementing memory in DNNs.
Revised: In compositional generalization tasks, both outperformed other baselines such as Universal Transformer and DNC, indicating that NAM is a more powerful mechanism for implementing memory in DNNs.

Original: Finally, it opens up further research possibilities in the fields of hierarchical data modeling, edge inference, and few-shot learning.
Revised: Finally, it opens up further research possibilities in the fields of hierarchical data modeling, edge inference, and few-shot learning.