Block-State Transformers
Authors: Jonathan Pilault
Mahan Fathi
Orhan Firat
Christopher Pal
Pierre-Luc Bacon
Ross Goroshin
Venue: NeurIPS 2023 poster (2023)
Keywords: State Space Models
Efficient Transformers
Long Range Language Modeling
Language Modeling
TL;DR: The Block-State Transformer combines State Space Models with attention, and outperforms and is more efficient over strong baselines on long sequences.

Abstract: 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 a more than *tenfold* increase in speed at the layer level compared to the Block-Recurrent Transformer when model parallelization is employed.

Block-State Transformers
Jonathan Pilault Mahan Fathi 123+
Orhan Firat
Christopher Pal24 Pierre-Luc Bacon 23
Ross Goroshin'
Google DeepMind
Mila
Université de Montréal
*Polytechnique Montréal
Abstract
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 plethora of tasks, in vision and audio; however, SSMs still lag Transformer performance Language Modeling tasks. 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-basee 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-Recurren Transformer when model parallelization is employed.
Introduction
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 NL.P and have mostly replaced Recurrent Neural Networks (RNNs). Several factors contribute this success, including computa- tional 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 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
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 hard NLP tasks such reasoning and question answering [2] 39 6), the demand for deploying even deeper and larger networks greater than ever before. An orthogonal scaling dimension, which could be potentially sequence. Despite the several advantages of Transformers over RNNs, it is still problematic to seale the input sequence length, again for both computational performance and quality reasons. Further, the Transformer's '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 (2). notoriously struggle long-input classification tasks 37 Vanilla transformers be unstable when trained
more consequential, is the size of the input
EEuual Contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

