Going Beyond Linear Transformers with Recurrent Fast Weight Programmers
Authors: Kazuki Irie
Imanol Schlag
Róbert Csordás
Jürgen Schmidhuber
Venue: NeurIPS 2021 Poster (2021)
Keywords: Transformers
memory augmented recurrent neural networks
fast weight programmers

Abstract: Transformers with linearised attention (''linear Transformers'') have demonstrated the practical scalability and effectiveness of outer product-based Fast Weight Programmers (FWPs) from the '90s. However, the original FWP formulation is more general than the one of linear Transformers: a slow neural network (NN) continually reprograms the weights of a fast NN with arbitrary architecture. In existing linear Transformers, both NNs are feedforward and consist of a single layer. Here we explore new variations by adding recurrence to the slow and fast nets. We evaluate our novel recurrent FWPs (RFWPs) on two synthetic algorithmic tasks (code execution and sequential ListOps), Wikitext-103 language models, and on the Atari 2600 2D game environment. Our models exhibit properties of Transformers and RNNs. In the reinforcement learning setting, we report large improvements over LSTM in several Atari games. Our code is public.

Going Beyond Linear Transformers with Recurrent Fast Weight Programmers
Kazuki Irie', Imanol Schlag', Róbert Csordás' .ürgen Schmidhuber
The Swiss AI Lab, IDSIA, University of Lugano (USI) & SUPSI, Lugano, Switzerland 2 King Abdullah University of Science and Technology (KAUST), Thuwal, Saudi Arabia (kazuki, imanol, robert, juergen)@idsia.de
Abstract
Transformers with linearised attention ("linear Transformers") have demonstrated the practical scalability and effectiveness of outer product-based Fast Weight Programmers (FWPs) from the '90s. However, the original FWP formulation more general than the one of linear Transformers: slow neural network (NN) continually reprograms the weights of a fast NN with arbitrary architecture. existing linear Transformers, both NNs are feedforward and consist of single layer. Here we explore new variations by adding recurrence to the slow and fast nets. We evaluate our novel recurrent FWPs (RFWPs) two synthetic algorithmic tasks (code execution and sequential ListOps), Wikitext- 103 language models, and the Atari 2600 2D game environment. Our models exhibit properties of Transformers and RNNs. In the reinforcement learning setting, we report large improvements over LSTM in several Atari games. Our code is public
Introduction
The Transformer has become of the most popular neural networks (NNs) for processing sequential data. Its success on neural machine translation quickly transferred to other problems in natural language processing, such as language modelling [2] question answering [4]. Recently, it has also been applied in other domains, including image processing S 6 or mathematical problem solving
Conceptually, the Transformer is deep feedforward NN that processes all elements of sequence in parallel: unlike in recurrent NNs (RNNs), the computations of layer for the entire sequence be packed into one big matrix multiplication. This scales well with the number of parallel processors.
Despite the benefits of parallelisation, major drawback of Transformers is that their computational complexity in time and space is quadratic in sequence length. Furthermore, the auto-regressive version 2 the focus of our work - the state size increases linearly with sequence length. This makes Transformers infeasible for auto-regressive settings dealing with very long or potentially infi- nite sequences, forcing practitioners truncate temporal contexts and ignore long-term dependen- cies beyond fixed-size time windows. Although recent work tries to address this issue [0 [T] this limitation makes applications of Transformers challenging, e.g., reinforcement learning (RL) in partially observable environments [3), which is still dominated by RNNs such as the Long Short-Term Memory (LSTM; (14) trained by policy gradients [5] [6 17 [8].
To scale Transformers to longer sequences, recent works have proposed to linearise the softmax in the self-attention computation and reorganise the latter in a sequential way [99. Such models include
Equal contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Katharopoulos et al.]s Linear Transformer (LT) [19]. Choromanski et al. Performer [20] and Peng al. [2] variant. They enjoy time and space complexities linear sequence length with states of constant size. While their performance on some tasks does not fully match the one of regular
Transformers [2 several improvements have already been proposed [21 23
(see our review in
Sec. S.c.2.2 which makes this Transformer family promising alternative.
Here we go one step further in advancing linear Transformer variants as powerful auto-regressive sequence processing models, adopting the perspective of "Fast Weight Programmers" (FWPs) [24] 25
26.. Recent work emphasised that linearised Transformers are essentially equivalent to outer product- based FWPs from the '90s (23): reviewed in Sec. 2 Here we explore this connection further and describe more powerful FWPs.
The original FWP 24 is two-NN system: a slow and a fast net, each with arbitrary architectures. The slow net learns generate rapid context-dependent weight modifications for the fast net. In the case of existing linear Transformer variants, the slow and fast nets are simple one layer feedforward NNs. Here we augment them with recurrent connections obtain recurrent FWPs (RFWPs). Recurrence enhances the model's theoretical power [27] and can help to solve tasks that naturally require recurrence as a part of the solution.
Our experiments on the language modelling dataset Wikitext-103 [28] show that our RFWPs are competitive compared to regular Transformers. We then study various properties of the proposed models on two synthetic algorithmic tasks: code execution (2 and sequential ListOps 30 Finally, it is straightforward to apply our models to RL problems as a drop-in replacement for LSTMs. Here our RFWPs obtain large improvements over LSTM baselines across many Atari 2600 2D game environments [1. Although LSTM still works better in few environments, show that
RFWPs generally improve by scaling them up.
The main contribution of this work is twofold: (1) from the perspective of FWPs, study novel powerful FWPs for sequence processing, demonstrating that NNs can easily learn to control NNs that are more complex than single feedforward layer, and (2) from the perspective of Transformer models, our RFWPs augment linear Transformers with recurrence, addressing general limitations of existing rego-regressive Transformer models.
2
Background on Fast Weight Programmers (FWPs)
Here review the general concept of FWPs, as well as two specific instances thereof: the linear Transformer 19 20 and the Delta Net 23
2.1 General Formulation
We refresh the concept of fast weight controllers or FWPs 24 25 using modern notation in sequence processing scenario. An FWP with trainable parameters slow sequentially transforms an input
sequence (0' with (!) E Rdd to an output sequence { 1 with y E Rdd of length T as
(²t) Fast q = ZlowNet((z(01) 1 ( (1 (e (0
1t) J= (1 astow)
y (t) FastNet((f)) ( (v (1))) (())
(2)
where ())), ( and ()) are initial variables. This is a system with two NNs called Fastlet and SlowNet in which the parameters t of FastNet are generated by SlowNet at each time step f. The weights of the fast net are fast in the sense that they may rapidly change at every step of the sequence while the weights of the slow net Oxow are slow because they can only change through gradient descent during training, remaining fixed afterwards" Eq. expresses a slow NN in its general form. The slow net can generate fast weights conditioned on various variables, depending on architectural choices
for the slow and fast NNs. In addition to the fast weights o( the slow net also generates or invents an input ( to be fed to the fast net (alternatively (2) can simply be (!)))) While the architectures of slow and fast nets are arbitrary, they are typically chosen to be differentiable such that the entire
FWP can be trained in an end-to-end manner using gradient descent. By interpreting the weights of
2 The fast net could also contain some additional slow weights; omit this possibility here.

---

an NN as a program [32, the slow net effectively learns to control, or program, the fast NN. Thus, the slow net a neural programmer of fast weights, and its parameter set stow embodies compressed information used to produce potentially infinite variations of context-dependent fast weights.
In many settings, it makes sense to generate the fast weights 0 incrementally in iterative fashion, where the SlowNet is further decomposed into two sub-parts:
SlowSubnet (x> (1) 1t j==11 (1 (0 1 (a 31 (2()) Jj=0 sslow)
(3)
UpdateRule (0(t-1) ())
(4)
fast
fast
where UpdateRule takes the fast weights (-- from the previous iteration to produce the new
fast weights O conditioned on 2(1), The update rule essentially the differentiable elementary
programming instruction used by the FWP. In the next section we review concrete examples of recent FWPs.
2.2 Linear Transformers as Fast Weight Programmers
fast is too large to be conveniently parameterised by an NN. Instead, it proposed in 1991 [24] to perform a rank-one update via the outer product of two vectors generated by the slow net. Two recent models directly correspond to such outer product-based
In general, the dimension of the fast weights (t)
FWPs: linear Transformers
19 and the Delta Net [23]
Linear Transformer. The "linear Transformer" 19 is class of Transformers where the softmax in the attention is linearised. This is achieved by replacing the softmax with kernel function o-then the self-attention can be rewritten as a basic outer product-based bbsed FWP 24 23 Previous works focused on different é maps with properties such as increased capacity [23 or guaranteed approximation of the softmax in the limit (20 21). For our purposes, the particular choice of ő is irrelevant and we simply assume Φ du Raher simplifying our equations below by writing k. q instead of (kk)), ((q)-
Using otherwise the same notation as above, for each new input ax!"), the output y is obtained by:
A (1) v (1) (()
(1 (., W
(5)
W W(-1) + (t) 2 k(()
(6)
y g()g(())
(7)
where the slow weight matrices WR E & and W. E dh are used to obtain the key k(()) E dey and the value v()) E Rdo, The key and value vectors are used to generate new weights via the outer product (t)) & (t() E IRdux A further simplification in the equations above is the omission of attention normalisation which has been experimentally shown to be unnecessary if the function produces normalised key and query vectors (23 In Eq. 6. the previous fast weight matrix ((111) E Xhuae is updated to yield W by adding the update term (ft) @ (t)), This corresponds the sum update rule or purely additive programming instruction. Here the fast NN is a simple linear transformation as in Eq. 7 which takes as input the query vector (()) E RRdyy generated by the slow weights W, E Rud Xda, Hence, in linear Transformers, the previous Eq. simplifies to: z(t) qit) SlowSubnet(z Ostow) with 2(t) (A)(1) (()).
Delta Net. The Delta Net 23
is obtained by replacing the purely additive programming instruction
(Eq. in the linear Transformer with the one akin to the delta rule (33 :
(()) (t) ( (t-1)) + (v(()
where g(t) E R is a fast parameter (learning rate) of the update rule generated by the slow net with weights Wa E Blxds and the sigmoid function a:
g o (1)
and (()) E Ed is generated as a function of the previous fast weights (t-1)) and the key k(())
( (1-1))
(10)
3

---

This update rule was introduced to address memory capacity problem affecting linear Trans- formers with the purely additive update rule (23). The corresponding Eq. 3 is: x0)) SlowSubnet (x)() W((+1); skow) with 2(1) = (lx(), ,()),()) 5()). Thus, unlike linear Transform- ers, the SlowNet in the Delta Net takes the previous fast weights ((+11) into account to generate
the new fast weight updates.
We typically use the multi-head version 1 of the computations above. After the projection (Eq. 5 the vectors (0) ((t), g(t)) are split into equally sized // sub-vectors, and the rest of the operations are conducted by H computational heads independently. The resulting output vectors from each head are concatenated to form the final output.
Other approaches. While our focus here is on outer product-based weight generation, which is an efficient method to handle high dimensional NN weights, there are also other approaches. For example, instead of generating new weight matrix, Hypernetworks [34 scale the rows of slow weight matrix with a generated vector of appropriate size. Weight compression to control fast weights in low dimensional compressed space has been also studied [35]. In the broad sense of context-
dependent weights [6 37 38), many concepts relate to FWPs: e.g. dynamic convolution [39 40 41
LambdaNetworks or dynamic plasticity 43
3 Fast Weight Programmers With Slow or Fast RNNs
The original formulation of FWPs reviewed in Sec. 2 more general than existing models presented in Sec. 2.2 In particular, both fast and slow networks in existing linear Transformers consist of single feedforward layer (Eqs. [ anddd 7 Here we present FWPs with recurrent fast nets in Sec. 3.] and FWPs with recurrent slow nets Sec. 3.2
3.1 Fast Network Extensions
In principle, any NN architecture be made fast. Its fast weight version is obtained by replacing the networks' weights with fast weights parameterised by an additional slow network. For example, consider a regular RNN layer with two weight matrices W and R:
h
( x (8) (1-))
(11)
fast weight version can be obtained by replacing and R with ( and R()) which are controlled as in Eq. 8 with all necessary variables generated by separate slow net at each time step i.
While this view illustrates the generality of FWPs, the angle under which approach these models is slightly different: we introduce recurrence as a way of augmenting existing linear Transformers.
Delta RNN. We obtain a fast weight RNN called Delta RNN by adding an additional recurrent term the feedforward fast net of the linear Transformer (Eq. 7
y
w ( + R ((y(()))
(12)
where Rdt) E IR an additional fast weight matrix which introduces recurrent connections. It also generated by the slow net using the delta update rule, similar to W(() in Eq. 8 but with additional slow weights. We apply an element-wise activation function ff to the previous output of the fast network (t-1) to obtain the recurrent query. The choice of activation function crucial here because, achieve stable model behaviour, the elements in key and query vectors should be positive and sum up to one when the delta update rule is used [23]. We use the softmax function
(ʃ = softmas in Eq. (2) to satisfy these conditions. An ablation study the choice of using Eq. 12 instead of the one similar to Eq. II can be found in Appendix A.2
Analogous to the Delta RNN, we also construct Delta LSTM with six fast weight matrices. The exact equations can be found in Appendix A.2
Alternative Feedforward Fast Nets. While the focus of this work on RNNs, there are also interesting fast feedforward models to be used in Eq. 7 which might result in stronger feedforward

---

baselines. For example, we can replace the single layer fast net of Eq. by a K K-layer deep network:
h (
l f(h for E [1..K] with hg ( q())
(13)
(8) (t) hk
(14)
where the slow network produces all K fast weights [ /==1 and query q( from single input 2²()) In light of the capacity limitation in linear Transformers [23]. this might introduce additional capacity without the need of larger representations, analogous to the trade-off in a multilayer
perceptron (MLP) between narrow & deep versus shallow & wide. We refer to this class of models as Delta MLPs. Again, for stable model behaviour with the delta rule, we apply the softmax activation the vectors be used query.
Another interesting approach is to use a Delta Net itself as a fast net, i.e., make the slow weights in the Delta Net fast (thus obtaining Delta Delta Net). Such model could in principle learn to adapt the way of generating fast weights depending on the context. While we plan to investigate the potential of such hierarchical FWPs in future work, we also include preliminary results of such model in our language modelling experiments (Sec. 4.1p. A discussion on the dimensionality of such model can also be found in Appendix A.3
We experimentally demonstrate that (slow) NNs can learn to control the weights of these rather complex fast networks (Sec.
3.2 Slow Network Extensions
In linear Transformers, the slow network is purely feedforward 5). It can be made recurrent at two different levels: within the slow network (i.e. the slow network computes weight updates based on its own previous outputs e.g., key, value, query vectors) or via the fast network by taking the fast net's previous output as an input. In our preliminary experiments, we found the former to be sub- optimal (at least in language modelling experiments). So we focus on the latter approach: we make the slow net in the Delta Net dependent on the previous output of the fast network. We refer to this model as the Recurrent Delta Net (RDN).
Recurrent Delta Net. We obtain the RDN by modifying the generation of key, value, and query vectors (Fq. (5) as well as the learning rate (Eq, (9) in the Delta Net. We add additional slow weights
(RR, R2 Ra dea R. E Rduuxda X and Rg R RIxdaa for recurrent connections which connect
the previous output of the fast net (1-1) (Eq. 7 to the new k(t), o(1), (()), and ((t) follows:
k(t)
Wi + R4 tanh (z (1-1)
(15)
(t
W. (', (1) + R. tanh ( (t-1))
(16)
q = W (1)
(t) + R, tanh (y (t-
(17)
g(t) = (W (1)) (t) + Rg tanh(y ((-1))
(18)
While the rest of the model remains as in the Delta Net, with these simple extra recurrent connections the model becomes proper RNN. The corresponding dependencies in Eq. 3 are: g(),g((
LLowSubnet(=!0 (t(()))(t))))))))) (t-1)), W(t-1); stow) with 2(1) (k, o(t) ,g0)) v(t)).
3.3 Related Models
All the RFWP models presented in Sec. 3.1 and 3.2 can be seen as a type of memory augmented recurrent neural networks 45 46) in the sense that they maintain two-dimensional fast weight states as a short-term memory, in addition to the standard one-dimensional RNN states.
There are also several previously proposed recurrent fast weight models. For example, Schmidhuber's recurrent FWP from 1993 [26] has been revisited by Ba et al. [47]. There, key and value vectors are not generated within the same time step, unlike in our models or in linear Transformers. The Fast Weight Memory (FWM) 48 is also a recurrent FWP: the slow net an LSTM and the fast net is a higher-order RNN. However, the FWM is single pair of slow and fast nets, and multi-layer version, as in the linear Transformer family, was not explored. Similarly, the Metalearned Neural Memory 49 uses an LSTM its slow net and a 3-layer MLP as its fast net but again limited to one pair.
5

---

Table 1: WikiText- 103 language model perplexity results with the small setting 21 23
For each model, its name, corresponding slow and fast networks, and weight update rule (Update) are specified. All models are trained and evaluated on the span of 256 tokens except for the models in the last two rows (+ full context) which are trained and evaluated without context truncation. Parameter count in millions. See Appendix A for further experimental details and results.
Name
Slow net
Update
Fast net
Valid
Test
#Prms
Transformer
33.0
34.1
44.0
Linear Transformer
Feedforward
Linear
37.1
38.3
44.0
Delta Net
delta
34.1
35.2
44.0
Delta MLP
Feedforward
delta
Deep MLP
35.8
36.8
44.3
Delta Delta Net
Delta Net
34.0
35.2
44.6
Delta RNN
RNN
33.8
35.0
44.6
Delta LSTM
LSTM
32.6
33.8
47.3
RDN
Recurrent
Linear
34.1
35.2
44.1
Delta RNN
31.8
32.8
44.6
full context
RDN
32.5
33.6
44.1
Others have investigated variants of RNNs with fast weights for toy synthetic retrieval tasks [50] 51]. In particular, Keller et al. [5]] augment the LSTM with a fast weight matrix in the cell update. In contrast, make all weights in the L.TTM fast and, importantly, our model specifications build upon the successful deep Transformer architecture using residual connections [52) (3), layer-norm [54). multiple attention heads and ferd-forward blocks [T]. Essentially, we replace the self-attention layers in the regular Transformers by the fast weight programmer operations described above.
4 Experiments
We conduct experiments in four different settings. We start by evaluating all models on a language modelling task (Sec. 4.1)) obtain a performance overview and discuss computational costs. Language modelling is an excellent task to evaluate sequence models. However, to highlight their different capabilities, we evaluate our models also algorithmic tasks. In fact, it is well-known that the actual capabilities of RNNs differ from one architecture to another (55). We are interested in discussing such differences. With that goal in mind, we conduct experiments on two synthetic algorithmic tasks, code execution (Sec. 4.2) and sequential ListOps (Sec. 4.3). which are designed to compare elementary sequence processing abilities of models. Finally, we apply our models to reinforcement learning in 2D game environments (Sec. 4.4) as replacement for LSTMs.
4.1 Language Modelling
We first evaluate discussed models on the generic language modelling task. This allows for obtaining a performance overview and reviewing the computational efficiency of different models. We use the Wikitext- dataset (28) and follow the small model setting similar what's 's used recent works by Peng et al. [21 and Schlag al. (3) This allows for training and evaluating different models with a reasonable amount of compute on this resource-demanding language modelling task.
Perplexity results. The results are shown in Table which also serves as a tabular summary recapitulating different models described in Sec. 2 and with various architectures for slow and fast nets, and two choices of update rule. The top block Table shows the performance of the baseline Transformer, Katharopoulos et al. I19's s Linear Transformer, and Schlag et al. 23 's Delta Net. The performance models presented in Sec. 3 can found in the middle block. First of all, the Delta MLP performs worse than the baseline Delta Net despite slight increase in parameter count (44.3 vs. 44.0 M). This supports the intuition that better to make the slow network aware of the outputs of intermediate layers to generate fast weights in deep network, instead of generating fast weights for all layers time. In all other models, the performance never degrades with the proposed architectural augmentation. The Delta Delta Net yields limited improvements; plan to study this

---

model in depth in future work. With the same amount of parameters (44.6 M), the Delta RNN yields greater improvements. Among the models presented here, the Delta LSTM variant exhibits the best performance. This shows that the slow network successfully controis the rather complex fast LSTM network, although also requires more parameters (47.3 M) than other models. Finally, the benefits of recurrent connections added to the baseline Delta Net do not directly translate into practical improvements in language modelling as demonstrated by the performance of RDN compared to the one of the baseline Delta Net. Importantly, given a constant memory size w.r.t. sequence length, it is straight-forward to train and evaluate our RNNs without context truncation (while still limiting the backpropagation span). Corresponding performances of Delta RNN and RDN are shown in the bottom part of Table [
they outperform the regular Transformer with a limited context (256 tokens).
While language modelling is useful as a sanity check (here for example, except for the Delta MLP, all models achieve reasonable performance), the task is too generic to identify certain important aspects of the models, such as real benefits of recurrence. Before we move on to trickier RI applications, will focus on studying such aspects using synthetic algorithmic tasks.
Sec. 4.2 and 4.3
Computational efficiency. The modifications we proposed in Sec. 3 introduce additional computa- tional costs to lincar Transformers/FWPs First of all, none of them affect the core complexity of linear Transformers: they all have a constant space and linear time complexity w.r.t. sequence length. However, the per-time-step computational costs differ a lot from one model to another, quantified here in terms of training speed using our implementation. All models are implemented using a custom CUDA kernel except the baseline Transformer for which we use regular Py PoTcc code [56]. Training speeds of LT and Delta Net in Table arre 66 K and 63 K words per second respectively (vs. 33 K for the baseline Transformer). The most expensive model is the Delta LSTM. This fast weight LSTM with tied input-forget gates has 6 weight matrices, and each of these are manipulated by separate delta rules. The corresponding speed is 14 K words per second, too slow for scaling to more experiments. In contrast, the speeds of Delta RNN and RDN remain reasonable: 41 K and 35 K words per second respectively. Therefore, the remaining experiments will focus on these two recurrent architectures which are promising and practical in terms of both performance and computational costs.
4.2 Code Execution Task: Learning to Maintain and Update Variable States
In code execution tasks 29. models are trained to sequentially read the input code provided as word- level text, and to predict the results of the corresponding code execution. We adopt the task setting from Fan et al. (57) with one conditional and three basic statements. We refer the readers to Appendix for a precise description of the task. This code execution task requires models to maintain the values of multiple variables, which has been shown to be difficult for relatively shallow Transformers with only feedforward connections (57).
The left block of Table 2 shows the results. Following again Fan et al. [57]. we control the task difficulty by modifying the number of variables (3 or 5). The model architectures are fixed: the LSTM has only one layer with 256 nodes and all Transformer variants have the same architecture with layers with a hidden size of 256 using 16 heads and an inner feedforward layer size of 1024.
We first note that the LSTM is the best performer for both difficulty levels, with the smallest performance drops through increasing the number of variables. In contrast to prior claims (57]. the LSTM is clearly capable of storing the values of multiple variables in its hidden and cell state vectors. With three variables. the regular Transformer already largely underperforms other models with a mutable memory: Delta Net, Delta RNN, and RDN. Linear Transformers completely fail at this task, likely due the memory capacity problem pointed out by Schlag al. (3 (see Appendix B.2 for further discussion). By increasing the number of variables to five, the baseline Transformers, Delta Net, and RDN become unstable as shown by high standard deviations w.r.t. the seed. The benefits of recurrent connections introduced in our RDN compared to the baseline Delta Net become more apparent (76.3 vs. 61.4%). In contrast, the Delta RNN remains stable and gives the best performance (85.1%) among Transformer variants, which shows the benefits of recurrence and in particular the regular RNN architecture in the fast net. To match the performance of I.STM on this task, however, these models need more layers (see Appendix B. 2
for more results).
7

---

Table 2: Test accuracies (%) with standard deviations on code execution (Code Exec) and sequential ListOps (Seq ListOps). The difficulty of the task controlled by the maximum number of possible variables (f variables) for code execution, and the list depth (10 or 15) for ListOps. For code execution with 5 variables, we report means over six seeds. In all other cases, the results are computed with three seeds. For more results, see Appendix B.2 (Code Exec) and (Seq ListOps).
Code Exec (# variables)
Seq ListOps (depth)
3
5
10
15
LSTM
99.0 + 0.1
93.2 + 6.1
88.5 + 2.9
24.4
+ 1.1
Transformer
71.8 2.6
35.4
28.2
79.1
0.9
75.3
0.4
Linear Transformer
0.0 + 0.0
0.0 + 0.0
64.0 + 0.3
64.4 + 0.4
Delta Net
90.7 2.7
61.4 20.0
85.7 1.8
77.6
1.4
Delta RNN
90.8 + 1.7
85.1 + 1.9
83.6 + 1.2
78.0
1.0
RDN
92.6 + 2.2
76.3 + 17.6
83.2 = 0.9
79.2 + 1.4
4.3 Sequential ListOps: Learning Hierarchical Structure and Computation
The ListOps task 30) is typical test for hierarchical structure learning, which requires list operation executions. We use a simple variant of ListOps whose detailed descriptions can be found in Appendix
B.4 For example, the list [MAX 6 [FIRST 2 3 ] 0 [MIN 4 7 1] ] is of depth two and the
expected output is While early research comparing self-attention
RNNs [58 has shown
some advantages of recurrence in hierarchical structure learning, more recent work (59] reports Transformers outperforming L.STMs on ListOps. According to Tay et al. [22, linear Transformer variants (LT and Performers) underperform other Transformer variants by large margin on ListOps.
The right block of Table shows results for two differrnt depths: 10 and 15. The model architectures are identical to those used in the code execution task (Sec. 4.2) At depth 10, we find LSTM to perform best, while mutable memory Transformer variants (Delta Net, Delta RNN, and RDN) outperform the regular and linear Transformers. At depth 15, the LSTM's performance drops drastically (to 24.4%)) while the differences between Transformer variants remain almost the same. We note that sequences are longer for the depth 15 problem (mean length of 185 tokens) than for the depth 10 version (mean length 98 tokens). This turns out to be difficult for the small 256-dimensional LSTM; see Appendix B.4 for the corresponding ablation study. The performance differences between the baseline Delta Net and the proposed Delta RNN and RDN are rather small for this task. Importantly, our models outperform both regular and linear Transformers on this task requiring hierarchical structure learning.
4.4 Reinforcement Learning 2D Game Environments
We finally evaluate the performance of our models as a direct replacement for the LSTM in reinforce- ment learning settings. In fact, only limited number of prior works have investigated Transformers
for RL. Parisotto et al. 12] and Rae et al. [11 evaluate them on the DMLab-30 [60] [1]. Parisotto
et al. 12 also evaluate them on Atari but in a multi-task setting (62). Others [57] [ use toy maze
environments. In contrast to Parisotto et al. [12's work, which presents multi-task Atari as side ex-
periment, we study the Transformer family of models on the standard Atari 2600 setting [31 63 64 by training game-specific agents.
Settings. We train an expert agent on each game separately with the Importance Weighted Actor- Learner Training Architecture (IMPALA) using the V-trace actor-critic setup (65] and entropy regularization [66] implemented in Torchbeast [67]. Our model follows the large architecture of Espeholt al. 655 which consists of a 1 15-layer residual convolutional NN with one 256-node LSTM layer which we replace by either the RDN (Sec. (.2) or the Delta RNN (Sec. 3.1 In line with the small LSTM used for Atari (only layer with 256 hidden nodes) we also configure a small RDN: layers with hidden size of 128 using 4 heads, and a feedforward dimension of 512. We find this small model to perform already surprisingly well. For the rest, we use the same hyperparameters as Espeholt et al. (65) which can be found in Appendix C
8

---

Relahwe Improsemmntt from LT
60
Ingrovimmnnss
1
Seorueee
Pang
B. Heur
Swaam
Amider
JanwA Band
Ofteert
Rabostek
Zazzon
Upin Down
Gipher
Kang 6
Aroutt
D. Abwck
Batewono
S. maaaing
Storgeest
Burgek
Pary
Reboteenk
Briiet
E
Staruanna
Arndar
Assech
King-Fu
somestiona
Aller
Batlienne
ordeet
Figure 1: Relative improvements in test scores obtained by the Recurrent Delta Net (RDN) compared to the Linear Trans- former (LT) after 50 M env. steps.
Figure 2: Relative improvements in test scores obtained by the Recurrent Delta Net (RDN) compared to the Linear Trans- former (I.T) after 200 M env. steps.
Main experiments. We evaluate our models in 20 environments. According to Mott et al. [68]. in about half of them, the LSTM outperforms the feedforward baselines-whiee whhic we confirm in our setting with 50 M steps (see Appendix C We report results at 50 M and 200 M environmental steps of training. Like Nair al. [69]. we run the trained agent for 30 test episodes. Here we repeat this evaluation five times to report the average score with a standard deviation. The following analysis focuses on the RDN (Sec. 3.2) compared to the regular linear Transformer and the LSTM. similar study of the Delta RNN, as well comparisons to more baselines, and the exact scores achieved by each model on each game can be found in Appendix C
In all our experiments above, we have shown that the Linear Transformer, i... Fast Weight Programmer with a purely additive update rule, consistently underperforms other models based on the delta rule. Here we confirm this trend once more. Figures [nd show the relative improvements of scores obtained by Recurrent Delta Net over those achieved by the linear Transformer on each game, respectively after 50 and 200 M interaction steps. The RDN matches or outperforms the Linear Transformer on all games except for two out of 20 games at both stages of training.
Figure 3 shows relative improvements of RDN over LSTM after 50 M interactions. In 12 games, the RDN yields improvements over LSTM, whereas in games, the LSTM performs better. In the
1.2
Reabre Improvements from the LSTM
$4.5
1.0
1.5
1.0
0.5
0.0
0.0
-05
-04
-1.
Konged
Boceest
Londomm Dom
2 Helld
Beeequns
Pont
Assauu
Gogher
O'Stet
deviter
Alon
Becoat
Copteer
Zacen
Konseuu
Pong
Biwaa
Rabetask
B invoseees
Wolder
Annitaa
D'Seet
Relative improvements in test scores obtained by 2-layer RDN compared to LSTM after 50 M env. steps.
Figure 3:
Figure Relative improvements in test scores obtained by 2-layer RDN compared to LSTM after 200 M env. steps.

---

remaining 5 games, both reach similar scores. Interestingly, this trend does not directly extrapolate to the 200 M case, which presented in Figure 4 With longer training, the LSTM surpasses the performance of the RDN in Baitlezone, Gopher, Seaquest and Zaxxon, while the RDN catches up in Up'N Down and Kung-Fu Master. Overall, there are games in which LSTM clearly outperforms RDN at 200 M steps, whereas in 9 games the result is the opposite.
On a side note, some the scores achieved by the RDN 200 M step are excellent: score of over 170 K and 980 K in Space Invader and Qe respectively beats the state-of-the-ard set by MuZero 70
and Agent57 62 However, a direct comparison not fair as we train game-specific agents.
Experiments with larger models. Given the results above, a natural question to ask is whether larger model size improves the RDN in games where the LSTT dominates. We focus on four such games: Battlezone, Berzerk, Gopher, and Seaquest (See Fig. 4. We double the model size to 3.4 parameters by increasing the number of layers to 4 and the hidden size to 256, with 8 heads. As shown in Table 3 larger RDN models reduce the gap to the LSTM (except in Berzerk). This indicates that further scaling RDN might be as promising as scaling regular Transformers in other domains.
Table 3: Performance of larger RDN in games where the LSTM dominates (200 steps).
Battlezone
Berzerk
Gopher
Seaquest
LSTM
24,873 + 1,240
1,150 + 92
124,914 + 22,422
12,643 + 1,627 4.373 1 504
RDN
10,980 1 1,104
348 1 17
86,008 1 11,815
RDN larger
28,273 1 5,333
346 1
118,273 1 14,872
14,601 1 712
5 Conclusion
Inspired by the formal equivalence of linear Transformers and certain traditional Fast Weight Pro- grammers (FWPs) from the early '90s, we propose various new linear Transformer variants with recurrent connections. Our novel Recurrent FWPs (RFWPs) outperform previous linear and regular Transformers on a code execution task and significantly improve over Transformers in a sequential ListOps task. On Wikitext- 103 in the "small" model setting, RFWPs compete well with the previous best linear Transformer variants for truncated contexts, and with full contexts, beat regular Trans- formers. Our RFWPs also be used as drop-in replacements for problems where RNNs are still dominant. In particular, we evaluate them reinforcement learning settings on 20 Atari 2600 en- vironments. They clearly outperform the regular Linear Transformer almost all environments. They also outperform the LSTM across many environments with a small model size and demonstrate promising scaling properties for larger models. Given the increasing interest in deploying Transform- ers in RL (2). in particular in the framework of Upside-Down RL
73 ZZ! our RFWP models
are particularly relevant: as RNNs, they conveniently handle long contexts with a constant memory size, while being powerful Transformer variants at the same time. Our work highlights the usefulness of the FWP framework from the '90s and its connection to modern architectures, opening promising avenues for further research into new classes of recurrent Transformers.
Acknowledgments and Disclosure of Funding
We thank Aleksandar Stanié and Sjoerd van Steenkiste for valuable comments on the first version of this paper. This research was partially funded by ERC Advanced grant no: 742870, project AlgoRNN, project NEUSYM. This work
and by Swiss National Science Foundation grant no: was partially supported by computational resources at the CSCS Swiss National Supercomputing Centre, project di 15. We thank NVIDIA Corporation for donating several DGX machines, and IBM for donating a Minsky machine.
References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Advances in Neural
10

---

Information Processing Systems (NIPS), pages 5998-6008 Long Beach, CA, USA, December 2017.
[2] Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character-level language modeling with deeper self-attention. In Proc. Conference on Artificial Intelligence (AAAI), pages 3159-3166, Honolulu, HI, USA, January 2019.
[3] Tom B Brown et al. Language models are few-shot learners. In Proc. Advances in Neural Information Processing Systems (NeurIPS), Virtual only, December 2020.
[4] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACL- HLT), pages 4171-4186, Minneapolis, MN, USA, June 2019.
[5] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In Inf. Conf. on Learning Representutions (ICLR), Virtual only, May 2021.
[6] Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai. Deformable DETR: Deformable transformers for end-to-end object detection. In Int. Conf. on Learning Representations (ICLR), Virtual only, May 2021.
[7] David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical reasoning abilities of neural models. In Int. Conf. Learning Representations (ICLR), New Orleans, LA, USA, May 2019.
[8] Imanol Schlag, Paul Smolensky, Roland Fernandez, Nebojsa Jojic, Jürgen Schmidhuber, and Jianfeng Gao. Enhancing the transformer with explicit relational encoding for math problem solving. Preprint PrXiv:1910.000 1, 2019.
[9] Francois Charton, Amaury Hayat, and Guillaume Lample. Learning advanced mathematical computations from examples. In Int. Conf. on Learning Representations (ICLR), Virtual only, May 2021.
[10] Zihang Dai, Zhilin Yang, Yiming Yang, William W Cohen, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context, In Proc. Association for Computational Linguistics (ACL), pages 2978-2988 Florence, Italy, July 2019.
[11] Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling. In Int. Conf. on Learning Representations (ICLR), Virtual only, April 2020.
[12] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing Transformers for reinforcement learning. In Proc. Int. Conf. on Machine Learning (ICML), pages 7487-7498, Virtual only, July 2020.
[13] Emilio Parisotto and Ruslan Salakhutdinov. Efficient transformers in reinforcement learning using actor-learner distillation. In Int. Conf. on Learning Representations (ICLR), Virtual only, May 2021.
[14] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 98): 1735-1780, 1997.
[15] Daan Wierstra, Alexander Förster, Jan Peters, and Jürgen Schmidhuber. Recurrent policy gradients. Logic Journal of IGPL., 18(2):620-634, 2010.
[16] Daan Wierstra, Alexander Förster, Jan Peters, and Jürgen Schmidhuber. Solving deep memory POMDPs with recurrent policy gradients. In Proc. Int. Conf. on Artificial Neural Networks (ICANN), pages 697-706, Porto, Portugal, September 2007.
11

---

[17] OpenAl et al. Dota 2 with large scale deep reinforcement learning. Preprint arXiv:1912.0668) 2019.
[18] Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, al. Grand- master level in StarCraft II using multi-agent reinforcement learning. Nature, 575(7782)3350-
354, 2019.
[19] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proc. Int. Conf. on Machine Learning (ICML), Virtual only, July 2020.
[20] 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. In Int. Conf. on Learning Representations (ICLR). Virtual only, 2021.
[21] Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In Int. Conf. on Learning Representations (ICLR), Virtual only, 2021.
[22] Yi Tay, Mostafa Dehghnn,, 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. In Int. Conf. on Learning Representutions (ICLR), Virtual only, May 2021.
[23] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers are secretly fast weight programmers. In Proc. Int. Conf. on Machine Learning (ICML), Virtual only, July 2021.
[24] Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to recurrent nets. Technical Report FKI-147-91, Institut für Informatik, Technische Universität München, March 1991.
[25] Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
[26] Jürgen Schmidhuber. Reducing the ratio between learning complexity and number of time varying variables in fully recurrent nets. In International Conference Artificial Neural Networks (ICANN), pages 460-463, Amsterdam, Netherlands, September 1993.
[27] Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8: 8:156-171, 2020.
[28] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In Int. Conf. on Learning Representations (ICLR), Toulon, France, April 2017.
[29] Wojeiech Zaremba and Ilya Sutskever. Learning to execute. Preprins arXiv: 1410.4615, 2014.
[30] Nikita Nangia and Samuel Bowman. ListOps: A diagnostic dataset for latent tree learning. In Proc. North American Chapter of the Association for Computational Linguistics (NAACL): Student Research Workshop, pages 92-99, New Orleans, LA, USA, June 2018.
[31] Mare G. Bellemare, Georg Ostrovski, Arthur Guez, Philip S. Thomas, and Rémi Munos. Increasing the action gap: New operators for reinforcement learning. In Proc. AAAI Conf. on Artificial Intelligence, pages 1476-1483, Phoenix, AZ, USA, February 2016. AAAI Press.
[32] Jürgen Schmidhuber. Making the world differentiable: On using fully recurrent self-supervised neural networks for dynamic reinforcement learning and planning in non-stationary environ- ments. Technical Report FKI-126-90, Institut für Informatik, Technische Universität München, 1990.
[33] Bernard Widrow and Marcian E. Hoff. Adaptive switching circuits. In Proc. IRE WESCON Convention Record, pages 96-104, Los Angeles, CA, USA, August 1960.
[34] David Ha, Andrew Dai, and Quoe Le. Hypernetworks. Int. Conf. on Learning Representa- tions (ICLR), Toulon, France, April 2017.
12

---

[35] Kazuki Irie and Jürgen Schmidhuber. Training and generating neural networks in compressed weight space. In Neural Compression: From Information Theory Applications Workshop, ICLR 2021, Virtual only, May 2021.
[36] Christoph von der Malsburg. The correlation theory of brain function. Internal Report 81-2, Goettingen: Department of Neurobiology, Max Planck Intitute for Biophysical Chemistry, 1981.
[37] Jerome A Feldman. Dynamic connections in neural networks. Biological cybernetics, 46(1), 27-39, 1982.
[38] James L McClelland. Putting knowledge in its place: A scheme for programming parallel processing structures on the fly. Cogmitive Science, 9(1):133-146, 1985.
[39] Benjamin Klein, Lior Wolf, and Yehuda Afek. A dynamic convolutional layer for short rangeweather prediction. In Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR), pages 4840-4848 Boston, MA, USA, June 2015.
[40] Hyconwoo Noh, Paul Hongsuck Sco, and Bohyung Han. Image question answering using convolutional neural network with dynamic parameter prediction. In Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR), pages 30-38, Las Vegas, NV, USA, 2016.
[41] Xu Jia, Bert De Brabandere, Tinne Tuytelaars, and Lue V Gool. Dynamic filter networks. In Proc. Advances in Neural Information Processing Systems (NIPS), pages 667-675, Barcelona, Spain, 2016.
[42] Irwan Bello. LambdaNetworks: Modeling long-range interactions without attention. In Int. Conf. on Learming Representations (ICLR), Virtual only, May 2021.
[43] Thomas Miconi, Kenneth Stanley, and Jeff Clune. Differentiable plasticity: training plastic neural networks with backpropagation. In Proc. Int. Conf. on Machine Learning (ICML), pages 3559-3568, Stockholm, Sweden, July 2018
[44] Thomas Miconi, Aditya Rawal, Jeff Clune, and Kenneth O. Stanley. Backpropamine: training self-modifying neural networks with differentiable neuromodulated plasticity. In Int. Conf. on Learning Representations (ICLR), New Orleans, LA, USA, May 2019.
[45] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural Turing Machines. arXiv:1410.5401, 2014.
Preprint
[46] Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwifiska, 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
[47] Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past. In Proc. Advances in Neural Information Processing Systems (NIPS), pages 4331-4339, Barcelona, Spain, December 2016.
[48] Imanol Schlag, Tsendsuren Munkhdalai, and Jürgen Schmidhuber. Learning associative infer- ence using fast weight memory. In Int. Conf. on Learning Representations (ICLR), Virtual only, May 2021.
[49] Tsendsuren Munkhdalai, Alessandro Sordoni, Tong Wang, and Adam Trischler. Metalearned neural memory. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 13310-13321, Vancouver, Canada, December 2019.
[50] Imanol Schlag and Jürgen Schmidhuber. Gated fast weights for on-the-fly neural program generation. In NIPS Metalearning Workshop, Long Beach, CA, USA, December 2017.
[51] T Anderson Keller, Sharath Nittur Sridhar, and Xin Wang. Fast weight long short-term memory. Preprint arXiv:1804.0651/ 2018
[52] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. Proc. European Conf. on Computer Vision (ECCV), pages 630-645, Amsterdam, Netherlands, October 2016.
13

---

[53] Rupesh K Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. In the Deep Learning workshop Int. Conf. on Machine Learning (ICML), Lille, France, July 2015.
[54] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. Preprint 2016.
[55] Gail Weiss, Yoav Goldberg, and Eran Yahav. On the practical computational power of finite precision mns for language recognition. Proc. Association for Computational Linguistics (ACL), pages 740-745, Melbourne, Australia, July 2018.
[56] Adam Paszke et al. PyTorch: An imperative style, high-performance deep learning library. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 8026-8037, Vancouver, Canada, December 2019.
[57] Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. Address- ing some limitations of Transformers with feedback memory. Preprint arXiv:2000.009402 2020.
[58] Ke Tran, Arianna Bisazza, and Christof Monz. The importance of being recurrent for modeling hierarchical structure. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP). pages 4731-4736, Brussels, Belgium, October 2018
[59] Kevin Lu, Aditya Grover, Pieter Abbeel, and Igor Mordatch. Pretrained transformers as universal computation engines. Preprint arXiv:2103.05247, 2021.
[60] Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Victor Valdés, Amir Sadik, et al. Deepmind lab. Preprint arXiv:1612.0380) 2016.
[61] Joel Z Leibo, Cyprien de Masson d' Autumme Daniel Zoran, David Amos, Charles Beattie, Keith Anderson, Antonio García Castañeda, Manuel Sanchez, Simon Green, Audrunas Gruslys, et al. Psychlab: psychology laboratory for deep reinforcement learning agents. Preprint arXiv: 18:1.8000000 1 16, 2018.
[62] Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the Atari human bench- mark. In Proc. Ini. Conf. on Machine Learning (ICML), pages 507-517, Virtual only, July 2020.
[63] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing Atari with deep reinforcement learning. In NIPS Deep Learning Workshop, Lake Tahoe, NV, USA, December 2013.
[64] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. Nature, 518(7540):529-53,, 2015.
[65] Lasse Espeholt, Hubert Soyer, Rémi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, Shane Legg, and Koray Kavukcuoglu. IMPALA: scalable distributed Re with importance weighted actor-learner architectures. In Proc. Ins. Conf. on Machine Learning (ICML), pages 1406-1415, Stockholm, Sweden, July 2018.
[66] Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforce- ment learning, In Proc. Int. Conf. on Machine Learning (ICML), pages 1928-1937 New York City, NY, USA, June 2016.
[67] Heinrich Küttler, Nantas Nardelli, Thibaut Lavril, Marco Selvatici, Viswanath Sivakumar, Tim Rocktäschel, and Edward Grefenstette. Torchbeast: A PyTorch platform for distributed RI. Preprint arXiv:1910.0355 2019.
[68] Alexander Mott, Daniel Zoran, Mike Chrzanowski, Daan Wierstra, and Danilo Jimenez Rezende. Towards interpretable reinforcement learning using attention augmented agents. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 12329-12333 12338, Vancouver, Canada, December 2019.
14

---

[69] Arun Nair, Praveen Srinivasan, Sam Blackwell, Cagdas Alcicek, Rory Fearon, Alessandro De Maria, Vedavyas Panneershelvam, Mustafa Suleyman, Charles Beattie, Stig Petersen, et al. Massively parallel methods for deep reinforcement learning. In Deep Learning Workshop, International Conference on Machine Learning (ICML., Lille, France, July 2015.
[70] Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering Atari, Go, Chess and Shogi by planning with a learned model. Nature, 588(7839)6004-099 2020.
[71] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch, Decision Transformer: Reinforcement learning via sequence modeling. Preprint 206x:2/066..1455 2021.
[72] Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as one big sequence modeling problem. Preprint arXiv:2106.02033 02039, 2021.
[73] Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards-just map them to actions. Preprint arXiv: 1912.02875, 2019.
[74] Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaskowski, and Jürgen Schmid- huber. Training agents using upside-down reinforcement learning. Preprint arXiv: 1912.02877, 2019.
[75] Stephen José Hanson. A stochastic version of the delta rule. Physica D: Nonlinear Phenomena, 42(1-3):265-272, 1990.
[76] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, 15(1):1929-1958, 2014.
[77] Noah Frazier-Logue and Stephen José Hanson. Dropout is special case of the stochastic delta rule: Faster and more accurate deep learning. Preprint 1808.0888800357 2018.
[78] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. Preprint arXiv:1412.6988, 2014.
[79] Felix A Gers, Jürgen Schmidhuber, and Fred Cummins. Learning to forget: Continual prediction with LSTM. Neural computation, 12(10))2451-2477 2000.
[80] Kazuki Irie, Albert Zeyer. Ralf Schlüter, and Hermann Ney. Language modeling with deep Transformers. In Proc. Interspeech, pages 3905-3909, Graz, Austria, September 2019.
[81] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The devil is in the detail: Simple tricks improve systematic generalization of transformers. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP), Punta Cana, Dominican Republic, November 2021.
[82] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The Neural Data Router: Adaptive control flow in Transformers improves systematic generalization. Preprint 110.0221 2021.
[83] Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5- RMSProp:
Divide the gradient by running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 4, 2012.
[84] Marlos C Machado, Marc G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
[85] Marcin Andrychowicz, Anton Raichuk, Piotr Stanezyk, Manu Orsini, Sertan Girgin, Raphaël Marinier, Léonard Hussenot, Matthieu Geist, Olivier Pietquin, Marcin Michalski. et al. What matters for on-policy deep actor-critic methods? A large scale study. In Int. Conf. on Learning Representations (ICLR), Virtual only, May 2021.
15


Reviews:
[{"review_id": "bDt3EhBz4B-", "reviewer": "Reviewer_7TUg", "summary": "The paper presents a new series of transformer architectures with novel linear attention mechanisms, all based on the formalism of fast weight programmers (FWPs). The fast weight programming perspective views attention as a \"slow\" network producing internal weight updates to a \"fast\" weight hyper-network, where an update to the fast weight matrix is predicted and accumulated at each time step in the input sequence. Therefore the fast weight network changes it's parameters during processing of a single sequence, while the slow network is updated using standard gradient descent on the loss function across sequences. Using this perspective of fast weights, the paper designs several novel attention replacements: (1) recurrent architectures for the fast weight network, (2) deep feedforward networks, and (3) a nested recurrent architecture, the Recurrent Delta Net (RDN), where both the fast and slow network include a recurrence. Results on language modeling, code execution and reinforcement learning domains reveal the benefits of approaching linearized attention under this formalism, with the RDN in particular showing advantages over even standard softmax-attention transformers. ", "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": "The paper is well-written and clear to follow. The formalism of fast weight programmers makes understanding the attention variants intuitive, and the design choices of the network architectures presented using this perspective follow naturally.  The experimental results seem to support the viability of the architectures, with similar or better results to both standard transformer and linear transformer baselines. Adding recurrence to the fast weight mechanism seems to be especially important on the code execution and list ops domains, which largely close the gap with the LSTM baseline.\n\nWhile the paper presents a formalism from which a few novel architectures are designed, there is a downside that the results do not always provide a clear advantage to one particular architecture over another. For example, Delta LSTM does best out of the presented architectures at language modeling, but is twice as slow as a regular transformer and still underperforms it at test perplexity. As another example, at code execution it is either the RDN or the Delta RNN which are best performing. This somewhat limits the significance of the results, since a practitioner can't immediately choose one architecture and might need to ablate over many of them. It could make the narrative of the experiments section cleaner if the authors chose to focus on a particular architecture, and demonstrated a clearer performance improvement under some chosen metrics (wall-clock training time, inference time, etc.).\n\nIt would also be interesting to see how these models perform at larger scale. I would be very interested to see the scaling properties of the recurrent attention architectures, although I understand that running these experiments could be expensive. \n\nAdditionally, is there an intuition on why any of the memory architectures on Atari would be particularly helpful? The results seem to support that, but a more thorough understanding of why exactly would be informative. Is this caused by increased parameter count, ease-of-optimization, or is the memory actually being used effectively in the policy?\n\nIn conclusion, I believe the paper has presented an interesting perspective on recurrent attention and, despite some issues with the focus of the results, I vote for acceptance. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have discussed the limitations of their model. The model also does not seem to have any significant potential for negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "aWjlQvZHIec", "reviewer": "Reviewer_x5eT", "summary": "The paper extends the previous line of work on interpreting linear transformer as fast weight programmers (FWP), and introduce recurrent variants of FWP. ", "questions": "", "limitations": "", "rating": 5, "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": "Originality:\n- The idea is not very novel in my opinion. It is a straightforward extension of the previous work, by changing either the slow network or the fast network from a feedforward neural network to a recurrent neural network.\n\nClarity:\n- The paper is written very nicely. When introducing the background material, it provides a unifying view of the subject. It helps me understand FWP much better and also the connection between FWP and linear transformer.\n\n\nQuality:\n- As the paper introduces a simple variant of the existing work, I would mainly hope that the empirical results to be strong enough to deserve a publication. However, the current results do not convince me at all. Mainly,\n   1. The results are not very strong. E.g., in Table 2, LSTM seems to beat all the variants proposed. \n   2. The comparisons of existing works are missing. In the wikiText dataset, it missed the results from [3], which attained better performance than any of the reported variants. For CodeExec task, it missed the results from [4], which also did better than the reported results. For atari, it missed the baselines such as [1] or [2]. Note that these references are already cited in the paper (some tasks are taken from those). So it seems the author knows these results but intentionally choose to not include these as baselines? \n- Over the past several years, many of the impressive results on realistic benchmarks were demonstrated without the use of recurrence. Hence this raises a question -- when is recurrence really needed? Part of the reason we prefer non-recurrence architecture is because of its easiness for scalability, and fast training speed, as one can parallelize computation across time. I would like to see a the necessity of recurrence on some practically relevant task.\n\nQuestion:\n- In the paper line 235-247 the authors discussed the computational efficiency of this model. But I am a bit confused by the measurement there. The computational complexity of vanilla transformer computes the embeddings for each token in the layer in one big matrix multiplication which can be highly parallelized in modern accelerators, but the recurrence would not allow computational complexity any lower than O(N), where N is the sequence length. But somehow the number of tokens per second is roughly the same. I wonder what's the reason here?\n  \n\nSignificance:\n- The work is a straightforward extension of the previous work. The results are not very convincing. Hence I wouldn't think the work is significant at its current stage.\n\nSummary:\n- I would encourage the authors to improve the empirical results to make it more convincing. In particular, I think the author should demonstrate on a practically relevant task that recurrence is necessary. \n\n\n[1] Mott et.al.: Towards interpretable reinforcement learning using attention augmented agents.\n\n[2] Parisotto et.al.: Stabilizing Transformers for reinforcement learning.\n\n[3] Peng et. al.: Random feature attention.\n\n[4] Fan et.al: Addressing some limitations of Transformers with feedback memory", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "NA", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "RvTjehwf0m", "reviewer": "Reviewer_piWh", "summary": "This paper studies linear attention variants within transformers from the perspective of fast weight programmers from the 90s.  In particular, it explores new more sophisticated variants that have recurrence on both the slow (fixed after training) and fast networks (generated) as opposed to previous variants such as Delta Net that uses one-layer feedforward networks for both.  The evaluation shows that recurrence leads to improved results compared to both standard and linear transformers.", "questions": "", "limitations": "", "rating": 7, "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": "The exploration of more sophisticated fast weight programmers for linear attention is an interesting and valuable direction since it enhances the understanding of the linear attention variants. Even though the connection to fast weight programmers itself is not new the proposed variants have sufficient novelty and increase the expressiveness in maintaining the internal state when used for attending to long sequences. \n\nThe technical exposition was quite clear and well-positioned with respect to prior work. The motivation for using recurrence becomes apparent given the simplicity of the networks used in previous variants. The study reveals that using a Delta Net with simple recurrence leads to best results and more sophisticated recurrence based on LSTMs removes the efficiency benefits in linear transformers. One interesting insight from this finding is that it is not possible to increase the expressiveness of linear attention beyond a certain point without being slower compared to the original transformer; albeit the model would still be scalable while the original transformer is not. \n \nThe evaluation is thorough, contains controlled comparisons, and uses a configuration similar configurations to prior work. The results support the claims in the paper and shed light on the strengths of different variants. I also found the details provided in the supplementary quite helpful. Perhaps one limitation is that the evaluation does not include any non-autoregressive tasks. \n\nQuestions:\n- Is the difference in LSTM performance on the code execution task from prior work due to the data creation process or some other factor?  \n \n- Can the proposed attention variants be applied to non-autoregressive tasks? Recent efficient variants are applicable to such settings, including linear transformers, Performer/RFA.   \n\n- The statement that fast weight variants maintain the standard one-dimensional state and two-dimensional fast weight states in Sec 3.2 was a bit confusing. Where do these dimension sizes come from? \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper discusses the limitations of the proposed models and explores their speed-quality tradeoffs in the evaluation section. I would suggest that the authors discuss the applicability of the proposed variants on non-autoregressive tasks and whether the proposed methods can handle them (and why not if they can't) so that the scope of this paper is clear to the readers. There is also a section dedicated to the potential negative societal impact. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5"}, "extra_text": {}}, {"review_id": "Lw3aJLbEWp0", "reviewer": "Reviewer_u3vG", "summary": "This work proposes modifications to the linear transformer paradigm that adds recurrent dependencies to past transformer embeddings. The resulting \"Recurrent Fast Weight Programmer\" model proposed is shown to outperform transformers on multiple sequence learning tasks. Furthermore, the proposed method remains competitive with RNNs in tasks where RNN architectures are still dominant over transformers, implying a best of both worlds feature in various sequence learning settings. ", "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": "The text presents several potential variations of the linear transformer formulation which incorporate various properties such as recurrence and differing types of interactions between fast and slow weights. But the motivation for such architectural modifications is unclear from the presentation of the paper. Why is the DeltaNet formulation necessary as compared to the Linear Transformer? What aspect of the DeltaNet formulation improves model capacity?\n\nOverall, the connections between linear transformers and Fast Weight Programmers are interesting because it may give rise to explanations of where the benefits and limitations exist among both styles of architectures. Furthermore, it may help explain what deviations could be expected in performance from the original/dominant Transformer. These differences are seen in the experiments of this work. But this work has little to no hypotheses and therefore no subsequent analysis of how these modifications improve/don't improve performance. For instance, how is the \\beta gating function behaving when given various difficulty levels of the code execution and sequence list tasks?\n\nPost Rebuttal: I have gone over the author responses and discussions from other reviewers and have decided to increase my score. My main concern is still the lack of exploration into the function of the additional modifications proposed in the RDN model. Without knowing what is systematically different between the proposed model after learning a task and other models, a never-ending paper stream of variations of linear transformers that do not necessarily perform any better than the original transformer formulation in the real world can result. But I agree with the author's responses that this may be too much to ask for in a paper submission.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Suggestions for improvement:\n\n-Introduction and specifically lines 36-45 begin discussion of the architectural modifications that are hard to understand without any formulation. (lines 42-43 seem contradictory)\n\n-Explain what are the scaling differences between Delta Delta Net and the size of the backpropagation graph if another meta-level of slow to fast weights are added?\n\n-Lines 186-190 need a better explanation for the formulations of the slow weights. More capacity and additional memory are too vague. Why formulate the fast weights with 1-D states and weights as 2-D states?\n\n-Since some of these models are not limited by the token span as much as vanilla transformers, why not try the models that can support it much longer sequence spans?\n\n-Lines 216-218 like earlier lines are too vague, an intuition for why the slow network requires additional input (and also untested). It seems like the lack of strong reasoning behind these architecture modifications makes conclusions drawn from the results ambiguous.\n\n-The surprisingly good performance of the LSTM on the code execution task is also not explained or analyzed. The lack of analysis here is unfortunate because this deviates from what is expected and is worth understanding.\n\n-The falloff in LSTM performance at higher depth (depth 15) is also an interesting observation that should be expanded upon.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper uses the connection between linear attention transforms and fast weight programmers to introduce a novel transformer architecture, recurrent fast weight programmers, which includes recurrence in the fast weights. 

The authors demonstrate competitive performance on three distinct benchmarks and provide interesting conceptual insights into the functioning of their new architecture.



The discussion focussed on the role of the recurrence (“Why is it needed?”) and novel baselines / ablations. In particular, the authors provided new experimental evidence showing that the recurrence indeed helps, even when compared to a feedforward NN with the same number of parameters. 

The clarifications and new evidence provided made one reviewer increase their score, while all other reviews remained unchanged. 



Given the positive reviews, the discussion and the overall contributions of the paper, I recommend this paper to be accepted.

Decision: Accept (Poster)