SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Perplexity of each model. The results for XL:2048, SLIDE:12L and BRECT:FIXED:SKIP are from [21] by converting  $\log_2$  of perplexity to raw perplexity. GSS-HYBRID-L performance was taken from [28]. Results with  $\pm$  are average scores and error bars of runs with three different random seeds. For a smaller computational budget, BST provides a small perplexity improvement compared to BRECT on PG19 and GitHub. For the same computational budget, BST outperforms GSS-HYBRID-L across datasets by  $1.5\%$  to  $4\%$ .
Model	eval seq. length	window length	number params	TPUv4 hours (k) PG19/arXiv/GitHub	PG19	arXiv	GitHub
SLIDE:12L	4096	512	190M	0.5 / 0.5 / 1.8	12.12	2.69	2.28
TRSF-XL:2048	2048	2048	190M	0.8 / 0.8 / 3.0	11.96	2.48	2.01
BRECT:FIXED:SKIP	4096	512	196M	0.8 / 0.8 / 3.0	11.55 ±1.1	2.36	2.04
BST:SH:S4	202M	0.5 / 0.5 / 1.8	11.57 ±1.1	2.51	2.14
BST:MH:S4	218M	0.8 / 0.8 / 1.8	11.60 ±1.1	2.52	2.15
BST:MF:S4	217M	0.5 / 0.5 / 1.8	11.63 ±1.2	2.48	2.07
BST:SH:UNSTRUCT	206M	0.5 / 0.5 / 1.8	11.52 ±1.1	2.49	2.09
BST:MF:UNSTRUCT	221M	0.5 / 0.5 / 1.8	11.56 ±1.2	2.44	2.03
GSS-HYBRID-L	4096	512	373M	0.8 / 0.8 / 1.8	10.52	2.51	1.88
BST:SH:S4-L	366M	0.8 / 0.8 / 1.8	10.47	2.49	1.86
BST:MF:S4-L	383M	0.8 / 0.8 / 1.8	10.52	2.46	1.84
BST:SH:UNSTRUCT-L	371M	0.8 / 0.8 / 1.8	10.37	2.46	1.85
BST:MF:UNSTRUCT-L	388M	0.8 / 0.8 / 1.8	10.42	2.41	1.83
benchmark for measuring progress on long-range language modeling performance. We report the "test" split evaluation performance.
arXiv dataset is a corpus containing scientific and technical articles on the subject of Mathematics [42]. The arXiv dataset contains latex source code as well as items such as theorems, citations, definitions that are referenced and discussed over long ranges of text. Using the same vocabulary as in [42] and [21] for a fair comparison, many special characters are broken up into small subwords. As a result, the number of tokens per paper in the arXiv dataset is approximately equal to the number of tokens per book in PG19. We report perplexity on "test" split.
GitHub dataset [42] is the largest of the three datasets and was assembled by extracting GitHub code repositories with open-source licences. Files were filtered to only contain the following programming languages: C, C++, Java, Python, Go and Typescript. While code files are relatively small, there are many import dependencies between each file. By traversing the directory tree and concatenating all code files along the path, a single document that preserves a repository's structure and dependencies is created. We report performance on the "validation" split.
For a fair comparison with the baselines, we keep the vocabularies consistent as used by [21] and [28]. Specifically, we used a pretrained T5 vocab with 32k tokens for PG19 [33] and LaMDA vocab with 32k tokens [39] for both arXiv and GitHub datasets. Due to the long training times and large number of experiments, we only provide error bars for the PG19  $\sim$ 200M parameter models by running our models with three different random seeds. BRECT:FIXED:SKIP error bars are from [21].
4.1 Comparing our Baselines and Models
We experiment three different types Block-State Transformer (BST) models: BST-SH, BST-MH and BST-MF as described in Section 3.3. Our models do not use global learned positional embeddings but encode positional awareness with an SSM at the first layer, right after the word embedding layer. We organize models into two groups: (i) fixed window size have either a 512 or a 2048 token training window size; and (ii) fixed parameter count have either a  $\sim 200\mathrm{M}$  or  $\sim 400\mathrm{M}$  total parameters. We run experiments with two types of SSMs:
BST:  $\{\mathrm{SH},\mathrm{MH},\mathrm{MF}\} :S4$  encode long context using a Structured State Space Model (S4) [16]. As described in Equation (3), S4 kernel matrix  $\mathbf{K}$  is compiled from matrices  $\mathbf{A}$ ,  $\mathbf{B}$  and  $\mathbf{C}$  and is independent of the length of the input evaluation sequence length. We show that the structured parameterization of  $\mathbf{K}$  allows our BST models to generalize to longer lengths. We refer the reader to
section 4.2 for results on length generalization. We only run one BST:MH using S4 since the model requires  $8\%$  more parameters while performing on par with the faster BST:SH variant. BST:MF also has  $8\%$  more parameters but performs better on arXiv and GitHub compared to SH. Interestingly, SH performs better than MF on the PG19, a dataset where local context is more important to predict the next token compared to arXiv and GitHub. We posit that this is likely due to the ability of the SH model to retrieve the most recent context captured by the SSM.
BST:  $\{\mathrm{SH,MF}\}:\mathrm{UNSTRUCT}$  are based on unstructured parameterized convolution filters, inspired by the Hyena Hierarchies [30] convolutional kernel. We exclude the utilization of the multiplicative gating mechanism employed in Hyena Hierarchies and solely apply the regularizations implemented on the parameterized kernel, denoted as  $\bar{\mathbf{K}}$  in Equation (4). This formulation has two important advantages over S4: (1) the  $\bar{\mathbf{K}}$  kernel does not need to be recompiled, allowing speedups when using multiple filters; (2)  $\bar{\mathbf{K}}$  has more free parameters because it is no longer restricted by A, B matrices in equation 3, potentially providing richer representations that can explain the improved perplexity scores over S4 variants. Nonetheless, UNSTRUCT kernel  $\bar{\mathbf{K}}$  relies on learned positional encoding which makes the method less extendable to larger length sequences at inference..
We compare the Block-State Transformer to four different baselines:
TRSF-XL:2048 [8] is a Transformer with a training window size of 2048. As expected, increasing the window size improves perplexity, especially on the arXiv and GitHub datasets. However, this model performs worse than BST:SH:HYENA on PG19 and is much slower, bottlenecked by the attention layer on higher sequence lengths.
SLIDE:12L [21] This model is almost identical to TRSF-XL:2048. It uses however a sliding window of size 512 over a segment of 4096 tokens. The sliding window is differentiable over two blocks, while TRSF-XL does not backpropagate through the cached keys and values from the previous window. This simple baseline is closest in terms of training speed to BST:SH. The perplexity scores show that integrating a representation of the past, as with BREC and BST, positively impacts LM performance.
BRECT:FIXED:SKIP [21] is the strongest performing and fastest Block-Recurrent Transformer architecture in [21]. This architecture is very similar to SlIDE:12L. There is however a sequential recurrent "skip" configuration, a simple linear layer gating mechanism that combines current block hidden representation with past information from the previous blocks.
GSS-HYBRID-L [28] is the closest SSM-Transformer hybrid model that was tested on long-range language modeling tasks. GSS-HYBRID-L is based on the Diagonal State Space (DSS) [16]. DSS and S4 are similar in performance and architecture, only differing on the initialization of the kernel K [15]. [16] further improves on DSS for LM tasks by introducing a Gated State Space version called GSS, which performs better on PG19, arXiv and GitHub. Unlike our method, GSS-HYBRID-L does not directly integrate SSMs states into the attention mechanism but only interleaves 32 GSS layers with Transformer layers. It must be noted that the GSS-HYBRID-L scores were obtained after grid searching over four learning rates  $\{6.4,3.2,1.6,0.8\} \times 10^{-3}$  and used a different learning rate and weight decay for the SSM layer and the Transformer layer to avoid training instabilities. In our experiment, we did not use grid search and used the same learning rate for all layers. BST results demonstrate that integrating SSM states into the Transformer attention provides larger benefits than interleaving SSM and attention layers as in GSS-HYBRID-L.
Fixed compute budget. As seen in Table 1, we track the exact amount of compute in TPUv4 hours that was spent training each model. The training TPUv4 hours for SLIDE:12L, TRSF-XL:2048, BRECT:FIXED:SKIP and GSS-HYBRID-L were taken from [28]. The TPUv4 hours metric measures the compute cost of training models. For our experiments, we align our training times with GSS-HYBRID-L for a fair comparison. Smaller parameter models all have 12 layers, 8 heads of size 128, embedding vectors of size 1024, an MLP with a hidden layer size of 4096 with ReLU activation functions. For larger BST models, we double the intermediate layer size from 4096 to 8192 and increase the number of attention heads to 12.
Training details We use the same training setup as [21] and we perform our experiments using the Meliad library in JAX/Flax [1, 17]. We use the Adam optimizer [25] and a batch size of 32
<sup>3</sup>https://github.com/google-research/meliad
and a sequence length  $L$  of 4k for training. Using a structured SSM's recurrence (such as S4) in the first layer allows us to extend the positional encoding to various lengths at inference. Smaller BST models have Block-State layer integrated in Transformer layers  $\{1, 7, 9\}$  and larger BST models at layers  $\{1, 5, 7, 9\}$ . Since our datasets contain long documents, it is possible to train on larger sequence lengths  $L$ . Training on 4k sequence lengths allows us to test length generalization since the convolution kernel  $\mathbf{K}$  in Equation (3) can be extended to any sequence length  $L$ . However, since we show in Section 4.2 that our model works well when extended to unseen lengths, we did not find it necessary to run expensive experiments with higher sequence lengths. For the MF model variants, we lower the SSM state dimension  $D$  by an additional factor of two to improve FFT efficiency. The state dimension reduction has negligible impact to perplexity. The MF models have  $S = 32$  filters while the larger MF models have  $S = 64$  filters.
4.2 Evaluating Length Generalization capabilities
We present our length generalization analysis and report perplexity in Figure 3. Our models and baselines all have  $\sim 400\mathrm{M}$  parameters, are trained on a sequence length of  $4\mathrm{k}$  and tested on sequences with lower and higher sequence lengths of  $\{512, 16\mathrm{k}, 65\mathrm{k}\}$ .
We notice that all models have similar perplexity for sequence lengths of 512. Both BST:SH:S4-L and GSS-HYBRID-L generalize well on 16k and 65k sequence lengths for PG19 and GitHub. For arXiv, GSS-HYBRID-L and BST:MF:UNSTRUCT-L perplexities increase drastically, potentially due to noise in the arXiv dataset (as indicated by variation in perplexity metric over time). [28] also reported that larger GSS models had difficulty generalizing to higher lengths. Interestingly, for arXiv again, BRECT:FIXED:SKIP-L performs very well at higher sequence lengths. We hypothesize that the Block-Recurrent model's access to the entire past during training, via a non-differentiable cache of representations across sequences, helps retain a "memory" of dependencies between key items in an arXiv article allowing the model to access past symbols, definitions, theorems or equations beyond the 4k training sequence length. We also note that BST:MF:UNSTRUCT-L and BRECT:FIXED:SKIP-L outperform other methods on PG19 up to a sequence length of 16K. Perplexity performance on PG19 is perhaps less reliant on long term relationships between tokens, which can explain the performance of models that have no explicit built-in mechanisms for length generalization.
The analysis also allows us to draw a clear distinction between structured and unstructured SSMs integrated in hybrid architectures. As previously mentioned in Section 3.1, SSMs such as DSS and S4 use a structured kernel  $\mathbf{K}$ , built from learned matrices  $\mathbf{A}$ ,  $\mathbf{B}$  and  $\mathbf{C}$  for any sequence length  $L$  in Equation 3. Since  $\mathbf{K}$  is extendable to any arbitrary sequence length  $L$ , both BST:SH:S4-L and GSS-HYBRID-L have a build-in mechanism for length generalization that the unstructured BST:MF:UNSTRUCT-L model does not. BST:MF:UNSTRUCT-L performs best on the training sequence of 4K and is on-par for 512 with perplexity increasing for unseen 16K and 65K sequence lengths. BST:SH:S4-L has by far the best perplexity for 65K sequence lengths on PG19, GitHub and arXiv. Similarly to [21], we also notice that perplexity improves when we extend context window (sequence length) for PG19 and GitHub.
Figure 3: Length Generalization for sequence lengths {512, 16k, 65k} on PG19 (left), arXiv (middle) and GitHub (right). BST:SH:S4-L generalizes better than other baselines, including GSS-HYBRID-L that uses GSS, a structured SSM. GSS-HYBRID-L numbers are from [28].
Figure 4: Left: The forward-pass computation time of a BST layer is compared against a layer of BRECT and SLIDE:12L. These experiments were executed on GPU, to demonstrate and exploit the parallelizability of BST layers. BST:SH is  $6 - 11\times$  faster than BRECT while BST:MH is  $3 - 4\times$  faster. Right: Perplexity of the trained models using different window lengths. The figure shows that increasing the training window length results, as expected, in better perplexity scores. We find however that both BST:MF:HYENA and BRECT:FIXED:SKIP are the least impacted by decreasing window lengths.
4.3 Efficiency
The improvement over Block-Recurrent Transformers, with time complexity of  $\mathcal{O}((W^2 + S^2 + 2SW) \cdot L / W) \approx \mathcal{O}(L \cdot W)$ , follows from the ability to run the Block Transformer's cells in parallel. The time complexity of the Block-State Transformer layer is comprised of the time complexity of the state space model sublayer,  $\mathcal{O}(D \cdot L \log L)$ , in addition to the time complexity required to execute the Transformer over the given context chunks (blocks) in parallel,  $\mathcal{O}(W^2)$ .
In spite of the superlinear growth of the SSM sublayer, our experiments indicate that significant performance improvements, up to a factor of 6, remain evident for sequences as long as 65k tokens, the point at which hardware saturation began to occur. When using a structured SSM, the computational complexity is closely tied to the internal memory state size of the SSM,  $N$  – specifics may vary depending on the exact type of the SSM. We set  $N = 16$  when reporting performance. Left side of Figure 4 shows the results of benchmarking the forward-pass of a Block-State Transformer layer on GPU. Our proposed layer runs almost  $6 - 11 \times$  faster than Block-Recurrent Transformers (including recurrent units), and yields comparable performance to a SLIDE:12L layer, i.e. BRECT without the recurrence. At 4k sequence length, which is mostly used during training, BRECT layer runs almost  $15 \times$  slower than SLIDE:12L with the same window size. We manage to reduce this gap to less than  $2 \times$  with BST layer. To reflect a realistic model, for these experiments we use a fixed window length of 128, an internal state size of 16 for the SSM, and 16 heads. Moreover, to highlight the performance gains that are only due to parallelization made possible by our framework, we use same embedding size as input to the SSM, which is 512. Note that we use the vanilla implementation of FFT and inverse FFT operations provided by JAX [1]. However, we believe that the speed of our method can be further improved with recent and faster hardware-specific I/O-aware implementations introduced in other auto-diff frameworks.
We have introduced a model that combines the attention mechanism of Transformers with the long-range memory mechanism, and parallelism afforded by State Space Models. We explored several memory state variants that make different trade-offs between redundancy and retrievability. Experiments show that our model can minimize perplexity on par with and often improves upon recent competing baselines, while achieving up to more than  $10 \times$  speedups at the layer level, provided there is hardware support to fully take advantage of parallelism. This is an appealing property for scaling up BST which makes the addition of SSMs into Transformers computationally appealing. We show that integrating SSM states into the Transformer attention provides larger benefits than simply interleaving SSM and attention layers. Finally, we show that the model generalizes to longer sequences than it was trained.
We would like to thank Caglar Gulcehre and Albert Gu for helpful discussions and support with the S4 codebase. We would also like to express our gratitude to Delesley Hutchins for providing valuable guidance throughout the project, as well as Xavier Garcia and Courtney Paquette for their careful review of the manuscript, where they identified and rectified several errors.
While BST's SSM layer allows the model to unroll and parallelize the recurrence that models long-term context between blocks of tokens, the SSM variants are reliant on efficient FFT operations. We have found that the FFT operation is an important speed bottleneck on TPUs that needs to be resolved to better scale BST to many layers and larger models. While we are still investigating the reasons, we found that JAX FFT was  $4 \times$  faster on GPUs. Further, new SSM variants such as S5 [34] bypass FFT operations using a binary associative operator<sup>4</sup>. Our implementation is modular enough that we can simply plug in S5 or use other FFT implementations.
One of our assumptions is that BST's SSM layer is able to capture the right long-term dependencies for each block. The SSM recurrence at step  $T = t$  provides a summarized representation of previous steps for  $T = 0$  to  $T = t$ . However, a single vector representation may not be powerful enough to support all important long-term dependencies. Despite the perplexity improvements on long-range language modeling tasks, this assumption needs to be tested on other long-range classification tasks such as Long Range Arena [37] as well. It is possible that our model can perform better if we feed to the attention layer  $k = W$  SSM representations that are chosen by a top- $k$  retrieval operation, similar to the one in Memorizing Transformer [42].
B More detailed comparisons with existing baselines
This section provides the reader with a more in-depth comparison with similar architectures. We cover BRECT [21] in Section B.1 and GSS-HYBRID [28] in Section B.2.
B.1 Comparison with Block Recurrent Transformer (BREC)
The Block Transformer sublayer (i.e. SlIDE:12L) processes keys and values from the previous window stored in a differentiable cache. This is implemented similarly to the sliding window attention pattern suggested in [21] and was originally introduced by Transformer-XL [8]. Using a causal mask, at every token inference step, the attention mechanism is applied to blocks of tokens of size  $W$  and is partially extended to the cached keys and values from the previous block with the sliding window. BREC T, as explained in [21], uses a non-differentiable cache that is carried from one sequence of size  $L$  to the next<sup>5</sup>. The last recurrent states of a sequence are stored in a non-differentiable cache and fed to the next training step on the following sequence in the document as a warm-start. We do not pass such a representation, since to compute the output of the convolution, we need access to the whole sequence. We believe that this is one advantage that BREC T has over our method, especially for very long examples that split into ordered sequences of length  $L$ , since the cache carried from one sequence to the next can provide very useful long-range information and (weak) access to the whole past. Since we need the whole sequence to compute SSM states, history beyond  $L$  may be lost in the process. We believe that BST can further be improved by adding non-differentiable sequence cache for very long documents.
While in other architectures, the history between blocks of tokens is not modeled, both BST and BREC T use a mechanism to model previous block context. The authors of BREC T experiment with various sequential gating mechanisms to condense the information from past blocks. With BST, we use SSM to provide context from previous blocks to the current block as explained in Section 3.2.
B.2 Comparison with the Transformer GSS-HYBRID
GSS-HYBRID [28] is a SSM-Transformer hybrid architecture that we first describe in Section 4.1. The architecture is significantly different from BST. GSS-HYBRID is primarily composed of Gated State Space (GSS) layers and has a few interleaved Transformer layers at every 4th layer starting with the 2nd layer. BST on the other hand is mainly composed of Block Transformer layers and has Block-State Transformer layers at positions \{1, 7, 9\} for the  $\sim 200\mathrm{M}$  model and  $\{1, 5, 7, 9\}$  for the  $\sim 400\mathrm{M}$  model. Our hybrid does not stack SSM and Transformer layers like the GSS-HYBRID but rather replaces the recurrence in BRECT with an SSM such as S4. In BST, the SSM generates states
4In JAX, this is equivalent to using jax.lax.associative_scan.
In our work and in [21], a document is split into multiple sequences of size  $L$  and each sequence is split into multiple blocks of size  $W$
for each Block Transformer representations and we then use cross-attention to mix the states and the self-attention outputs. The authors in [28] initially built GSS, a gated version of DSS [16], to (1) reduce SSM parameter dimensions, (2) stabilize training of the SSM and (3) allow better length generalization. However, when experimenting with SSMs such as S4 or DSS, we found that the gating was not necessary to achieve all three objectives stated above. We decided that using GSS's Gated Attention Unit [20] was therefore not needed when integrating SSM states into the attention mechanism. We also reiterate that the authors in [28] used hyperparameter search to get the best performance while we did not.
C Scaling Experiments
Figure 5: Scaling properties on PG-19.
In this section, we compare how BST scales compared to Transformer-XL with  $4 \times$  the window size and BREC. In Figure 5, we see that at lower scales, from 80M to 200M, BREC and BST have very similar performances. Beyond 200M, the perplexity performance percentage gap between BREC and BST increases from  $2.5\%$  at 200M parameters to  $4.0\%$  at 1.3B parameters. The perplexity performance percentage gap between BREC and TRSF-XL is even more pronounced as it starts at  $7.6\%$  at 200M parameters to  $10.6\%$  at 1.3B parameters.
D Long Range Arena Experiments
MODEL	LISTOPTS	TEXT	RETRIEVAL	IMAGE	PATHFINDER	PATH-X	AVG
Transformer	36.37	64.27	57.46	42.44	71.40	X	53.66
Linear Trans.	16.13	65.90	53.09	42.34	75.30	X	50.46
Reformer	37.27	56.10	53.40	38.07	68.50	X	50.56
Performer	18.01	65.40	53.82	42.77	77.05	X	51.18
BigBird	36.05	64.02	59.29	40.83	74.87	X	54.17
Mega	63.14	90.43	91.25	90.44	96.01	97.98	88.21
S4D	60.47	86.18	89.46	88.19	93.06	91.95	84.89
S4	59.60	86.82	90.90	88.65	94.20	96.35	86.09
S5	62.15	89.32	91.40	88.00	95.33	98.58	87.46
Methods with chunked input sequences
BRECT:FIXED:SKIP	37.29	66.14	58.76	50.41	76.33	75.89	60.80
MEGA-CHUNK	58.76	90.19	90.97	85.80	94.41	93.81	85.66
BST:SH:S4 (ours)	61.49	87.63	90.51	91.07	95.75	95.28	86.96
Table 2: Performance on Long-Range Arena (LRA). For a fair comparison, we adjust the number of layers and model dimensions on each task so that BST and BRECT have similar number of parameters with S4 and MEGA-CHUNK. BRECT results are from our own runs and all other baselines are from published results.
While the main focus of our research was to demonstrate that hybrid Transformer-SSM models are efficient and perform well on long context autoregressive LM, we also evaluate our method on standard classification task where long range dependencies in a sequence are important to capture. In Table 2, we present our results on the Long Range Arena (LRA) benchmark [38] which incorporates three different modalities including text, images, and mathematical expressions. The LRA dataset also tests models on various sequence lengths from 1K to 16K.