block output token embeddings
output embedding
Block Trunstormer Layer
block Input token token embeddings
linear
self- amt
lineut
Block TRE
Block TRE
Block TRE
V.Q
projection
MILP
linear
Cross-
Q
altn
State-Spare
linear
leamedd conteat
K,
Block Transtomer
tokes embeedding
Block-State Transformer layer. The BST-SI TH layer is illustrated on the left, and includes state space model (SSM, in green) and Block Transformers (in red). For demonstration purposes the sequence is divided into blocks in the picture. The details of the Block Transformer sublayer are on the right. *TRF Transformer.
Figure 1:
on long sequences [26 and token importance is concentrated in a local receptive field of around 50 tokens around the current time step 355
An emerging body of research suggests that State Space Models (SSMs) can
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 (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 [ 12] SSMs have recently outperformed Transformers 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
long-range dependency benchmarks by large
Recent findings suggest that Transformers and SSMs are complementary models for the purpose of language modeling (88. In this work, we propose architecture that integrates 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 process long input sequences, while still incorporating an attention mechanism to predict next tokens. BST is fully parallclizable, scales much longer sequences, and offers a 10x speedup compared comparable Transformer- based layers.
In every BST layer, SSM takes the entire sequence input and maps it into "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 then cross-attended to the corresponding block of
21 The block of input token embeddings context states; see Figure [ Note that by introducing SSMs as a means of contextualization. we completely remove the need for sequential recurrences and we are able run our hybrid SSM- Transformer layer fully in parallel. The resulting runtime complexity can be expressed as the sum of ((W2) + OLL 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 major improvement over (LL) of computation is available. Moreover, due to hardware imposed restrictions, the runtime complexity of SSM on a full sequence comparable that Block Transformer on block 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 evident by observing that the bottom-moo two lines on the left of Figure are almost overlapping.
Transformer, so long as hardware support parallel

---

2
Related Work
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 range dependencies, beyond the length 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 (BRECT) [21] uses recurrent memory mechanism to extend the theoretical context length of the Transformer. In the recurrent unit of the BRECT 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, maintain context states. We also conduct a more extensive exploration of maintaining and updating context states. Earlier works that augment transformers with non-differentiable external memory include the Memorizing Transformer 42). Transformer-Xe S was an early work that combined recurrent memory with Transformers. Our work can be
linear state space models instead of recurrent cells to
continued evolution of those models incorporating state-of-the-ar recurrent memory models inspired by SSMs.
State space models can be considered as linear RNNs 12] 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, cnabling Transformers to leverage theoretically infinite context.
This simplicity facilitates their analysis
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 language. Nevertheless, SSMs are able to capture longer range dependencies than Transformers in both theory and practice, while also being highly parallelizable
1 We therefore elect to combine the best aspects of SSMs and Transformers into single model. The idea of communication across blocks, similar GSS (28), was later implemented by MEGA [27., through an Exponentially Moving Average (EMA) update rule instead of SSM However, both GSS and MEGA use 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 rross-attention, allowing to plug any out-of-the--ool SSMs or attention layers. Further, we investigate three ways to mix SSM signals with attention outlined in Section 3.3
3 Method
We consider the problem of next token prediction via decoder-only language model. This seemingly simple pretext task has led to spectacular progress in language understanding 19. 2 29]. During training, the decoder takes in a sequence of length L of tokens embeddings and tasked generate the next token at every step in the sequence.
We start by brief review of SSMs that are essential for understanding the Block-State Transformer layer (3.T). Our full Block-State Transformer architecture is outlined in Section 3.22 Section 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 structured
initialization of the convolutional kernel by unrolling linear time-invariant (LTT) dynamical system of the following form:
Zs' Az Bus
(I)
y/s Ca + Dus
"The authors in (27 show mathematical form of EMA that has state transition and also derive convolution kernel efficiently compute EMA similarly to $4.
3

---

The system is parameterized by a state matrix A E RNXN, vectors B E RNXI, CERIXN and D RI the SSM maps 1-D input signal ug, to 1-D output signal ys- Internally, the SSM projects the input signal to N-D representation state the C matrix. The term Dug 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, yr- can be computed as a discrete convolution, by realizing that the recurrence can be explicitly unrolled:
before mapping it down to scalar using
Let
I- == 0
Σ CA'B uk-j.
Uk
==0
The CA*B entries are collected to create the SSM kernel K E R2, and the convolution could be expressed as:
(CB, CAB, ... CA BB)
& Σ 2=
yk =
K U = K u
Given input sequence u E R2, possible to compute the output y E R4 sequentially through the recurrence in Equation (T. 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 in parallel. This is done via Fast Fourier Transform (FFT) [7]. provided we have already computed K.
Additional inductive biases have been imposed on SSMs by analytically deriving closed-form expressions for the matrices A and B using the HiPPO framework [12]. In this framework, the state H, represents the coefficients of polynomials that approximate the sequence Ut-
Explicitly Parameterized Filters In contrast structured kernels, one can parameterize the convolution kernel, as trainable weights and optimize them, K RR- However, this would result in poor performance unless certain types of regularization are applied to the kernel. [1]] simply makes use of squashing the kernel weights, and subsequently applying a smoothing technique. Trainable kernels are also used in attention-fie free alternative models to Transformers, such Hyena (30), which involves exponentially decaying the weights along the kernel:
K, -at (FFN PositionalEnccoinng\tt
(4)
where K, entry in the filter at location and FFN is feed-forward network used for decoupling the parameter count from the seuqnece 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 tokens, is sampled from longer document. The tokens are then embedded and fed the model. Our model consists of 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 operate similarly (BRECT) cell. The SSM sublayer takes the sequence of token embeddings from the previous layer as input, and produces sequence of the same length L as the output.
Block-Recurrent Transformer
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 number of "context states," S, from the context sequence, and we set S & 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 fed to the Block Transformer, in place of what was referred to as "recurrent state Transformer 211. The subsequent operations, shown on the right side of
vectors" in

---

Block TRF
Block TRE
Black TRF
Block TRF
TRRE
TRF
w
H
CONTEST
State-Spice Model
State-Space Model
isput token embelkinng
input token enbeldding
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 [ 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 then concatenated with that of self-attention over the input embeddings, followed by simple projection.
In addition the ability of SSMs 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 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 BRECT This representation ensures retrievability 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 introducing retrievability at the cost of redundancy. The shape of the output of single SSM layer is (B x L x D), where is the batch size, T, 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 context tensor of shape (B S x D x H), where S is the number of context states; we usually set S < L and S = W similar to
Transformers (BRECT).
We now discuss the three different approaches that we evaluate to generate 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 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. this case we set S W, and we pick the window of SSM outputs that correspond to the current block, and
5

---

triangular mask is used to implement causal masking of context states. This approach is shown in Figure
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 [C1. C2.... CH] matrices, to produce context states that are fed to the attention heads. This enables the SSM to extract comple- mentary features from the summarized history. The conceptual difference is that the C matrix, from Equation [ has direct access to the full memory state of the SSM (rx),, 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
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 done by convolving the sequence of embeddings with different kernel/filters. The output of each convolution operation, corresponding to specific filter, is a tensor of shape (B X 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 (B x S x D) tensor. Feed forward networks are then used to lift this tensor to different heads, (B x D x 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 no longer consists of adjacent SSM states. Since the context is taken from the entries of the previous window, cross- attsntattnn 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 retrievability 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 cross-attention. However, in the cases where the context states correspond to different time-steps along the sequence, namely Single- (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 add global positional bias to the token embeddings, and use a T5-style relative 32 instead, as the SSM does also encode positional information into the context.
position bias
Down-sampling Consistent with findings in (28], we find FFT operations to be the main source of bottleneck when training SSMs TPUs. We project the input embeddings to oower-dimensional space, that is 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.
4 Results
Our results are presented in Table I We conduct experiments with BST on three different datasets, PG19, arXiv and GitHub, allowing to test our method suite of varying documents lengths composed of English texts, latex scientific articles and source code.
PG19 dataset is from 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 popular

---

Table 1: Perplexity each model. The results for XL:2048, SLIDE:12L and BRECTHXEDD:SII are by converting log2 of perplexity to raw perplexity. GSS-HYBRID-LL performance was taken Results with + are average scores and error bars of runs with three different random seeds. For smaller computational budget, BST provides small perplexity improvement compared to BRECT on PG19 and GitHub. For the same computational budget, BST outperforms GSS- -SS-HYLRID--L
from
21
from 28
across datasets by 1.5% to 4%.
eval seq. window number params
TPUv4 hours (k)
PG19
arXiv
GitHub
Model
length
length
PG19/rcii//itHHbb
SLIDE:12L
4096
512
190M
0.5 / 0.5 1.8
12.12
2.69
2.28
TRSP-XL:2048
2048
2048
190M
0.8 / 0.8 3.0
11.96
2.48
2.01
BREDSSAE
4096
512
196M
0.8 / 3.0
11.55 +1.1
2.36
2.04
BST:SH::4
202M
0.5 /
11.57
+1.1
2.51
2.14
BST:MH:S4
218M
0.8 1.8
11.60 #1.1
2.52
2.15
BST:MF:S4
217M
0.5
11.63
#1.2
2.48
2.07
BST:SH: UNSTRUCT
206M
0.5
11.52
#1.1
2.49
2.09
BSNSMU:
221M
0.5/
11.56 +1.2
2.44
2.03
GSS-HYBRID-A
4096
512
373M
0.8
10.52
2.51
1.88
BST-
366M
0.8 0.8
10.47
2.49
1.86
BST:MF:S4-L
383M
0.8 1.8
10.52
2.46
1.84
BST:SIE:USTTRUCTL BST..F::USTTRUCTL
371
0.8 /
10.37
2.46
1.85
388M
0.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 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 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 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 many import dependencies between each file. By traversing the directory tree and concatenating all code files along the path, single document that preserves repository's 's structure and dependencies is created. We report performance on the "validation" split.
For fair comparison with the baselines, we keep the vocabularies consistent as used by (2] 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 200M parameter models by running our models with three different random seeds. BRE error bars from 21
4.1 Comparing our Baselines and Models
We experiment three different types Block-State Transformer (BST) models: BST-SH, BST- and BST-MF as described Section 3.3 Our models do use global Iearned 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 ~200 or 4-00M total parameters. We run experiments with two types of SSMs:
BST: MF): encode long context using a Structured State Space Model (S4) [6]. As described in Equation 33) $4 kernel matrix K is compiled from matrices A, B and C and is independent of the length of the input evaluation sequence length. We show that the structured parameterization of K allows our BST models to generalize to longer lengths. We refer the reader to
7

---

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, dataset where local context is more important to predict the next token compared to Xiviv 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.
are based of unstructured parameterized convolution filters, inspired by the Hyena Hierarchies (00 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 K in Equation 44) This formulation has two important advantages over $4: (1) the K kernel does not need to be recompiled, allowing speedups when using multiple filters; (2) 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 K relies 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:
ReSF-XL:2048 g is 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: on PG19 and much slower, bottlenecked by the attention layer on higher sequence lengths.
SLIDE: 12L [11] This model is almost identical to TRSF-XL:2048. It uses however a sliding window of size 512 over segment of 4096 tokens. The sliding window is differentiable over two blocks, while TRSF-XI XL does not backpropagate through the cached keys and values from the previous window. This simple baseline closest in terms of training speed to BST:SH. The perplexity scores show that integrating a representation of the past, as with BRECT and BST, positively impacts LM performance.
BRECT:AXED:SKIP 21 is the strongest performing and fastest
Transformer ar- chitecture in (11] This architecture is very similar to SLIDE: 1 21.. There is however sequential recurrent "skip" configuration, 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 S.MM-Trasffomer hybrid model that was tested on long-range language modeling tasks. GSS-HYBRIDS is based the Diagonal State Space (DSS) [16]. DSS and $4 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 Gated State Space version called GSS, which performs better on PG19, arXiv and GitHub. Unlike our method, GSS- -SS-HH 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 BSS-HYBRID-L, scores were obtained after grid searching over four learning rates (6.4.3.2.1.6,0.33 x 10 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:: 21., TRS2-XI.:2048, BRECT.HIXE::SA. and ESS-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 -SS-IDBRID- L for a fair comparison. Smaller parameter models all have 12 layers, heads of size 128, embedding vectors of size 1024, an MLP with 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 experiments using the Meliad library in JAX/Flax 1 [7 We use the Adam optimizer (25 and a batch size of 32
8

---

and 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 possible to train on larger sequence lengths L. Training on 4k sequence lengths allows us to test length generalization since the convolution kernel 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 additional factor of two 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.
at layers (1, 5, 7, 9). Since our datasets contain long documents,
4.2 Evaluating Length Generalization capabilities
We present our length generalization analysis and report perplexity Figure 3 Our models and baselines all have 4400M parameters, are trained on sequence length of 4k and tested on sequences with lower and higher sequence lengths of (512, 16k, 65k).
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 L and SSSSSRF.USTTUCC-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:AXEDD:SKI--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.MFFNNTTRLIIT and 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 $4 use a structured kernel K, built from learned matrices A, and C for any sequence length L, in Equation Since K is extendable to any arbitrary sequence length J., both BST:SH:S4- and GYRI-LL have build-in mechanism for length generalization that the unstructured BST.MFF:USTTRCCTLL model does not. BST:MF::NSTRRUCT 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 PG19, GitHub and arXiv. Similarly to (21]. we also notice that perplexity improves when we extend context window (sequence length) for PG19 and GitHub.
BST:5H:54-L
BST.MPunntructt
GSS Hybrid
BRocTrriked:kipp
0.000
GISADD
Sequence Length
Sequence Length
Sequence Length
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. SSS--YBRID-LL numbers are from

---

Benchmarking Block-State Transformer
Perplexity by Window Length
SRecT: skip
S4T:MI:54
(ours)
13.0
MFT Mn:uuttuuct
BST=SH:54
S4::SH:54 (ours)
Stide 122L.
-lybril
S1ide:12L
3.0x
Rec: fixed: skip
6.6x
100
4.0x
3.0k
Sequance Langth
Window Langth
The forward-pass pass computation time of BST layer is compared against layer BRECT and SLIDE: 12I.. These experiments were executed on GPU, to demonstrate and exploit the parallelizability of BST layers. BST:SH is 6-11x x faster than BRECT while BST:MH is 3-4x 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 AST:MF:HYEN: and BRECT. the least impacted by decreasing window lengths.
Figure 4: Left:
4.3 Efficiency
The improvement over
Transformers, with time complexity of O(W? + 52 2SW) - L/W) = O(L. 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, O(D L log L), in addition to the time complexity required to execute the Transformer over the given context chunks (blocks) in parallel, O(W)).
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 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 F shows the results of benchmarking the forward-pass pass of a Block-State Transformer layer on GPU. Our proposed layer runs almost 6-1 X faster than Recurree 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 15x slower than SLIDE: 12L with the same window size. We manage to reduce this gap to less than 2x with BST layer. To reflect a realistic model. for these experiments 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 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 [I]. However, we believe that the speed of our method can be further improved with recent and faster hardware-speciffe 1/O-aware implementations introduced in other auto-diff frameworks.
use fixed window length
5 Conclusion
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 more than 10x 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 was trained.
10

---

Acknowledgments
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.
References
[1] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NummP programs, 2018.
[2] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. CoRR, abs/2005. 14165, 2020.
[3] T.S. Chihara. An Introduction to Orthogonal Polynomials. Dover Books Mathematics. Dover Publications, 2011.
[4] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904. 2019.
[5] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Be- langer, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. CoRR, abs/2009. 14794, 2020.
[6] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du Ben Hutchinson, Reiner Pope, James Bradbury. Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang. Brennan Sacta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-HHllstern Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022.
[7] James Cooley and John Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of Computation, 19(90):2973311, 1965.
[8] Zihang Dai, Zhilin Yang, Yiming Yang. Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov, Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint 1901.0090 2019.
[9] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/l ftbs/1810.. 810.04805, 2018.
[10] Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023.
[11] Daniel Y. Fu, Elliot L. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efficient long convolutions for sequence modeling, 2023.

---

[12] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Re. Hippo: Recurrent memory with optimal polynomial projections, 2020.
[13] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[14] Albert Gu, Karan Goel, and Christopher Ré, Efficiently modeling long sequences with structured state spaces, 2022.
[15] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models, 2022.
[16] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces, 2022.
[17] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2023.
[18) Sepp Hochreiter. The vanishing gradient problem during learning recurrent neural nets and problem solutions. International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems, 6(2):107-116 1998.
[19] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation. 1755:1735-1780, 1997.
[20] Weizhe Hua. Zihang Dai. Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 9099-9117. PMLR, 17-23 Jul 2022.
[21] DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block- recurrent transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[22] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5156-5166 5165. PMLR, 13-18 Jul 2020.
[23] Urvashi Khandelwal, He He, Peng Qi, and Dan Jurafsky. Sharp nearby, fuzzy far away: How neural language models use context. CoRR, 005/1805.04622 2018
[24] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), San Diega, CA, USA, 2015.
[25] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, ICLR (Poster), 2015.
[26] Conglong Li, Minjia Zhang, and Yuxiong He. The stability-effccinnyy dilemma: Investigating sequence length warmup for training GPT models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[27] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig. Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention, 2023.
(28] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur, Long range language modeling via gated state spaces. In The Eleventh International Conference on Learning Representations, 2023.
[29] OpenAl. Gpt-4 technical report, 2023,
12

---

[30] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models, 2023.
[31] Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling. In 8th International Confer- ence on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[32] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. CoRR, abs/1910. 10683, 2019.
[33] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21:140:1-140:67, 2020.
[34] Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling, 2023.
[35] Sandeep Subramanian, Ronan Collobert, Marc' Aurelio Ranzato, and Y-Lan Boureau. Multi- scale transformer language models. CoRR, abs/2005.0053 00581, 2020.
[36] Yi Tay, Dara Bahri, Liu Yang. Donald Metzler, and Da-Cheng Juan. Sparse Sinkhorn attention. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pages 9438-9447. PMLR, 13-18 Jul 2020.
[37] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers, 2020.
[38] Yi Tay, Mostafa Dehghnni, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In International Conference on Learning Representations, 2021.
[39] Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du YaGuang Li, Hongrae Lee, Huaixiu Steven Zheng, Amin Ghafouri, Marcelo Menegali, Yanping Huang. Maxim Krikun, Dmitry Lepikhin, James Qin, Dehao Chen, Yuanzhong Xu, Zhifeng Chen, Adam Roberts, Maarten Bosma, Yanqi Zhou, Chung-Ching Chang. Igor Krivokon, Will Rusch, Mare Pickett. Kathleen S. Meier-Hellstern, Meredith Ringel Morris, Tulsee Doshi, Renelito Delos Santos, Toju Duke, Johnny Soraker, Ben Zevenbergen, Vinodkumar Prabhakaran, Mark Diaz, Ben Hutchinson, Kristen Olson, Alejandra Molina, Erin Hoffman-John, Josh Lee, Lora Aroyo, Ravi Rajakumar, Alena Butryna, Matthew Lamm, Viktoriya Kuzmina, Joe Fenton, Aaron Cohen, Rachel Bernstein, Ray Kurzweil, Blaise Aguera-Arcas Claire Cui, Marian Croak, Ed H. Chi, and Quoc Le. Lamda: Language models for dialog applications. CoRR, abs/2201. .08239, 2022.
[40] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[41] Sinong Wang. Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04766. 2020.
[42] Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers, In International Conference on Learning Representations, 2022."
[43] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33, 2020.
13

---

A Limitations
While BST's SSM layer allows the model 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 4x faster on GPUs. Further, new SSM variants such as SS [34] bypass FFT operations using a binary associative operato 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 step T 6 provides 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 we feed to the attention layer k = W SSM representations that are chosen by 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 2T in Section B.I and GSS-HYBRID 28 in Section B.2
B.1
Comparison with Block Recurrent Transformer (BRECT)
The Block Transformer sublayer (i.e SLIDE: 12L) processes keys and values from the previous window stored in differentiable cache. This is implemented similarly to the sliding window attention pattern suggested in [21] and was originally introduced by Transformer-Xi [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. BRECT, as explained in [21], uses a non non-differentiable cache that is carried from one sequence of size L to the nex The last recurrent states of a sequence are stored in a non -on-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 BRECT 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 improved by adding non-differentiable sequence cache for very long documents.
While other architectures, the history between blocks of tokens is not modeled, both BST and BRECT use mechanism to model previous block context. The authors of BRECT 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
28 is a SSM-Transformer hybrid architecture that we first describe in Section 4.1 The architecture is significantly different from BST. GSS-HYARID 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 ~200M model and {1, 5, 7, 9) for the 4400M model. Our hybrid does not stack SSM and Transformer layers like the GSS-HYBA HYBRID but rather replaces the recurrence in BRECT with an SSM such as $4. In BST, the SSM generates states
GSS-HYBRID
'In JAX, this is equivalent to using jar. asw.ciative scon. "In our work and in (21). document is split into multiple sequences of size L and each sequence is split into multiple blocks of size W
14

---

for each Block Transformer representations and we then use cross cross- attention to mix the states and the self-attention outputs. The authors in [28] initially built GSS, a gated version of DSS [T6], 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
88T.aHunstud
Figure 5:
Scaling properties PG-19.
14
Yellow: (BST:SH:UNSTRUCT)
12-layer Block-State Transformer.
Nepott
Red: (REC:HXED:SKIIP)
12-layer Block-Recurrent Transformer.
Blue: (TRSF-XL XL-2048))
230M
13-layer Transformer-Xi
Number parers
In this section, we compare how BST scales compared to Transformer-Xi XI. with 4x the window size that lower scales, from 80M to 200M, BRECT and BST have very
and BRECT. In F similar performances. Beyond 200M, the perplexity performance percentage gap between BRECT and BST increases from 2.5% at 200M paramaters 4.0% 1.3B parameters. The perplexity performance percentage gap between BRECT and TRSF-XX XL is even more pronounced as it starts at 7.6% at 200M parameters 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
50.46 X
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
$4D
60.47
86.18
89.46
88.19
93.06
91.95
84.89
$4
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:ANED:SKII
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::4 (ours)
61.49
87.63
90.51
91.07
95.75
95.28
86.96
Table Performance on Long-Range Arena (LRRA). For a fair comparison, we adjust number 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.
15

---

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 sequence are important to capture. In T we present our results on the Long Range Arena (LRA) benchmark [88] which incorporates three different modalities including text, images, and mathematical expressions. The LRA dataset also tests models on various sequence lengths from IK to 16K.
SS is composed of four BST layers (no BRT layers are interlcaved) and two S4 layers on top. We use the same standard block length of 512 for BST and BRT. However, we train BST and BRT on the full sequences (up to 16K for Path-X). We use AdamW as our optimizer [24] with warmup for the learning rate, where we start from value of le-T and increase the learning rate linearly up specified value {le 2e 4e 3 for the first 10% of training. This is followed by cosine annealing for the rest of training down to value of le-7 All layers are bidirectional, including the $4 layer in BST:SH:S4 as described in 13] Our weight decay is chosen from (0, 0.05, 0.1, 0.15) and our dropout is chosen from [0, 0.1) Except for Path-X experiments, we use weight decays E {0.03, 0.05, 0.07] for all parameters except S4D matrices A and B. Also, for Path-X, the initialization range our discretization time step for PathX is decreased from (Amins A Amax)= (0.001, 0.1) to (Amin. Amax) (0.0001, 0.01).
Our results on LRA are very promissing and show that, compared to other state-of the art methods that chunk sequences into blocks, BST is able to model long range dependencies. For example, BST outperforms MEGA-CHUNK [27] on four out of six LRA tasks and by 1.5% on the average score. However, BST still needs to improve (perhaps by extending the block size) to catch up to MEGA (without chunks).
E
Ablation Studies
In the following section, we perform ablations to investigate (1) the placement a single SSM layer in Table 3 in the overall architecture, (2) the effects of the number of SSM layers added in Table and (3) the size of the SSM state in Table 5 For the ablations, we use the ~200 parameter BST:SH:S4, since it is the fastest model, and assess various configurations PG19.
Table 3: single BST at various layer index.
Table Multiple BST layers at various locations.
Table 5: Increasing BST's S4 model
state size D.
Layer index Perplexity
Num layers Perplexity
State Size Perplexity Step Time
12.41
11.69
8
11.95
x0.7
7
11.92
3
11.57
16
11.57
x1.0
11.88
11.21
32
11.55
x1.8
12
12.03
5
11.20
64
11.54
x3.2
In Table 3 we experiment adding a single BST layer at layer indices 3, 6, 9, 12. We notice that a single BST layer with state size D = 16 located closer the middle of the whole Block Transformer stack, at index = 9, has the greatest effect on perplexity. This finding is inline with findings in prior work [42 22
In Table we test adding multiple BST layers yields improvements on performance. We start with BST layers with state size D 16 at indices 0, 9. We follow by adding another BST layer at index 7 for a total of three BST layers and then another at index 5, followed by another at index 12. Adding more BST layers lowers perplexity. However, the results seem plateau at 5 BST layers. We note also that there is a 3.5% training step time increase for each added layer.
In Table we train our models with different state sizes D. For the state size ablation, we use three BST layers indices 0, 7,9. 9. We find that increasing D improves perplexity to the detriment of training speed (step time). For this reason, chose D 16 for Table BST results.
16

---

JAX Implementation of BST
Pscudocode contains function that implements convolution of multiple filters over the same imput
sequence using FFT and inverse FFT operations. Pscudocodes 2 and 3 respectivel implement context state collection of BST variants: Single-Head (SH), Multi-Head (MH) and Multi-Filter (MF).
Finally, Pseudocode runs the Block Transformer sublayer in parallel by feeding the context states their corresponding block.
"" "nUstructtre filters and convolutions. "..
import jax
from jax import numpy jnp from einops import rearrange
Len_tlegthh 512
# (w)
seq_length 4096
# (1)
def get unnstrutttccchanelss Returrs trainable filters and biases.
Args:
channels: number of filters.
Returns:
h: filter of shape @seq_length, channels, dim)
b: bias of shape (channels, dim)
.**
jnp. linspace(o.o, 1.0, seq_length) jnp.exp(- alpha t) dense(positionll_emm(> b get__iiaee (i_biiss( return h, b
def multichannel_ cetiihaaneeLccoooouutninCCCC h, b): *Mulltccaanne convolution function.
Args:
u: input of shape (seq_length dim)
h: filters of shape (seg_hength, channels, dim)
b: bias of shape (channels, dim)
***
rearrange(h,, "l -> d 1")
fft ift_size = seq_length 2
u_f fff,fft. n=fft_sizel h_f fnp.fft.rffthh, n=fft_size)
y jnp. fnt.fftttrffthh.. x_f, n=fft size, norm="forward"
:seg_length]
# (c, d, 1) # (c, d, 1)
y x b[..., None]
x
y rearrange(y, "c d - cdc" return
Pscudocode 1: Unstructured filters and convolutions.
""*Context state collection for BST-SH variant.
num heads 8 # (h) # (s)
num_ sum_states 32
# (SH): Single-Head def contextt [H__context_sttes(())) .** Singleee-Hedd Context collection, ...
h, b =
17

---

y_1 mu muttichamee__cnvvlutiinnn h, b) y_1: (t, d, 1)
lift multiple heads
y_h = dense(y_1) y_h: (1, d, h)
pottxt_states splitt
y_h, seq_length // win_length, axis=o)
return context_states (L/w, w, d, h)
Pseudocode 2: Context state collection for BST-SH variants.
"Context state collection for BST-MH variant."
(MH): Multi-Head
def MH MM_context_states((
** M*ultii-Hadd Context Collection." ..*
h, b get_filters__eeeeeeeeeee_tsse___ (
y_h muttichamee_convouutiinnnu h, b) y_h: (L, d, h)
context_stated split( y_h, seq_length // win_length, axis=0)
return context_states (l/w, w, d, h)
Pseudocode 3: Context state collection for BST-MH variants.
"Context state collection for BST-MF variant."" "**
# (MF): Multi-Filter
def MF cerontext_state((((
-- "Multi-Filteer Context Collection. ... h, b) y_s: (1, d, s)
h, b =
y_s
context_states jnp.split(
y_s, seq_length // win_tength, axis=0)
# context_states: (L/w, w, d, s)
collect the last context states
context_states rearrange( context_states
santext_states[:. -1, ...] # (L/w, d, s)
context_states, "lw d -> lw d")
# shift context states corresponding to windows
context_stater gnpprool(contexx_ssttee 1, axis=1)
# replace the initial window with trainable weights init init_context get_init sunitttooneetnnumsseeeeeeeee context_stateslo) init_context
(d, s)
lift to multiple heads
context_states dense(context_states
return context_states (1/w, s, d, h)
Pseudocode 4: Context state collection for BST-MF variants.
"""Block-State Transformer Layer.
Block Transformers are non-recurrent and parallelizable. Block_trrasforreet .nnmoeBeeet...eeeeeeeeet......l
def BST(u):
"""Block-Stat" Transformer Layer.
18

---

global MF # True if Multi-Filter, False otherwise (SH/MH)
split inputs into windows (1/w, w, d)
jnp.split(u, seq_length // win_length, axis=o)
collect context states from SSM outputs
context_states (S5/////Ii_ccontext__tttttttee
pass the contexts in place of recurrent states
y block_trand transsorrrfffr((
token gooke__embeddingsaat
cecurrent_staaee-conteek
use_cros att. attaaattnncssaaassseeee MF, pee_cross_positiooallemmeee context IDs
)
return rearrange(y, "lw w -> (lw w) d") # (1, d)
Pseudocode 5: Block-State Transformer Layer.
19


Reviews:
[{"review_id": "zSmwWdbhPA", "reviewer": "Reviewer_V1Sd", "summary": "The authors present a novel architectural framework called the Block-State Transformer (BST), which integrates State Space models and Block-Recurrent Transformers to create a competitive autoregressive language model capable of effectively processing lengthy sequences. The input sequence is passed through a State Space model like S4, and the output of which is later in Block-recurrent Transformer as a replacement of the recurrent state vectors. To obtain the final output, the input embeddings are divided into fixed-sized windows and processed in parallel by a series of Block-Recurrent Transformers. Due to the usage of the S4 output as recurrent state vectors within the Block-Recurrent Transformers, the absence of recurrence allows for parallel computation. The authors propose three distinct integration modes, which differ in terms of how the S4 output is integrated within the recurrent state vectors. To evaluate the performance of BST, the authors compare it against four baseline models: Transformer-XL, Slide, Block-Recurrent Transformer, and a hybrid Gated State Space model. The comparison is conducted across three diverse datasets, namely PG19, arXiv Math, and GitHub. BST demonstrates slight perplexity improvements in the PG19 and GitHub datasets. Additionally, the authors present ablation studies on various parameters, including SSM layer placement, the number of SMM layers, and the state dimensionality of SMM.", "questions": "Broadly, understanding BST capabilities in other settings would help enrich the paper results. See weaknesses.", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "+ **Strong Presented Results**: Authors presented results on competitive benchmarks across reasonable prior baselines such as Transformer-XL, Slide, Block-Recurrent Transformers etc. and do outperform them across several tasks.\n\n+ **Computational efficiency**: The proposed method is able to provide a huge improvement in terms of computational efficiency over models like Block-Recurrent Transformers by parallelization.\n\n+ **Interesting combination of prior ideas**: By using the parallelizable nature of the SSMs the authors were able to introduce parallelization to Block-Recurrent Transformers thereby achieving computation efficiency. ", "weaknesses": "- **Incomplete Related works**: The authors' treatment of related works, particularly in the context of models combining Transformers and S4 appears to be lacking. It would have been beneficial for the authors to provide a more comprehensive discussion on existing models that incorporate both Transformers and S4. This would have allowed for a deeper exploration of the advancements and limitations of these models, highlighting the unique contributions of the proposed Block-State Transformer (BST). In addition to that, there could be more details on SSM development in related works since S4, S5, and S4D were mentioned in the paper later.\n\n- **Missing Preliminaries on S4/S5** : A more detailed description of the S4 models within the method section would have been beneficial, particularly regarding the computation of the kernel since the complexity of the S4 model is the major part of the BST model complexity. Specifically, the computation of the kernel is not trivial if you want to keep overall L log(L) complexity and it relies on the form of the A and B matrices. A more comprehensive exposition regarding the computational aspects of the S4 models is deemed necessary for a thorough understanding of the subject matter.\n\n- **Additional benchmarking** : As the authors themselves admit in the limitations section, there are further results required, especially on well benchmarked domains such as the Long Range Arena and also other long-term datasets to provide convincing evidence of BST performance. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your valuable and constructive comments. We agree that Block-State Transformer is a novel architecture that shows strong results and computational efficiency. We think that we have only scratched the surface of possibilities with this interesting combination of ideas. We have conducted more experiments that can be found in the **1-page PDF** (see general comments to access PDF), some of which are related to your review, namely benchmarking BRT and BST on Long-Range Arena. Additional experiments also include scaling results and more baselines on language tasks. We believe that these supplementary results address most of your  feedback on proving that we can harness the strengths of S4 and Transformer in long-range classification and language modeling tasks.\n \n___\n\n#### Extending Related works  \nWe have indeed included a more comprehensive review of related work in the latest version of our paper, while also emphasizing the key differences. Our discussion will include S4 kernel computation, GSS-Hybrid, Hybrid-H3, and Mega. We have already outlined some of the core contributions in [response to Reviewer 9Xt3](https://openreview.net/forum?id=XRTxIBs2eu&noteId=wjrwjedtwd), and we invite you to review them for further insights.\n\n---\n\n#### Covering Preliminaries\n\nWe have incorporated your feedback by expanding the discussion in the efficiency section in the latest version of the paper. To summarize, the efficiency of our models is largely afforded by standard SSM implementations, which are due to:\n\n* Imposing Diagonal Plus Low-Rank (DPLR) structure on the $A$, $B$, and $C$ matrices, rematerializing the convolutional kernel can be carried out in tandem with the convolution operation with $\\mathcal{O}(L \\log L)$ complexity, as described in S4. \n* By replacing RNNs (BRT) with SSMs (BST), BST transformer blocks can run in parallel instead of sequentially as in BRT.\n\nFurther, by demonstrating two vastly different SSM as means of contextualization, we show that our proposed architecture/framework remains agnostic to the specific type of SSM used, making the advancement of SSMs orthogonal to our work. That being said, and as highlighted in your review, since the computational efficiency of BST layers _can be_ dictated by the SSM, it is worth discussing them in more detail. \n\n---\n\n#### Additional benchmarking\n\nPlease see **Table 2 in the 1-page PDF** for additional experiments on Long-Range Arena (LRA). \n\nOur motivating hypothesis was that replacing recurrence with a more efficient yet powerful component for capturing long-range dependencies would improve performance on language modeling tasks. As we openly acknowledged in the limitations section, we recognized the need for more concrete evidence to support this hypothesis. Consequently, in response to your request, we conducted additional experiments on Long-Range Arena (LRA) (please see **Table 2 of 1-page PDF**), demonstrating that BST indeed surpasses BRT (and other methods that chunk inputs such as Mega-chunk) in capturing long-range dependencies.\n\n---\n\nThank you for reading our response!\n\n\n"}, {"review_id": "NGhCi4TyYc", "reviewer": "Reviewer_9Xt3", "summary": "This paper proposes block-state transformers, a method to combine state space models with transformers for language modeling. The paper evaluates block-state transformers on PG19 and arxiv math and finds promising results.", "questions": "How does BST compare to the architectures listed in the weaknesses section? It is important to compare against the original architectures that inspired components of BST, as well as modern standard Transformers.", "limitations": "The paper would be stronger with more discussion of limitations.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Combining state space models and Transformers is an interesting idea worth exploring. The presentation of the paper is clear. The explanation of state space models, which can be quite complex, is very clear. The evaluation hits the right notes in terms of the major questions to ask.", "weaknesses": "The evaluation is missing many recent works combining state space models and attention in various ways. The claim that SSMs do not match Transformers on language has not been true for a while. Most of these methods were released significantly before the NeurIPS deadline and are critical to compare against for evaluation.\n\n* Mega [1] combines attention and state space models.\n* BiGS [2] is a new SSM-based architecture that matches Transformers in language.\n* H3 [3] combines SSMs and attention in alternate layers.\n* Hyena [4] removes attention completely and replaces it with a convolution-based layer (similar to an SSM).\n\nConfusingly, many of these works are cited in the paper - and ideas from the papers are used extensively in the methods proposed (e.g., \"BST:{SH,MF}\" uses the structure from H3 and Hyena without comparing against those architectures as baselines). Using the ideas from these papers without comparing against them makes it difficult to understand how this method compares against previous methods and where the performance improvement comes from.\n\nPerformance is also hard to evaluate compared to standard models such as Transformers (GPT-Neo [5], Pythia suite [6]). TransformerXL is an older model that is not trained as well as modern Transformer-based LLMs.\n\n[1] https://arxiv.org/abs/2209.10655\n[2] https://arxiv.org/abs/2212.10544\n[3] https://arxiv.org/abs/2212.14052\n[4] https://arxiv.org/abs/2302.10866\n[5] https://github.com/EleutherAI/gpt-neo\n[6] https://github.com/EleutherAI/pythia", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We hope that we have addressed most of your concerns with the additional experiments and comparisons in the **1-page PDF**. \n\nOur new experiments show:\n\n1. Large improvements on PG-19 compared to Hyena and Hybrid-H3. Specifically, integrating Hyena and attention in `BST:SH:unstruct` allows our model to surpass standalone Hyena ***by ~1.4 perplexity points***. See **Table 1 in 1-page PDF**.\n\n2. We show that we surpass all Transformer variants that you have mentioned and other methods that chunk inputs (allowing similar speed-ups to our BST) such as Mega-chunk. Specifically, on Long-Range Arena (LRA), `BST:SH:S4` performs better than: \n    - S4 on 5 out of 6 tasks and on average by 0.9% points\n    - Mega-chunk on 4 out of 6 tasks and on average by 1.3% points\n\nFurthermore, we respond to your specific questions and comments below.\n\n---\n\n> “Mega [1] combines attention and state space models.”\n\nThanks for bringing Mega to our attention, we will review and cite it in our related works section. \n\nWe compare our model to Mega on Long-Range Arena (LRA), where we see that our model achieves comparable performance to Mega/Mega-chunk. However on language modeling, according to Table 1 in [1], **Mega** (252M) achieves _18.07_ perplexity on WikiText103 that is roughly the same performance, _18.50_ perplexity, achieved by **Hybrid-H3** [3] (125M, half the size of Mega). ***`BST:SH:unstruct` outperforms Hybrid-H3 by 3.0 perplexity points*** on the PG-19 long text language modeling benchmark, see **Table 1 in the 1-page PDF**. Therefore, although we were unable to directly compare our model against Mega on Language Modeling in the rebuttal phase, we have reason to believe that our model will outperform Mega when using the same number of parameters. \n\n--- \n\n> \"BiGS [2] is a new SSM-based architecture that matches Transformers in language.\"\n\nRegarding BiGS, we consider this work to be less relevant compared to the other papers mentioned. The reason is that BiGS is designed as a bidirectional model, making it more suitable for serving as an encoder or for natural language understanding tasks, rather than auto-regressive language modeling or language generation tasks. In contrast, our focus is on decoder-only language models, which have different requirements and objectives. Further, BiGS is not evaluated on Long-Range Arena which makes it difficult to compare against.\n\n---\n\n> \"H3 [3] combines SSMs and attention in alternate layers.\"\n> \"Hyena [4] removes attention completely and replaces it with a convolution-based layer (similar to an SSM).\"\n\nWe have compared BST to H3-Hybrid and Hyena on PG-19, results can be found in **Table 1 of the 1-page PDF**. We have trained both BST and BRT models under a fixed parameter setting using the same tokenizer (GPT2) and vocab size, for a fair comparison. BST remains to be state-of-the-art at this scale. Specifically, integrating Hyena and attention in `BST:SH:unstruct` allows our model to surpass standalone Hyena by ~2 perplexity points. \n\n---\n\n> \"Confusingly, many of these works are cited in the paper - and ideas from the papers are used extensively in the methods proposed (e.g., \"BST:{SH,MF}\" uses the structure from H3 and Hyena without comparing against those architectures as baselines).\"\n\nPlease see **Table 1 in 1-page PDF** for a direct comparison against H3 and Hyena where our model outperforms both under a fixed parameter count setting.  \n\n---\n\n> \"Performance is also hard to evaluate compared to standard models such as Transformers (GPT-Neo [5], Pythia suite [6]).\"\n\nAs we opted to implement our models using JAX, we were unable to utilize other open-source PyTorch-based codebases like GPT-Neo and Pythia by the rebuttal deadline. However, we should have such comparisons before the camera-ready deadline. It is crucial to mention that any enhancements made to the attention layer are independent/orthogonal to our model which focuses on enabling off-the-shelf transformer-attention layers to model long-range dependencies.  \n\nFurthermore, we want to note that since GPT-Neo and Pythia employ pipeline parallelism, we anticipate similar speed improvements when applying SSMs since they can be parallelized in a similar way. As part of our ongoing work, we are actively developing a PyTorch implementation for BST and BRT. These implementations will enable us to reliably compare performance with other PyTorch models and explore their potential benefits in conjunction with our proposed method. Please note that we do evaluate performance against linear transformer, reformer, performer and BigBird in **Table 2 of 1-pagePDF**.\n"}, {"review_id": "x9dY4AUgLU", "reviewer": "Reviewer_XDNb", "summary": "State space models (SSMs) perform well on modeling long-range dependencies with good efficiency scaling, but on language modeling, transformers still outperforms SSMs. This paper tries to combine the best of both worlds and proposes a hybrid model, Block-State Transformer, which combine SSMs’ capacity on long range modeling and Transformer’s ability on modeling local context. The input sequence is split to multiple smaller segments. For each segment, transformer layer will do a self attention on this token embeddings and cross attention to the output of SSMs. Their experiments show that on language modeling, their approach achieve reasonable speedup with comparable performance to Transformers. ", "questions": "The main concerns are listed above. There are two more questions:\n1. On line 235, the 6,966,499 English language words seems to be a typo according to dataset statistics on https://github.com/deepmind/pg19. \n2. What are the latency of one forward step and what about memory consumption? ", "limitations": "Yes. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed combination of SSMs and transformers allow the model to exploit advantages of two powerful methods while avoiding their drawbacks. \n2. The SSMs used in the proposed method can be swapped to different SSMs making it possible to enjoy the advancement on SSMs field. \n3. The proposed method give similar performance on language modeling compared to Transformers. \n", "weaknesses": "1. There was already existing work on SSMs that achieve similar performance on language modeling (https://arxiv.org/abs/2212.14052) compared to Transformers. The authors should include a discussion and comparison to the relevant work. \n2. The evaluations are performed on language modeling for 4096 length sequences. On this setting, there are a lot of strong transformer baselines with efficient self-attention designs. It would be good if the authors can provide an empirical comparison with these baselines. \n3. Perplexity is only one indicator of how the language models performance. To get more precise understanding of performance, it would be good to include a comparison on downstream tasks. \n4. Code is not available. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your generally positive review. We have taken your valuable feedback into account to improve our current version of the paper. A more comprehensive review of related works, including H3, will be provided. Additionally, we have included a number of experiments in the **1-page PDF**, to further assess long-range modeling capabilities (LRA) and scaling properties of our model. Moreover, we have expanded our baselines on language to also include Hyena, a more recent and general framework that subsumes H3 and GSS. \n\nOur additional experiments show:\n- Large improvements on PG-19 compared to Hyena and Hybrid H3. Specifically, integrating parameterized convolutions and attention in `BST:SH:unstruct` allows our model to **surpass Hyena by 1.4 perplexity points**. See **Table 1**.\n- ***+30%-35% point improvement on Long-Range Arena (LRA)*** over recent and strong Transformer baselines. See **Table 2**.\n\n___\n\n> W1. \"There was already existing work on SSMs that achieve similar performance on language modeling (https://arxiv.org/abs/2212.14052) compared to Transformers.\"\n\nWe have added Hybrid-H3 and Hyena as additional baselines in the **Table 1 of the 1-Page PDF**. \n\nHyena Hierarchy (https://arxiv.org/pdf/2302.10866.pdf) is a language model that draws inspiration from GSS and H3, and captures both of these methods within a more generic framework. In all these approaches, attention and retrieval are conceptually simulated by element-wise multiplication of a sequence of tokens, and its corresponding contextualized (via SSM or parameterized convolutional kernels) counterpart. Hyena, outperforms both GSS and H3 on tasks such as associative recall and in language modeling, particularly on the WikiText103 dataset. \n\nWe have trained both BST and BRT models under a fixed parameter setting using the same tokenizer (GPT2) and vocab size, for a fair comparison. BST achieves SoTA performance on this dataset. Results can be found in **Table 1 of the 1-Page PDF**. We hope this addresses your concern regarding benchmarking SSM-inspired language models.\n\n---\n\n> W2. \"The evaluations are performed on language modeling for 4096 length sequences. On this setting, there are a lot of strong transformer baselines with efficient self-attention designs.\"\n\nTo the best of our knowledge, GSS and BST were the state-of-the-art on the datasets we use in our work, outperforming other efficient Transformers, e.g. linear implementation etc. More importantly, other Transformer-based architectures do not generalize to sequence lengths not seen during training, whereas our method does (See **Figure 4 in Appendix B**). Even with relative positional embeddings, Transformers cannot reliably go beyond 3x the trained sequence length (see https://arxiv.org/abs/2305.19466), let alone 16x (at a sequence length of 65K for example). That being said, if the reviewer can point us to specific works, we are well-positioned to include additional baselines in the camera-ready version of the paper. Further, as seen in **Table 2 of the 1-Page PDF**, our method greatly outperforms very recent and efficient Transformers such as Linear Transformer, Reformer, Performer and BigBird.\n\n---\n\n> W3. \"Perplexity is only one indicator of how the language models performance.\"\n\nAlthough there is no unanimous consensus, most practitioners in the field generally agree that the performance on downstream tasks seems to be well correlated with perplexity for LLMs (https://arxiv.org/pdf/2210.14199.pdf, https://aclanthology.org/2021.emnlp-main.478.pdf). Like other works in this field, we focus on developing decoder-only models that achieve lower perplexity. Nevertheless, we acknowledge that this is an important step in developing large language models and believe BST can also serve as a powerful encoder model which may be evaluated on downstream tasks. Because BST uses off-the-shelf transformer models augmented with SSM states as context, our approach can be easily adapted into existing LLM codebases. While no one has evaluated SSM or hybrid-SSM pretraining performance on downstream tasks yet, we look forward to doing that with BST in a follow-up project.\n\n---\n\n> W4. \"Code is not available.\"\n\nAs part of our ongoing work, we are working closely with the maintainers of the Block-Recurrent Transformer codebase to integrate our implementation of BST into the repository (https://github.com/google-research/meliad). In the meantime, we have provided the JAX pseudo-code for all of our variants in **Appendix E** in the supplementary materials. \n\n---\n\n> Q1. \"On line 235, the 6,966,499 English language words seems to be a typo...\"\n\nThank you for pointing this out. We have fixed this error in the latest version. \n\n---\n\n> Q2. \"What are the latency of one forward step and what about memory consumption?\"\n\nThe computational and space complexity of a BST layer consists of that of the transformer blocks and the SSM sublayer. We discuss this topic in more detail in the Efficiency section of the paper. In the left-side plot of Figure 3 in the paper, the y-axis represents the latency for a forward pass of our model when executed on an NVIDIA V100 GPU.\n\nWe assume that your question refers to the auto-regressive “forward step” at test time. The exact forward step latency may vary slightly depending on the BST variant used. In Single-Head and Multi-Head variants, every token generation step is immediately followed by feeding that token back into the SSM and adding it to the context in $\\mathcal{O}(1)$ operations – via the RNN view of SSMs. On the other hand, when the Multi-Filter variant is used, contextualizing can be postponed to when all the tokens corresponding to the current block are generated. The new block of tokens are then added to the context in one go, using the recurrent view of SSMs in $\\mathcal{O}(W)$ operations. The token generation done within the Transformer is similar to BRT, $\\mathcal{O}(W)$, since every token attends to $W$ previous tokens. "}, {"review_id": "YXABaY0See", "reviewer": "Reviewer_rmDJ", "summary": "This paper focuses on combining two efficient techniques for long-range modeling: state-space models (global contextualization) and block-recurrent transformers (local contextualization). In particular, they propose two different approaches, the first uses SSMs to output contexts for multiple heads (multi-head), and the second concatenates the last entries from the previous window to form a combined context state (multi-filter). Evaluation is performed on three language modeling datasets that outperform block-recurrent transformers in perplexity and is much faster when compared layer-wise. ", "questions": "Q1: Evaluating methods that model long context typically involves more tasks than language modeling. What is the level of confidence that proposed models perform to other tasks such as long context classification and text generation? \n\nQ2:  It would be interesting to measure the time it takes for every method to converge even within a fixed time budget. Have all methods converged in the fixed training budget in Table 1 and do they have any differences worth discussing?\n\n", "limitations": "Discussion about the limitations of the proposed method would be useful, I'd suggest talking about scaling behavior, performance on conventional accelerators, and generalizability to long-context tasks. ", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "S1. Exploring different ways to combine SSMs and block-recurrent models to improve efficiency is a compelling direction.  SSMs offer a parallelizable way to capture long-term information and avoid sequential computation in block-recurrent models. The results of this study should be of interest to researchers that study architectures that capture local and global information.\n\nS2. The evaluation even though it focuses mainly on comparison with block-recurrent models and SSMs on three language modeling tasks, it is thoroughly described and well-executed. ", "weaknesses": "W1. Even though the paper is mainly empirically driven, the delivery lacks a comprehensive and diverse set of evaluations to demonstrate the effectiveness and limitations of the method.  \n\nW2. Experiments are targetting language modeling on three tasks but there is no experiment that measures the long-range capabilities of the model. There are several long-context classification benchmarks that the authors can use in addition to language modeling: LRA [1], MuLD [2], and CAB [3]. \n\nW3. The method design makes specific assumptions about the hardware to be employed and bases its evaluation on it; e.g. efficiency comparisons are made per layer. It's not explored to what extent the benefit remains when comparing training time/speed vs quality for the whole model and evaluation is performed on typical accelerators. That reduces the practical impact in my view. \n\nW4. A study regarding the behavior of the model when increasing the model size is missing. Scaling aspects are important to consider when making claims about outperforming transformers. \n\n[1] https://arxiv.org/pdf/2011.04006.pdf\n[2] https://arxiv.org/pdf/2202.07362.pdf \n[3] https://arxiv.org/pdf/2210.07661.pdf\n ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful feedback. In response to your comments and concerns, we have conducted additional experiments, which we have included in the **1-page PDF**. These experiments address the points you have raised and also explore areas that you alluded to in your review, namely scaling aspects and performance on Long-Range Arena (LRA), outside the language modeling domain. \n\n---\n\n#### Re W1 and W2: Extending to areas other than language/perplexity\n\nAs we acknowledged in the limitations section (Appendix D) of our submission, we recognize the significance of conducting experiments to assess the long-range capabilities of our proposed model. These new experiments are designed to test if the improved performance on language modeling can be attributed to our model’s ability to infer long-range dependencies. To test this hypothesis rigorously, we have performed comparisons between variants of our proposed model, BST, and its recurrent predecessor, BRT, as well as several other baselines that utilize SSMs. We use the Long-Range Arena (LRA) benchmark as the primary testbed. Kindly refer to **Table 1 of the 1-page PDF**. The results demonstrate that, in line with our expectations, SSMs enhance long-range modeling capabilities compared to BRT. \n\n---\n\n#### Re W3: Efficiency of our model\n\nIn the interest of convenience while maintaining an exploratory approach and ensuring a fair comparison to former baselines on the dataset, we intentionally chose not to utilize model parallelism, despite the fact that our model is well-suited for it. Unlike BRT, whose recurrent layers cannot be parallelized, our model’s per-layer speedups (Figure 3) will trivially translate to full model speedups provided that there is sufficient compute resources to support parallelism. **Appendix C, Table 3** shows that replacing transformer layers with BST layers monotonically improves perplexity performance. This evidence, combined with the per-layer speedup results, demonstrate the performance and speedup potential of our full model at scale. \n\n---\n\n#### Re W4: Scaling capabilities \n\nPlease see the additional scaling experiments in the **1-Page PDF, Figure 1**. Scaling experiments relate parameter count (80M to 1.3B) to the perplexity performance. We show a 0.1% relative perplexity improvement at 80M parameters to a 3.8% relative perplexity improvement at 1.3B parameters over an equivalently large Block-Recurrent Transformer (BRT) model. Additionally, ablation studies in **Appendix C Tables 2-4** relate the placement, capacity, and the number of the BST layers to performance. \n\n---\n\n#### Re Q1: LRA\n> \"What is the level of confidence that proposed models perform to other tasks such as long context classification and text generation?\"\n\nWe have benchmarked BST and BRT on Long-Range Arena in **Table 2 of the 1-page PDF**. Results show that BST is indeed better at capturing long-range relations, which aligns with previous results (e.g. S4 vs Transformers) and our initial motivating intuition. \n\n---\n\n#### Re Q2: Training/Convergence Time\n> \"Have all methods converged in the fixed training budget in Table 1 and do they have any differences worth discussing?\"\n\nThis is an interesting question. Our experiments show that BRT and BST both had room to continue improving and our validation set perplexity continued to diminish on all datasets. While our experimental setups were designed to align with prior works (GSS & BRT), this is a reasonable question and we are well positioned to include these additional experiments in the camera-ready version of the paper."}, {"review_id": "FQVSgIqyZf", "reviewer": "Reviewer_4BCd", "summary": "The authors present a new long-range transformer architecture by incorporating SSMs. This novel model outperforms several established baselines, such as Transformer XL, Block Recurrent Transformer, and Sliding Window Transformer, in terms of cost-effectiveness trade-off for tasks involving long-document or code modeling.", "questions": "1) Did the authors observe any instable training process? And again, would it become instable when scaling up?\n", "limitations": "See weakness", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1) Well-motivated. Long-range modeling is becoming increasingly important for LLM community.\n2) Good results on language modeling (PG19, arXiv, Github)\n3) The writing is clear and effectively conveys the ideas and findings.\n4) The model design is intuitive and well-reasoned. The inclusion of both local full attention and linear components to handle long sequences is a sensible approach.", "weaknesses": "1) The scale is too small. For language modeling, based on the success of LLM, we always expect good scalability. This paper only conduct experiments up to 380M params, which are far from many emergent abilities threshold. When scaling up, many inductive bias would become useless[1]. \n2) More insightful experiments beyond language modeling are required. For instance, as the authors mentioned in the limitation section, what about the results on Long Range Arena? \n3) Any case study about how this model captures long-range dependency? Why this model is indeed better? It seems that putting one efficient attention layer with linear or sub-linear complexity before self-attention should work similarly.\n4) Some strong baselines like CoLT5[2] are missing.\n\n[1] Scaling Laws vs Model Architectures: How does Inductive Bias Influence Scaling?\n[2] CoLT5: Faster Long-Range Transformers with Conditional Computation\n\nI enjoyed reading this idea but I believe the missed experiments above, especially (1), would be highly desirable. Without a set of experiments about scaling the model up, I cannot agree this paper is useful enough.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback. We have taken your comments and concerns into careful consideration and conducted additional experiments to address them. These experiments have been included in the **1-page PDF**, focusing on scaling aspects and performance in areas beyond language tasks, namely experiments that assess long-range modeling capabilities and scaling properties of BST. We would also like to draw the reviewer’s attention to the results in **Appendix C** which also provide insights into the scaling properties of our model. \n\n---\n\n> \"The scale is too small.\"\n\n> \"This paper only conduct experiments up to 380M params.\"\n\n> \"When scaling up, many inductive bias would become useless[1].\"\n\nWe acknowledge the growing significance of scalability in language models, and in response, we have incorporated additional scaling experiments in relation to the number of parameters in **Figure 1 of 1-page PDF**, where we scale BST from 80M to 1.3B parameters. We show a 0.1% relative perplexity improvement at 80M parameters, to a 3.8% relative perplexity improvement at 1.3B parameters over an equivalently large Block-Recurrent Transformer (BRT). These experiments demonstrate favorable scaling properties of our proposed model on perplexity. However in Table 1 of the paper all models have the same number of trainable parameters for a fair comparison. \n\nAblation Studies in **Appendix C**, also suggest favorable scaling properties of our model. We provide experiments showing that performance uniformly improves as we add more contextualizing layers, i.e. BST layers, into the architecture. Moreover, an equally significant alternative scaling property pertains to the length of the sequence, and our model plays a direct role in addressing and contributing to this aspect, especially at inference time. See **Appendix B** in the supplementary material, on Evaluating Length Generalization capabilities.  \n\n---\n\n> \"what about the results on Long Range Arena?\"\n\n> \"Any case study about how this model captures long-range dependency? Why this model is indeed better?\"\n\nWe thank the reviewer for suggesting experiments on Long-Range Arena (LRA), the results of which can be found in **Table 2 of the 1-page PDF**. Because the model’s ability to model long-range dependencies is not the only factor that influences perplexity, this experiment was necessary to directly demonstrate that our model captures long range dependencies better than the baseline (BRT). Further a fair comparison with Mega-chunk shows that we surpass on 4 out 6 tasks and on average the latest and strongest “chunked input” baseline on LRA.\n\n---\n\n> \"It seems that putting one efficient attention layer with linear or sub-linear complexity before self-attention should work similarly.\"\n\nAs seen in Table 2 of the 1-page PDF, SSMs such as S4D, S4 and S5 are 30%-35% points higher on Long-Range Arena (LRA) compared Transformers with linear and sublinear attention (Linear Transformer, Performer, BigBird). It is therefore unlikely that we would gain any benefit from replacing the BST or SSM layers in hybrid models with linear or sub-linear attention layers. \nFurther, transformer memories are fundamentally limited by their context length - $L$ (with $L^2$ complexity), whereas SSMs (like RNNs) can encode information indefinitely in their latent states. For this reason SSMs strongly outperform (linear) Transformers on long-range modeling tasks. Moreover, at inference time, Transformers completely lack the ability to generalize to unseen sequence lengths during training, unlike structured SSMs (See **Appendix B**).\n\n---\n\n> \"Some strong baselines like CoLT5[2] are missing.\"\n\nWe were not aware of this baseline at the time of writing and submitting the paper. According to the arXiv publication date, this paper was made public on _March 17th_, which indicates that it is concurrent work that emerged during a similar timeframe as our research. This new work does not report perplexity on our targeted tasks nor LRA. We have cited CoLT5 and will nonetheless attempt to replicate and test CoLT5 on LRA and PG-19.\n\n---\n\n> \"Did the authors observe any unstable training process? And again, would it become instable when scaling up?\"\n\nNo, the model trains satably without the need of additional tricks. \n"}]

Author rebuttal:
We would like to thank all the reviewers for their insightful comments. We believe that we addressed the vast majority of reviewer’s concerns by conducting additional experiments, found in the attached **1-page PDF attached to this message**, and responding to reviewer’s individual questions. Our additional experiments include: 


#### **Scaling Properties of Block-State Transformer (BST)** [_Figure 1_]

- Showing the competitive scaling properties of Block-State Transformer (BST), going from a 0.1% relative perplexity improvement at 80M parameters to a ***3.8% relative perplexity improvement at 1.3B parameters*** over an equivalently large Block-Recurrent Transformer (BRT).


#### **Comparisons to Hybrid-H3 and Hyena** [_Table 1_]

- Demonstrating that our model achieves ***superior perplexity*** compared to Hyna & Hybrid:H3, respectively by 1.4 and 3.0 points, under standard evaluation conditions (fixed parameter count setting). 


#### **Comparisons to Transformer variants and Mega on Long-Range Arena (LRA)**  [_Table 2_]

- Achieving a substantial lead over the Block-Recurrent Transformer (BRT). Further, BST outperforms Transformer variants **by 30%-35% points on average**. Further, in a fair comparison with other methods that chunk input sequences, we see that BST outperforms Mega-chunk **on 4 out of 6 tasks** and by 1.3% point on average.

Our strong results on Long-Range Arena (LRA) demonstrate that BST significantly outperforms BRT on LRA, providing further evidence that the gains on language modeling tasks can indeed be attributed to our model’s ability to capture long-range dependencies more effectively while improving computational efficiency. 

---

We have included the responses to the reviewers in their rebuttal sections individually. We would like to thank all reviewers again for their time and attention. 

---

Find the PDF containing additional results below.


Meta-review: This paper introduces the Block-State Transformer (BST), a novel hybrid layer that combines state space models (SSMs) for long-range contextualization and Block Transformer sublayer for short-term sequence representation. The authors propose three parallelizable variants for integrating SSMs and block-wise attention. The reviewers are overall positive towards the contribution of this paper. I would recommend to accept this paper.
Decision: Accept (poster)