Block-Recurrent Transformers
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We introduce the Block-Recurrent Transformer, which applies a transformer layerin a recurrent fashion along a sequence, and has linear complexity with respect tosequence length.
Our recurrent cell operates on blocks of tokens rather than singletokens during training, and leverages parallel computation within a block in orderto make efﬁcient use of accelerator hardware.
The cell itself is strikingly simple.
Itis merely a transformer layer: it uses self-attention and cross-attention to efﬁcientlycompute a recurrent function over a large set of state vectors and tokens.
Our designwas inspired in part by LSTM cells, and it uses LSTM-style gates, but it scalesthe typical LSTM cell up by several orders of magnitude.
Our implementationof recurrence has the same cost in both computation time and parameter countas a conventional transformer layer, but offers dramatically improved perplexityin language modeling tasks over very long sequences.
Our model out-performsa long-range Transformer XL baseline by a wide margin, while running twice asfast.
We demonstrate its effectiveness on PG19 (books), arXiv papers, and GitHubsource code.
Our code has been released as open source.
Introduction
Transformers have mostly replaced recurrent neural networks (RNNs), such as LSTMs [1], on tasksthat involve sequential data, most notably in the domain of natural language processing.
There areseveral reasons for their success.
First, transformers process all elements of the sequence in parallel,and are thus more efﬁcient to train on modern accelerator hardware.
In contrast, an RNN must processtokens sequentially which leads to slow step times during training and large batch sizes in order tofully saturate GPUs or TPUs.
Second, an RNN must summarize and compress the entire previous sequence into a single statevector which is passed from one token to the next.
The size of the state vector limits the amount ofinformation that the RNN can encode about the previous sequence.
In contrast, a transformer canattend directly to past tokens, and does not suffer from this limitation.
Third, attention operates effectively over longer distances.
The forget gate in an LSTM discardsinformation moving forward, and causes vanishing gradients during backpropagation.
In practice,this means that LSTMs struggle to send a clear signal over more than a few hundred tokens, far lessthan the typical size of the attention window in a transformer [2].
Despite these advantages, transformers also have a disadvantage.
The computational complexityof self-attention is quadratic with respect to the sequence length, which is a limiting factor whenattempting to process long documents, such as books, technical articles, or source code repositories.
Moreover, a transformer has no memory of past context; any tokens that it cannot attend to are“invisible” to the model.
input token embeddings output token embeddings c u rr e n t s t a t e v ec t
o r s n e x t s t a t e v
ec t
o r s linear Q sv linear K e V e Q ev s s cross-attn self-attn linear projection +
MLP + input token embeddings c u rr e n t s t a t e v ec t
o r s linear K e V e linear K s V s Q sh linear Q eh cross-attn self-attn
M L P li n ea r p r
o
j ec ti on gate gate + learnedstate
IDs + learnedstate
IDs output token embeddings n e x t s t a t e v ec t
o r s
In this work, we describe an architecture which combines the beneﬁts of attention and recurrence.
Like previous implementations of recurrence, our architecture constructs and maintains a ﬁxed-sizestate, which summarizes the sequence that the model has seen thus far.
However, our implementationof recurrence differs from previous work in several important aspects which together address thethree limitations mentioned above.
Instead of processing the sequence one token at a time, our recurrent cell operates on blocks oftokens ; see Figure 1.
Within a block, all tokens are processed in parallel, at least during training.
Therecurrent cell likewise operates on a block of state vectors rather than a single vector.
This meansthat the size of the recurrent state is orders of magnitude larger than in an LSTM, which dramaticallyimproves the model’s capacity to capture the past.
Processing the sequence in blocks also helpspropagate information and gradients over longer distances, because the number of recurrent steps (andthus the number of times that the forget gate is applied) is orders of magnitude smaller.
We show thatthe Block-Recurrent Transformer can remember information over distances of 60k tokens or more.
The recurrent cell itself is strikingly simple.
For the most part, it consists of an ordinary transformerlayer applied in a recurrent fashion along the sequence length.
There are a few tricks that arenecessary to stabilize training ; see Sections 3.2 and 3.4 for details.
The cost of recurrence, in termsof both computation time and parameter count, is essentially the same as simply adding one morelayer to our transformer baseline.
We demonstrate empirically that adding a single recurrent layerresults in a much larger improvement in perplexity on multiple datasets than adding a conventionaltransformer layer, while training time and memory use are equivalent.
Moreover, our recurrent cellis very easy to implement because it largely makes use of existing transformer code.
Thus, ourtechnique is a cheap and cheerful way to improve language modeling perplexity on long sequences.
Related Work
The quadratic cost of attention is well known in the literature, and a great deal of work has been doneon efﬁcient long-range attention mechanisms; see [3, 4] for recent surveys.
Sparse strategies such as
Big Bird [5], Routing Transformers [6], and Reformer [7] select only a subset of tokens to attend to.
Hierarchical mechanisms [8] combine multiple tokens into phrases or sentences to reduce sequencelength.
Expire-span [9] learns to prune far-away tokens that the model has labelled as “unimportant”.
Memorizing transformers [10] replace dense attention with k -nearest-neighbor lookup.
Yet another approach is to reduce the sequence length by pooling, averaging, or compressing it insome way.
Hierarchical 1D attention [11], and Combiner [12] apply pooling or averaging over tokensat longer distances.
Linformer [13] applies a linear transformation to the key and value matricesto reduce the sequence length.
Compressive transformers [14] and funnel transformers [15] applyadditional learned compression layers to compress the sequence. [Figure]
The equation for attention is (roughly) softmax ( QK T ) V where Q , K , and V are the query, key,and value matrices of the attention layer.
If the softmax operation is removed from this equation orsomehow “linearized”, the equation can be rearranged as Q ( K T V ) , where ( K T V ) can be computedincrementally (i.e., in a recurrent fashion) as a cumulative sum over the sequence [16].
Linearizedattention thus has linear rather than quadratic complexity with respect to sequence length.
Followingthis line of reasoning, there have been several proposals that approximate the softmax [17, 18] orreplace it
Linear transformers are related to earlier work on fast weight programmers [19][21], and can be extended with other forms of recurrence [22].
Our work differs from all of the above mechanisms, because we rely only on standard dense attentionwith softmax.
A few other lines of research have combined the transformer architecture with recurrence in someway.
The feedback transformer [23] allows lower layers to attend to the output of the topmost layer.
Feedback has minimal cost at inference time, but it is unfortunately very slow to train because tokensmust be processed sequentially.
Simple Recurrent Units [24, 25] use a recurrence function that doesnot involve matrix multiplication, and is consequently much faster.
RNMT+ combines RNNs andtransformers in an encoder/decoder architecture to improve on translation tasks [26].
“Sandwichmodels” alternate between transformer and RNN layers and out-perform both transformers and RNNs on tasks involving source code [27].
The R-Transformer introduces an additional local RNNwhich can be computed in parallel in order to better model sequential structure [28].
The perceiverarchitecture [29] is somewhat similar to ours; it also applies a transformer layer in an iterative fashion.
To the best of our knowledge, the idea of performing recurrence on blocks of tokens is underexplored.
In the context of translation, [30] operates on sentences rather than tokens.
Staircase Attention [31]also operates on blocks of tokens; each layer takes, as input, the outputs of the same layer from theprevious block.
Method
The Block-Recurrent Transformer is based on sliding-window attention [32], which is an extensionof ideas from Transformer-XL [33].
A long document, such as a book, consists of a sequence of tokens. Due to memory limitations, itis usually not possible to ﬁt the entire sequence into device memory. Thus, the sequence is dividedinto segments of length N ( N = 4096 in our experiments), which are processed sequentially over anumber of training steps.
Each training step processes one segment.
The sliding window attention pattern is illustrated in Figure 2. Given a segment of N tokens, thesliding window applies a causal mask in which each token can only attend to the W previous tokens, where W is the window size ( W = 512 in our experiments). Because of the causal mask, mostentries of the N × N attention matrix are masked out (assuming that W << N ).
Thus, the attentioncomputation can be optimized by breaking it into smaller tiles along the diagonal.
The sequence of Ntokens is subdivided into blocks of size W , and each block attends locally to itself and to the previousblock, so the size of each local attention matrix is W × 2 W .
Using this mechanism, attention isquadratic with respect to the window size W , but linear with respect to the segment length N .
Borrowing an idea from Transformer-XL, the keys and values from the last block in each segmentare stored in a non-differentiable cache for use on the next training step. By using the cache, theﬁrst block in the next segment can attend to the last block in the previous segment, which extendsthe sliding window to cover the entire (book-length) sequence. The cache implements a form oftruncated backpropagation through time [34] over long documents.
Note that if N = W , then sliding window attention will behave exactly like Transformer-XL; it willprocess and cache one segment (i.e. one block) per training step.
Setting N >> W does not changethe context length of attention, but it allows gradients to backpropagate across multiple blocks duringtraining; we show that the improved differentiability provides a modest beneﬁt to perplexity over
Transformer-XL. See Appendix A for more details. 3.1 Recurrent Cell A Block-Recurrent Transformer layer extends the sliding-window attention mechanism by adding aset of recurrent states, which are updated at the end of each block of W tokens.
Our design for therecurrent cell is illustrated in Figure 1, which depicts the operations done within a single block of theinput sequence.
The recurrent cell receives two tensors as inputs: a set of W token embeddings, where W is the blocksize, and a set of S “current state” vectors.
The cell produces two tensors as outputs: a set of Woutput embeddings, as well as a set of S “next state” vectors.
We denote the function going frominput token embeddings to output token embeddings as the vertical direction, and the function goingfrom the current state vectors to the next state vectors as the horizontal direction.
The number of statevectors S and the window size W are independent hyperparameters, but we set S = W = [-
-]512 in ourexperiments to simplify comparisons against baselines.
The vertical direction of the cell is an ordinary transformer layer with an additional cross-attentionoperation, much like a decoder layer in a standard encoder-decoder architecture [35].
It does selfattention over the input tokens, and cross-attends to the recurrent states.
Unlike a typical decoderlayer, we do self-attention and cross-attention in parallel.
The results of both forms of attention areconcatenated together and fed into a linear projection.
The horizontal direction of the cell mirrors the forward direction, except that it performs selfattention over the current state vectors, and cross-attends to the input tokens.
The recurrent directionalso replaces the residual connections with gates, which allows the model to “forget”, an ability thatis important for algorithmic tasks [36], or when processing long documents, where it has been centralto the success of LSTMs [37].
Note that the presence of gates is the reason why self-attention and cross-attention are done in parallel.
Doing them sequentially, as is standard practice, would introduce a third gate in the horizontaldirection, which led to worse perplexity in our experiments.
Recurrence is integrated with the sliding window attention mechanism.
Although not shown in Figure1, each cell also receives keys and values from the previous block as input, these are concatenatedwith ( K e , V e ) from the current block in order to implement sliding-window attention.
A Block-Recurrent Transformer layer processes the blocks within a segment sequentially by stackingrecurrent cells horizontally, with the “next states” output of the previous cell feeding into the “currentstates” input of the next cell.
In code, this is implemented as a simple for-loop over blocks. Multiplelayers can also be stacked vertically in the usual fashion.
Our experiments use a single recurrentlayer, sandwiched between a number of non-recurrent layers that use sliding-window attention.
The ﬁnal set of state vectors from the last block in the segment are cached, along with the keys andvalues, and used as the initial state for the ﬁrst block on the next training step.
Every layer in thestack (both recurrent and non-recurrent) has its own cache.
Sharing of keys and values.
Keys and values are shared between the vertical and horizontaldirections.
One set of keys and values ( K e , V e ) are computed from the input token embeddings, andanother set of keys and values ( K s , V s ) are computed from the recurrent state vectors.
Queries arenot shared, so there are four separate sets of queries: Q ve and Q vs in the vertical direction, and Q hsand Q he in the horizontal direction.
3.2 State IDs and Position Bias
With a large number of state vectors, the total size of the recurrent state is far larger than that of an
LSTM.
However, the same weights (projection matrices and MLP) are applied to each state vector.
Without some way to differentiate the states, the model will compute the same result for each statevector, thus negating any advantage from having multiple states.
To prevent this failure mode, we adda set of learned “state IDs” to the state vectors before computing the keys, values, and queries.
These“state IDs” allow each state vector to consistently issue different queries against the input sequence,and against other states.
State IDs are identical to learned position embeddings; we use a differentname because there’s no notion of “position” between states.
We do not add global position embeddings to the tokens, because global position embeddings don’twork well for long sequences [33].
Instead, we add a T5-style relative position bias [38] to the selfattention matrix in the vertical direction.
(Although similar, T5 relative positions differ slightly fromthe relative positions used in the Transformer-XL paper [33].)
When the recurrent states cross-attendto input tokens, there is no position bias, because the relative distance between “state” and “token” isundeﬁned.
We also normalize queries and keys as described in [39]; we found that normalization improved thestability of Transformer-XL when used with a relative position bias.
3.3 Gate Type
We experimented with two different gating mechanisms for the recurrent cell.
Each state vector hasits own gate, but all state vectors are updated in parallel, using the equations below.
Fixed gate.
The ﬁxed gate uses a learned convex combination, similar to highway networks [40].
[Equation]
where W ( z ) is a trainable weight matrix, b ( z ) and b ( g ) are trainable bias vectors, σ is the sigmoidfunction, c t is the cell state for the current block (i.e., the state for the block at index t in the sequenceof blocks), (cid:12) is the element-wise multiplication, and h t is the current input to the gate.
In our model,h t is either the output of attention, in which case W ( z ) is the linear projection that feeds into the gate,or h t is the output of the hidden layer of the MLP, in which case W ( z ) is the ﬁnal layer of the MLP.
Note that the bias b ( g ) is a simple learned vector, so the value of g does not depend on either thecurrent state c t , or the input h t , and thus remains constant (i.e., ﬁxed) after training.
The ﬁxed gateessentially implements an exponential moving average over previous blocks.
LSTM gate.
The LSTM gate uses the standard combination of input and forget gates:
[Equation]
[Equation]
[Equation] [Equation] [Equation] where W ( z ) , W ( i ) , W ( f ) are trainable weight matrices, and b ( z ) , b ( i ) , b ( f ) are trainable bias vectors.
The LSTM gate is strictly
3.4 Gate Initialization and Training Stability
We observed that training stability is quite sensitive to how the gates are initialized.
Recurrence hasa failure mode where the model learns to completely ignore the recurrent state, in which case itsperformance reverts to that of the non-recurrent transformer.
Moreover, this situation appears to be alocal optimum; once the model has reached this point, it does not recover.
We stabilize training byinitializing the weights and bias to small but non-zero values, and adding a constant -1 and +1 to theinput and forget gates to bias them to “remember”.
See Appendix B for details.
3.5 Gate Conﬁguration
We experimented with three different gate conﬁgurations.
Dual.
The dual gate conﬁguration is the one shown in Figure 1, in which both of the residualconnections in the cell are replaced with gates.
The disadvantage of this conﬁguration is that thereare two gates, both of which can forget.
Single.
The single gate conﬁguration removes the linear projection and the gate that is attached toit.
Instead, the concatenation of self-attention and cross-attention is fed directly into the MLP.
Skip.
The skip conﬁguration removes the MLP and the gate that is attached to it.
This conﬁgurationis similar to the single-gate version, except that it is strictly weaker.
Instead of a two layer MLP witha very large hidden layer, it uses a linear projection with no nonlinearity.
3.6 Placement of Recurrence and Computation Cost
Single recurrent layer.
The basic version of the Block-Recurrent Transformer uses a single recurrent layer sandwiched between a number of non-recurrent transformer layers with sliding attention.
We use a 12-layer model with recurrence on layer 10.
All layers have a Transformer-XL-style cache.
Cost of recurrence. During training, the 12-layer Block-Recurrent Transformer has almost exactlythe same computation cost, in both parameters and FLOPS, as a 13-layer model without recurrence.
The two are equivalent because the recurrent cell does almost the same operations as a conventionaltransformer layer, merely in the horizontal instead of the vertical direction.
The inference cost for autoregressive decoding is also nearly identical, for the same reason.
Recurrence adds an additional attention operation per token, the cost of which is the same as self-attentionin a 13th layer.
Results
We tested the Block-Recurrent Transformer on three different data sets of long documents: PG19,arXiv, and GitHub.
The PG19 dataset [41] contains full-length books written prior to 1919 fromproject Gutenberg.
The arXiv dataset [10] is a corpus of technical papers downloaded via thearXiv Bulk Data Access 2 , and ﬁltered to include only articles labeled as “Mathematics” and whose
L A TEX source is available.
The GitHub dataset [10] is a corpus of source code from different GitHubrepositories with open-source licenses.
All of the ﬁles in each GitHub repository are concatenatedtogether to make one long document.
The task is auto-regressive language modeling, where the goal is to predict the next token in thesequence.
We report bits-per-token numbers (i.e. log 2 perplexity; lower is better) for all models.
Further training details for each dataset can be found in Appendix C.
4.1 Baselines
We compare the Block-Recurrent Transformer to ﬁve different baselines.
The ﬁrst baseline, XL:512 ,establishes a reference point against which various other improvements can be compared. It’s a [Table]
Feedback:lstm:single  1.
Feedback:fixed:skip 1. Memorizing Trans. 32k  1. - 12-layer Transformer-XL model with a window size of 512, and 150 million parameters.
It has 8heads of size 128, embedding vectors of size 1024, an MLP with a hidden layer of size 4096, and therelu nonlinearity.
It uses a Transformer-XL style cache, but no sliding window, so the segment lengthis the same as the window size, i.e., it is trained on segments of 512 tokens.
XL:1024 and XL:2048 are similar, but have window sizes of 1024 and 2048, respectively.
As expected,increasing the window size improves perplexity, especially on the arXiv data set.
However, these twomodels still have worse perplexity than the recurrent model, as well as being much slower.
Slide:12L is a 12-layer transformer with a window size of 512, but uses a sliding window over asegment of 4096 tokens.
This model is almost identical to XL:512 ; the only difference is that thesliding window is differentiable over multiple blocks, while the Transformer-XL cache is not.
Slide:13L adds a 13th layer, and is directly comparable to the recurrent models in terms of both computation cost (FLOPS or step-time), number of parameters, and segment length.
Notice that addinganother layer with more parameters yields a much smaller improvement than adding recurrence.
Relative cost.
All ﬁve baselines, and all 6 recurrent models, have roughly the same number ofparameters: between 151 million (12 layer) and 164 million (13 layer or recurrent).
The trainingspeed (i.e. step time) of each model is shown in Table 1 (lower is better).
Because the raw step timedepends on hardware and compiler, we report numbers relative to the Slide:13L baseline.
Batch Size.
We adjust the batch size so that each model processes the same number of tokens (andthus the same amount of training data) per training step.
Thus, XL:512 (segment length 512) runs at abatch size of 256 (8 per replica), while Slide:12L (segment length 4096) runs at a batch size of 32 (1per replica) on PG19.
4.2 Beneﬁt of Recurrence
We compare the 5 baselines to all six gate conﬁgurations for the Block-Recurrent Transformer.
The recurrent model reliably outperforms all ﬁve baselines.
The best overall conﬁguration is Rec:fixed:skip , which outperforms the others in 3 out of 4 cases, and comes within the margin oferror in the remaining case.
This is especially notable because it is also the fastest conﬁguration,having a slightly lower step time and fewer parameters than Slide:13L , because it does not have FLOPs are the same between the two models at a given parameter count.
At larger sizes, adding recurrence is equivalent to doubling the number of parameters.
Details in Appendix F. the MLP.
It is better than the 13-layer baseline by a wide margin, and it is even better than the Transformer-XL model with a window size of 2048, which runs over 2 times slower.
The other gate conﬁgurations also outperform the 13-layer baseline, but their relative ranking variesaccording to the dataset.
Despite being theoretically more powerful, the LSTM gate tends to lag behindthe fixed gate in all of our experiments.
Scaling up.
Figure 3 shows the effect of adding recurrence as the transformer model is scaled upand down in size.
We trained six different models on PG19, ranging in size from 40M parameters to1.3B parameters.
For the four smaller models, we compare a 12-layer Block-Recurrent Transformeragainst a 13-layer Transformer-XL baseline, while for the two larger models, we compare a 24-layer
Block-Recurrent Transformer, with recurrence at layers 10 and 20, against a 26-layer Transformer-XLbaseline.
This experiment used a cosine-decay learning rate as described in [42], and a custom 32k
SentencePiece vocabulary [43].
More details are in Appendix F.
Our experiments show that recurrence provides a consistent beneﬁt across all scales.
The relativeimprovement actually seems to increase with the number of parameters; at larger sizes recurrenceprovides a beneﬁt which is greater than doubling the number of parameters .
4.3 Ablations
Multiple recurrent layers.
Adding two recurrent layers right next to each other in the stack (layers9 and 10) did not improve model perplexity.
Adding two layers widely separated in the stack (layers 4and 10) did provide an improvement, but the improvement was no better than simply adding anothernon-recurrent layer to the stack. Previous work on Memorizing Transformers [10] showed a similareffect.
In our qualitative study, we saw that the model seems to use recurrence primary for long-rangename lookups, much like memory.
We conclude that one layer of recurrence is sufﬁcient for themodel to extract most of the beneﬁts, although we did use two layers for our largest models.
Number of recurrent state vectors.
We trained the model with differing numbers of state vectors,from 128 to 2048.
Increasing the number of states makes a small but measurable improvement up to1024, but the model does worse with 2048 (see Appendix D).
We hypothesize that the model hastrouble learning to use the recurrent state effectively if the state space grows too large.
Reducing window size.
Reducing the size of the sliding window dramatically reduces perplexityfor Transformer-XL, because it reduces the amount of context that the transformer is able to attend to.
Reducing the size of the window in a recurrent transformer has a smaller effect, because the modelcan use recurrence to compensate (see Appendix D).
4.4 Block feedback
Inspired by the feedback transformer [23], which allows all layers to attend to the topmost layer,we implemented a variation in which every layer of the transformer (not just the recurrent one) can [Table] cross-attend to the state vectors in the recurrent layer.
This variation further improves perplexity, butat a cost; step time increased by approximately 35-40%, and the additional queries also increase thenumber of parameters.
Results are shown in Table 1, and further described in Appendix E.
4.5 Comparisons against prior published work
The PG19 test set contains 6,966,499 words [14], which are broken into 10,229,476 tokens usinga SentencePiece vocabulary, trained on PG19.
Our 24-layer 1.3B parameter model achieves 3.22bits per token, and thus achieves a new state of the art word-level perplexity of 26.50 (Table 2).
However, we note that raw perplexity numbers are not necessarily a meaningful way to comparearchitectures, because they depend on numerous other factors, such as the number of parameters,vocabulary, learning rate schedule, batch size, etc.; a more detailed discussion is in Appendix C.2.
We were able to run a fair comparison (identical vocabulary, conﬁguration, and hyperparameters)of the Block-Recurrent Transformer against the Memorizing Transformer [10], with a memory ofsize 32k (Table 1).
The memorizing transformer is constructed similarly to our model; it has onelayer which has been augmented with a mechanism that gives it the ability to attend over much longerdistances.
We ﬁnd that recurrence performs nearly as well as the memorizing transformer, but trainsalmost twice as fast.
However, there are many ways of implementing approximate k -nearest-neighborlookup, so relative speed will be highly implementation-dependent; our implementation runs on TPU,and does not use custom CUDA kernels.
4.6 Qualitative analysis
Prior work on long-context transformers [41, 10] has found that attention at long ranges is typicallyused to look up proper names, such as characters or places.
We performed a qualitative analysisin an attempt to determine whether our model is using recurrence in the same way.
We selected 5books at random from the PG19 test set, ran both the Block-Recurrent Transformer and the 13-layer
Transformer-XL on each book, and then compared the cross-entropy loss for all tokens.
We sortedthe results, and examined the top 4 tokens from each book with the greatest difference: the tokens forwhich the predictions of the recurrent model have the largest improvement over the baseline.
In 17/20 cases, the recurrent model predicted a proper name, usually with relatively high probability,that Transformer-XL was unable to predict.
In 2 cases it predicted a chapter title (having previouslyseen the table of contents), and in the last case, it predicted a foreign-language word that was uniqueto that book.
In 19/20 cases, the predicted word was nowhere within the attention window, so it musthave been stored within the recurrent state (details in the appendix, Section G).
In a second study, we compared the recurrent model, running normally, against a variation in whichthe recurrent state is cleared at the end of each 4096-token segment, instead of being cached.
Clearingthe state degrades the model’s ability to predict dependencies at a longer range than the segmentlength; typical mispredictions once again included proper names and chapter titles.
Interestingly, thisstudy also showed that the recurrent model is able to remember the title and author of a book (whichis part of the Gutenberg boilerplate at the beginning and end of each book) across the entire lengthof the book – more than 60,000 tokens.
See Appendix G.1.
A further quantitative comparison of the per-token cross-entropy between Transformer-XL and the
Block-Recurrent Transformer is
Discussion
Our implementation of recurrence was inspired by the way that humans seem to process longsequences.
When a human reads a novel, they do not attempt to remember every single word in thebook.
Instead, a human reader will construct a mental model, or knowledge graph, which summarizesthe story thus far, i.e., the names of the main characters, the relationships between them, and anymajor plot points.
When a human reads a paragraph of text, they will parse the information in theparagraph, process and interpret the information using background knowledge from their mentalmodel, and ﬁnally update their mental model with new information.
Our recurrent architectureloosely mimics this process.
It takes a block of text, and parses it by running it through a conventionaltransformer stack.
Tokens in the text attend to the recurrent states (i.e. the mental model), and thestates, in turn, are updated by attending to the text.
Based on our qualitative analysis, it seems that the model is, in fact, using the recurrent state tosummarize some of the information about frequently occurring characters and places.
However, itdoes not seem to be doing much complex reasoning, as evidenced by the fact that our best performingmodel is the fixed:skip conﬁguration.
This conﬁguration does not use a complex LSTM-style gate,which chooses to remember or forget based on its current state and inputs; instead, it simply computesan exponential moving average, not unlike some other forms of long-range approximate attention.
Moreover, the skip conﬁguration cuts out the large MLP from the recurrent transformer layer.
In avanilla transformer, removing the MLP from all layers would severely degrade the model [45]; thoselarge MLPs are computing something important.
In a recurrent layer, removing the MLP makes littledifference; it does not seem to be computing anything useful.
We conclude that training the recurrentlayer to make full use of its capabilities for knowledge extraction and summarization will requirefurther advances.
5.1 Ethics
The potential negative social impacts from this work are similar to any other advance in languagemodelling.
Large language models could potentially be used to create disinformation and fake news,power malicious chatbots, or generate spam.
The Block-Recurrent Transformer can potentially createlonger documents than was previously feasible, thus expanding the range of applications in whichthese negative impacts could occur.
The best way to mitigate these risks is to train models that canreason about text, and ﬂag misinformation or malicious content.
Conclusion
We have shown that when training language models on long documents, the Block-Recurrent
Transformer provides a greater beneﬁt at lower cost than scaling up the transformer model inother ways.
Adding recurrence to a single layer has roughly the same cost as adding an additionalnon-recurrent layer, but results in a much larger improvement to perplexity.
We have also shown thatrecurrence provides a larger beneﬁt than simply increasing the window size of attention, or increasingthe number of parameters.
Our medium-sized model has lower perplexity than a Transformer-XLmodel with 4 times the window size, but runs twice as fast, and our larger model outperforms a
Transformer-XL model with twice the number of parameters .
Furthermore, in contrast to some other recently proposed transformer variants, the Recurrent Transformer is very easy to implement, since it consists mostly of ordinary transformer components and RNN gates. No custom CUDA kernels are required.
Evaluating block-recurrent transformers on downstream tasks is an important direction for futurework.
We believe that the Block-Recurrent Transformer will be most useful in situations that requirelong-range context; examples of potential applications include writing book reports, summarizing longnews articles, code completion, or question/answering over book-length works.
There are are a number of new and emerging benchmarks that test long-range performance [46, 47, 3].
Previous studieshave found a strong correlation between language modeling and diverse downstream tasks [48, 49]. Despite our initial successes, we also believe that the recurrent architecture that we present herehas not yet achieved its full potential, and there are opportunities for future research and furtherimprovements in this area.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the
paper’scontributions and scope?
[Yes
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Section 5.1.(d)
Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
Yes ] We plan to open-source the code itself, although we have not yet done so
We have given a detailed description of architecture, hyper-parameters, and trainingmethodology.
Our main results are for PG19, which is a publicly available dataset.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix C.
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] Given the large number of experiments, we did not have the computationalresources to run every experiment multiple times, but we do run experiments overmultiple datasets.
For the main headline numbers on PG19-tokens, we ran the primaryexperiments three times each with different initial seeds and with dataset shufﬂing.
Error bars are given in Table 1.
Note that actual measured error bars are somewhatlower than reported in Table 1; we round results to the nearest 0.01, and thus round theerror up to match the precision of the reported results.
See Appendix C for details.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of
GPUs , internal cluster, or cloud provider)?
[Yes] See Appendix C.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes] See Appendix C.1.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[Yes] See Appendix C.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[Yes] See Appendix C