PaperlayerView exact revision on arXiv

Source-linked AI summary

Block-Recurrent Transformers

DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, Behnam Neyshabur

arXiv:2203.07852v3cs.LGcs.AIcs.NE

TL;DR

Long-document transformers face quadratic attention cost and lack memory beyond their attention window. The Block-Recurrent Transformer adds blockwise recurrent state updates using a transformer-like cell with parallel within-block computation. It improves perplexity over Transformer-XL baselines while maintaining comparable layer-level cost and faster execution.

  • Problem

    Transformers lack memory beyond their attention window and incur quadratic self-attention cost on long documents.

  • Method

    The Block-Recurrent Transformer applies a transformer layer recurrently across blocks of tokens and maintains a large set of recurrent state vectors with attention and gates.

  • Results

    The model provides larger perplexity gains than adding conventional layers, outperforms Transformer-XL baselines, and runs twice as fast in the reported comparison.

  • Takeaways & Limitations

    Blockwise recurrence offers a low-cost way to improve language modeling on very long sequences while retaining parallel computation within blocks.

  • Takeaways & Limitations

    The recurrent architecture has not reached its full potential, and further advances are needed for knowledge extraction and summarization.

Abstract

from arXiv · show

We introduce the Block-Recurrent Transformer, which applies a transformer layer in a recurrent fashion along a sequence, and has linear complexity with respect to sequence length. Our recurrent cell operates on blocks of tokens rather than single tokens during training, and leverages parallel computation within a block in order to make efficient use of accelerator hardware. The cell itself is strikingly simple. It is merely a transformer layer: it uses self-attention and cross-attention to efficiently compute a recurrent function over a large set of state vectors and tokens. Our design was inspired in part by LSTM cells, and it uses LSTM-style gates, but it scales the typical LSTM cell up by several orders of magnitude. Our implementation of recurrence has the same cost in both computation time and parameter count as a conventional transformer layer, but offers dramatically improved perplexity in language modeling tasks over very long sequences. Our model out-performs a long-range Transformer XL baseline by a wide margin, while running twice as fast. We demonstrate its effectiveness on PG19 (books), arXiv papers, and GitHub source code. Our code has been released as open source.

1 Introduction

Transformers handle long-range dependencies efficiently but incur quadratic attention cost and lack memory beyond their attention window. The Block-Recurrent Transformer combines attention with blockwise recurrence to enlarge recurrent memory while retaining parallel training and transformer-like cost.

  • Quadratic self-attention cost and invisible out-of-window tokens limit transformers on long documents.
  • Block-Recurrent Transformer recurrence combines fixed-size sequence state with attention while addressing the limitations of prior recurrent approaches.
  • Blockwise recurrence processes tokens and states in parallel within each block, enlarging recurrent capacity by orders of magnitude and extending information propagation.
  • The recurrent cell costs about one conventional transformer layer in computation and parameters but improves language-modeling perplexity more than adding a conventional layer.

2 Related Work

Prior work reduces long-range attention cost through sparsity, pooling, compression, linearization, or recurrence. The Block-Recurrent Transformer instead performs recurrence over token blocks while retaining standard dense softmax attention.

  • Efficient attention methods reduce long-range cost by selecting tokens, pooling or compressing sequences, or replacing dense attention with approximate lookup or linearized computation.
  • The model builds on sliding-window attention, which extends Transformer-XL-style context handling.
  • The Block-Recurrent Transformer differs from linearized-attention approaches by relying only on standard dense attention with softmax.
  • Prior recurrence-based transformer variants include feedback, simple recurrent, encoder-decoder, and block-oriented mechanisms, but block recurrence remains underexplored.

3 Method

The method combines sliding-window attention with a recurrent cell that updates many state vectors per token block. Its design preserves linear segment-length scaling, parallel within-block computation, and approximately transformer-layer cost.

  • Sliding-window attention is linear in segment length because each W-token block attends only to itself and the previous block.
  • Cached keys and values connect consecutive segments, while larger segments than windows allow gradients to propagate across multiple blocks beyond Transformer-XL.
  • 3.1 Recurrent Cell: The recurrent cell maps W token embeddings and S current states to W output embeddings and S next states, with S and W independently configurable.
  • 3.1 Recurrent Cell: Its vertical direction uses parallel self-attention over tokens and cross-attention to states, while its horizontal direction attends over states and cross-attends to tokens.
  • 3.1 Recurrent Cell: Gated recurrent connections replace residual connections, and parallel attention avoids an additional gate that produced worse perplexity.
  • State IDs differentiate recurrent vectors, while relative position bias is applied to token self-attention rather than global position embeddings.
  • A fixed learned gate implements an exponential moving average over previous blocks.
  • 3.6 Placement of Recurrence and Computation Cost: A 12-layer Block-Recurrent Transformer has almost exactly the computation and parameter cost of a 13-layer Transformer-XL, with nearly identical autoregressive inference cost.

