Compressive Transformers for Long-Range Sequence Modelling

1 INTRODUCTION . Humans have a remarkable ability to remember information over long time horizons . When reading a book , we build up a compressed representation of the past narrative , such as the characters and events that have built up the story so far . We can do this even if they are separated by thousands of words from the current text , or long stretches of time between readings . During daily life , we make use of memories at varying time-scales : from locating the car keys , placed in the morning , to recalling the name of an old friend from decades ago . These feats of memorisation are not achieved by storing every sensory glimpse throughout one ’ s lifetime , but via lossy compression . We aggressively select , filter , or integrate input stimuli based on factors of surprise , perceived danger , or repetition — amongst other signals ( Richards and Frankland , 2017 ) . Memory systems in artificial neural networks began with very compact representations of the past . Recurrent neural networks ( RNNs , Rumelhart et al . ( 1986 ) ) learn to represent the history of observations in a compressed state vector . The state is compressed because it uses far less space than the history of observations — the model only preserving information that is pertinent to the optimization of the loss . The LSTM ( Hochreiter and Schmidhuber , 1997 ) is perhaps the most ubiquitous RNN variant ; it uses learned gates on its state vector to determine what information is stored or forgotten from memory . However since the LSTM , there has been great benefit discovered in not bottlenecking all historical information in the state , but instead in keeping past activations around in an external memory and attending to them . The Transformer ( Vaswani et al. , 2017 ) is a sequence model which stores the hidden activation of every time-step , and integrates this information using an attention operator ( Bahdanau et al. , 2014 ) . The Transformer will thus represent the past with a tensor ( depth × memory size × dimension ) of past observations that is , in practice , an order of magnitude larger than an LSTM ’ s hidden state . With this granular memory , the Transformer has brought about a step-change in state-of-the-art performance , within machine translation ( Vaswani et al. , 2017 ) , language modelling ( Dai et al. , 2019 ; Shoeybi et al. , 2019 ) , video captioning ( Zhou et al. , 2018 ) , and a multitude of language understanding benchmarks ( Devlin et al. , 2018 ; Yang et al. , 2019 ) amongst others . One drawback in storing everything is the computational cost of attending to every time-step and the storage cost of preserving this large memory . Several works have focused on reducing the computational cost of attention with sparse access mechanisms ( Rae et al. , 2016 ; Child et al. , 2019 ; ∗Authors contributed equally , † DeepMind , London , UK . ‡ CoMPLEX , Computer Science , University College London , UK . Please direct correspondence to { jwrae , apotapenko } @ google.com . Sukhbaatar et al. , 2019 ; Lample et al. , 2019 ) . However sparse attention does not solve the storage problem , and often requires custom sparse kernels for efficient implementation . Instead we look back to the notion of compactly representing the past . We show this can be built with simple dense linear-algebra components , such as convolutions , and can reduce both the space and compute cost of our models . We propose the Compressive Transformer , a simple extension to the Transformer which maps past hidden activations ( memories ) to a smaller set of compressed representations ( compressed memories ) . The Compressive Transformer uses the same attention mechanism over its set of memories and compressed memories , learning to query both its short-term granular memory and longer-term coarse memory . We observe this improves the modelling of text , achieving state-of-the-art results in character-based language modelling — 0.97 bpc on Enwik8 from the Hutter Prize ( Hutter , 2012 ) — and word-level language modelling — 17.1 perplexity on WikiText-103 ( Merity et al. , 2016 ) . Specifically , we see the Compressive Transformer improves the modelling of rare words . We show the Compressive Transformer works not only for language , but can also model the waveform of high-frequency speech with a trend of lower likelihood than the TransformerXL and Wavenet ( Oord et al. , 2016 ) when trained over 400,000 steps . We also show the Compressive Transformer can be used as a memory component within an RL agent , IMPALA ( Espeholt et al. , 2018 ) , and can successfully compress and make use of past observations . Furthermore we present a new book-level language-modelling benchmark PG-19 , extracted from texts in Project Gutenberg1 , to further promote the direction of long-context sequence modelling . This is over double the size of existing LM benchmarks and contains text with much longer contexts . 2 RELATED WORK . There have been a variety of recent attempts to extend the range of attention , particularly in the Transformer , or to replace the attention operation with something less expensive . Wu et al . ( 2019 ) show that a convolution-like operator that runs in linear time can actually exceed the performance of the quadratic-time self-attention layer in the Transformer at sentence-to-sentence translation and sentence-level language modelling . However such a mechanism inhibits the flow of information across a large number of time-steps for a given layer , and has not shown to be beneficial for longrange sequence modelling . Dai et al . ( 2019 ) propose the TransformerXL , which keeps past activations around in memory . They also propose a novel relative positional embedding scheme which they see outperforms the Transformer ’ s original absolute positional system . Our model incorporates both of these ideas , the use of a memory to preserve prior activations and their relative positional embedding scheme . The Sparse Transformer ( Child et al. , 2019 ) uses fixed sparse attention masks to attend to roughly√ n locations in memory . This approach still requires keeping all memories around during training , however with careful re-materialization of activations and custom kernels , the authors are able to train the model with a reasonable budget of memory and compute . When run on Enwik8 , the much larger attention window of 8 , 000 improves model performance , but overall it does not significantly outperform a simpler TransformerXL with a much smaller attention window . The use of dynamic attention spans is explored in Sukhbaatar et al . ( 2019 ) . Different attention heads can learn to have shorter or longer spans of attention — and they observe this achieves state-ofthe-art in character-based language modelling . This idea could easily be combined with our contribution — a compressive memory . However an efficient implementation is not possible on current dense-linear-algebra accelerators , such as Google ’ s TPUs , due to the need for dynamic and sparse computation . Our approach builds on simple dense linear algebra components , such as convolutions . 3 MODEL . We present the Compressive Transformer , a long-range sequence model which compacts past activations into a compressed memory2 . The Compressive Transformer is a variant of the Transformer 1Project Gutenberg : https : //www.gutenberg.org/ 2A TF implementation can be found in Sonnet : https : //github.com/deepmind/sonnet ( Vaswani et al. , 2017 ) , a deep residual network which only uses attention to propagate information over time ( namely multi-head attention ) . We build on the ideas of the TransformerXL ( Dai et al. , 2019 ) which maintains a memory of past activations at each layer to preserve a longer history of context . The TransformerXL discards past activations when they become sufficiently old ( controlled by the size of the memory ) . The key principle of the Compressive Transformer is to compress these old memories , instead of discarding them , and store them in an additional compressed memory . 3.1 DESCRIPTION . We define nm and ncm to be the number of respective memory and compressive memory slots in the model per layer . The overall input sequence S = x1 , x2 , . . . , x|s| represents input observations ( e.g . tokens from a book ) . These are split into fixed-size windows of size ns for the model to process in parallel . The model observes x = xt , . . . , xt+ns at time t , which we refer to as the sequence ( e.g . in Figure 1 ) . As the model moves to the next sequence , its ns hidden activations are pushed into a fixed-sized FIFO memory ( like the TransformerXL ) of size nm . The oldest ns activations in memory are evicted , but unlike the TransformerXL we do not discard them . Instead we apply a compression operation , fc : Rns×d → Rb ns c c×d , mapping the ns oldest memories to bnsc c compressed memories which we then store in a secondary FIFO compressed memory of size ncm . d denotes the hidden size of activations and c refers to the compression rate , a higher value indicates more coarse-grained compressed memories . The overall temporal range of the model becomes l× ( ns + nm + c ∗ ncm ) , where l is the number of layers — as discussed in Supplementary Section A . The full architecture is described in Algorithm 1 . Algorithm 1 Compressive Transformer At time zero 1 : m0 ← 0 // Initialize memory to zeros ( l × nm × d ) 2 : cm0 ← 0 // Initialize compressed memory to zeros ( l × ncm × d ) At time t 3 : h ( 1 ) ← xWemb // Embed input sequence ( ns × d ) 4 : for layer i = 1 , 2 , . . . , l do 5 : mem ( i ) ← concat ( cm ( i ) t , m ( i ) t ) // ( ( ncm + nm ) × d ) 6 : ã ( i ) ← multihead attention ( i ) ( h ( i ) , mem ( i ) t ) // MHA over both mem types ( ns × d ) 7 : a ( i ) ← layer norm ( ã ( i ) + h ( i ) ) // Regular skip + layernorm ( ncm × d ) 8 : old mem ( i ) ←m ( i ) t [ : ns ] // Oldest memories to be forgotten ( ns × d ) 9 : new cm ( i ) ← f ( i ) c ( old mem ( i ) ) // Compress oldest memories by factor c ( bnsc c × d ) 10 : m ( i ) t+1 ← concat ( m ( i ) t , h ( i ) ) [ −nm : ] // Update memory ( nm × d ) 11 : cm ( i ) t ← concat ( cm ( i ) t , new cm ( i ) ) [ −ncm : ] // Update compressed memory ( ncm × d ) 12 : h ( i+1 ) ← layer norm ( mlp ( i ) ( a ( i ) ) + a ( i ) ) // Mixing MLP ( ns × d ) Algorithm 2 Attention-Reconstruction Loss 1 : Lattn ← 0 2 : for layer i = 1 , 2 , . . . , l do 3 : h ( i ) ← stop gradient ( h ( i ) ) // Stop compression grads from passing ... 4 : old mem ( i ) ← stop gradient ( old mem ( i ) ) // ... into transformer network . 5 : Q , K , V← stop gradient ( attention params at layer i ) // Re-use attention weight matrices . 6 : def attn ( h , m ) ← σ ( ( hQ ) ( mK ) ) ( mV ) // Use content-based attention ( no relative ) . 7 : new cm ( i ) ← f ( i ) c ( old mem ( i ) ) // Compression network ( to be optimized ) . 8 : Lattn ← Lattn + ||attn ( h ( i ) , old mem ( i ) ) − attn ( h ( i ) , new cm ( i ) ) ||2

Summary: This paper proposes a way to compress past hidden states for modeling long sequences. Attention is used to query the compressed representation. The authors introduce several methods for compression such as convolution, pooling etc. The outcome is a versatile model that enables long-range sequence modeling, achieving strong results on not only language model tasks but also RL and speech. For testing and evaluating the modeling of really long context sequence modeling, the authors introduce PG-19, a new benchmark based on Project Gutenberg narratives.