Addressing Some Limitations of Transformers with Feedback Memory

1 INTRODUCTION . In recent years , the Transformer architecture ( Vaswani et al. , 2017 ) has brought large improvements to a wide range of Natural Language Processing tasks such as machine translation , sentence representation ( Devlin et al. , 2019 ) , and summarization ( Edunov et al. , 2019 ) . Transformers are also successfully used as an autoregressive model on sequential tasks such as language modeling ( Dai et al. , 2019 ; Rae et al. , 2020 ) and reinforcement learning ( Parisotto et al. , 2019 ) . Unlike more traditional recurrent architectures such as RNNs and LSTMs , the Transformer architecture processes a sequence in parallel in an order-invariant way . Techniques such as position embeddings ( Sukhbaatar et al. , 2015 ; Shaw et al. , 2018 ) and attention masking are required to capture input order information . In this work , we focus on several limitations of the Transformer architecture as an autoregressive model and present a straightforward solution — Feedback memory . These limitations and our proposed solution target sequential token prediction tasks , such as language modeling or other auto-regressive generative tasks . The feedforward nature of Transformers makes them efficient on modern hardware , but restricts the Transformer from taking full advantage of the input ’ s sequential property . In particular , the current hidden representation of a Transformer only accesses the past representations of lower layers , even though higher level representations of the past have already been computed as an autoregressive model . At generation , the Transformer generates only one token at a time , so it could access these representations for better performance , but does not exploit these at training time due to parallelization . However , if these past higher level representations could be used at training time , they would enrich future lower level representations , enabling shallower models to have the same representation power . Another inherent limitation of Transformers on sequential tasks is the lack of recursive computation ( Dehghani et al. , 2018 ) , and the number of transformations possible on the input is bounded by the model depth . Such disadvantages have impact on tasks that require careful tracking of a world state or modeling hierarchical structures ( Tran et al. , 2018 ; Hahn , 2020 ) . On the other hand , while RNNs can maintain an internal state for an unbounded time while accumulating more computations upon it , the size of this internal state is limited by the dimension of the hidden state . In this work , we propose a novel autoregressive model , the Feedback Transformer , that makes all previous hidden representations accessible to the computation of a representation at any depth — the model feeds back previous computations to itself . The feedback allows the model to perform recursive computation , building stronger representations iteratively upon previous states . To achieve this , we modify self-attention to attend to higher level representations rather than lower ones . As shown in Figure 1 , the Feedback Transformer merges the hidden states from all layers into a single vector for every time step and stores them in a memory . Instead of self-attention , all subsequent layers attend to this memory , which means every previously computed representation is accessible by all future layers , mediated by the memory . This allows Feedback Transformers to recursively compute and transform an input as many times as the input length , which is something Transformers can not achieve . While RNNs can perform recursive computation , the amount of information that Feedback Transformers can maintain is not limited by the number of layers . There are computational benefits to this straightforward modification . First , it uses less memory because all the layers share a single Feedback memory , thus reducing the memory size by L times , where L is the number of layers . There is also less computation because we share the key and value projections during attention computation , which increases the speed of the attention over the Feedback Memory . Further , the GPU memory usage is reduced due to the memory sharing — the overall model is 2x smaller — allowing the batch size to be increased for computational efficiency . During inference , the increased batch size contributes to substantially faster decoding speeds . In summary , our main contributions are : ( 1 ) The Feedback Transformer architecture , which completely changes the way a Transformer works to access available higher level representations immediately . ( 2 ) We show the Feedback Transformer can achieve state of the art results with smaller , shallower models that have faster decoding speed and smaller memory footprint . ( 3 ) The Feedback Transformer uses substantially less memory during training and inference time . 2 RELATED WORK . Several previous works have analyzed the limitations of Transformer architectures , such as the inability to process input sequentially ( Dehghani et al. , 2018 ) or represent hierarchical structure ( Tran et al. , 2018 ) . Hahn ( 2020 ) demonstrate that Transformers can not model structures involving bounded recursion , such as closing parentheses . Pérez et al . ( 2019 ) study Transformers in the context of Turing machines , where they must produce unbounded numbers of decoding steps . Various work in probing Transformers identified several limitations where Transformers may not have the computational capacity of recurrent architecture like an LSTM ( Hahn , 2020 ) . From the architectural perspective , our work shares similarities with recurrent networks augmented with external shared memories ( Graves et al. , 2014 ; Joulin & Mikolov , 2015 ; Sukhbaatar et al. , 2015 ) . For example , the stack augmented RNN of Joulin & Mikolov ( 2015 ) adds an external memory to a recurrent network to keep long term dependencies . Closer to our work , the Neural Turing Machine of Graves et al . ( 2014 ) models an unconstrained memory that resembles the self-attention layer of a Transformer . Further improvements to recurrent networks , such as the Gated Feedback RNN ( Chung et al. , 2015 ) , are based on better controlling signal from different layers and extended to feedback through multiple pathways ( Jin et al. , 2017 ) . These works are built on recurrent networks with additional components to store long term dependencies . Other works have studied modifications to the Transformer architecture by enriching its structure with components inspired by recurrent networks . For example , Wang et al . ( 2019 ) propose adding a local recurrent sublayer to the Transformer layer to remove the need for position embeddings in the multi-head self-attention layers . Universal Transformer ( Dehghani et al. , 2018 ) share the parameters between the layers of a Transformer , leading a recurrent network in depth . Hao et al . ( 2019 ) and Chen et al . ( 2018 ) augment Transformers with a second , recurrent encoder . As opposed to our work , these prior investigations do not change the computational path in a Transformer to reduce the discrepancy between the training and inference time . Closer to our work , Merity ( 2019 ) proposes adding a self-attention layer on top of the past outputs from an LSTM cell . However , this approach keeps the recurrent and the self-attention mechanisms decoupled , as opposed to ours which makes the attention mechanism recurrent . In particular , the LSTM layer of Merity ( 2019 ) still intrinsically has a bottleneck corresponding to the dimension of the hidden layer . 3 METHOD . In this section , we propose the Feedback Transformer , which provides capacity to build richer representations of each timestep t of a sequential modeling task . 3.1 TRANSFORMER ARCHITECTURES . We briefly describe the Transformer ( Vaswani et al. , 2017 ) . Each layer is composed of a multihead self-attention sublayer ( Attn ) followed by a feedforward sublayer ( FF ) , and each sublayer is followed by an add-norm operation that combines a skip-connection ( He et al. , 2016 ) and layer normalization ( Lei Ba et al. , 2016 ) . The l-th layer of a Transformer processes an input sequence of vectors Xl = ( xl1 , . . . , x l t ) into a sequence of vectors of the same length . First , the self-attention sublayer computes a representation for each time step t by taking its related input vector xt along with its past context , { xlt−τ , ... , xlt−1 } : zlt = Attn ( x l t , { xlt−τ , . . . , xlt−1 } ) . Within the self-attention sublayer , xlt is used to form query vectors while its context is used to compute key and value vectors , forming a memory of the past information . Then the feedforward sublayer processes each vector zlt independently , i.e. , x l+1 t = FF ( z l t ) . The Transformer layer transforms its input sequence into an output sequence Xl+1 = FF ( Attn ( Xl ) ) . In practice , a block of steps { xlt−M+1 , . . . , xlt } is computed in parallel during training , where M can be seen as the backpropagation through time ( BPTT ) length . This makes training Transformers efficient on hardware such as GPUs . However , to operate on sequences of unbounded length , Transformers require modifications such as caching and relative position embeddings ( Dai et al. , 2019 ; Sukhbaatar et al. , 2019 ) . 3.2 LIMITATIONS OF TRANSFORMERS . Previous work has analyzed the impact of several limitations of the Transformer architecture , such as the inability to track long sequences and process hierarchical inputs ( Hahn , 2020 ) . In this work , we focus on two major limitations of Transformer architectures . Limited Access to Higher Level Representations . Layer by layer , Transformers build more abstract , high level representations of the input sequence . At each layer , the representations for the input sequence are treated in parallel . As a consequence , a Transformer does not leverage the highest level representations from the past to compute the current representation , even though these highest level representations have already been computed for autoregressive models . Maintaining a Belief State . Many sequential tasks require models to maintain an internal state for two main purposes . First , internal states act as memory for recalling past inputs , where Transformers excel because their internal state xlt is directly accessible to future steps through self-attention . The second role of an internal state is to act as a belief state that tracks the world state that is not directly observable in inputs . For example , when inputs are actions taken on a Markov Decision Process , an internal state can apply those changes to the current belief state and correctly predict the outcome . As a feedforward model , Transformer have inherent limitations in this area — only a fixed number of transformations can be applied to its internal states . Since both Attn and FF sublayers contain a fixed number of transformations and there are L layers of them , the total number of transformations between the input and output is limited by the depth . This means Transformers can not maintain an internal state for long time if it has to be frequently updated . 3.3 FEEDBACK TRANSFORMER . We propose to change the Transformer architecture by using the most abstract representations from the past directly as inputs for the current timestep . This means that the model does not form its representation in parallel , but sequentially token by token . More precisely , we replace the context inputs to attention modules with memory vectors that are computed over the past , i.e. , zlt = Attn ( x l t , { mt−τ , . . . , mt−1 } ) , where memory vectors mt are computed by summing the representations of all layers at time step t : mt = L∑ l=0 Softmax ( wl ) xlt , ( 1 ) where wl are learnable scalar parameters . Note these scalars are the only new parameters introduced by our change , with all else the same as the standard Transformer . Here l = 0 corresponds to token embeddings . The weighting of different layers by a softmax output gives the model more flexibility as it can average them or select one of them . This modification of the self-attention input adapts the computation of the Transformer from parallel to sequential , summarized in Figure 2 . Indeed , it provides the ability to formulate the representation xlt+1 based on past representations from any layer l ′ , while in a standard Transformer this is only true for l′ < l. This change can be viewed as exposing all previous computations to all future computations , providing better representations of the input . Such capacity would allow much shallower models to capture the same level of abstraction as a deeper architecture . This has several practical advantages , as more shallow models have reduced memory footprint and increased decoding speed . An alternative view of such an architecture modification is providing the capacity for recursive computation — outputs from a sublayer can feed back to the same sublayer through the memory . The model can then maintain an internal state for unbounded time . This is a clear advantage over Transformers , in which a submodule never looks at its own output . While an RNN can also repeat its computation on its internal state , its internal state has a limited capacity determined by the number of layers and their hidden dimension . In contrast , the internal state of a Feedback Transformer is its whole memory , which can grow with the input length . This allows the model to keep track of a large number of things within its internal state . While our modification requires sequential computation , we significantly improve training speed by sharing the key and value projections W lk and W l v across all layers . This sharing reduces computation because we need to compute key and value vectors only once instead of computing them per layer klt = kt = Wkmt v l t = vt = Wvmt . For the same reason , the memory footprint is smaller than a standard Transformer because only one set of kt , vt needs to be stored . To be more precise , the memory requirement for processing a single token is reduced from O ( L× T ) to O ( T ) , where L is the number of layers and T is the context size . Further , the reduced memory usage allows the batch size to be increased to recover some of the lost parallelism , which improves training speed . Thus , the Feedback Transformer is not much slower compared to the standard Transformer . Note that the same sharing of projections will not make the standard Transformer efficient because those projections are applied to different representations at each layer ( the key and value vectors will not the same for all layers ) . Lastly , we note that the sequential nature of the Feedback Transformer does not affect the performance during generation where one needs to compute one step at a time anyway . The same is true for online reinforcement learning where the input must be processed sequentially even during training . Task / Model Accuracy ( % ) Copy Char Seq Transformer 59.1 6.2 Feedback Transformer 76.2 23.6 Reverse Char Seq Transformer 50.2 5.9 Feedback Transformer 74.8 29.2 Counting Len 50 Len 1K Transformer 99.6 82.4 Feedback Transformer 99.7 95.3 Random Walk Transformer 68 Feedback Transformer 100 Algorithmic Task 3 vars 5 vars Transformer 4L 33.7 37.5 Transformer 8L 47.4 29.1 LSTM 82.8 32.1 Feedback Trans . 4L 99.1 92.6 Table 1 : Results on toy tasks . Char is character accuracy , Seq is sequence accuracy . 20 40 60 80 100 Memory size 0.5 0.6 0.7 0.8 0.9 S u cc es s ra te Transformer Feedback Transformer Figure 3 : Results on the Corridor task . The Transformer degrades as the memory size decreases , but the Feedback Transformer maintains performance .

Summary: This paper modifies transformers with feedback memory. Specifically, for each timestep, it merges hidden representations of all layers into a high-level single vector and stores it in memory. For the current timestep, it attends past memory vectors. The authors claim that in this way, low layers of the current timestep can utilize high-level representations of past timesteps. The authors show that the proposed models with shallow layers can achieve stronger performance than comparable transformers. However, it seems that the models need a much longer time to train.