Neural Attention Memory
Authors: Hyoungwook Nam
Seung Byum Seo
Venue: Submitted to ICLR 2023 (2023)
Keywords: Neuro-symbolic AI
Transformer
Memory-augmented neural network
compositional generalization

Abstract: Scaled dot-product attention has become the essence of state-of-the-art deep neural networks for various machine learning tasks. Though its ubiquitous accomplishments, it is inefficient for long sequence tasks and problematic for tasks requiring memory states such as compositional generalization. We propose a novel perspective of the attention mechanism by reinventing it as a memory architecture for neural networks, namely Neural Attention Memory (NAM). 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. Indeed, we show that our normalized outer-product attention mechanism is mathematically equivalent to the conventional attention mechanism. Then, we evaluate a NAM-based Transformer on long-range arena tasks and demonstrate its efficiency and efficacy. Finally, we propose two NAM-based memory-augmented neural networks, namely Long Short-Term Attention Memory (LSAM) and NAM Turing Machine (NAM-TM), and test their compositional generalization capability using four different tasks. 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. The experimental results show that the proposed models outperform traditional Transformer and LSTM, as well as DNC. NAM opens up possibilities in diverse machine learning research problems, including hierarchical data modeling, efficient edge inference, and few-shot learning.

Under review as a conference paper at ICLR 2023
NEURAL ATTENTION MEMORY
Anonymous authors
Paper under double-blind review
ABSTRACT
Scaled dot-product attention has become the essence of state-of-the-art deep neural networks for various machine learning tasks. Though its ubiquitous accomplish- ments, it is inefficient for long sequence tasks and problematic for tasks requiring memory states such as compositional generalization. We propose a novel perspec- tive of the attention mechanism by reinventing neural networks, namely Neural Attention Memory (NAM). 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 by adding outer products of value and unit key vectors, and reads it by multiplying the matrix with unit query vector. We define read and write primitives of NAM and mathematically prove their functionalities. One benefit of NAM is that it can be basis for efficient linear attention, namely normalized outer-product attention. We evaluate a NAM-baseed Transformer on long-range tasks and demonstrate NAM's efficiency and efficacy. Most importantly, NAM provides building blocks for memory-augmente neural networks. We propose two NAM-augmented neural networks, namely Long Short-Term Attention Mem- ory (LSAM) and NAM Turing Machine (NAM-TM), and test their compositional generalization capabilities using four different tasks. LSAM replaces LSTM's long-term cell state with NAM memory matrix and NAM-TM implements a Turing tape structure using NAM read/write primitives. The experiments show that they have better computational power than Transformer and I.STM, as well as DNC. NAM opens up possibilities diverse research problems, including hicrarchical data modeling, efficient edge inference, and few-shot learning.
a memory architecture for
INTRODUCTION
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 process- ing (Devlin et al., 2018). multi-modal task (Li et al., 2019), and graph data processing (Hamilton let al. 2017). Specifically, the Transformers using the self-attention method have replaced recurrent neural networks (RNN) by outperforming them in most of the tasks. Despite its success, there exist limitations the mechanism. First, it needs the information of the entire sequence compute one attention so that its computational complexity becomes quadratic to the length of the sequence. Hence. is inefficient for long sequence tasks (Tay et al. 2020) or edge inference environments (Tambe et al. 2020). Also, its stateless design enables efficient parallelism but makes it impossible to solve tasks that require memory states. Hence, Transformers fail to generalize the rules that require inductive bias (Dehghani et al., 2018) or compositional generalization (Lake & Baroni 2018).
There have been studies designing neural networks with external memory to solve algorithmic tasks where Transformers fail. These memory-augmentee neural networks (MANN) design differentiable read/write functions that can trained by backpropagation. 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). They outperform generic neural networks in synthetic algorithmic tasks but considered impractical due to their complexities and inefficiencies.
In this work, we re-invent the attention mechanism memory architecture for neural networks, namely neural attention memory (NAM). NAM's design objective is to build simple, efficient, yet powerful external memory which also incorporates the attention mechanism. Following the same

---

