SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jonathan Pilault $^{124*}$  Mahan Fathi $^{123*}$   
Orhan First $^{1}$  Christopher Pal $^{24}$  Pierre-Luc Bacon $^{23}$  Ross Goroshin $^{1}$ $^{1}$ Google DeepMind  $^{2}$ Mila  $^{3}$ Université de Montréal  $^{4}$ Polytechnique Montréal
State space models (SSMs) have shown impressive results on tasks that require modeling long-range dependencies and efficiently scale to long sequences owing to their subquadratic runtime complexity. Originally designed for continuous signals, SSMs have shown superior performance on a plethora of tasks, in vision and audio; however, SSMs still lag Transformer performance in Language Modeling tasks. In this work, we propose a hybrid layer named Block-State Transformer (BST), that internally combines an SSM sublayer for long-range contextualization, and a Block Transformer sublayer for short-term representation of sequences. We study three different, and completely parallelizable, variants that integrate SSMs and block-wise attention. We show that our model outperforms similar Transformer-based architectures on language modeling perplexity and generalizes to longer sequences. In addition, the Block-State Transformer demonstrates more than tenfold increase in speed at the layer level compared to the Block-Recurrent Transformer when model parallelization is employed.
Transformers have shown impressive performance on a wide range of natural language processing (NLP) tasks. While they have been primarily used for language modeling the Transformer architecture [40] has also been successfully applied to other tasks outside of the NLP and have mostly replaced Recurrent Neural Networks (RNNs). Several factors contribute to this success, including computational efficiency and architectural inductive biases that are well-suited for training on natural language tasks at scale. On the computational upside, Transformers are able to process tokens of a given input sequence in parallel, making the most of modern accelerator hardware. Moreover, the attention mechanism enables Transformers to find relationships in longer sequences by providing ready access to all the extracted information from past tokens when inferring the next token. Compared to RNNs and LSTMs [19], the benefits of self-attention are two-fold: (i) the capacity of what could be stored and directly accessible as context is drastically increased, and (ii) training on longer sequences is more stable [18, 23].
Given the remarkable achievements of Transformers in language modeling tasks, and their improved performance at scale on hard NLP tasks such as reasoning and question answering [2, 39, 6], the demand for deploying even deeper and larger networks is greater than ever before. An orthogonal scaling dimension, which could be potentially even more consequential, is the size of the input sequence. Despite the several advantages of Transformers over RNNs, it is still problematic to scale the input sequence length, again for both computational performance and quality reasons. Further, the Transformer's runtime is quadratic with respect to the input sequence length, which makes training these models increasingly expensive. Furthermore, Transformers with attention, that is local [8], sparse [4, 43, 36], low-rank approximated [41] or linearized via kernel methods [5, 22], notoriously struggle on long-input classification tasks [37]. Vanilla transformers can be unstable when trained
Figure 1: Block-State Transformer layer. The BST-SH layer is illustrated on the left, and includes a state space model (SSM, in green) and Block Transformers (in red). For demonstration purposes the sequence is divided into 3 blocks in the picture. The details of the Block Transformer sublayer are on the right. *TRF = Transformer.
on long sequences [26] and token importance is concentrated in a local receptive field of around 50 tokens around the current time step [35].
An emerging body of research suggests that State Space Models (SSMs) can serve as an alternative to Transformers because they are able to capture dependencies in extremely long sequences, while being more computationally efficient and parallelizable [14]. While still falling into the category of autoregressive sequence models, the underlying linear time-invariant dynamical system of SSMs allows the efficient processing of sequences using parallelizable convolution operators with the Fast Fourier Transform (FFT) [7], with  $\mathcal{O}(L\log L)$  complexity, where  $L$  is the length of the sequence. Moreover, retention of past information over long sequences, up to thousands of steps, can be ensured by deriving recurrent update rules by borrowing ideas from online function approximation [3, 12]. SSMs have recently outperformed Transformers on long-range dependency benchmarks by a large margin [37]. Despite their success on long-range classification tasks, SSMs have not yet completely matched Transformers as an off-the-shelf sequence model for general language modeling tasks [10].
Recent findings suggest that Transformers and SSMs are complementary models for the purpose of language modeling [28]. In this work, we propose an architecture that integrates a strong local attention-based inductive bias with the long-term context modeling abilities of SSMs into a single layer, that we call Block-State Transformer (BST). Our model is able to process long input sequences, while still incorporating an attention mechanism to predict next tokens. BST is fully parallelizable, scales to much longer sequences, and offers a  $10 \times$  speedup compared to comparable Transformer-based layers.
In every BST layer, an SSM takes the entire sequence as input and maps it into a "context" sequence of the same length. The SSM sublayer takes advantage of FFT-based convolutions. This sequence of context is then divided into blocks of equal size, i.e. window length  $(W)$ , and each context block is then fed to a Block Transformer layer, that attends to the subsequences of size  $W$  as defined in [21]. The block of input token embeddings are then cross-attended to the corresponding block of context states; see Figure 1. Note that by introducing SSMs as a means of contextualization, we completely remove the need for sequential recurrences and we are able to run our hybrid SSM-Transformer layer fully in parallel. The resulting runtime complexity can be expressed as the sum of  $\mathcal{O}(W^2) + \mathcal{O}(L\log L)$ , where the first term represents the time complexity of the Transformer sublayer, while the second term represents the time complexity of the SSM sublayer. This is a major improvement over  $\mathcal{O}(LW)$  of Block-Recurrent Transformer, so long as hardware to support parallel computation is available. Moreover, due to hardware imposed restrictions, the runtime complexity of the SSM on a full sequence is comparable to that of Block Transformer on a block of tokens, which further implies the absence of a speed bottleneck in the BST layer, empirically validated for sequences containing hundreds of thousand of tokens. This is evident by observing that the bottom-most two lines on the left of Figure 4 are almost overlapping.
This work is primarily related to two branches of recent research: (i) combining local attention with recurrent networks in order to extend their capacity to capture long-range dependencies, beyond the length of the attention window size, and (ii) State Space Models (SSMs) which describe sequences via linear dynamical systems whose outputs can be computed in parallel. Block-Recurrent Transformer (BREC) [21] uses a recurrent memory mechanism to extend the theoretical context length of the Transformer. In the recurrent unit of the BREC cell, the updates made to the "recurrent state vectors," are extracted by employing a cross-attention mechanism over a block-window of input token embeddings. Different from their work, we use linear state space models instead of recurrent cells to maintain context states. We also conduct a more extensive exploration of maintaining and updating context states. Earlier works that augment transformers with a non-differentiable external memory include the Memorizing Transformer [42]. Transformer-XL [8] was an early work that combined recurrent memory with Transformers. Our work can be seen as a continued evolution of those models incorporating state-of-the-art recurrent memory models inspired by SSMs.
State space models can be considered as linear RNNs [12]. This simplicity facilitates their analysis and even enables analytical derivation of recurrent weights for optimally representing arbitrarily long sequences. The linear property also allows the recurrence to be unrolled and parallelized during training and inference [14]. Our work combines these state-of-the-art models, enabling Transformers to leverage theoretically infinite context.
Other works have attempted to replace Transformers, and their attention mechanism with SSMs [28, 27, 10, 30], however despite recent progress, the performance achieved by the Transformer architecture remains unparalleled in language. Nevertheless, SSMs are able to capture longer range dependencies than Transformers in both theory and practice, while also being highly parallelizable [7, 11]. We therefore elect to combine the best aspects of SSMs and Transformers into a single model. The idea of communication across blocks, similar to GSS [28], was later implemented by MEGA [27], through an Exponentially Moving Average (EMA) update rule instead of  $\mathrm{SSM}s^2$ . However, both GSS and MEGA use a single-head Gated Attention Unit (GAU) [20]. MEGA further mixes layer inputs, GAU outputs and EMA outputs via two gating mechanisms. Our method uses a simpler architecture to mix signals from local attention and SSM outputs via cross-attention, allowing us to plug any out-of-the-box SSMs or attention layers. Further, we investigate three ways to mix SSM signals with attention as outlined in Section 3.3.
We consider the problem of next token prediction via a decoder-only language model. This seemingly simple pretext task has led to spectacular progress in language understanding [9, 2, 29]. During training, the decoder takes in a sequence of length  $L$  of tokens embeddings and is tasked to generate the next token at every step in the sequence.
We start by a brief review of SSMs that are essential for understanding the Block-State Transformer layer (3.1). Our full Block-State Transformer architecture is outlined in Section 3.2. Section 3.3 describes three approaches for integrating SSM states into the attention mechanism. Important implementation details are described in Section 3.4.
3.1 State Space Preliminaries
State space models can be divided into two categories:
State Spaces: Structured Kernels S4 [14], S5 [34], S4D [15], DSS [16], follow a structured initialization of the convolutional kernel by unrolling a linear time-invariant (LTI) dynamical system of the following form:
$$
x _ {k} = \mathbf {A} x _ {k - 1} + \mathbf {B} u _ {k}, \tag {1}
$$
$$
y _ {k} = \mathbf {C} x _ {k} + \mathbf {D} u _ {k}.
$$
2The authors in [27] show a mathematical form of EMA that has a state transition and also derive a convolution kernel to efficiently compute EMA similarly to S4.
The system is parameterized by a state matrix  $\mathbf{A} \in \mathbb{R}^{N \times N}$ , vectors  $\mathbf{B} \in \mathbb{R}^{N \times 1}$ ,  $\mathbf{C} \in \mathbb{R}^{1 \times N}$ , and  $\mathbf{D} \in \mathbb{R}^{1 \times 1}$ , the SSM maps a 1-D input signal  $u_{k}$ , to a 1-D output signal  $y_{k}$ . Internally, the SSM projects the input signal to an  $N$ -D representation state  $x_{k}$ , before mapping it down to a scalar using the  $\mathbf{C}$  matrix. The term  $\mathbf{D}u_{k}$  can be thought of as a skip connection and will be omitted for the remainder of the discussion for convenience. The output of the above recurrent equation,  $y_{k}$ , can be computed as a discrete convolution, by realizing that the recurrence can be explicitly unrolled:
$$
\begin{array}{c c} \text {L e t} & x _ {- 1} := \vec {0}, \end{array}
$$
$$
y _ {k} = \sum_ {j = 0} ^ {k} \mathbf {C A} ^ {j} \mathbf {B} \cdot u _ {k - j}. \tag {2}
$$
The  $\mathbf{CA}^k\mathbf{B}$  entries are collected to create the SSM kernel  $\mathbf{K}\in \mathbb{R}^{L}$ , and the convolution could be expressed as:
$$
\mathbf {K} = (\mathbf {C B}, \mathbf {C A B}, \dots , \mathbf {C A} ^ {L - 1} \mathbf {B}),
$$
$$
y _ {k} = \sum_ {j = 0} ^ {k} \mathbf {K} _ {j} \cdot u _ {k - j}, \quad y = \mathbf {K} * u. \tag {3}
$$
Given an input sequence  $u \in \mathbb{R}^L$ , it is possible to compute the output  $y \in \mathbb{R}^L$  sequentially through the recurrence in Equation (1). While this property is useful for autoregressive decoding, sequential computation is prohibitively slow to train with long inputs and, instead, the convolution from the Equation (3) can be used to compute all elements of  $y$  in parallel. This is done via Fast Fourier Transform (FFT) [7], provided we have already computed  $\mathbf{K}$ .
Additional inductive biases have been imposed on SSMs by analytically deriving closed-form expressions for the matrices  $\mathbf{A}$  and  $\mathbf{B}$  using the HiPPO framework [12]. In this framework, the state  $x_{t}$  represents the coefficients of polynomials that approximate the sequence  $u_{t}$ .
Explicitly Parameterized Filters In contrast to structured kernels, one can parameterize the convolution kernel, as trainable weights and optimize them,  $\bar{\mathbf{K}}\in \mathbb{R}^{L}$ . However, this would result in poor performance unless certain types of regularization are applied to the kernel. [11] simply makes use of squashing the kernel weights, and subsequently applying a smoothing technique. Trainable kernels are also used in attention-free alternative models to Transformers, such as Hyena [30], which involves exponentially decaying the weights along the kernel:
$$
\bar {\mathbf {K}} _ {t} = e ^ {- \alpha t} \cdot (\text {F F N} \circ \text {P o s i t i o n a l E n c o d i n g}) (t), \tag {4}
$$
where  $\bar{\mathbf{K}}_t$  is an entry in the filter at location  $t$ , and FFN is a feed-forward network used for decoupling the parameter count from the seuqnec length.
3.2 Block-State Transformer (BST) Layer
We now introduce the Block-State Transformer layer, which combines SSMs with Block Transformers. At each training iteration, a sequence of  $L$  tokens, is sampled from a longer document. The tokens are then embedded and fed to the model. Our model consists of a stack of Block-State Transformer layers. Each BST layer optionally includes an SSM sublayer that is responsible for providing long-range context to the Block Transformer layer, which operates similarly to a Block-Recurrent Transformer (BREC) cell. The SSM sublayer takes the sequence of token embeddings from the previous layer as input, and produces a sequence of the same length  $L$  as the output.
The output of the SSM is contextually encoded, meaning that entries at every time-step, potentially include information about all the time steps preceding elements in the sequence. We collect a number of "context states,"  $S$ , from the context sequence, and we set  $S \ll L$ . In order to prevent the model from accessing future information, we only allow the model to access context states that precede the current token. Various ways to gather context states from the context sequence are discussed in section 3.3 in detail.
The context states are fed to the Block Transformer, in place of what was referred to as "recurrent state vectors" in Block-Recurrent Transformer [21]. The subsequent operations, shown on the right side of
Figure 2: Summarizing our approaches. The left side shows the cases where the SSM is required to output Multi-Head (MH) contexts. On the right Multi-Filter (MF) approach is depicted where the last entries from the previous window are concatenated into a set of context states of size  $S$ . Dashed lines represent the current block.
Figure 1, are kept unaltered, except that we no longer need to run the recurrent unit of the BRECT cell since we are maintaining the context via an SSM. In addition to the context states, the Block Transformer also receives a block/window of length  $W$  of token embeddings as input, which are cross-attended to the context states. The output of the cross-attention operation is then concatenated with that of self-attention over the input embeddings, followed by a simple projection.
In addition to the ability of SSMs to retain information over longer time horizons compared to Transformers and RNNs, using the SSM to maintain context states as a replacement for recurrent cells makes for a more computationally efficient layer. Removing recurrence by integrating SSMs into Transformer layers, allows the Block-State Transformer layer to be fully parallelizable, whereas the Block-Recurrent architecture processes blocks of tokens sequentially using a for loop.
3.3 Context States
Although the latest SSM output technically contains information about the entire sequence, retrieving individual tokens from only the final state may not be feasible. To compensate, we concatenate a sequence of states, corresponding to the latest block of tokens. This is also analogous to the approach taken by BREC. This representation ensures retrievalability and ease of access, through redundancy. It is redundant because adjacent states are highly correlated, however this also makes it possible to easily recover the current block of tokens, if necessary.
In our approach, the context states are constructed from the output of the SSM and fed to the attention heads of the Transformer. These context states can be constructed in various ways. To guide these design decisions we consider each of the below proposed schemes as introducing *retrievability* at the cost of redundancy. The shape of the output of a single SSM layer is  $(B \times L \times D)$ , where  $B$  is the batch size,  $L$  is the number of the tokens processed, and  $D$  is the embedding dimension. When doing cross-attention in the Transformer cell with  $H$  different heads, this tensor needs to be transformed into a context tensor of shape  $(B \times S \times D \times H)$ , where  $S$  is the number of context states; we usually set  $S \ll L$  and  $S = W$  similar to Block-Recurrent Transformers (BRECT).
We now discuss the three different approaches that we evaluate to generate a context tensor for each block sequence:
SH: Single-Head The first approach constructs the context tensor by sequentially concatenating the  $S$  states from the SSM with a single filter (each of size  $D$ ). Note that because the SSM captures information from preceding blocks, the context state also captures information about blocks that preceded the current block. The resulting context vector is highly retrievable and redundant, as defined above. As in typical Transformers, fully connected layers are used to project each context vector to  $H$  different heads of size  $D$ . Note that in the cross-attention operation, context states that correspond to future tokens from the current block need to be causally masked out. In this case we set  $S = W$ , and we pick the window of SSM outputs that correspond to the current block, and a
triangular mask is used to implement causal masking of context states. This approach is shown in Figure 1.
MH: Multi-Head This approach differs from Single-Head (SH) in that here the SSM is tasked to generate a separate output for different heads. We use separate  $\left[\mathbf{C}_1,\mathbf{C}_2,\dots,\mathbf{C}_H\right]$  matrices, to produce context states that are fed to the attention heads. This enables the SSM to extract complementary features from the summarized history. The conceptual difference is that the C matrix, from Equation (1), has direct access to the full memory state of the SSM  $(x_{k})$ , that in theory could be thought of as a compact representation of the history, before it gets mapped down to a scalar. The Multi-Head (MH) approach is illustrated on the left side of Figure 2. Because the  $H$  different C matrices may extract complementary information, the context vector constructed by this method is theoretically less redundant compared to the single-head method described above.
MF: Multi-Filter In this approach the SSM sublayer produces  $S$  context states, which we set to be independent from  $W$ . This is done by convolving the sequence of embeddings with  $S$  different kernels/filters. The output of each convolution operation, corresponding to a specific filter, is a tensor of shape  $(B\times L\times D)$ . After convolving the input with all the filters, the context states of size  $D$  that correspond to the last token from the previous window are stacked together to make a  $(B\times S\times D)$  tensor. Feed forward networks are then used to lift this tensor to different heads,  $(B\times S\times D\times H)$ . Different from the previous two approaches, the context is formed by taking only the last  $S$  context states, from the previous window, outputted by the  $S$  SSMs. The context is less redundant because it no longer consists of adjacent SSM states. Since the context is taken from the entries of the previous window, cross-attention masking is no longer required, as shown on the right of Figure 2.
The memory states of the Multi-Filter (MF) approach is least redundant, while Multi-Head (MH) strikes a middle ground, and Single-Head (SH) has the most redundancy. The incorporation of redundancy in these approaches aims to facilitate retrievalability of the most recent context captured by the SSM, albeit at the expense of potentially inefficient utilization of the network capacity. The last approach attains highest utilization, as the cross-attention is done in the space of unique features extracted by specialized filters.
3.4 Implementation Details
Context IDs & Positional Embedding To allow distinction between the entries supplied to the attention mechanism, a positional embedding is commonly added to the inputs. When using the Multi-Filter (MF) approach, the collected context states correspond to different features extracted from the sequence, hence we add a set of unique learned "context IDs" to the context states, before using them as input to cross-attention. However, in the cases where the context states correspond to different time-steps along the sequence, namely Single-Head (SH) and Multi-Head (MH) approaches, inherent positional encoding is incorporated into the context states, due to the incremental nature of convolutions; as such, we find the addition of context IDs to be unnecessary. We also realize that we do not need to add global positional bias to the token embeddings, and use a T5-style relative position bias [32] instead, as the SSM does also encode positional information into the context.
Down-sampling Consistent with findings in [28], we find FFT operations to be the main source of bottleneck when training SSMs on TPUs. We project the input embeddings to a lower-dimensional space, that is a quarter of embedding size in our experiments, this reduces the required total number of FFTs by a factor of 4. The output of the SSM, i.e. the context states, are later lifted to the original embedding size before being passed to the Block Transformer.
Our results are presented in Table 1. We conduct experiments with BST on three different datasets, PG19, arXiv and GitHub, allowing us to test our method on a suite of varying documents lengths composed of English texts, latex scientific articles and source code.
PG19 dataset is from a large collection of full-length books from Project Gutenberg [31]. All extracted 28,602 books were published prior to 1919 and contain 6,966,499 English language words. When tokenized, each PG19 book has between 50k-100k tokens. PG19 has become a popular