4 Results

Across long-document language modeling experiments, Block-Recurrent Transformers consistently outperform Transformer-XL baselines at comparable or lower cost, with benefits increasing at larger scales.

  • Overall results: The recurrent models outperform all five baselines, and Rec:fixed:skip is best in 3 of 4 cases while also running slightly faster than Slide:13L.It also uses fewer parameters because it omits the MLP.
  • Baselines: Rec:fixed:skip outperforms the 13-layer baseline by a wide margin and beats XL:2048, which runs more than twice slower.Table 1 reports average bits-per-token, where lower is better, at comparable computational cost.
  • Ablations: A single recurrent layer is sufficient for most benefits: adjacent recurrent layers do not help, while widely separated layers perform no better than adding a non-recurrent layer.The authors nevertheless use two recurrent layers in their largest models.
  • Ablations: Increasing recurrent states improves performance up to 1024 states, but 2048 states degrade performance.The authors hypothesize that very large state spaces are harder to use effectively.
  • Ablations: Recurrence compensates for smaller attention windows more effectively than it does for Transformer-XL, whose perplexity worsens substantially when its window is reduced.Reducing the recurrent transformer's window has a smaller effect because recurrence can preserve context.
  • Published comparison: On PG19, the 1.3B-parameter model achieves 3.22 bits per token and word-level perplexity of 26.50.The authors describe this as a new state of the art, while cautioning that raw perplexity depends on multiple experimental factors.
  • Published comparison: Block-Recurrence performs nearly as well as Memorizing Transformer on arXiv, matches it on PG19, and trains almost twice as fast.The comparison uses identical vocabulary, configuration, and hyperparameters.

5 Discussion

The recurrent state appears to summarize recurring entities and places rather than perform complex reasoning, while the paper identifies both model limitations and broader misuse risks.

  • Model behavior: Qualitative analysis suggests the recurrent state summarizes frequently occurring characters and places, but the best model uses a simple exponential moving average rather than complex gating.The best-performing fixed:skip configuration does not use an LSTM-style remember-or-forget gate.
  • Limitations: Removing the recurrent layer’s MLP has little effect, suggesting that the layer is not using that component effectively for knowledge extraction and summarization.The authors conclude that further advances are needed to make fuller use of the recurrent layer’s capabilities.
  • Ethics: The model could extend language modeling to longer documents, which may broaden applications but also expand opportunities for disinformation, malicious chatbots, and spam.The paper describes these as potential negative social impacts shared with other language-model advances.

6 Conclusion

The paper concludes that recurrence improves long-document language modeling more efficiently than increasing attention windows or parameter counts, while emphasizing implementation simplicity and unresolved evaluation needs.

  • Conclusion: Adding recurrence to one layer costs roughly one additional non-recurrent layer but improves perplexity more than increasing parameters or attention-window size.The medium model beats Transformer-XL with four times the window and runs twice as fast; the larger model beats a Transformer-XL model with twice the parameters.
  • Conclusion: The recurrent transformer is easy to implement because it mostly uses ordinary transformer components and RNN gates, requires no custom CUDA kernels, and is open source.The implementation claim concerns the released Recurrent Transformer code.
  • Future work: The paper identifies downstream-task evaluation as an important direction, especially for applications requiring long-range context such as book reports, long-article summarization, code completion, and book-length question answering.The cited examples are proposed applications rather than completed evaluations in this work.
  • Future work: The authors state that the recurrent architecture has not yet reached its full potential and remains open to further research and improvement.

Appendix A Further Analysis.

The appendix analyzes context length, receptive fields, recurrence, and memory, showing how block recurrence supports long-range information use while retaining linear sequence-length complexity.

  • A vanilla transformer's average context length is N/2, whereas sliding-window attention provides a fixed W-token context at every position.
  • For sliding-window models, the theoretical receptive field is W · L, but using that additional context requires multiple attention hops and is harder to learn.
  • The block-recurrent transformer's effective context length appears large in practice, with accurate predictions observed across distances exceeding 60k tokens.
  • The recurrent layer has complexity O((W^2 + S^2 + 2SW) · N/W), which is linear in segment length N for fixed W and S.
  • On PG19, Slide:13L outperforms XL:2048 despite its shorter window, while XL:2048 performs better on arXiv, suggesting dataset-dependent attention requirements.
  • Recurrence summarizes text into recurrent states, whereas kNN memory performs direct lookups; the former may better capture subtle long-range information, while the latter favors precise facts.

Appendix B Gate Initialization and Training Stability