Under review as a conference paper at ICLR 2023
juery-key-valuu structure of attention, NAM stores key-value pairs to memory matrix via additively writing their outer- anter-products. Reading the memory matrix is simply done by multiplying the matrix with unit query vector. We provide mathematical formulation for the read/write primitives, and make theoretical analyses showing that these read and write primitives can replace attention.
One big benefit of NAM is that it can perform attention in more efficient way. By sacrificing the erasure capability of the NAM write operation, we can design an efficient and parallel attention mechanism, namely normalized outer-product attention. This special variant of NAM almost equivalent to linear attention Katharopoulos et al. 2020), enjoying the same linear computational complexity to the sequence length. We evaluate NAM-based efficient Transformer in long-range arena TTay et al.) (020) tasks. Its efficacy is on par with the base Transformer and Linear Transformer, implying that NAM can an efficient alternative to the scaled dot-product attention.
The bigger value of NAM is that its read and write primitives can be building blocks for augmenting memory structures to deep neural networks. Using NAM read/write primitives, we design two memory-angmenteed neural networks (MANN), namely Long Short-term Attention Memory (LSAM) and NAM Turing Machine (NAM-TM). LSAM is a generic RNN architecture that replaces LSTM's long-term cell state with a memory matrix. Instead of additively writing vector cell state, LSAM reads and writes the memory matrix using NAM primitives. The design combines strengths of attention and RNN while maintaining the same computational complexity LSTM. NAM-TM is MANN for algorithmic tasks, leveraging Turing tape structure. A tape has read and write heads accessing the memory with NAM read/write primitives. They move along the tape with four actions: NO-OP, LEFT, RIGHT, and JUMP. The actions are implemented as differentiable functions enable end-to-end training with backpropagation.
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). 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. The evaluation results show that their computational powers are superior to other baselines, including Universal Transformer (Dehghani et al., 2018) and DNC (Graves et al. (016). While the generic LSAM model consistently outperforms the others, NAM-TM shows even better results at algorithmic tasks. The results indicate that NAM is powerful method to implement memory in neural networks.
The efficient, simple, and flexible structure of NAM opens up new possibilities in multiple machine learning research fields. One straightforward application is leveraging NAM's efficiency for edge inference environment. Another possibility is using NAM for hierarchical data modeling by general- izing NAM with tensor products. Moreover, memorization of input-output mapping using NAM can be solution for one-shot and few-shot learning.
The main contributions of this work are as follows:
We re-invent the attention mechanism as a memory architecture for neural networks, namely neural attention memory (NAM).
We present mathematical basis for NAM read/write primitives, and give theoretical proofs that NAM is equivalent to attention in certain conditions.
We show that NAM can construct an efficient Transformer for long-range sequence tasks.
We propose two mumory-augmented neural network designs of LSAM and NAM-TM and show their capabilities in compositional generalization tasks.
2 BACKGROUND
2.1 SCALED DOT-PRODUCT ATTENTION
Attention mechanisms of deep neural networks (Bahdanau et al. 2014 Luong et al., 2015) provide differentiable methods of selectively attending items from variable-length sequence. 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 representa- tions using the scores. Among the variations, scaled dot product attention /Vaswani et al., 2017) has been the most successful. For each token, there are key vector and value vector associated to it.
2

---

Under review as a conference paper at ICLR 2023
Given query vector, the scores determined by the scaled dot- dot-product of the query and the keys. Then the output is computed by weighted sum of softmaxed scores and the value vectors.
The self-attention mechanism based the scaled dot-product attention has proven to be very powerful, replacing the needs of RNNs. Since attention reaches every element a sequence in an O(1) path, 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). However, its stateless and parallel architecture also bring multiple limitations. First, the computational cost quadratically increases with the sequence length, making it inefficient for long-rang contexts (Beltagy et al.
2020 Choromanski et al., 2020; Kitaev al., 2020 Katharopoulos et al. 2020) and edge inference environments (Tambe et al. (0020). Also, the memory- -lessry-esssee architecture lacks inductive bias (Dehghani et al. 2018), making impossible to generalize inductive algorithmic rules (Kim et al. (021). There are researches to resolve those issues, but simple and practical solution is yet to be found.
2.2 MEMORY-AUGMENTED NEURAL NETWORK
In theory, RNNs are proven to as powerful as Turing machines (Hyötyniemi 1996). However, they fail to learn algorithmic patterns that require at least pushdown automata in practice (Nam et al. 2019). Therefore, there have been efforts to augment external memory architecture to neural networks, as known as eeemory-augmeee
neural networks (MANN) (Grefenstette et al. [2015 Joulin & Mikolov 2015 Graves et al. 2014; 2016). The main challenge for MANNs is to design differentiable read and write functions that can be trained via back- poopppropagttion. Some MANNs use attention mechanism for differentiable read/write functions. For instance, Neural Turing Machine (NTM) (Graves et al.) 2014) and Differentiable Neural Computer (DNC) (Graves al. 2016) leverage attention mechanism for implementing content-bases addressing. 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 (NN) cost. Hence, MANNs often become inefficient and complex so that they are considered impractical outside of the algorithmic task domain.
2.3 TRANSFORMER AND MEMORY
Transformers' attention mechanism can be understood as an external memory where all hidden activations are stored and selectively read. Since this method's storage footprint scales up with the sequence, there have been studies to limit the space complexity and scale to longer context by using explicit memory architecture to transformers. Transformer XL (Dai et al., 2019) has attempted to overcome the limited context length of Transformers by using segment-level recurrence and relative positional encoding. Compressive Transformer (Rae et al., 2019) and co-former (Martins et al. 2022) further have extended the idea by compressing the recurrent segment into smaller memory unbounded continuous memory, respectively. While these works effectively address the scalability problem of Transformers, they have extended their ideas to generic memory architecture that can be freely read and written.
and using
3 NEURAL ATTENTION MEMORY
3.1 NAM READ AND WRITE PRIMITIVES
The main idea of NAM is implementing attention mechanism via matrix multiplication of a memory matrix M E Rdd a unit query vector E IRd and read probability 0 S pr S 1. Hereby d, and dy are feature dimensions of value and key vectors respectively. Then, the read operation (RD) of NAM computes the read vector r E Idd as follows.
RD g.pr) P.
The memory matrix M is written by adding outer products of unit key vectors k; and value vectors Vi- The write operation W R updates the memory matrix with a write probability Pve and an erase probability Pe below.
M W M. k U. Pw, Pe) = M + ok Makk
This write operation guarantees that reading with the same key vector yields the most recently written value. Such a property be mathematically proven as below.
3

