Block-Recurrent Transformers

TL;DR

The paper introduces the Block-Recurrent Transformer, a model that combines recurrence with transformer layers to efficiently process very long sequences, achieving better perplexity and faster speed than existing models.

Contribution

It presents a novel recurrent transformer architecture that operates on blocks of tokens with linear complexity, improving long-sequence language modeling performance.

Findings

Outperforms Transformer XL on long sequences

Runs twice as fast as baseline models

Achieves lower perplexity on language modeling tasks

Abstract

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.