Going Beyond Linear Transformers with Recurrent Fast Weight Programmers

1 Introduction . The Transformer [ 1 ] has become one 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 , 3 ] or question answering [ 4 ] . Recently , it has also been applied in other domains , including image processing [ 5 , 6 ] or mathematical problem solving [ 7 , 8 , 9 ] . Conceptually , the Transformer is a deep feedforward NN that processes all elements of a sequence in parallel : unlike in recurrent NNs ( RNNs ) , the computations of a layer for the entire sequence can be packed into one big matrix multiplication . This scales well with the number of parallel processors . Despite the benefits of parallelisation , a major drawback of Transformers is that their computational complexity in time and space is quadratic in sequence length . Furthermore , in the auto-regressive version [ 1 , 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 infinite sequences , forcing practitioners to truncate temporal contexts and ignore long-term dependencies beyond fixed-size time windows . Although recent work tries to address this issue [ 10 , 11 ] , this limitation makes some applications of Transformers challenging , e.g. , reinforcement learning ( RL ) in partially observable environments [ 12 , 13 ] , which is still dominated by RNNs such as the Long Short-Term Memory ( LSTM ; [ 14 ] ) trained by policy gradients [ 15 , 16 , 17 , 18 ] . 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 [ 19 ] . Such models include ∗Equal contribution . 1https : //github.com/IDSIA/recurrent-fwp 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Katharopoulos et al. ’ s Linear Transformer ( LT ) [ 19 ] , Choromanski et al. ’ s Performer [ 20 ] and Peng et al . [ 21 ] ’ s variant . They enjoy time and space complexities linear in sequence length with states of constant size . While their performance on some tasks does not fully match the one of regular Transformers [ 22 ] , several improvements have already been proposed [ 21 , 23 ] ( see our review in Sec . 2.2 ) which makes this Transformer family a 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 productbased 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 a two-NN system : a slow and a fast net , each with arbitrary architectures . The slow net learns to 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 to 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 [ 29 ] 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 [ 31 ] . Although LSTM still works better in a few environments , we show that our RFWPs generally improve by scaling them up . The main contribution of this work is twofold : ( 1 ) from the perspective of FWPs , we study novel powerful FWPs for sequence processing , demonstrating that NNs can easily learn to control NNs that are more complex than a single feedforward layer , and ( 2 ) from the perspective of Transformer models , our RFWPs augment linear Transformers with recurrence , addressing general limitations of existing auto-regressive Transformer models . 2 Background on Fast Weight Programmers ( FWPs ) . Here we 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 a sequence processing scenario . An FWP with trainable parameters θslow sequentially transforms an input sequence { x ( t ) } Tt=1 with x ( t ) ∈ Rdin to an output sequence { y ( t ) } Tt=1 with y ( t ) ∈ Rdout of length T as θ ( t ) fast , q ( t ) = SlowNet ( { x ( j ) } tj=1 , { y ( j ) } t−1j=0 , { θ ( j ) fast } t−1 j=0 , { q ( j ) } t−1j=0 ; θslow ) ( 1 ) y ( t ) = FastNet ( { q ( j ) } tj=1 , { y ( j ) } t−1j=0 ; θ ( t ) fast ) ( 2 ) where y ( 0 ) , θ ( 0 ) fast , and q ( 0 ) are initial variables . This is a system with two NNs called FastNet and SlowNet in which the parameters θ ( t ) fast of FastNet are generated by SlowNet at each time step t. 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 θslow are slow because they can only change through gradient descent during training , remaining fixed afterwards2 . Eq . 1 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 θ ( t ) fast , the slow net also generates or invents an input q ( t ) to be fed to the fast net ( alternatively q ( t ) can simply be x ( t ) ) . 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 2The fast net could also contain some additional slow weights ; we 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 is a neural programmer of fast weights , and its parameter set θslow 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 θ ( t ) fast incrementally in an iterative fashion , where the SlowNet is further decomposed into two sub-parts : z ( t ) , q ( t ) = SlowSubnet ( { x ( j ) } tj=1 , { y ( j ) } t−1j=0 , { θ ( j ) fast } t−1 j=0 , { q ( j ) } t−1j=0 , { z ( j ) } t−1j=0 ; θslow ) ( 3 ) θ ( t ) fast = UpdateRule ( θ ( t−1 ) fast , z ( t ) ) ( 4 ) where UpdateRule takes the fast weights θ ( t−1 ) fast from the previous iteration to produce the new fast weights θ ( t ) fast conditioned on z ( t ) . The update rule is 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 . In general , the dimension of the fast weights θ ( t ) fast is too large to be conveniently parameterised by an NN . Instead , it was 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 FWPs : linear Transformers [ 19 ] and the Delta Net [ 23 ] . Linear Transformer . The “ linear Transformer ” [ 19 ] is a class of Transformers where the softmax in the attention is linearised . This is achieved by replacing the softmax with a kernel function φ—then the self-attention can be rewritten as a basic outer product-based 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 φ : Rdkey → Rdkey , simplifying our equations below by writing k , q instead of φ ( k ) , φ ( q ) . Using otherwise the same notation as above , for each new input x ( t ) , the output y ( t ) is obtained by : k ( t ) , v ( t ) , q ( t ) = Wkx ( t ) , Wvx ( t ) , Wqx ( t ) ( 5 ) W ( t ) = W ( t−1 ) + v ( t ) ⊗ k ( t ) ( 6 ) y ( t ) = W ( t ) q ( t ) ( 7 ) where the slow weight matrices Wk ∈ Rdkey×din and Wv ∈ Rdout×din are used to obtain the key k ( t ) ∈ Rdkey and the value v ( t ) ∈ Rdout . The key and value vectors are used to generate new weights via the outer product v ( t ) ⊗ k ( t ) ∈ Rdout×dkey . 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 matrixW ( t−1 ) ∈ Rdout×dkey is updated to yieldW ( t ) by adding the update term v ( t ) ⊗ k ( t ) . This corresponds to 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 q ( t ) ∈ Rdkey generated by the slow weights Wq ∈ Rdkey×din . Hence , in linear Transformers , the previous Eq . 3 simplifies to : z ( t ) , q ( t ) = SlowSubnet ( x ( t ) ; θslow ) with z ( t ) = ( k ( t ) , v ( t ) ) . Delta Net . The Delta Net [ 23 ] is obtained by replacing the purely additive programming instruction ( Eq . 6 ) in the linear Transformer with the one akin to the delta rule [ 33 ] : W ( t ) = W ( t−1 ) + β ( t ) ( v ( t ) − v̄ ( t ) ) ⊗ k ( t ) ( 8 ) where β ( t ) ∈ R is a fast parameter ( learning rate ) of the update rule generated by the slow net with weightsWβ ∈ R1×din and the sigmoid function σ : β ( t ) = σ ( Wβx ( t ) ) ( 9 ) and v̄ ( t ) ∈ Rdout is generated as a function of the previous fast weightsW ( t−1 ) and the key k ( t ) v̄ ( t ) = W ( t−1 ) k ( t ) . ( 10 ) This update rule was introduced to address a memory capacity problem affecting linear Transformers with the purely additive update rule [ 23 ] . The corresponding Eq . 3 is : z ( t ) , q ( t ) = SlowSubnet ( x ( t ) , W ( t−1 ) ; θslow ) with z ( t ) = ( k ( t ) , v ( t ) , β ( t ) , v̄ ( t ) ) . Thus , unlike linear Transformers , the SlowNet in the Delta Net takes the previous fast weightsW ( t−1 ) 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 k ( t ) , v ( t ) , q ( t ) are split into equally sized H 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 a new weight matrix , Hypernetworks [ 34 ] scale the rows of a slow weight matrix with a generated vector of appropriate size . Weight compression to control fast weights in a low dimensional compressed space has been also studied [ 35 ] . In the broad sense of contextdependent weights [ 36 , 37 , 38 ] , many concepts relate to FWPs : e.g . dynamic convolution [ 39 , 40 , 41 ] , LambdaNetworks [ 42 ] , or dynamic plasticity [ 43 , 44 ] .

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.