Training stability depends strongly on gate initialization because recurrence can be ignored by the model, leaving performance equivalent to a non-recurrent transformer.

  • Poor gate initialization can cause the model to ignore recurrent states permanently, reverting performance to that of the non-recurrent transformer.
  • The authors attribute this failure mode to the recurrent transition learning more slowly than direct token attention, making recurrent states initially less informative.
  • Gate initialization must account for the optimizer, because Adafactor can produce very small updates for bias terms initialized at zero.
  • The implementation uses small non-zero gate biases and adds -1 and +1 to input and forget gates to initially favor remembering.
  • With this initialization strategy, the recurrent cell reliably learns to use its recurrent state.

Appendix C Training Details

The appendix documents dataset and optimization choices, emphasizing that long-document benchmarks are selected for long-range evaluation and that published perplexities require tightly controlled comparisons.

  • GitHub results are noisy because the dataset spans many programming languages and coding styles, so the experiments use a batch size four times larger than PG19's.
  • The experiments do not provide error bars for every condition, although PG19-token headline experiments used three runs and measured error bars of 0.002–0.007.
  • ArXiv and GitHub cannot necessarily be redistributed as public datasets because their authors retain copyright, and access was obtained privately.
  • PG19 was chosen for long-range experiments because it contains book-length works, is larger than older benchmarks, and is publicly available.
  • Perplexity comparisons across published models are not meaningful unless vocabulary, learning rate schedule, batch size, training steps, optimizer, and related variables are controlled.

Appendix D Window Size and Number of Recurrent States

The appendix evaluates window size, recurrent-state count, feedback, and model scaling, finding measurable design trade-offs and substantial computational costs for feedback.

  • Reducing window size worsens perplexity for both recurrent and Transformer-XL models, but the recurrent model incurs a smaller penalty.
  • Adding recurrent states improves performance modestly through 1024 states, but performance worsens at 2048 states with a window size of 512.
  • Feedback increases step time by approximately 35–40% and adds parameters because all transformer layers receive cross-attention modules.
  • Block feedback improves perplexity in most cases, with especially large gains for the LSTM gate, although the effect depends on dataset and gate configuration.
  • Large PG19 models generally perform better, but their training costs can prevent researchers from reproducing or surpassing the reported results.
  • The scaling study spans 40M to 1.3B parameters and varies embedding, MLP, attention-head, head-size, and layer dimensions.

Appendix G Qualitative Analysis Results

The qualitative analyses show that the Block-Recurrent Transformer’s advantages over Transformer-XL are concentrated in recalling long-range names, titles, captions, and other book-specific information beyond the attention window.

  • On average, the recurrent model performed slightly better than Transformer-XL, although it was not better on every individual token.Figure 4 plots the per-token NLL difference for one randomly selected book.
  • Most tokens where the Block-Recurrent Transformer improved over Transformer-XL were proper names that did not occur within the previous 512 tokens.The same pattern persisted when the search expanded to the top 40 tokens for each book.
  • The recurrent model also recalled chapter titles and illustration captions from earlier table-of-contents occurrences, often beyond the 4096-token segment.The qualitative analysis found this pattern across the selected books, not only in isolated examples.
  • The model could remember a book’s title and author across a distance of 60,000 tokens or more through its recurrent state.This information appeared repeatedly near the beginning and end of the books, while the relevant token was outside the local context.

Appendix I Vocabulary Ablation Experiments

The appendix examines how tokenizer and vocabulary choices affect perplexity measurements, showing that larger PG19-trained vocabularies change tokenization statistics and bits per token. It also explains why WLP and BPC comparisons require consistent tokenization and vocabulary.

  • Comparison validity: BPC is potentially simpler and fairer across datasets, but inconsistent whitespace handling prevents reliable character counts across many SentencePiece vocabularies.Tokenizers may normalize, merge, or strip whitespace, affecting whether whitespace and newlines are represented in the model’s predictions.
  • Comparison validity: WLP comparisons can be misleading when models use different vocabularies, because vocabulary choice can produce significant differences unrelated to the model innovation being judged.The authors argue that tokenizer and vocabulary should be treated as part of the model.
  • Perplexity metrics: WLP rescales average token loss using the ratio of tokens to words, so its value depends on the chosen tokenization scheme and vocabulary.For PG19, the word count is fixed at 6,966,499, while the token count varies with tokenization.
  • Mechanisms: Vocabulary size affects model behavior through embedding capacity, token length, context coverage, and the amount of training data processed per step.Larger vocabularies can produce longer tokens, allowing a fixed token context to span more text and processing more data per training step.
  • Experimental setup: The ablation trains Rec:fixed:skip for 500k steps with a 4096-token segment while varying only vocabulary, including byte-level, T5, LaMDA, and PG19-trained SentencePiece vocabularies.SentencePiece vocabulary sizes range from 512 to 128k tokens.
Loading 2203.07852v3…