---

Under review as a conference paper at ICLR 2023
Theorem 1. If lo is unit vector and M' = W,R,,,,,,,,1 RD(M', M.M,,, k,
Proof. M'k Mk + v k- Mkk k= =
The verbal explanations of this theorem are as follows. First, MI = RD(M, MMKk k. yields the value in associated to l so that -Pekkk' erases out the associated value from M. Then, adding Po
writes the new value on top of
3.2 NORMALIZED OUTER-PRODUCT PRODUCT ATTENTION
At each time step, the computational and space complexity of NAM are both O(d,da). Since de and dę are per-head dimensions, can be further reduced by the factor of the number of the heads H, to O (d.... (dx/ H)) = O(d,dz/ H). In deep learning setups, these are identical to those of other neural network layers. On the other hand, scaled dot-product attention's s compute and space complexities are OSS(d, +x.).. if the sequence has S tokens. This opens up the opportunity for using NAM as an efficient linear attention mechanism for long-range sequential tasks or edge inference environments. Furthermore, the outer-product operation vk" for NAM is often more efficient than dot-products in modern parallel HW since its compute complexity (dd) much higher than the amount of memory access required, (d).
However, the sequential nature W can make NAM inefficient in parallel GPU and NPU architec- tures. This is because the erasure part Pe Mkk of W R. depends the current state of M. Hence, we can make W R parallel by setting the erase probability P. to zeros. The reads can also become parallel if we make it bidirectional by reading the same final memory matrix M E, Pur which is sum of outer product from every time step E 1....,S. S. Reading this memory matrix can act an attention mechanism, namely normalized outer-product attention, as the following theorem holds.
(t) ((() ((t)T,
Theorem 2. If 1 ((()) are orthonormal, then E S,
KO), 1) - Pw +
Proof. Mk E, 0 (0) pw (t) as & (00) is when and otherwise.
In practice, such hard guarantee not possible since is hard to make all the keys orthonormal. However, normalized outer- outer-produce attention can act as soft attention mechanism if the keys are well separated. By removing the erasure process, trades off the computational power of NAM for the maximal efficiency and parallelism. In the following sections, we show efficiency and efficacy of NAM as an attention mechanism in long-range sequential tasks (Section 4 and computational power of NAM in compositional generalization tasks (Section 5
EFFICIENT TRANSFORMER WITH NAM
4.1 LINEAR TRANSFORMER USING NAM
Although NAM is theoretically capable of replacing scaled dot-product attention, there exists danger of information loss due the limited capacity of the memory M. Meanwhile, such an information loss is not an issue for Transformers because they utilize hidden activations from the entire sequence. Hence we need empirical evidences that NAM is an effective alternative of attention. Hereby we implement an efficient linear bidirectional transformer using normalized outer-product attention, namely NAM H.MM-Trnsformer. As proof-of-concepp we prove the efficacy and efficiency of NAM by at long-range sequence tasks.
evaluating
Nrm-Transformrr is the most basic way of leveraging normalized outer-product attention, by setting the read/write probabilities to 1. Given queries Q E Sxd, keys K E Sxd, and values V E R Sxd. of a sequence with tokens, TAM--ransformer" self-attention is computed as follows.
Altm AM (Q,K, V) (VT (VT(())))(()
Hereby a is a unit-vector normalization function applied to each row vector of K. Computationally, this is almost identical to the self-attentio of Linear Transformer (Katharopoulos al. 2020). The
4

---

Under review as a conference paper at ICLR 2023
differences from Linear Transformer to ours are: 1) unit vector normalization instead of ELU kernel function, 2) so that we do not need to compute the causal masking factor Z (Katharopoulos et al. 2020). That is, one can think Linear Transformer a special variant of NAM, sacrificing some computational capabilities for parallelism and efficiency.
4.2 LONG-RANGE TASK EVALUATION
Table Accuracy (%) and relative training speedup comparison of the three Transformer models in listops (Listops), text classification (Text), pixel-level image classification (Image) tasks.
Listops
Text
Image
Model
Ace
Speedup
Acc
Speedup
Acc
Speedup
Transformer
29.75
1
57.28
1
40.16
1
Linear Transformer
36.5
2.59
64.27
2.35
38.76
10.02
NAM--ransformer (Ours)
36.1
2.66
63.49
2.44
37.18
10.07
We test our NAM-Transformer long-range arena tasks (Tay et al., 2020) that have been used to compare efficient Transformer architectures. Since the previous work has already proven that Linear Transformer is as capable others, conduct partial comparison of ours to the original Transformer and Linear Transformer. Table [elee[ssee shows that NAM can be an efficient alternative to Transformers. 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 some cases. We see bigger speedups the image classification task because the models have much smaller per-head dimension than others (16 vs 64). Recall that the computational complexities of both linear transformer and NAM transformer are quadratic the per-head dimensions. Hence, the image classification result shows that smaller per-head dimension can bring bigger speedups for NAM, but may result in lesser capacity, Overall, the normalized outer-prodee onter-product attention as powerful as the scaled dot- hot-product attention while enjoying greater computational efficiency. In other words, NAM can be an efficient and effective alternative for attention mechanism.
Setup We use the code base of the original benchmark and implemented our NAM encoder by modifying the Linear Transformer implementation. The three models share the same hyperparameters and the only differences come from the attention algorithms. The evaluations 20.04 system with RTX 3080. The results differ from the original work because we used smaller batch sizes due to the limited VRAM capacity. All other experimental setups, including hyperparameters. are identical to the original benchmark (Tay et al. 2020). Details are available in the source code included as supplementary material.
run the Ubuntu
NAM AUGMENTED NEURAL NETWORKS
In this section, we design two types of nemory-augmeente neural networks (MANN) based on NAM, namely Long Short-term Attention Memory (LSAM) and NAM Turing Machine (NAM-TM). LSAM a generic recurrent neural network architecture derived from LSTM (Hochreiter Schmidhuberi 1997). LSAM replaces the long-term cell state of LSTM with NAM memory matrix. NAM Turing Machine a MANN design for algorithmic tasks inspired by Neural Turing Machine (Graves et al.. 2014). Its read and write heads can move along the tape with four actions: left, right, no-op, and jump, Implementations of the heads are based on NAM read/write primitives. We evaluate the models with tasks of number sequence prediction, sequence reduction, and SCAN. We test their compositional generalization capability by splitting the data based on the sequence lengths.
5.1 LONG SHORT-TERM ATTENTION MEMORY
Long Short-term Memory (LSTTM (Hochreiter & Schmidhuber 1997) leverages two recurrent state vectors: the short-term hidden state and the long-term cell state. To mitigate the problems
5

---

Under review as a conference paper at ICLR 2023
of vanishing/exploding gradients, the cell state is additively updated using forget and input gates. Then, the output gate selectively reads the cell state to determine the hidden state. Long Short-term Attention Memory (LSAM) follows the same principle. Instead of using the vector cell state, it leverages the memory matrix Mt E IRd which is also additively updated using the NAM write primitive. The hidden state h2 E R is retrieved by reading the memory matrix. Given the input Ie Rd, the update rule M,, h4 LS (Iz, M,-1, is defined follows.
[q: : k2 ve] W Wakular : he-1] 1 bohr
< Pr, Pw > o (Wr h + brw)
M₁ W RVM(MM (((k),, (e), Ubs Pers PwrP
h+ (DDMD.,,,C p.))
Hereby is the concatenate operator, (.) is the sigmoid function, and Woke: Wrw: Byku, brw are trainable weights and biases. Although the memory matrix M, has much higher capacity than vector cell state, the computational complexity of LSAM is identical to that of LSTM. This is because W R and RD have complexity of (dd) which is identical to that of matrix-vector multiplication of the weights and the states. Like Transformers, we can design multi-headed LSAM by concatenating per-head states MM and h4. Bidirectional LSAM is also possible by splitting the multiple heads into two directions. The backward heads are updated in the opposite direction by the rule of Mi, hi = LSAM (Iz, M/41, 1142))
The LSAM architecture combines the strengths of recurrent neural networks (RNN) and Transformers. Since LSAM follows the RNN design so that it enjoys strengths of RNNs such as low computational cost for inference and recurrent inductive bias. Additionally, it benefits from the strengths of attention because reading and writing the memory matrix natively incorporates the attention mechanism.
5.2
NAM TURING MACHINE
Neural Turing Machine (NTM) (Graves et al. 2014) is one the early neural networks that implement external memory structure with differentiable read and write methods. It is basis of Differentiable 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. Their external memory matrix is accessed by read and write heads using differentiable attention mechanisms.
Hereby we design NAM Turing Machine (NAM-TM) which adopts the design principles of NTM and DNC. The main idea of NAM-TM is to treat the tape state T = [U1, U2, Un, O.... E R x. as memory matrix as follows.
T [11, U2, ... Uns 0 O,... = Σ
mi
T
(v. CERREGRR
Hereby e; are standard basis vectors < O, 0, ...1, ...,, 0 > of L where L is the size of the tape. This tape state now be accessed by using NAM read/write primitives.
NAM-TM is a differentiable function that takes the tape state T, read and write heads Hr, H., E RR4, and the input vector x E Fd and produces the read output R [R along with the updated tape and head states T", Hi, H'..
H, = (TM H,, x)
The read and write heads are positional vectors to attend the memory matrix for reading and writing the states. At each time step, the positions can be updated by four actions: LEFT, RIGHT, NO-OP, and JUMP. They are controlled by a controller neural network mincontrol(c)) which emits read and write probabilities Pr, action probabilities Pright, Pleft, Pnoops Pjump for each head and unit jump query vector fjump E Rdd Given the controller outputs and the value W.r to write, reading and writing the memory are conducted by NAM primitives as follows.
R = RD(T, H,,Pr)
T' = W R(T, Hw: U. Pw: Pw)

---

Under review as a conference paper at ICLR 2023
Then, each head is updated the next position based the action probabilities. LEFT and RIGHT actions can be performed by the differentiable roll function, which is a linear transformation R'. - R mapping e, e++1- (HLEFT, HRIGHT) = (rolt- roll(H))
The jump position H is determined by reading the transpose of key tape K with the unit query vector Ljump- The key tape written the same way as the tape T, but it stores the corresponding unit key vector derived using the weight matrix Wx E Raxdd and the unit vector normalization p(.)..
K' WR(K, Hw, ((( Pies Pw)
HJUMP = RD(K'T <9jump>?
One understand the transpose of the key tape KT =E; eikk as jump table. Reading K with a key vector K; returns the corresponding position vector 6, the keys are orthonormal. Finally, the next head position HI' updated as weighted sum of the positions.
Proop x H + Pteft x HLEFT + Pright X HRIGHT + Pjump x HUMMP
While all of the computations parameters depend on the value of L. That is, a NAM-TM TM trained on certain length L can be applicd to any tape length I' without modification nor re-training. Theoretically, it be extended to infiniteddimensional Hilbert spaces.
technically done with a fixed tape length L, none of the trainable
There are multiple strengths in NAM-TM design compared to the memory structures of NTM and DNC. First, unlike NTM and DNC, NAM-TM's building blocks are simple and computationally efficient. Second, NAM-TM's addressing mechanism is based the query-key-value attention mechanism of NAM, which is more powerful than the content- based attention mechanism used in NTM and DNC. Finally, NAM-TM's design is flexible in that easy to add/remove transition rules of the read/write heads. For example, the JUMP transition rule is optional in that Turing machine only requires LEFT and RIGHT transitions in theory. One can also add another transition rule for the head positions if the rule can be defined with differentiable functions.
5.3 COMPOSITIONAL GENERALIZATION TASKS
We test the computational powers of LSAM and NAM-TM using three types of algorithmic tasks in compositional generalization setups. 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. can test the compositional generalization capability by testing/validating the models with the longer decimal numbers that are never observed during the training stage. In this setup. many models often suffer from drastic fall of test/validation accuracy, due lack of inductive bias (Kim al. 202T). We use two representative sequences from NSP: Fibonacci (Fib) and Palindrome (Palin). The two tasks require the generalization of digital addition and sequence reversal rules, respectively.
Next is a sequence reduction (Reduce) task with simple rule: given the sequence of digits, the target the reduced sequence by skipping zeros. This is that can be easily solved by proper Turing machine. We also use similar compositional generalization setup by testing/validating using the longer target sequences that are longer than any target sequences the training dataset.
For NSP and Reduce tasks, we use training datasets with d 1 10 decimal digit sequences in little-endian order. Then we validate/test the models with two validation sets (ID, OD-easy) and test set (OD-hard) for each task. An in-distribution (ID) validation set consists of d 5... 10- digit sequences, and an out-of-distribution validation set (OD-easy) consists of d = 11 13-digit sequences. The harder out-o--distribution test set (OD-hard) consists of d 14 16-digit sequences, challenging the models with generalization to longer contexts. A training dataset has 25600 samples, and each validation/test set has 2048 samples.
The last task is SCAN (Lake & Baroni 2018) task which has familiarized the concept of compositional zero-shot generalization. SCAN consists of simplified natural language input sequences and the
7

---

Under review as a conference paper at ICLR 2023
corresponding output action sequences. Among many data split methods, train/test split by length to be consistent with the other tasks. Since SCAN only has a limited number of sequences, only have two datasets of train and test.
Fibonacci
Palindrome
Reduce
Inpui
Target
Figure 1: Input and output sequence examples of the tasks. The Fibonacci sequence is given in the little-endian order.
We format the problems as masked sequence completion problems as shown in Figure [ An input sequence consists of input tokens followed by masks, and output sequence consists of target tokens following the masks. Since the input and output sequences are 1:1 matched, we can compare structures like DNC. We choose four baseline models to compare: a bidirectional Transformer encoder (TF) (Devlin et al., 2018). an 2-layer LSTM model with attention (Bahdanau al. 2014), Universal Transformer (Dehghani et 2018). and a Differentiable Neural Computer (DNC). The TF model follows the architecture of BERT medium and the hyperparameters of the other models are adjusted to have similar parameter counts. Our 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. As an ablation study, also evaluate NAM-TM design without the JUMP transition (No Jmp). All experiments are run with PyTorch 1.10 on the system with Ubuntu 20.04 and RTX 3080. Each experiment takes less than five hours to run 200 training epochs. The implementation details and hyperparameters can be found at source code the supplementary materials.
to baseline models without peuueence-t-sequeene
5.4 EVALUATION RESULT
Table 2: Sequence accuracy (%) comparison on the compositional generalization tasks.
Model
TF
LSTM
UT
DNC
LSAM
NAM-TM
No Jmp
Parameters
28. 28.7M
31.5M
26.0M
40.2M
25.8M
23.3M
22.2M
ID
100%
100%
100%
100%
100%
100%
100%
Palin
OD-easy 65.6%
100%
99.2%
100%
100%
100%
100%
OD-hard
19.0%
100%
5.2%
100%
100%
100%
100%
ID 98.9%
46.2%
99.1%
46.9%
100%
97.4%
40.4%
Fib
OD OD-casy
0.4%
8.3%
21.3%
1.8%
39.9%
89.7%
19.8%
OD-hard
0.0%
0.0%
0.1%
0.0%
2.9%
71.5%
1.1%
ID
99.6%
99.5%
98.7%
100%
99.9%
100%
100%
Reduce
OD-easy 0.2%
91.5%
15.7%
95.7%
96.8%
100 %
100 %
OD-hard
0.0%
63.6%
0.0%
60.1% %
77.6%
100 %
100 %
Train
99.8%
99.9%
99.6%
99.3%
99.5%
99.6%
99.9%
SCAN
Test
0.0
14.5%
4.5%
8.6%
14.9%
11.2%
9.6%
Table 2lolwssswwws the sequence accuracy comparison of the models in compositional generalization tasks. We report the sequence accuracies of the models, where a wrong prediction of one token is counted as failed prediction of the entire sequence. To avoid cherry-picking, the epochs of the best OD-easy validation accuracies are presented for the evaluation results. However, we present the results of the best test accuracy for SCAN tasks because there are only two available datasets (Train, Test).
Although LSAM's architecture is not specifically designed for algorithmic tasks, LSAM performs consistently better than the other baselines. Surprisingly, it performs better than DNC, which is a MANN model designed for such algorithmic problems. This is not a matter of model capacity
güttuh.com/heeeeellaaeSCCCC
8

---

Under review as a conference paper at ICLR 2023
because LSAM has a slightly smaller parameter count. The results imply that the computational power of NAM memory architecture superior to that of DNC's memory architecture.
As expected, NAM-TM performs significantly better in algorithmic tasks (Palin, Fib, and Reduce) possibly due to its specialized architecture. Especially, it finds casier to solve OD-hard problems, whereas the other models experience steep performance decline. A potential explanation is that the action-based positional transitions provide robustness in long-context cases. However, NAM-TM is not the best-performing model for the SCAN task.
NAM-TM remains effective at palindrome and reduction tasks even without the JUMP transition (No Jmp). Hence, LEFT and RIGHT transitions seem be enough for emulating simple Turing machines, but more complex transition rules can augment the computational power of NAM-TM. This suggests extending NAM-TM to variety of tasks by augmenting specialized transition actions.
6 FURTHER APPLICATIONS
Since NAM offers simple and flexible building blocks for memory architecture, there are infinite number of potential applications. Hereby we suggest few possible extensions of NAM. While we do not evaluate the ideas in this work, we leave them future work.
Hierarchical data modeling Tensor product, the key operation behind NAM, is not limited to two-dimensional outer product. Any dimensional tensors T; can construct memory tensor M by performing sum of tensor products M = Σ; T; @ ki with the unit key vectors k; For example, document tensor D can be constructed by sentence matrices S,, each of which is also sum of outer products of word embeddings and word- level keys. Nested attention such hierarchical memory can be conducted by two inner products with sentence-leve unit query vector 4s and word-level unit query vector Jw-
of tensor products of sentence-level keys and
Efficient edge inference While Transformers are very successful in many MIL tasks, deploying such models for edge inference has been challenging task since Transformer's computation and memory cost per each time step varies with the sequence length. However, NAM's cost does not depend on the sequence length at all. Also, the outer product operation is more compute intensive than the dot products, making NAM more friendly to high- -ighhhhoughhutt accelerators. Hence, NAM can be an efficient Transformer alternative for edge inference.
Few-shot learning Given a input-output vector pair (x,, y). is possible to conduct one-shot learning of weight matrix W by W R(W, y/|z|, >/|=| 1. 1). This is because Wx is guaranteed to return y by Theorem[]] Therefore, we can implement one-shot or few-shot learning as memorization to NAM. This aligns well with the human behavior since we conduct few-shot learning by memorizing, not by repetitive training.
CONCLUSION
We proposed redesign of attention mechanism to construct a differentiable memory for neural networks, namely neural attention memory (NAM). 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. Then we can read it by multiplying the matrix with unit query vector. The first strength of NAM is that its computational complexity does not rely on the sequence length. 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. Next, NAM can be powerful basis for constructing MANN models. We designed two NAM-based MANNs: LSAM for generic sequential tasks and NAM-TM for algorithmic tasks. In compositional generalization tasks, both outperformed other baselines such as Universal Transformer and DNC, indicating that NAM is more powerful mechanism for implementing memory in DNNs. Finally, it opens up further research possibilities in the fields of hierarchical data modeling, edge inference, and few-shot learning.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.047. 2014.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.051 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quae V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv:1901.02860, 2019.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals. Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.0338 100.00399... 2018.
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint 140.541410.5401, 2014.
Alex Graves. Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwinska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 538(7626): 471-476, 2016.
Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, and Phil Blunsom. Learning to transduce with unbounded memory. Advances neural information processing systems, 28, 2015.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30, Curran Asso- ciates, Inc., 2017. URL
www.:iin..........o.e.e.e.. wwwrips...... http:pps.ccppppppeeeeee wttpp//ppppppp///........ Paperr.ppppppp.... pdf
Sepp Hochreiter. The vanishing gradient problem during learning recurrent neural nets and problem solutions. International Journal of Uncertainty, Fuzziness and Knowledge-Bald -assd Systems, 6(02): 107-116, 1998.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1780- 1997.
Heikki Hyötyniemi. Turing machines recurrent neural networks. Proceedings of step, 96, 1996.
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. Advances in neural information processing systems, 28, 2015.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference Machine Learning, pp. $156-5165. PMLR, 2020.
Segwang Kim, Hyoungwook Nam, Joonyoung Kim, and Kyomin Jung, Neural sequence-to to-grid mod- ule for learning symbolic nules. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 8173-8171, 2021.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04455 2020.
10

---

Under review as a conference paper at ICLR 2023
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sesquencee---sqqueee recurrent networks. In International conference on machine learning, pp- PMLR, 2018.
Liunian Harold Li, Mark Yatskar, Da Yin, Cho-Jui Hsieh, and Kai-Wei Chang. Visualbert: A simple and performant baseline for vision and language, 2019. URL http:::/aar....ogg//////
1908. 03557
Minh-Thang Luong, Hieu Pham, and Christopher D Manning. Effective approaches attention-based neural machine translation. arXiv preprint arXiv:I508..44 2015.
Pedro Henrique Martins, Zita Marinho, and André FT Martins. -oomer: Infinite memory transformer- former: Infinite memory transformer. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 5468 5485, 2022.
Hyoungwook Nam, Segwang Kim, and Kyomin Jung. Number sequence prediction problems for evaluating computational powers of neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 4626-4633 2019.
Jack Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy Lillicrap. Compressive trans- formers for long-range sequence modelling. International Conference Learning Representations, 2019.
Thierry Tambe, Coleman Hooper, Lillian Pentecost, Tianyu Jia, En-Yu Yang, Marco Donato, Victor Sanh, Paul N. Whatmough, Alexander M. Rush, David Brooks, and Gu- Yeon Wei. Edgebert: Sentence-level energy optimizations for latency-awar multi-task nlp inference, 2020. URL
Dttps://axxx..orrr/////////11/1
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:2011.0400 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Iliia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017


Reviews:
[{"review_id": "zfKvWTvo_3D", "reviewer": "Reviewer_g5E5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes neural attention memory (NAM), an alternative to the standard attention mechanisms, which suffer from well known limitations: their quadratic dependency with respect to the sequence length, and thus their inefficiency when dealing with long sequences, as well as their difficulty to solve tasks that require memory states.\n\nNAM follows the same query-key-value structure of scaled dot-product attention; it first writes a memory matrix by adding outer products of key-value pairs, and then reads it by multiplying the memory matrix with a unit query vector.  NAM attention reduces the computational complexity to linear with respect to the sequence length. \n\nThey perform experiments on long-range arena tasks, comparing a transformer with NAM self-attention with a vanilla transformer and with the Linear Transformer of Katharopoulos et al. (2020). They further experiment with their memory augmented neural network designs (LSAM and NAM-TM) on compositional generalization tasks like number sequence prediction, sequence reduction, and SCAN. \n", "summary_of_the_review": "The paper is well written and addresses an important problem. The formulation is interesting and might be used for other applications that were not explored in this paper. See the Weakness section for further details on what I think the paper would benefit from. ", "strength_and_weaknesses": "Strengths:\n\n* The paper is well written and is easy to follow. \n* The motivation is very clear from the beginning: trying to address two limitations of current attention mechanisms (inefficiency when dealing with long sequences and stateless design).\n\nWeaknesses:\n* The discussion about memory-augmented neural networks (MANN) in $\\S2.2$. seems limited. It would be beneficial to discuss the similarities/differences of your proposal and recent works that try to augment the transformer architecture with a memory network (e.g., [1], [2]). The same applies for the experiments in $\\S5$. \n* The writing of $\\S4.1$ could be improved. In particular, I find the explanation of NAM-based attention a bit confusing and unstructured. Also, I don’t think the connection between $\\S3$ and $\\S4$ is very clear. \n* Several transformer architectures that try to efficiently model long sequences have been proposed in the last couple of years. I believe that your experimental results would be stronger if you compared NAM to some of these – the linear transformer is just one. \n\nMinor Comments/Questions:\n\n* I would slightly change the sentence “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.”, because attention may be spread/not concentrated on one item only. \n* Typo in “algorithmic patterns that require more that pushdown automata in practice”\n* Is something missing in “Note that $V^\\top(p_w \\odot \\mu(K))$ is one of the special cases of constructing the memory matrix by setting the erase probabilities to zeros W R operations.”?\n* Typo (*hyperparameters*) in “The implementation details and hypereparameters can be found at the source code in the supplementary materials.” Also, it would be better to have the hyperparameters listed in, e.g., an appendix, and not only as part of the code.\n* Can you further comment on “However, NAM-TM is not the best-performing model for the SCAN task.”?\n* Typo in “For example, a document tensor $D$ can constructed by sum”\n* Typo in “NAM’s cost does not depend the sequence\n\n\n\n[1] Compressive Transformers for Long-Range Sequence Modelling, Rae et al., ICLR 2020\n\n[2] $\\infty$-former: Infinite Memory Transformer, Martins et al., ACL 2022\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "USAJpgIwmp9", "reviewer": "Reviewer_83W7", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work proposes a method called neural attention memory. It constructs the attention output leveraging a memory matrix. The computational cost is linear to the sequence length (vs. quadratic in the conventional attention mechanism). The method is evaluated for long-range sequence tasks. It is also used to derived variants of LSTM and neural tuning machine. \n\n\n", "summary_of_the_review": "The proposed method is very limited in its novelty in comparison with Katharopoulos et al 2020. The empirical results do not justify the merits of the method among other alternatives. ", "strength_and_weaknesses": "Strength\nEfficient attention is a topic of general interests to the community. The presentation of the proposed method is clear. It is shown that the method can be applied to different settings such as LSTMs and NTMs.\n\nWeakness\n1. The novelty in compression with the linear attention work (Katharopoulos et al 2020) is very limited. In fact the mathematical forms are almost the same. The difference is that this work uses \"unit vector normalization\" so that they do not need to compute the causal masking factor Z. It is unclear to me what benefit, if any, the unit vector normalization provides.\n2. In the long-range task evaluation section, the proposed method performs worse than Katharopoulos et al 2020. Also, other linear transformer works (discussed in Section 2.1) should be compared with too. \n3. One more comment: although the proposed approach is linear to the sequence length, the writing process is a sequential operation. In practice, this can be slow on GPUs/TPUs in comparison with conventional attention which can be fully parallelized. It would be nice to discuss the implementation as well as the break even sequence length with efficient GPU or TPU implementations. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "JOjsc68fPNw", "reviewer": "Reviewer_ek6u", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a memory-based attention mechanism. It leverages the outer product of the key and query to write into a memory matrix. To retrieve information from the memory, the query does a dot product with the memory matrix. The benefit of this attention mechanism is that the computation complexity is reduced from quadradic to linear. Experiment results show the efficiency of the proposed method. The author further proposes NAM-based LSTM and NTM, both models achieve stronger performance compared to their vanilla version while maintaining a similar computation complexity.", "summary_of_the_review": "Overall the paper proposes a new outer product-based memory mechanism, that has an incremental novelty. The author also leverages the proposed method to extend recurrent neural networks such as LSTM and NTM. The compositional generalization results are interesting, but the detail of the experiments requires further clarification.\n", "strength_and_weaknesses": "Strength:\n1. The proposed NAM transformer is fairly efficient compared to the standard transformer model and slightly faster than the linear transformer model.\n2. The newly proposed LSAM and NAM TM models expand the memory capacity of the original LSTM and NTM model and achieve strong compositional generalization.\n\nWeakness\n1. The NAM is not substantially different from the linear transformer and achieves lower performance compare to the linear transformer while the speedup is similar.\n2. The performance of LSAM and NAM TM still falls behind the SOTA models in the SCAN task.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "BQ7U34axX74", "reviewer": "Reviewer_Ln3M", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes neural attention memory (NAM) as an efficient memory-augmented attention mechanism. The authors propose a well-motivated read and write mechanism that is efficient and convenient for hardware implementations. Then, they show a self-attention mechanism inspired by NAM. The resulting NAM Transformer is comparable to Linear Transformer (direct competitor). The authors also design an LSTM-like and an NTM-like architectures that use NAM as attention mechanisms. The authors demonstrate that the NAM-augmented architectures are superior to direct competitors, such as LSTM, NTM, DCN, UT and TF on simple algorithmic tasks, such as Fibonnaci, Palindrom, Reduce and SCAN. Sometimes the proposed methods generalize very well on OOD tasks.", "summary_of_the_review": "The paper is well executed and the constructions are meaningful, so I recommend weakly the paper to be accepted. To improve the paper: it would be nice to experiment with read and write into the self-attention mechanism. Further improvements might come from trying some simple few-shot learning tasks. ", "strength_and_weaknesses": "Strengths:\n\n1. Simple and working idea. \n\n2. Meaningful implementation that can be applied in hardware efficiently, and it could help few-shot learning as well.\n\nWeaknesses:\n\n1. The gains are only marginal compared to the Linear Transformer (see Table 1).\n\n2. It seems to me that the Self-attention NAM model is a bit disconnected from Attention NAM. Namely, currently posed, there are no read/ write operations in SelfAttn_NAM. Could you elaborate a bit further, or potentially test some SelfAttn mechanisms with read / write?\n\n3. It would be nice to see experiments on the domains you suggested trying, e.g. few-shot learning.\n\n4. Minor:\n\n* I think you mean \"bigger speadups at image classification\" in place of  \"bigger speadups at text classification\" when discussing Table 1. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject