Block-State Transformer

🎯 Pitch

A hybrid architecture combining state space models with block-wise attention achieves more than a tenfold speedup over Block-Recurrent Transformers under model parallelism, all while maintaining or improving language modeling perplexity. Crucially, the model only generalizes to sequences far beyond the training length when using a structured SSM kernel.

1. Executive Summary

This paper introduces the Block-State Transformer (BST), a hybrid architecture that combines a State Space Model (SSM) sublayer for long-range contextualization with a Block Transformer sublayer for local, short-term sequence representation — for example, using an S4-based SSM to capture global context across an entire sequence while a windowed attention mechanism processes local blocks of tokens in parallel. The authors study three fully parallelizable variants — Single-Head (SH), Multi-Head (MH), and Multi-Filter (MF) — that differ in how they construct and feed context states from the SSM into the Transformer’s cross-attention, trading off redundancy against retrievability. Evaluated on language modeling across PG19, arXiv, and GitHub datasets at ~200M and ~400M parameter scales, BST matches or modestly improves upon the Block-Recurrent Transformer (BRECT) and the GSS-HYBRID architecture in perplexity, while achieving more than a tenfold increase in layer-level speed compared to BRECT when model parallelization is employed (reducing a ~15× slowdown relative to a sliding-window baseline to less than 2×). The strongest length-generalization variant — BST:SH:S4-L — maintains the best perplexity on sequences extended to 65K tokens, establishing that SSM-based context states can substitute for recurrent state vectors in hybrid Transformer architectures and generalize beyond the training length only when a structured SSM kernel is used.

2. Context and Motivation

The Core Problem: Scaling Sequence Length in Transformer Language Models

The fundamental tension this paper addresses is the quadratic compute cost of self-attention relative to sequence length. In a standard Transformer, every token attends to every other token, yielding an $\mathcal{O}(L^2)$ time and memory complexity where $L$ is the sequence length. This means that doubling the context window quadruples the attention cost, making it prohibitively expensive to train Transformers on sequences beyond a few thousand tokens. Meanwhile, many real-world language modeling tasks — books, code repositories, scientific papers — contain dependencies that span tens or hundreds of thousands of tokens. A model limited to, say, 2,048 or 4,096 tokens simply cannot capture cross-chapter plot developments in a novel or cross-file import dependencies in a codebase.

This is not merely a matter of convenience or cost. It is a quality ceiling. As the authors note in Section 1, vanilla Transformers can be unstable when trained on long sequences, and token importance is concentrated in a local receptive field of roughly 50 tokens around the current time step. This means that even if you could afford to train on longer sequences, the attention mechanism itself may not efficiently use that longer context — it tends to focus on nearby tokens and struggles to surface relevant information from thousands of steps ago. The demand for deploying deeper and larger networks is growing, but scaling the sequence length dimension remains an open, arguably even more consequential, challenge than scaling model depth or width.

Why This Matters: The Long-Range Dependency Bottleneck

The practical importance of long-range sequence modeling extends across multiple domains:

Language modeling on long documents: Books (PG19), scientific papers (arXiv), and code repositories (GitHub) contain dependencies that span far beyond a typical 4K training window. A model that cannot access earlier chapters, earlier definitions, or earlier function declarations is fundamentally limited in what it can predict.
Multi-modal and structured data: Beyond text, tasks like processing long audio streams, video understanding, or reasoning over knowledge bases require models that can track relationships across vast temporal or structural distances.
Inference-time generation: When generating long outputs (stories, code files, dialogue histories), the model needs to remain coherent over thousands of tokens — a requirement that exceeds the training window length of most deployed models.

The authors frame this as an orthogonal scaling dimension to model size (Section 1): scaling sequence length could be "potentially even more consequential" than scaling parameter count or layer depth, because it unlocks qualitatively different capabilities (understanding entire books, reasoning over entire codebases) rather than just incrementally improving performance on existing tasks.

Two Families of Prior Approaches — and Their Shortcomings

Prior work on long-range sequence modeling largely falls into two camps: efficient attention variants and state space models (SSMs). Each has made progress but each has clear limitations that this paper seeks to overcome.

Efficient Attention Approximations

A large body of work attempts to reduce the $\mathcal{O}(L^2)$ cost of self-attention while preserving its essential properties. The paper catalogs several approaches in Section 1: local or sliding-window attention (Dai et al., 2019; Transformer-XL), sparse attention (Child et al., 2019; Zaheer et al., 2020), low-rank approximations (Wang et al., 2020; Linformer), and kernel-based linearization (Choromanski et al., 2020; Performer; Katharopoulos et al., 2020). These methods succeed in reducing the asymptotic complexity, but the paper identifies a critical weakness: they "notoriously struggle on long-input classification tasks" (Section 1, citing Tay et al., 2020). The gap between efficiency and effectiveness is real — making attention cheaper often means making it less capable of actually finding and using distant information.

A more specific variant is the Block-Recurrent Transformer (BRECT) (Hutchins et al., 2022), which this paper treats as its most direct predecessor. BRECT processes sequences in fixed-size blocks (e.g., windows of 512 tokens) and uses a recurrent mechanism — a learned gating function — to pass information between consecutive blocks. The attention within each block is $\mathcal{O}(W^2)$ , where $W$ is the window size, and the recurrence across blocks adds $\mathcal{O}(L \cdot W)$ cost overall. This is a significant improvement over $\mathcal{O}(L^2)$ , but it introduces a sequential bottleneck: each block must wait for the recurrent state from the previous block, making the layer inherently sequential. As the authors emphasize (Section 1, Section 3.2), this sequential dependency prevents full parallelization of the forward pass — a major limitation on modern hardware that thrives on parallel computation. The paper's Figure 4 (left) quantifies this: a BRECT layer runs roughly 15× slower than a simple sliding-window Transformer layer (SLIDE:12L) at a 4K sequence length with a window size of 128.

State Space Models (SSMs)

SSMs represent a fundamentally different approach to sequence modeling. Rather than building an explicit attention matrix, they model sequences as the output of a linear time-invariant (LTI) dynamical system (Equations 1–3 in the paper). A latent state $x_k$ evolves through a transition matrix $\mathbf{A}$ and input projection $\mathbf{B}$ , and is mapped to an output via $\mathbf{C}$ . Crucially, this recurrence can be unrolled into a convolution and computed efficiently using the Fast Fourier Transform (FFT) in $\mathcal{O}(L \log L)$ time. Moreover, by deriving $\mathbf{A}$ and $\mathbf{B}$ from the HiPPO framework (Gu et al., 2020), SSMs can be designed to optimally capture long-range dependencies — in theory, much better than Transformers.

In practice, SSMs have demonstrated impressive results. The paper cites their success on long-range dependency benchmarks like the Long Range Arena (LRA; Tay et al., 2020), where models like S4 (Gu et al., 2022) and S5 (Smith et al., 2023) have "outperformed Transformers by a large margin" (Section 1). They are also computationally attractive: the convolution formulation allows parallel training, and the recurrent formulation enables efficient autoregressive inference.

However, SSMs have a critical limitation that motivates this entire paper: they have "not yet completely matched Transformers as an off-the-shelf sequence model for general language modeling tasks" (Section 1). Despite their theoretical advantages for long sequences, when it comes to the core task of next-token prediction on standard language modeling benchmarks, Transformers still hold the lead. The authors cite Fu et al. (2023; H3) as evidence that this gap persists. The reasons are not fully spelled out in the paper, but the implication is that the Transformer's attention mechanism — with its ability to precisely compare specific token pairs — provides an inductive bias that is particularly well-suited to the local linguistic structure of natural language. SSMs, which compress all history into a fixed-size state vector, may lose the fine-grained token-level information that attention excels at capturing.

Hybrid Approaches: GSS-HYBRID and MEGA

Recognizing the complementary strengths of Transformers and SSMs, two prior works attempted to combine them:

GSS-HYBRID (Mehta et al., 2023) interleaves Gated State Space (GSS) layers with standard Transformer layers — for example, placing a Transformer layer every 4th layer in a stack otherwise composed of SSM layers. The SSM handles the long-range context, and the Transformer layers provide local attention. However, the paper notes two issues: (a) GSS-HYBRID required grid-searching over learning rates and using different learning rates for SSM and Transformer layers to avoid training instability, suggesting that simply stacking the two layer types is not seamless; (b) more fundamentally, the SSM and attention mechanisms operate in separate layers — they never directly interact. The SSM produces a hidden state that feeds into the next layer, but there is no mechanism for the attention to explicitly query the SSM's compressed history.

MEGA (Ma et al., 2023) takes a different approach, using an Exponentially Moving Average (EMA) — which the authors note can be seen as a simple form of state space model — and mixing its outputs with a Gated Attention Unit (GAU) through explicit gating mechanisms. This represents a tighter integration than GSS-HYBRID, but it is tied to a specific attention variant (GAU with a single head) and a specific gating scheme. The authors argue this makes MEGA less modular — you cannot easily swap in a different SSM or a different attention mechanism.

The Unresolved Gap

The landscape prior to this paper can be summarized as follows:

Architecture	Long-range	Local precision	Parallelism	Modularity
Standard Transformer	No ( $\mathcal{O}(L^2)$ )	Yes (full attention)	Yes	Yes
Efficient attention variants	Partial	Degraded	Mostly yes	Yes
BRECT	Yes (recurrence)	Yes (block attention)	No (sequential)	Moderate
Pure SSMs (S4, DSS)	Yes (theoretically optimal)	Lagging on LM	Yes (convolution)	N/A
GSS-HYBRID	Yes	Yes (separate layers)	Yes	No (interleaved, not integrated)
MEGA	Yes	Yes (GAU)	Yes	No (tied to specific components)

The gap is clear: no existing architecture simultaneously provides the local precision of attention, the long-range modeling of SSMs, full parallelizability, and modularity to swap components. BRECT offers the right functional combination but is sequential and slow. GSS-HYBRID and MEGA are parallel but either keep SSM and attention separate or tie them to specific architectural choices.

How This Paper Positions Itself

The Block-State Transformer (BST) is designed to fill exactly this gap. The key architectural insight, stated in Section 3.2, is to integrate SSM states directly into the attention mechanism via cross-attention, rather than keeping them in separate layers (GSS-HYBRID) or using them only as gating signals (MEGA). Specifically:

An SSM sublayer processes the entire input sequence to produce a sequence of "context states" that encode long-range history.
These context states are fed into a Block Transformer's cross-attention, where the local block of tokens can explicitly attend to and retrieve information from the compressed history.
The Block Transformer cells operate in parallel across all blocks because the SSM has already computed all context states upfront — there is no sequential recurrence between blocks.

This design simultaneously addresses all four desiderata:

Long-range modeling: The SSM captures dependencies across the full sequence (and beyond, for structured SSMs).
Local precision: The Block Transformer's self-attention within each window provides fine-grained token comparisons.
Full parallelizability: By removing the recurrent cell of BRECT and replacing it with an SSM whose convolution can be computed for all time steps at once, the entire layer can run in parallel. The paper reports a 6–11× speedup over BRECT at the layer level (Section 4.3, Figure 4 left).
Modularity: The cross-attention interface is clean — any SSM (S4, DSS, unstructured Hyena-style kernels) can generate context states, and any attention variant can consume them. The paper's three variants (SH, MH, MF) demonstrate this modularity by exploring different ways to construct the context states without changing the core architecture.

The paper also positions its three context-state variants — Single-Head, Multi-Head, and Multi-Filter — as exploring a deliberate tradeoff between redundancy and retrievability (Section 3.3). SH provides the most redundant (and therefore most retrievable) context by including all adjacent SSM states from the current window. MH reduces redundancy by having the SSM produce separate features per attention head. MF minimizes redundancy by using separate filters that each produce a single context state from the last position of the previous window, but sacrifices some retrievability of recent information. This taxonomy gives practitioners a design knob that can be tuned based on whether local (SH, for tasks like PG19 where local context dominates) or structured long-range memory (MF, for tasks like GitHub with cross-file dependencies) is more important.

Finally, an important subtlety in the paper's positioning concerns length generalization (Section 4.2). The authors distinguish between structured SSMs (like S4), whose convolution kernel $\mathbf{K}$ is parameterized by fixed-size matrices $\mathbf{A}$ , $\mathbf{B}$ , $\mathbf{C}$ and can be extended to arbitrary sequence lengths without recompilation, and unstructured SSMs (like the Hyena-inspired variant), whose kernel is parameterized directly as a weight vector of length $L$ and cannot be trivially extended. This distinction proves critical: only the structured variants (BST:SH:S4) generalize to 65K sequences, while the unstructured variants (which perform best at the training length of 4K) degrade at unseen lengths. The paper thus positions structured SSMs as the necessary choice for deployment scenarios that require evaluation on sequences longer than those seen during training — a practical consideration that prior hybrid architectures did not systematically analyze.

3. Technical Approach

3.1 Reader Orientation

The Block-State Transformer (BST) is a hybrid neural network layer that you can stack to build a language model — it takes a sequence of token embeddings as input and produces a sequence of context-aware output embeddings, just like a standard Transformer layer, but with a fundamentally different internal structure. The problem it solves is the sequential bottleneck in existing architectures that combine local attention with long-range memory: previous approaches like the Block-Recurrent Transformer force each block of tokens to wait for the previous block's recurrent state before processing can begin, making the forward pass inherently serial. BST eliminates this by replacing the recurrent cell with a State Space Model whose convolution can be computed for all time steps simultaneously, enabling every block to process in parallel while still receiving a compressed representation of the entire preceding sequence through cross-attention.

3.2 Big-Picture Architecture (Diagram in Words)

The BST layer has four major components that operate in a specific sequence:

State Space Model (SSM) sublayer — receives the entire input sequence of length $L$ and produces a "context sequence" of the same length $L$ , where each position contains a compressed representation of all preceding tokens. This is computed via Fast Fourier Transform (FFT) convolution in $\mathcal{O}(L \log L)$ time.
Context state collector — takes the SSM's output sequence and, depending on the variant (SH, MH, or MF), extracts or reorganizes a smaller set of $S$ "context states" per block. These context states serve as the long-range memory that each block's attention mechanism can query. Crucially, the collection strategy determines the tradeoff between redundancy (adjacent states carrying overlapping information) and retrievability (ease of recovering recent token-level detail).
Block Transformer sublayer — processes the input sequence divided into blocks of window size $W$ . For each block, a self-attention mechanism captures local token relationships within the block, and a cross-attention mechanism queries the context states to retrieve relevant long-range information. The block cells run fully in parallel because the SSM has already pre-computed all context states.
Output projection — concatenates the self-attention and cross-attention outputs and projects them back to the model dimension.

The total time complexity is $\mathcal{O}(W^2) + \mathcal{O}(L \log L)$ , where the first term is the per-block attention cost and the second term is the SSM convolution cost. This contrasts with BRECT's $\mathcal{O}(L \cdot W)$ , which includes a sequential factor that prevents parallelization.

3.3 Roadmap for the Deep Dive

The technical breakdown proceeds in the following order, which mirrors the flow of computation through the layer and builds understanding from the underlying mathematical machinery to the architectural integration:

First, State Space Model fundamentals (Section 3.1 expanded): The LTI dynamical system, its recurrent and convolutional forms, the HiPPO framework, and the distinction between structured and unstructured kernels. This is the computational engine that makes long-range context possible without sequential recurrence.
Second, the BST layer structure (Section 3.2 expanded): How the SSM sublayer connects to the Block Transformer sublayer, what replaces BRECT's recurrent cell, and the cross-attention interface that enables full parallelization.
Third, the three context state variants (Section 3.3 expanded): Single-Head, Multi-Head, and Multi-Filter — their construction algorithms, their different tradeoffs along the redundancy-retrievability axis, and the implications for causal masking and positional encoding.
Fourth, implementation details (Section 3.4 expanded): Context IDs, positional embedding choices, the down-sampling strategy that reduces FFT cost by 4×, and specific hyperparameter configurations.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a new architecture paper whose core idea is that replacing the sequential recurrent cell in a Block-Recurrent Transformer with a parallelizable State Space Model — and connecting it via cross-attention rather than gating — yields a layer that is both faster and equally or more performant, while enabling length generalization when structured SSM kernels are used.

State Space Model Fundamentals: The Computational Engine

The Linear Time-Invariant (LTI) Dynamical System

The BST layer's ability to compress long sequences into context states rests on the mathematical machinery of state space models. An SSM is defined by a linear time-invariant dynamical system with the following discrete-time recurrence:

\begin{aligned} x_k &= \mathbf{A}x_{k-1} + \mathbf{B}u_k, \\ y_k &= \mathbf{C}x_k + \mathbf{D}u_k. \end{aligned}

where $\mathbf{A} \in \mathbb{R}^{N \times N}$ is the state transition matrix that determines how the internal state evolves from one time step to the next, $\mathbf{B} \in \mathbb{R}^{N \times 1}$ is the input projection vector that maps the 1-D input $u_k$ into the $N$ -dimensional state space, $\mathbf{C} \in \mathbb{R}^{1 \times N}$ is the output projection vector that reads out a scalar from the $N$ -dimensional state, $\mathbf{D} \in \mathbb{R}^{1 \times 1}$ is a skip-connection term (treated as optional and omitted from subsequent discussion), $u_k$ is the scalar input at time step $k$ , $x_k \in \mathbb{R}^N$ is the internal state vector at step $k$ (a compressed representation of the history $u_0, u_1, \dots, u_k$ ), and $y_k$ is the scalar output at time step $k$ .

What it computes: At each time step $k$ , the system takes the current input $u_k$ and combines it with the previous state $x_{k-1}$ (transformed by $\mathbf{A}$ ) to produce a new state $x_k$ . This state is then read out through $\mathbf{C}$ to produce the output $y_k$ , with an optional skip connection $\mathbf{D}u_k$ adding the input directly to the output. Operationally, you can think of $x_k$ as a running summary of everything seen so far — the system compresses the entire history $\{u_0, \dots, u_k\}$ into a fixed-size vector of dimension $N$ .

Why this form: The linearity of the system is the key property that enables everything that follows. Because the state evolves linearly, the recurrence can be unrolled into an explicit convolution (shown below), and the convolution can be computed in parallel via FFT. A non-linear recurrence would not permit this unrolling and would require sequential computation. The choice of $\mathbf{A}$ , $\mathbf{B}$ , and $\mathbf{C}$ as learnable parameters (or analytically derived parameters, in the case of HiPPO) determines what aspects of the history the state $x_k$ preserves — the HiPPO framework, discussed below, provides a principled way to choose these matrices so that the state optimally represents the history in a polynomial basis.

Unrolling the Recurrence into a Convolution

The critical insight that makes SSMs computationally efficient is that the recurrence in Equation (1) can be explicitly unrolled into a discrete convolution. Starting from the initial condition $x_{-1} = \vec{0}$ , successive substitution yields:

y_k = \sum_{j=0}^{k} \mathbf{C} \mathbf{A}^j \mathbf{B} \cdot u_{k-j}.

where $j$ indexes how many steps back in time we are looking, $\mathbf{C} \mathbf{A}^j \mathbf{B}$ is a scalar (since $\mathbf{C}$ is $1 \times N$ , $\mathbf{A}^j$ is $N \times N$ , and $\mathbf{B}$ is $N \times 1$ ) that represents the impulse response of the system — how much the input from $j$ steps ago influences the current output.

What it computes: This equation expresses the output $y_k$ as a weighted sum of all past inputs $u_{k-j}$ , where the weights are given by $\mathbf{C} \mathbf{A}^j \mathbf{B}$ . The weight for an input from $j$ steps ago decays (or oscillates, or grows) according to the $j$ -th power of the state transition matrix $\mathbf{A}$ , projected through $\mathbf{C}$ and $\mathbf{B}$ . This is exactly a linear time-invariant filter — the system's response to an impulse at time $k-j$ is determined solely by the lag $j$ , not by the absolute time $k$ .

Why this form: This unrolling transforms a sequential computation (recurrence) into a parallel one (convolution). Once you have the convolution formulation, you can define a kernel $\mathbf{K} \in \mathbb{R}^L$ as:

\mathbf{K} = (\mathbf{C} \mathbf{B}, \mathbf{C} \mathbf{A} \mathbf{B}, \dots, \mathbf{C} \mathbf{A}^{L-1} \mathbf{B}).

Then the entire output sequence $y$ can be computed in one shot as $y = \mathbf{K} * u$ , where $*$ denotes convolution. The Fast Fourier Transform (FFT) computes this convolution in $\mathcal{O}(L \log L)$ time rather than the $\mathcal{O}(L^2)$ that a naive convolution would require, or the $\mathcal{O}(L \cdot N^2)$ that the recurrence would require if computed sequentially (since each step involves an $N \times N$ matrix multiplication). For large $L$ , the $\log L$ factor is the difference between practical training and infeasible training.

The HiPPO Framework: Structured State Initialization

The matrices $\mathbf{A}$ and $\mathbf{B}$ are not arbitrary — SSMs that achieve strong long-range dependency modeling use the HiPPO framework (Gu et al., 2020) to derive them analytically. The HiPPO framework poses the following question: given a continuous input signal $u(t)$ observed up to time $t$ , what is the optimal way to compress it into a fixed-dimensional state vector $x(t)$ such that $x(t)$ can be used to reconstruct the history? The answer involves projecting the signal onto a set of orthogonal polynomials (Legendre, Laguerre, or Chebyshev) and tracking the coefficients of this projection over time.

The result is that $\mathbf{A}$ and $\mathbf{B}$ have closed-form expressions that depend on the chosen polynomial basis. For example, the HiPPO-LegS (scaled Legendre) matrices ensure that the state $x_k$ represents the coefficients of the optimal polynomial approximation of the history $\{u_0, \dots, u_k\}$ . This is what gives structured SSMs like S4 their theoretical guarantee of capturing long-range dependencies — the state is explicitly designed to preserve information uniformly across the entire history, rather than exponentially forgetting old inputs as a vanilla RNN would.

Why this matters for BST: The paper uses S4 (Gu et al., 2022) and its diagonal variant DSS (Gupta et al., 2022) as the structured SSM backends. These models take the HiPPO-derived $\mathbf{A}$ and $\mathbf{B}$ matrices, apply additional parameterizations to make them diagonal (which simplifies the computation of $\mathbf{A}^j$ ), and expose $\mathbf{C}$ as a learnable output matrix. The key practical property for BST is that the kernel $\mathbf{K}$ is length-independent — you can compile $\mathbf{K}$ from $\mathbf{A}, \mathbf{B}, \mathbf{C}$ for any sequence length $L$ without retraining, because $\mathbf{A}, \mathbf{B}, \mathbf{C}$ have fixed sizes independent of $L$ . This is what enables the length generalization results in Section 4.2.

Unstructured (Explicitly Parameterized) Filters: The Hyena-Inspired Alternative

In contrast to structured SSMs, the paper also experiments with unstructured filters, inspired by the Hyena hierarchy (Poli et al., 2023). Here, the convolution kernel is parameterized directly as a trainable weight vector, without deriving it from an underlying state space:

\bar{\mathbf{K}}_t = e^{-\alpha t} \cdot (\text{FFN} \circ \text{PositionalEncoding})(t).

where $t \in \{0, 1, \dots, L-1\}$ indexes the position in the kernel, $\alpha$ is a learned decay rate that exponentially suppresses weights for larger lags (implementing a recency bias), $\text{FFN}$ is a small feed-forward network (a multi-layer perceptron), and $\text{PositionalEncoding}(t)$ is a sinusoidal or learned positional encoding that gives the FFN a sense of absolute position.

What it computes: For each lag $t$ , the FFN takes a positional encoding of $t$ as input and outputs a $D$ -dimensional vector. This vector is multiplied by the exponential decay $e^{-\alpha t}$ to ensure that weights for very distant past inputs are suppressed. The collection of all $\bar{\mathbf{K}}_t$ for $t = 0, \dots, L-1$ forms the convolution kernel $\bar{\mathbf{K}} \in \mathbb{R}^{L \times D}$ , which is then used in the same FFT-based convolution as the structured kernel.

Why this form: Separating the kernel parameter count from the sequence length via the FFN (which has a fixed number of parameters regardless of $L$ ) makes the model more parameter-efficient than directly learning $L \times D$ independent weights. The exponential decay provides a simple but effective regularization that enforces a recency bias — something that the structured S4 kernel achieves implicitly through its HiPPO parameterization. The paper notes two practical advantages over S4: (1) the kernel does not need to be "recompiled" from $\mathbf{A}, \mathbf{B}, \mathbf{C}$ for different lengths (though this is also somewhat true for S4), and (2) the unstructured kernel has "more free parameters" because it is no longer constrained to lie in the space of kernels realizable by the low-rank factorization $\mathbf{C} \mathbf{A}^j \mathbf{B}$ , potentially providing richer representations.

The critical tradeoff: The unstructured kernel $\bar{\mathbf{K}}$ is parameterized for a fixed maximum length $L$ (the training sequence length). Unlike the structured kernel, it cannot be extended to arbitrary lengths at inference time — the positional encoding and FFN were trained for positions $\{0, \dots, L-1\}$ , and extending beyond $L$ requires either interpolation or retraining. This is why the unstructured variants (BST:SH:UNSTRUCT, BST:MF:UNSTRUCT) perform best at the training length of 4K but degrade at 16K and 65K in Figure 3, while the structured variants (BST:SH:S4) generalize well.

The Block-State Transformer Layer: Integration of SSM and Attention

What the BST Layer Replaces

To understand the BST layer, it helps to understand what it replaces. The Block-Recurrent Transformer (BRECT) cell operates as follows:

Divide the input sequence into blocks of size $W$ .
Process block 1: self-attention within block 1, cross-attention to an initial recurrent state (usually zeros).
Compute a new recurrent state from block 1's outputs via a learned gating function.
Process block 2: self-attention within block 2, cross-attention to the recurrent state from block 3.
Compute a new recurrent state from block 2's outputs.
Repeat for all blocks.

Steps 2 through 6 form a sequential chain — block $i$ cannot begin until the recurrent state from block $i-1$ is available. This for-loop over blocks is what makes BRECT slow (15× slower than a sliding-window Transformer without recurrence, per Figure 4).

The BST layer replaces this with:

Run the entire input sequence (all $L$ tokens) through the SSM convolution, producing a context sequence of length $L$ — all in parallel via FFT.
Divide both the input sequence and the context sequence into blocks of size $W$ .
For each block, collect the relevant context states from the SSM output (the collection method depends on the variant: SH, MH, or MF).
Process each block in parallel: self-attention within the block, cross-attention to its assigned context states.
Concatenate self-attention and cross-attention outputs, project back to the model dimension.

The SSM has already encoded the entire history into the context sequence in step 1, so there is no sequential dependency between blocks in steps 3–5. This is what enables full parallelization of the forward pass.

The Cross-Attention Interface: How Blocks Query Long-Range Context

The cross-attention mechanism within each Block Transformer cell is the bridge between the SSM's compressed representation and the Transformer's token-level precision. Given:

A block of input token embeddings: $E_{\text{block}} \in \mathbb{R}^{W \times D}$ (window size $W$ , embedding dimension $D$ ).
Context states for that block: $C \in \mathbb{R}^{S \times D \times H}$ (number of context states $S$ , embedding dimension per head $D$ , number of heads $H$ — with $S$ typically set equal to $W$ for SH and MH variants, and an independent small number like 32 or 64 for MF).

The cross-attention operation computes, for each attention head:

$\text{CrossAttn}(E_{\text{block}}, C) = \text{softmax}\left(\frac{Q(E_{\text{block}}) \cdot K(C)^T}{\sqrt{d_k}}\right) \cdot V(C)$

where $Q$ projects the token embeddings to queries (shape $W \times d_k$ ), $K$ projects the context states to keys (shape $S \times d_k$ ), and $V$ projects the context states to values (shape $S \times d_v$ ). The softmax is computed over the $S$ dimension, so each token in the block computes a weighted average of the context states, where the weights are determined by how well the token's query matches each context state's key.

What it computes: Each token in the current block produces a query vector — "what information do I need from the past?" — and compares it against key vectors derived from the context states — "what information does each context state contain?" The attention weights determine which context states are most relevant, and the weighted sum of value vectors produces a context-aware representation for each token. This is the same mechanism as standard Transformer cross-attention, but with the crucial difference that the keys and values come from the SSM's compressed history rather than from a separate encoder.

Why cross-attention (rather than gating): The paper explicitly contrasts this with GSS-HYBRID and MEGA. GSS-HYBRID keeps SSM and attention in separate layers — the SSM's output feeds into the next layer's input, but attention never directly queries the SSM's state. MEGA uses gating mechanisms to mix SSM outputs with attention outputs, which is more integrated than GSS-HYBRID but still doesn't allow attention to selectively retrieve from the compressed history. BST's cross-attention interface allows each attention head to learn what to retrieve from the context states — the query-key matching provides a content-addressable memory over the compressed history. This is a more flexible and powerful interface than gating, which applies a fixed (learned but input-independent) mixture of SSM and attention outputs.

The self-attention and cross-attention outputs are concatenated (not added or gated) and then projected back to dimension $D$ via a learned linear transformation. This preserves both the local token-token interactions (from self-attention) and the long-range context retrieval (from cross-attention) as separate, non-interfering signals that the projection can learn to combine optimally.

Causal Masking in Cross-Attention

A subtler detail concerns causality. The SSM's convolution is computed over the entire sequence, so the context state at position $t$ contains information about tokens at positions $0, 1, \dots, t$ . When a token at position $t$ within a block attends to context states, it must not attend to context states that correspond to future positions within the same block (since those context states encode information about tokens the model hasn't seen yet during autoregressive generation).

For the Single-Head (SH) and Multi-Head (MH) variants, the context states for a block are taken from the SSM outputs corresponding to the same block's positions. Since the SSM outputs at positions $t, t+1, \dots, t+W-1$ are arranged sequentially as the $S = W$ context states, a triangular causal mask is applied during cross-attention: a token at the $i$ -th position within the block can only attend to context states $0, 1, \dots, i$ , not to context states $i+1, \dots, W-1$ . This is identical to the causal masking in standard self-attention.

For the Multi-Filter (MF) variant, the context states are taken from the last position of the previous window, not from the current window. Therefore, all context states encode information that strictly precedes the current block — there is no risk of attending to future information, and no causal mask is needed in cross-attention. This is one of the practical advantages of the MF approach: simpler attention computation without masking overhead.

Context State Construction: The Three Variants

The three BST variants — Single-Head (SH), Multi-Head (MH), and Multi-Filter (MF) — differ solely in how they construct the context tensor from the SSM's output sequence. The rest of the architecture (SSM convolution, Block Transformer, cross-attention) is identical across variants. Understanding these variants requires understanding the shape transformations involved.

The SSM output shape. For a single SSM layer processing a sequence of $L$ tokens with embedding dimension $D$ , the output is a tensor of shape $(B \times L \times D)$ , where $B$ is the batch size. This is $L$ vectors of dimension $D$ , each encoding the history up to that position. To use this output in cross-attention with $H$ attention heads, the context tensor must have shape $(B \times S \times D \times H)$ for each block, where $S$ is the number of context states per block (typically $S = W$ for SH and MH; $S$ is an independent hyperparameter like 32 or 64 for MF).

Single-Head (SH): Maximum Redundancy, Maximum Retrievability

The SH variant is the simplest and most directly analogous to BRECT's recurrent states. The construction proceeds as follows:

Run a single SSM: The SSM has one output channel (one $\mathbf{C}$ matrix), convolving a single filter over the input sequence to produce a context sequence $y_1 \in \mathbb{R}^{L \times D \times 1}$ — that is, one context vector of dimension $D$ per position. This is implemented by setting channels=1 in the convolution function (Pseudocode 2).
Lift to multiple heads: A learned dense (fully-connected) layer maps the single context vector at each position from dimension $D$ to dimension $D \times H$ , producing $y_h \in \mathbb{R}^{L \times D \times H}$ . This is the same operation that a standard Transformer uses to project inputs to multiple heads, but applied to the SSM outputs rather than token embeddings.
Split into blocks: The $L$ positions are divided into $L/W$ blocks of size $W$ along the time axis, producing context states of shape $(L/W \times W \times D \times H)$ . For each block, $S = W$ context states are available.
Apply causal mask: During cross-attention, a triangular mask ensures token $i$ in the block only attends to context states $0, \dots, i$ .

Redundancy analysis: Because adjacent positions in the SSM output are highly correlated (the SSM state evolves smoothly), the $W$ context states contain substantial overlap — position $t$ and position $t+1$ encode nearly the same history, differing only by the contribution of the single token at $t+1$ . This is redundant in the information-theoretic sense: the $W$ vectors do not represent $W$ independent pieces of information about the past. However, this redundancy makes the context highly retrievable: if a token needs to know about a specific recent position, its exact context state is in the set (in the correct sequential order), and the attention mechanism can simply learn to attend to the corresponding position.

When SH performs best: On PG19 (books), where local context (recent sentences and paragraphs) is most predictive of the next token, SH outperforms MF (Table 1). The redundancy of adjacent states makes it easy to recover exact recent information — the model can attend to the specific context state that corresponds to the relevant prior token.

Multi-Head (MH): Reduced Redundancy per Head, Complementary Features

The MH variant modifies step 1 of the SH procedure:

Run an SSM with $H$ separate output channels: Instead of a single $\mathbf{C}$ matrix, the SSM has $H$ different $\mathbf{C}$ matrices $[C_1, C_2, \dots, C_H]$ , each producing a separate output channel. This means the convolution runs with channels=num_heads, producing $y_H \in \mathbb{R}^{L \times D \times H}$ directly — no dense projection is needed (Pseudocode 3).
Split into blocks: Same as SH — each block gets $W$ context states, each already shaped for $H$ heads.
Apply causal mask: Same as SH.

The conceptual difference from SH: In SH, the SSM produces one compressed representation of the history (via one $\mathbf{C}$ ), and the multi-head projection simply copies this same information to all heads (with different learned projections, so heads can emphasize different aspects of the same underlying representation). In MH, the SSM's internal state $x_k \in \mathbb{R}^N$ is a rich, $N$ -dimensional representation of the entire history. Each $\mathbf{C}_i$ reads out a different linear projection of this rich state. Since the $\mathbf{C}_i$ matrices can be learned to extract complementary features, the $H$ context vectors at each position potentially contain more diverse information about the history than the SH variant, where all heads ultimately derive from the same scalar readout.

Redundancy-reduction analysis: The context is less redundant than SH because the $H$ readout matrices can specialize. For example, $\mathbf{C}_1$ might learn to extract syntactic structure, $\mathbf{C}_2$ might extract semantic content, $\mathbf{C}_3$ might extract named entities, etc. (These are hypothetical specializations — what the $\mathbf{C}_i$ actually learn is determined by the training objective.) However, the context is still redundant along the time axis — adjacent positions still encode highly similar information, just through multiple complementary lenses.

Practical note: The paper reports that MH requires approximately 8% more parameters than SH (Table 1: BST:MH:S4 at 218M vs. BST:SH:S4 at 202M) and performs "on par with the faster BST:SH variant." Given the added parameter cost and negligible perplexity difference, the paper does not strongly advocate for MH over SH, but includes it to demonstrate the design space.

Multi-Filter (MF): Minimum Redundancy, Independent Context States

The MF variant represents the most significant departure from BRECT's design. The construction is as follows:

Run $S$ separate SSM filters (or one SSM with $S$ output channels): The convolution runs with channels=num_states (e.g., $S = 32$ ), producing $y_S \in \mathbb{R}^{L \times D \times S}$ (Pseudocode 4). Each of the $S$ channels is the result of convolving the input with a different filter — effectively, $S$ independent SSMs, each learning to extract a different feature from the history.
Take only the LAST position from the PREVIOUS window: For each block, instead of collecting all $W$ context states from the block's own positions (as in SH/MH), the MF variant collects only the single last context state from the preceding window. Specifically, the context states for block $b$ are the $S$ output values at position $W \cdot b - 1$ (the last token of block $b-1$ ), yielding a tensor of shape $(S \times D)$ . This is done for each block.
Shift and initialize: The context states are shifted by one block position (so block $b$ gets the states from block $b-1$ ), and the first block gets learned initial context vectors (since it has no preceding block).
Lift to multiple heads: A dense layer projects the $S \times D$ context tensor to $S \times D \times H$ , producing the final context states for cross-attention.
No causal mask needed: Since all context states come from a position that precedes the entire current block, cross-attention can be unmasked.

Redundancy analysis: The MF context is the least redundant of the three variants. There is no overlap along the time axis (only one position per block is used), and the $S$ different filters are trained to extract complementary information (similar to how different attention heads learn different aspects). The context is a compact, low-redundancy summary of the history.

The retrievability cost: By discarding all but the last position of the previous window, MF sacrifices the ability to easily recover information about specific recent tokens. If the model needs to know exactly what the 10th token in the previous block was, that information is not directly available as a context state — it has been compressed through the SSM's state and mixed with information from all other tokens. The $S$ filters must learn to preserve enough detail that the attention mechanism can retrieve what it needs.

When MF performs best: On arXiv and GitHub (Table 1), where long-range structural dependencies (cross-references, import statements, function definitions) are more important than exact recent word order, MF outperforms SH. The $S$ independent filters can each specialize in tracking different types of long-range dependencies (e.g., one filter tracks variable definitions, another tracks function signatures, another tracks class hierarchies), and the attention mechanism can query the relevant filter's state when needed.

The Redundancy-Retrievability Tradeoff, Formalized

The paper frames the three variants as points on a spectrum:

SH: High redundancy (adjacent SSM states overlap significantly), high retrievability (easy to recover specific recent information), but potentially wasteful use of the $S$ context "slots" (storing highly correlated information).
MH: Medium redundancy (multiple readouts of the same underlying SSM state provide complementary views, but adjacent positions still overlap), medium retrievability.
MF: Low redundancy (only one position per block, each filter independent), medium-to-low retrievability (compressed summary may lose fine-grained detail), but efficient use of the context slots (each slot stores a different feature).

The "correct" choice depends on the data distribution. The paper's empirical results suggest that local-context-heavy datasets (PG19) favor SH's retrievability, while datasets with long-range structural dependencies (arXiv, GitHub) can benefit from MF's efficient feature extraction. This is a design knob for practitioners: if your task requires the model to frequently refer to exact recent wordings, use SH; if it requires tracking many different types of long-range dependencies, use MF.

Implementation Details: Making It Work in Practice

The paper describes several implementation choices that are essential for making the BST layer practical, particularly on TPU hardware.

Context IDs and Positional Embeddings

The paper makes a deliberate departure from standard Transformer practice by not using global learned positional embeddings on the token embeddings. Instead, it relies on two sources of positional information:

The SSM itself encodes positional information into the context states. Because the convolution kernel is applied across the time dimension, the SSM output at position $t$ is a function of inputs at positions $0, \dots, t$ with weights that depend on the lag — this inherently captures sequence order. The paper argues that this is sufficient for the context states and that additional positional embeddings on token embeddings are unnecessary.
A T5-style relative position bias (Raffel et al., 2020) is used in the self-attention mechanism within each Block Transformer. Relative position bias adds a learned scalar bias to the attention logits based on the relative distance between the query and key positions, rather than the absolute position. This is more compatible with windowed attention (where absolute positions within the window don't correspond to absolute positions in the document) and generalizes better to unseen window sizes.

Context IDs for MF variant: In the SH and MH variants, the context states correspond to sequential positions within the current block, so the SSM's inherent positional encoding provides ordering information. In the MF variant, however, the $S$ context states correspond to different filters, not different time positions — they have no natural ordering. To allow the attention mechanism to distinguish between the different filters, the paper adds learned "context IDs" — a set of $S$ unique trainable vectors that are added to the context states before cross-attention. These serve the same role as positional embeddings but for the filter dimension rather than the time dimension: they give each filter a persistent identity that the attention mechanism can learn to associate with specific types of information.

Why no context IDs for SH and MH: The paper explicitly states that for SH and MH, "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." In other words, the fact that context state $i$ in SH/MH corresponds to time step $i$ within the block already provides ordering information, and adding explicit position IDs would be redundant.

Down-Sampling: Reducing FFT Cost by 4×

The paper identifies FFT operations as "the main source of bottleneck when training SSMs on TPUs" (Section 3.4). To address this, the input embeddings to the SSM sublayer are projected to a lower-dimensional space — specifically, to one quarter of the embedding dimension ( $D/4$ ). The SSM convolution then operates in this reduced dimension, and the output is projected back to the full dimension $D$ before being passed to the Block Transformer.

What it computes: A learned linear projection (a dense layer) maps the input from $\mathbb{R}^{L \times D}$ to $\mathbb{R}^{L \times D/4}$ . The SSM convolution runs on this lower-dimensional signal, producing output of shape $\mathbb{R}^{L \times D/4}$ . A second learned linear projection maps back to $\mathbb{R}^{L \times D}$ . This reduces the total number of FFTs by a factor of 4, since the FFT operates on each channel independently.

Why this is effective: The FFT's computational cost scales with the embedding dimension — halving the dimension roughly halves the FFT cost (for the transform itself; the convolution multiplication scales similarly). The paper reports that this dimension reduction has "negligible impact to perplexity" (for the MF variant, the state dimension is further reduced by an additional factor of two, again with negligible impact). This suggests that the SSM's compressed representation is highly redundant across the embedding dimension — the information can be represented in a much lower-dimensional space without loss, and the projection back to the full dimension recovers any lost expressivity through learned up-sampling.

A practical note on MF: For the Multi-Filter variant, the paper reduces the SSM state dimension by an additional factor of two (so $D/8$ instead of $D/4$ ) to "improve FFT efficiency." The fact that MF still performs competitively with this further reduction (and even outperforms SH on arXiv and GitHub) suggests that the multiple independent filters compensate for the reduced per-filter capacity.

Training Configuration and Model Architecture Details

The paper provides specific architectural hyperparameters (Section 4.1, "Training details"):

Smaller models (~200M parameters): 12 layers total, with BST layers (containing both SSM and Block Transformer) at layers {1, 7, 9}. The remaining 9 layers are standard Block Transformer layers without SSM (i.e., sliding-window attention only). 8 attention heads of dimension 128, embedding dimension 1024, MLP hidden dimension 4096 with ReLU activation.
Larger models (~400M parameters): The intermediate MLP size is doubled from 4096 to 8192, the number of attention heads is increased from 8 to 12, and BST layers are placed at {1, 5, 7, 9} (four BST layers instead of three).
The SSM is always placed at the first layer (layer index 1), right after the word embedding layer. This ensures that even the earliest layers have access to long-range context. The ablation in Table 3 confirms that a single BST layer at index 9 (closer to the middle) has the greatest effect on perplexity, but the first layer placement ensures context propagates through the entire stack.
Training uses the Adam optimizer with a batch size of 32 and a sequence length of 4K tokens. This sequence length (4096) is chosen as a standard long-range benchmark and is a multiple of the window size $W = 512$ , yielding 8 blocks per sequence.
For the MF variant, the number of filters $S$ is set to 32 for smaller models and 64 for larger models. The SSM state dimension $D$ is halved compared to SH/MH (as discussed above, for FFT efficiency).
For structured SSMs (S4): The internal state size $N$ is set to 16. Table 5 (Appendix E) shows that increasing $N$ from 16 to 32 improves perplexity from 11.57 to 11.55 (on PG19, ~200M model) but increases step time by a factor of 1.8×; $N = 64$ gives 11.54 perplexity but a 3.2× step time increase. The paper chooses $N = 16$ as the efficiency-performance sweet spot.
For the larger BST-L models compared against GSS-HYBRID-L, the training budget is matched at 0.8K TPUv4 hours for all three datasets, ensuring a fair FLOPs-matched comparison.

Layer Placement Ablation Insights

The paper's ablation studies (Appendix E) provide several insights into the design:

Single BST layer placement (Table 3): A single BST layer placed at index 9 gives the best perplexity (11.88) compared to index 3 (12.41), index 7 (11.92), or index 12 (12.03). This aligns with findings from Memorizing Transformer and Block-Recurrent Transformer that mid-to-late layers benefit most from access to long-range context — early layers focus on local linguistic features where long-range context is less critical, while deeper layers that integrate information for prediction benefit from knowing the broader context.
Multiple BST layers (Table 4): Adding more BST layers consistently improves perplexity — from 11.69 with 2 layers to 11.57 with 3 layers to 11.21 with 4 layers — but the improvement plateaus at 5 layers (11.20). Each additional layer adds a "3.5% training step time increase," so the paper selects 3 layers as the practical tradeoff.
The first layer is always a BST layer. This is a conscious design choice: by placing an SSM at the very beginning, positional information from the SSM's convolution is available to the entire stack. This replaces the need for global learned positional embeddings, as the SSM's inherent temporal encoding propagates forward.

Comparison with Prior Integration Strategies: Why Cross-Attention Wins

The paper's central architectural claim is that integrating SSM states into Transformer attention via cross-attention is more effective than interleaving SSM and attention layers (as in GSS-HYBRID). The evidence for this is in Table 1: BST:SH:S4-L and BST:MF:UNSTRUCT-L both match or outperform GSS-HYBRID-L on all three datasets while using a comparable parameter count and identical training compute budget (0.8K TPUv4 hours).

The mechanism explanation: In GSS-HYBRID, the SSM layers process the sequence and pass their outputs to the next layer (which may be another SSM or a Transformer). The Transformer layers never explicitly "query" the SSM's compressed history — they just receive whatever the SSM chose to output as the next layer's input. In BST, the cross-attention mechanism gives each token the ability to actively retrieve from the SSM's compressed history based on its own content. A token representing the word "it" can produce a query that attends to context states containing recent noun phrases; a token representing a function call can query for context states containing the function's definition. This content-addressable retrieval is exactly what attention excels at, and BST gives the attention mechanism direct access to the long-range memory rather than forcing information to pass through multiple layers of transformations.

The training stability advantage: The paper notes that GSS-HYBRID "used a different learning rate and weight decay for the SSM layer and the Transformer layer to avoid training instabilities" and required grid search over four learning rates. BST did not require per-layer learning rate tuning — the same learning rate was used for all layers. This suggests that the cross-attention interface provides a more stable gradient flow between the SSM and Transformer components, possibly because the cross-attention's query-key softmax normalizes the interaction, whereas direct interleaving can lead to mismatched scales between SSM and attention outputs.

4. Key Insights and Innovations

Innovation 1: Cross-Attention as a First-Class Integration Interface Between SSMs and Transformers

The dominant assumption in prior hybrid architectures — exemplified by GSS-HYBRID (Mehta et al., 2023) — was that State Space Models and Transformers communicate most naturally by interleaving their layers: an SSM layer processes the sequence and passes its output to the next layer, which might be a Transformer, and the two modalities never interact within the same computational step. This is architecturally clean but functionally weak — the Transformer layers never explicitly query the SSM's compressed history; they simply receive whatever representation the SSM chose to output. The interaction is passive, mediated entirely by the forward pass through the layer stack.

MEGA (Ma et al., 2023) tightened this integration by using gating mechanisms to mix SSM and attention outputs within a layer, but its design was tied to a specific attention variant (single-head Gated Attention Unit) and a specific gating scheme, sacrificing modularity.

BST makes a conceptually distinct move: it treats the SSM's output as a memory bank that attention can actively query through cross-attention. The SSM produces context states (compressed representations of the history), and the Block Transformer's cross-attention mechanism computes query-key affinities between each token in the current block and each context state. This is not gating — it is content-addressable retrieval. A token representing "it" can produce a query that selectively attends to context states encoding recent noun phrases; a token representing a function call can query for context states encoding the function's definition. The interaction is active and content-dependent rather than passive and layer-sequential.

Why does this matter beyond the architectural detail? Because it reframes the SSM-Transformer relationship from pipeline (SSM → Transformer → SSM → ...) to tool-use (Transformer uses SSM as a memory). The Transformer's attention mechanism — the very component that gives it local precision — is now also the mechanism by which it accesses long-range context. This unification is conceptually elegant: attention serves double duty as both the local token-comparison engine and the long-range retrieval engine. The SSM's role becomes analogous to a compressed, queryable representation of the past rather than an independent sequence processor.

The practical significance of this reframing extends beyond performance. Because the interface is simply cross-attention (a standard building block), the BST layer is modular in a way that GSS-HYBRID and MEGA are not: you can swap in any SSM backend (S4, DSS, S5, Hyena, or future SSMs) without changing the cross-attention interface, and you can swap in any attention variant without changing the SSM. The paper's three context-state variants (SH, MH, MF) are themselves a demonstration of this modularity — each uses a different method to construct the memory that cross-attention queries, but the interface remains identical. This clean separation of concerns is a design principle that future hybrid architectures can inherit.

Evidence: Table 1 shows that BST:SH:S4-L and BST:MF:UNSTRUCT-L match or outperform GSS-HYBRID-L on all three datasets (PG19, arXiv, GitHub) at a comparable parameter count and identical training budget (0.8K TPUv4 hours), despite GSS-HYBRID using a separate, grid-searched learning rate schedule per layer type. More tellingly, the paper notes that BST required no such per-layer tuning — the cross-attention interface appears to provide more stable gradient flow, possibly because the softmax normalization in attention naturally handles scale mismatches between SSM and Transformer outputs.

Innovation 2: The Redundancy-Retrievability Tradeoff as a Design Axis for Memory Construction

Prior work on augmenting Transformers with memory — from Transformer-XL's cached hidden states to BRECT's recurrent state vectors to the Memorizing Transformer's kNN retrieval — treated the memory representation as a fixed design choice: you pick a mechanism, and the memory representation is whatever that mechanism produces. There was no systematic language for comparing what kind of memory different mechanisms construct, or for understanding why one mechanism might work better than another on a given task.

BST introduces a conceptual taxonomy for memory construction in hybrid sequence models: the tradeoff between redundancy (how much overlapping information adjacent memory entries carry) and retrievability (how easy it is for the attention mechanism to recover specific information from the memory). This taxonomy is not merely descriptive — it is predictive of which variant will perform well on which data distribution.

The three BST variants are designed as points on this spectrum:

Single-Head (SH) maximizes redundancy (adjacent SSM states within a window overlap substantially) and therefore maximizes retrievability — if you need to know what the 10th token in the current block was, its exact context state is in the memory, in order. This favors datasets where local context dominates prediction, such as PG19 books, where the next word depends heavily on recent sentences and paragraphs.
Multi-Head (MH) reduces redundancy by having the SSM produce separate readouts per attention head, extracting complementary features from the same underlying state. Adjacent positions still overlap, so time-axis redundancy remains, but feature-axis redundancy decreases. This represents a middle ground — useful when the model benefits from diverse features but still needs time-addressable memory.
Multi-Filter (MF) minimizes redundancy by using independent SSM filters, each producing a single context state from the last position of the previous window. There is no time-axis overlap (only one position per block), and the filters are trained to extract complementary information. This sacrifices retrievability of specific recent tokens (they have been compressed through the SSM state) in exchange for efficient use of the memory slots — each slot stores a different type of information. This favors datasets where long-range structural dependencies matter more than exact recent wordings, such as arXiv (cross-references, theorem citations) and GitHub (import statements, function signatures).

What makes this a genuine conceptual contribution rather than just "we tried three things and one worked"? The taxonomy is actionable: it tells a practitioner how to choose based on their data's dependency structure, not what worked on our benchmarks. The empirical pattern in Table 1 validates the taxonomy's predictions: SH outperforms MF on PG19 (local-context-dominated), while MF outperforms SH on arXiv and GitHub (long-range-structure-dominated). The MH variant — theoretically the middle ground — performs on par with SH but requires 8% more parameters, consistent with the prediction that reducing feature-axis redundancy without reducing time-axis redundancy yields diminishing returns when time-axis redundancy is the dominant factor.

This reframes memory design from a mechanism-choice problem ("should I use a recurrent cell, an SSM, or kNN retrieval?") to a representation-quality problem ("what properties should my memory representation have for my data?"). The specific implementations (SH, MH, MF) are instantiations of the principle, but the principle itself — that redundancy and retrievability are in tension and that the optimal tradeoff depends on the task — is a reusable insight for any future architecture that feeds compressed history into attention.

Evidence: Table 1 shows BST:SH:S4 outperforming BST:MF:S4 on PG19 (11.57 vs. 11.63) while BST:MF:S4/UNSTRUCT outperforms SH on arXiv (2.48/2.44 vs. 2.51/2.49) and GitHub (2.07/2.03 vs. 2.14/2.09). This cross-task reversal is the key pattern — the same architecture choice yields opposite rankings depending on the data, confirming that the redundancy-retrievability axis captures a real task-dependent tradeoff rather than a universal quality ranking.

Innovation 3: Structured vs. Unstructured Kernels as the Boundary Condition for Length Generalization

A persistent challenge in long-range sequence modeling is length generalization — the ability of a model trained on sequences of length $L$ to perform well on sequences of length $L' \gg L$ at inference time without retraining. Prior hybrid architectures either did not systematically study this (BRECT's recurrent gating was trained at fixed window sizes and tested only on modest extensions) or reported mixed results (GSS-HYBRID's authors noted that larger models "had difficulty generalizing to higher lengths").

BST provides what is, to the authors' knowledge, the clearest diagnostic separation of what enables length generalization in hybrid SSM-Transformer architectures. The key finding is not that structured SSMs like S4 generalize better than unstructured Hyena-style kernels — that much is predictable from the mathematics of how the kernels are constructed. The innovation is in demonstrating that this mathematical property survives integration into a hybrid architecture with attention and that it is the dominant factor determining whether the full hybrid model generalizes, overwhelming any effects of the attention mechanism or the context-state construction method.

The mechanism is straightforward but its empirical validation is significant: structured kernels (S4/DSS) are parameterized by fixed-size matrices $\mathbf{A}$ , $\mathbf{B}$ , $\mathbf{C}$ independent of sequence length, so the kernel $\mathbf{K}$ can be compiled for any $L'$ at inference time. Unstructured kernels (Hyena-inspired) are parameterized directly as a weight vector of length $L$ , with positional encodings that assume positions $\{0, \dots, L-1\}$ , so extending beyond $L$ requires ad-hoc interpolation. When these kernels are embedded inside a BST layer — surrounded by attention, cross-attention, and feed-forward networks — the structured variants generalize and the unstructured variants do not.

Why is this a conceptual contribution rather than an expected result? Because it was not obvious that the SSM's length generalization would propagate through the cross-attention interface to the model's overall language modeling performance. The cross-attention mechanism could have introduced its own length-dependent biases — the relative position bias in self-attention, the learned context IDs in MF, the way context states are split into blocks — that would dominate at longer lengths and swamp the SSM's generalization capability. The fact that BST:SH:S4-L generalizes well to 65K tokens (Figure 3) while BST:MF:UNSTRUCT-L degrades shows that the SSM kernel's length-generalization property is preserved and dominant in the hybrid architecture. This is a non-trivial architectural property: the cross-attention interface is transparent to the SSM's length generalization.

The diagnostic significance is amplified by the pattern inversion at the training length: unstructured variants perform best at 4K (where their richer parameterization pays off), but structured variants win at 65K (where their mathematical extendability becomes essential). This creates a practitioner's decision rule: if you need length generalization, use structured kernels even though they may underperform at the training length; if you only ever evaluate at the training length, unstructured kernels may give you better results for the same compute.

Evidence: Figure 3 shows BST:SH:S4-L maintaining or improving perplexity from 4K to 16K to 65K on PG19 and GitHub, while BST:MF:UNSTRUCT-L degrades on the same extension. On arXiv, BST:SH:S4-L degrades less severely than BST:MF:UNSTRUCT-L. The paper explicitly states: "BST:SH:S4-L has by far the best perplexity for 65K sequence lengths on PG19, GitHub and arXiv." This is not just "structured SSMs generalize" — it is "the full hybrid architecture inherits the SSM's generalization property, and this property dominates the overall model behavior at long lengths."

Innovation 4: Eliminating the Sequential Bottleneck Without Sacrificing Locality — The Speed-Accuracy Frontier Shift

The Block-Recurrent Transformer (BRECT) established that combining local block-wise attention with a recurrent inter-block memory yields strong language modeling performance on long documents. But BRECT's design imposed an architectural penalty: the recurrent cell that passes information between blocks creates a sequential dependency — each block must wait for the previous block's recurrent state before it can begin processing. At a 4K sequence length with a window size of 128, the paper reports that a BRECT layer runs approximately 15× slower than an equivalent sliding-window Transformer without recurrence (SLIDE:12L; Section 4.3, Figure 4 left). This speed penalty makes BRECT impractical for many deployment scenarios, despite its accuracy advantages.

The question BST answers is: can you eliminate the sequential bottleneck while preserving the accuracy benefits of inter-block memory? The answer is not obvious a priori. The SSM could have been too lossy a compression mechanism — replacing BRECT's explicit recurrent states (which are produced by attention over the previous block's tokens, a high-fidelity but expensive operation) with an SSM's convolution-based states (which compress the entire history through a fixed-size state vector) could have sacrificed too much local precision. Or the cross-attention interface could have been less effective than BRECT's learned gating at integrating long-range context. Or the FFT overhead could have eaten up the parallelization gains, particularly on TPUs where the paper notes FFT is a bottleneck.

BST's result is a frontier shift: it simultaneously improves speed and matches or improves accuracy. On the speed axis, the paper reports a 6–11× speedup over BRECT at the layer level on GPU (Figure 4 left), reducing the gap from the non-recurrent baseline (SLIDE) from ~15× to less than 2×. This is not merely an engineering optimization — it is an architectural property that emerges from replacing the sequential recurrence with a parallelizable convolution. On the accuracy axis, Table 1 shows BST matching or modestly outperforming BRECT on all three datasets at comparable parameter counts and training budgets: BST:SH:UNSTRUCT achieves 11.52 vs. BRECT:FIXED:SKIP's 11.55 on PG19; BST:MF:UNSTRUCT achieves 2.03 vs. 2.04 on GitHub.

The intellectual significance is that this establishes a new Pareto-optimal point on the speed-accuracy frontier for hybrid recurrent-attention architectures. Prior to BST, the choice was between: (a) fast but local-only (SLIDE:12L, which lacks inter-block memory and has worse perplexity), (b) accurate but slow (BRECT, which has inter-block memory but is 15× slower), and (c) fast and long-range but less accurate on language modeling (pure SSMs, which lag Transformers on next-token prediction). BST occupies a previously empty region: fast (nearly as fast as SLIDE), long-range (SSM context encodes full history), and accurate on language modeling (matching or exceeding BRECT). This is not an incremental improvement along one axis — it is opening a new region of the design space that was previously inaccessible.

Evidence: Figure 4 (left) directly benchmarks layer forward-pass time on GPU: BST:SH and BST:MH are 6–11× faster than BRECT at a 4K sequence length with a window size of 128. Table 1 shows that this speed advantage does not come at the cost of perplexity — BST variants either match or slightly outperform BRECT within the same training budget. The paper also reports (Section 4.3) that the speed advantage persists up to 65K tokens, the point at which hardware saturation begins.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper evaluates on three long-document language modeling benchmarks: PG19 (28,602 full-length books from Project Gutenberg published before 1919; 6,966,499 English words; each book is 50K–100K tokens when tokenized; test split reported), arXiv (a corpus of scientific and technical mathematics articles containing LaTeX source code — theorems, citations, definitions that are referenced over long ranges; test split reported), and GitHub (code repositories with open-source licenses, filtered to C, C++, Java, Python, Go, and TypeScript, where files are concatenated along directory tree paths to create a single document preserving repository structure and dependencies; validation split reported). These three datasets were chosen because they span natural language prose (PG19), structured scientific text with cross-references (arXiv), and code with import/file dependencies (GitHub), providing varying distributions of local-vs-long-range dependency structures.
Base model(s). All BST models and baselines use decoder-only Transformer architectures at two scales: ~200M parameters (12 layers, 8 attention heads of dimension 128, embedding dimension 1024, MLP hidden dimension 4096 with ReLU activation) and ~400M parameters ("-L" models: MLP hidden dimension doubled to 8192, attention heads increased to 12). The models are trained from scratch; no pretrained weights are used. The vocabularies are matched to prior work for fair comparison: a pretrained T5 vocabulary with 32K tokens for PG19 (Raffel et al., 2020), and a LaMDA vocabulary with 32K tokens (Thoppilan et al., 2022) for both arXiv and GitHub. The sequence length is 4K tokens during training, with a window size of 512 tokens (8 blocks per sequence). Importantly, the authors do not use global learned positional embeddings, instead relying on the SSM's inherent positional encoding and T5-style relative position bias in self-attention.
Metrics. The primary evaluation metric is language modeling perplexity (lower is better), computed as the exponential of the cross-entropy loss on the held-out evaluation split. For the length generalization experiments in Figure 3, perplexity is reported at evaluation sequence lengths of {512, 16K, 65K} tokens, even though all models were trained only at 4K. An additional metric in the efficiency comparison (Figure 4, left) is layer-level forward-pass wall-clock time on GPU, measured to quantify the parallelism advantage of BST over BRECT. The paper also tracks total training compute as TPUv4 hours for each model-dataset pair (Table 1), enabling FLOPs-matched comparisons.
Baselines. The paper compares against four architectures, all drawn from prior work and reproduced under matched conditions where possible:
- TRSF-XL:2048 (Transformer-XL; Dai et al., 2019): A standard Transformer with a training window size of 2048 tokens and segment-level recurrence via cached hidden states. This serves as the "larger window, standard attention" baseline, representing the conventional approach to extending context.
- SLIDE:12L (from Hutchins et al., 2022): A sliding-window Transformer with window size 512 over a segment of 4096 tokens, where the sliding window is differentiable over two consecutive blocks. This is the "fastest, no-inter-block-memory" baseline — it has the same computational structure as BST without the SSM context, measuring the pure perplexity benefit of adding long-range memory.
- BRECT:FIXED:SKIP (Block-Recurrent Transformer; Hutchins et al., 2022): The strongest and fastest BRECT variant from the original paper, using a "skip" gating configuration where a simple linear layer combines the current block's hidden representation with past information from previous blocks. This is the most direct architectural predecessor — it has inter-block memory (via recurrence) and local block attention, but processes blocks sequentially. It represents the accuracy target BST aims to match or exceed while being faster.
- GSS-HYBRID-L (Gated State Space Hybrid; Mehta et al., 2023): The closest SSM-Transformer hybrid from prior work, with approximately 373M parameters. GSS-HYBRID interleaves 32 Gated State Space (GSS) layers with Transformer layers (every 4th layer starting from layer 2). This represents the "interleaved SSM and attention" approach, contrasting with BST's "integrated cross-attention" approach. The paper notes that GSS-HYBRID's perplexity scores were obtained after grid-searching over four learning rates and using separate learning rates for SSM and Transformer layers.
Generation budget / compute accounting. The paper uses TPUv4 hours (thousands of hours) as the compute budget metric, reported in Table 1 for each model on each dataset. Training times are aligned across models for fair comparison: the ~200M parameter models are trained at a budget of 0.5K TPUv4 hours for PG19 and arXiv, and 1.8K for GitHub (a larger dataset). The larger ~400M models are trained at 0.8K TPUv4 hours for all three datasets. The TRSF-XL:2048 baseline is trained at 0.8K/0.8K/3.0K TPUv4 hours due to its larger window size increasing attention cost. The paper also benchmarks layer-level speed (Figure 4, left) by measuring forward-pass wall-clock time on GPU for a single layer, fixing the window size at 128, embedding dimension at 512, SSM internal state size at 16, and 16 attention heads — this isolates architectural speed from training-time overhead.
Cross-validation / statistical protocol. The paper runs the ~200M BST models on PG19 with three different random seeds and reports the average perplexity with error bars (standard deviation ±1.1–1.2 in Table 1), establishing variance estimates for this configuration. For arXiv and GitHub, and for all larger models, only single runs are reported due to long training times and the large number of experiments. The BRECT baseline error bars on PG19 (±1.1) are taken from the original paper (Hutchins et al., 2022). The paper does not perform cross-validation or hyperparameter sweeps for BST — the same learning rate is used for all layers (unlike GSS-HYBRID, which required grid search), and the training recipe is fixed from BRECT with modifications only to the architectural components under study. There is no evaluation on a separate held-out "test" set beyond the standard dataset splits (PG19 test, arXiv test, GitHub validation).

Main Quantitative Results

The paper organizes its primary results around three axes of evaluation: (1) language modeling perplexity comparisons at matched compute budgets (Table 1), (2) length generalization to unseen sequence lengths (Figure 3), and (3) layer-level speed benchmarks (Figure 4, left). Each axis tests a different claim: accuracy competitiveness, extrapolation capability, and parallelism gains, respectively.

Language Modeling Perplexity at Matched Compute Budgets (Table 1)

~200M parameter models, trained at 0.5K TPUv4 hours (PG19/arXiv) or 1.8K TPUv4 hours (GitHub).

The headline finding is that BST variants match or modestly outperform BRECT on all three datasets while being trained at the same or smaller compute budget, and substantially outperform SLIDE:12L (which lacks inter-block memory), confirming that the SSM-based context states provide a modeling benefit comparable to BRECT's recurrent states.

On PG19, the best BST variant is BST:SH:UNSTRUCT at 11.52 perplexity (±1.1), compared to BRECT:FIXED:SKIP at 11.55 (±1.1) and SLIDE:12L at 12.12. The difference between BST and BRECT is within one standard deviation, indicating essentially equivalent performance. However, BST achieves this at a lower training budget: 0.5K TPUv4 hours vs. 0.8K for BRECT. The structured S4 variants — BST:SH:S4 (11.57), BST:MH:S4 (11.60), BST:MF:S4 (11.63) — all cluster tightly around BRECT's perplexity. The unstructured variants (SH and MF) produce the two best perplexities on PG19 (11.52 and 11.56), consistent with the paper's hypothesis that unstructured kernels have "more free parameters" and provide richer representations at the training length.

The pattern flips on arXiv and GitHub, where MF variants outperform SH. On arXiv, BRECT:FIXED:SKIP achieves the best perplexity at 2.36, but BST:MF:UNSTRUCT is close at 2.44, followed by BST:MF:S4 at 2.48. The SH variants lag: BST:SH:S4 at 2.51 and BST:SH:UNSTRUCT at 2.49. The SLIDE baseline (2.69) is substantially worse, and TRSF-XL:2048 (2.48) matches BST:MF:S4 but requires a larger training budget (0.8K vs. 0.5K TPUv4 hours). This reversal — SH wins on PG19, MF wins on arXiv and GitHub — is the primary evidence for the paper's redundancy-retrievability tradeoff hypothesis: local-context-dominated tasks favor SH's high retrievability, while long-range-structure-dominated tasks favor MF's efficient, complementary feature extraction.

On GitHub, BST:MF:UNSTRUCT achieves the best perplexity among the ~200M models at 2.03, essentially tied with BRECT:FIXED:SKIP at 2.04 and slightly better than BST:SH:UNSTRUCT at 2.09. The TRSF-XL:2048 baseline (2.01) performs slightly better than all BRIECT/BST variants, but at a significantly higher training cost (3.0K TPUv4 hours vs. 1.8K for BST). BST:MF:S4 attains 2.07, and BST:SH:S4 falls further behind at 2.14. The same MF > SH pattern holds as on arXiv.

~400M parameter models ("-L" variants), trained at 0.8K TPUv4 hours across all three datasets.

The larger models are compared primarily against GSS-HYBRID-L (373M parameters, also trained at 0.8K TPUv4 hours). The key finding across all three datasets: BST-L variants match or outperform GSS-HYBRID-L while using the same compute budget and not requiring per-layer learning rate tuning.

On PG19, BST:SH:UNSTRUCT-L achieves the best perplexity at 10.37, outperforming GSS-HYBRID-L at 10.52 and BST:SH:S4-L at 10.47. BST:MF:UNSTRUCT-L (10.42) and BST:MF:S4-L (10.52) also match GSS-HYBRID-L. The gap of ~0.10–0.15 perplexity points is small but consistent across the SH unstructured variant.

On arXiv, BST:MF:UNSTRUCT-L achieves the best perplexity at 2.41, compared to GSS-HYBRID-L at 2.51 and BST:SH:S4-L at 2.49. This is a ~4% relative improvement over GSS-HYBRID-L. BST:MF:S4-L reaches 2.46, and BST:SH:UNSTRUCT-L reaches 2.46 — both also outperform GSS-HYBRID-L. This is the largest relative gap in BST's favor across the three datasets.

On GitHub, BST:MF:UNSTRUCT-L achieves 1.83, followed by BST:MF:S4-L at 1.84 and BST:SH:UNSTRUCT-L at 1.85. GSS-HYBRID-L is at 1.88. The differences are small (0.03–0.05 perplexity points) but consistently favor BST across all three BST-L variants. The paper emphasizes that these gains come despite GSS-HYBRID benefiting from "grid searching over four learning rates" and using "a different learning rate and weight decay for the SSM layer and the Transformer layer," while BST used the same learning rate for all layers without grid search.

Perplexity ranking consistency. A notable pattern across Table 1: BST:MF:UNSTRUCT appears in the top two perplexities for five of the six model-scale/dataset combinations (PG19 ~200M, PG19 ~400M, arXiv ~400M, GitHub ~200M, GitHub ~400M), with BST:SH:UNSTRUCT taking the top spot for PG19 ~200M and ~400M. The unstructured kernel variants consistently outperform their S4-structured counterparts at the training length of 4K across nearly all configurations. This is a practically important finding: if length generalization is not required, unstructured kernels provide better perplexity per unit of training compute.

Length Generalization to Unseen Sequence Lengths (Figure 3)

The critical experiment in Figure 3 evaluates all models at evaluation sequence lengths of {512, 16K, 65K} tokens after being trained only at 4K. The models compared are the ~400M parameter variants: BST:SH:S4-L, BST:MF:UNSTRUCT-L, GSS-HYBRID-L, and BRECT:FIXED:SKIP-L.

The headline finding: BST:SH:S4-L has "by far the best perplexity for 65K sequence lengths on PG19, GitHub and arXiv," demonstrating that structured SSM kernels enable length generalization in the full hybrid architecture, while unstructured kernels do not.

On PG19 (Figure 3, left), all models perform similarly at 512 tokens (around 10–12 perplexity). At 16K, BST:MF:UNSTRUCT-L and BRECT:FIXED:SKIP-L perform best (approximately 10–11 perplexity), with BST:SH:S4-L slightly behind. At 65K, BST:SH:S4-L pulls ahead decisively, dropping to approximately 9.5 perplexity (lower is better, and perplexity improves with longer context on PG19 — a phenomenon also noted in BRECT), while BST:MF:UNSTRUCT-L rises sharply to approximately 13–14 perplexity. GSS-HYBRID-L also generalizes well to 65K but at a higher perplexity (~11) than BST:SH:S4-L.

On arXiv (Figure 3, middle), the pattern diverges sharply. All models degrade at longer lengths (perplexity increases). At 65K, BST:SH:S4-L shows the least degradation (rising from ~2.5 at 4K to ~3.0–3.5 at 65K), while GSS-HYBRID-L degrades more severely (rising to ~4.0–4.5) and BST:MF:UNSTRUCT-L degrades substantially (rising to ~5.0). BRECT:FIXED:SKIP-L interestingly generalizes well, performing best or second-best at 65K, which the authors attribute to BRECT's "access to the entire past during training via a non-differentiable cache of representations across sequences" — the non-differentiable cache from one 4K training sequence to the next provides weak access to history beyond the 4K window, acting as an implicit length-generalization mechanism. The paper notes that GSS-HYBRID's degradation on arXiv is consistent with the original GSS paper's report that "larger GSS models had difficulty generalizing to higher lengths."

On GitHub (Figure 3, right), similar to PG19, perplexity improves with longer context for some models. BST:SH:S4-L improves from ~2.0 at 4K to ~1.8 at 65K, the best overall. BST:MF:UNSTRUCT-L also improves initially (to ~1.9 at 16K) but degrades slightly at 65K back to ~2.0. GSS-HYBRID-L improves modestly to ~1.9 at 65K. BRECT:FIXED:SKIP-L is the second-best at 65K.

The paper draws a "clear distinction between structured and unstructured SSMs integrated in hybrid architectures." The structured S4 kernel — parameterized by fixed-size matrices A, B, C and extendable to any length — provides a "built-in mechanism for length generalization that the unstructured BST:MF:UNSTRUCT-L model does not." The unstructured kernel, parameterized by a weight vector of length L with learned positional encodings, cannot be meaningfully extended beyond its training length. The key diagnostic: BST:MF:UNSTRUCT-L "performs best on the training sequence of 4K and is on-par for 512 with perplexity increasing for unseen 16K and 65K sequence lengths," while BST:SH:S4-L's performance continues to improve or holds steady.

A subtle observation in the arXiv results: the authors hypothesize that BRECT's surprising length generalization on arXiv is due to the non-differentiable sequence cache — "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." They note that BST lacks this cross-sequence cache (since the SSM convolution needs the whole sequence, history beyond the current 4K chunk is lost), and suggest this as a potential area for improving BST on very long documents.

Layer-Level Speed Benchmarking (Figure 4, Left)

Figure 4 (left) benchmarks the forward-pass wall-clock time of a single BST layer against a BRECT layer (including recurrent units) and a SLIDE:12L layer (no recurrence), on GPU with a fixed window size of 128, SSM internal state size of 16, 16 attention heads, and an embedding dimension of 512.

The headline finding: BST:SH is 6–11× faster than BRECT, while BST:MH is 3–4× faster, across sequence lengths from 1K to 65K tokens. At the training length of 4K, "BRECT layer runs almost 15× slower than SLIDE:12L with the same window size," and BST reduces this gap "to less than 2×." The BST:SH and SLIDE:12L curves are nearly overlapping in the plot (the paper states "the bottom-most two lines on the left of Figure 4 are almost overlapping"), confirming that the SSM sublayer's O(L log L) FFT cost is not a significant bottleneck relative to the block attention cost, even at 65K tokens.

The speed advantage persists across the full range of sequence lengths. At 65K, "hardware saturation began to occur," limiting further scaling, but "significant performance improvements, up to a factor of 6, remain evident for sequences as long as 65k tokens." BST:MH is slower than BST:SH because it runs the SSM convolution with H separate output channels (one per attention head), increasing the FFT dimension; BST:MF is not shown on this plot but would be expected to fall somewhere between SH and MH since it runs S separate filters.

The paper notes that these benchmarks use "the vanilla implementation of FFT and inverse FFT operations provided by JAX," suggesting that further speedups are possible with "recent and faster hardware-specific I/O-aware implementations." The authors also flag that the FFT operation is "an important speed bottleneck on TPUs" (4× faster on GPUs in their experience), though these benchmarks are on GPU.

Combined Accuracy-Speed Frontier

The paper does not plot this explicitly, but the combination of Table 1 and Figure 4 (left) establishes a new Pareto frontier for hybrid recurrent-attention architectures. BRECT sits at high accuracy but low speed (15× slower than SLIDE). SLIDE sits at high speed but lower accuracy (worse perplexity by 0.5–1.0 points on PG19). BST occupies a previously inaccessible region: accuracy matching BRECT with speed approaching SLIDE (within 2×). This is the paper's central practical contribution — it demonstrates that the sequential bottleneck is not inherent to architectures that combine local attention with long-range memory; it is an artifact of the specific recurrence mechanism, and replacing it with an SSM recovers parallelism without sacrificing modeling quality.

The paper also demonstrates a scaling trend (Appendix C, Figure 5): the perplexity gap between BST and BRECT widens with model scale (from 2.5% at 200M parameters to 4.0% at 1.3B parameters on PG19), suggesting that the SSM-based context states may scale more favorably than recurrent gating mechanisms. However, this is based on a single dataset (PG19) and only two data points beyond 200M, so the scaling trend should be interpreted cautiously.

Ablation Studies and Robustness Checks

The paper includes ablation studies in Appendices D and E, plus additional comparisons throughout the main results that serve as implicit ablations.

Layer placement of a single BST layer (Appendix E, Table 3): A single BST:SH:S4 layer (state size D=16) is placed at different positions in the 12-layer stack. A BST layer at index 9 gives the best perplexity (11.88) on PG19, compared to index 3 (12.41), index 7 (11.92), and index 12 (12.03). The finding that mid-to-late layers benefit most from long-range context is consistent with prior work (Memorizing Transformer, BRECT), but the optimal position being layer 9 (out of 12) — very late in the stack — is notable. The paper states this is "inline with findings in prior work [42, 21]" but does not explore why the very last layer (index 12) underperforms index 9, which might suggest that the final layer benefits more from local precision for immediate prediction than from long-range context.

Number of BST layers (Appendix E, Table 4): Increasing the number of BST layers (with state size D=16) from 2 to 5 monotonically improves perplexity on PG19: 2 layers at 11.69, 3 layers at 11.57, 4 layers at 11.21, 5 layers at 11.20. The improvement plateaus at 4–5 layers (11.21 to 11.20 is negligible). Each additional layer adds approximately 3.5% training step time. The paper selects 3 layers for the main experiments as a practical tradeoff, but the results suggest that 4 layers provide a meaningful improvement (11.57 → 11.21, a 0.36 perplexity drop) for only a modest cost increase. This tradeoff is not fully explored in the paper — the main results use 3 BST layers for ~200M models and 4 for ~400M models, so the scaling axis (more BST layers vs. more parameters) is partially confounded.

SSM state size D (Appendix E, Table 5): Increasing the S4 state dimension D from 8 to 64 improves perplexity on PG19: D=8 gives 11.95, D=16 gives 11.57, D=32 gives 11.55, D=64 gives 11.54. The improvement saturates quickly after D=16, while step time grows significantly: D=16 is the baseline (1.0×), D=32 is 1.8×, D=64 is 3.2×. The paper chooses D=16 as the efficiency-performance sweet spot. Interestingly, D=8 is 0.7× step time with a relatively small perplexity penalty (11.95 vs. 11.57, a 3.3% relative increase), which might be a preferable tradeoff for very latency-sensitive deployments, though only one seed is reported.

Structured (S4) vs. unstructured (UNSTRUCT/Hyena-inspired) kernels (implicit in Table 1 and Figure 3): This comparison runs throughout the paper as a thread. At the training length of 4K, unstructured kernels consistently outperform structured S4 kernels on all three datasets at both model scales (Table 1), with the gap being particularly pronounced on PG19 (~200M: 11.52 vs. 11.57; ~400M: 10.37 vs. 10.47). The paper attributes this to the unstructured kernel having "more free parameters" than the low-rank factorization of the structured kernel. However, at evaluation lengths of 16K and 65K (Figure 3), the structured variants generalize while the unstructured variants degrade. This establishes a clear decision boundary: unstructured kernels are preferable if evaluation lengths match training lengths; structured kernels are necessary for length generalization.

Single-Head (SH) vs. Multi-Filter (MF) vs. Multi-Head (MH) context states (Table 1): This is the central architectural ablation and reveals task-dependent behavior. On PG19, SH consistently outperforms MF (e.g., ~200M S4: 11.57 vs. 11.63; ~400M UNSTRUCT: 10.37 vs. 10.42), while on arXiv and GitHub, MF consistently outperforms SH (e.g., ~400M UNSTRUCT: arXiv 2.41 vs. 2.46; GitHub 1.83 vs. 1.85). The MH variant performs similarly to SH on PG19 (11.60 vs. 11.57) but requires 8% more parameters (218M vs. 202M), making it strictly less parameter-efficient. The MF variant also requires 8% more parameters in the S4 configuration (217M vs. 202M) but its accuracy advantage on arXiv and GitHub justifies the cost. The task-dependent reversal is the core evidence for the redundancy-retrievability tradeoff — SH's redundancy helps when local context matters, while MF's efficiency helps when diverse long-range features matter.

Down-sampling the SSM input (Section 3.4): The paper reports that projecting SSM inputs to one-quarter of the embedding dimension (D/4) has "negligible impact to perplexity" while reducing the total required number of FFTs by a factor of 4. For the MF variant, the state dimension is further reduced by an additional factor of two (to D/8), also with "negligible impact." This result is stated narratively without a supporting table, making it difficult to assess the exact magnitude of "negligible." A table comparing perplexity with and without down-sampling would strengthen this claim.

Context IDs for MF (Section 3.4): The paper states that learned context IDs are necessary for the MF variant (since the context states correspond to different filters rather than sequential positions, they lack inherent ordering) and unnecessary for SH and MH (where the SSM's temporal encoding provides positional information). No ablation is provided showing perplexity with vs. without context IDs, so this claim is based on design reasoning rather than empirical validation.

Long Range Arena (LRA) classification (Appendix D, Table 2): While not an ablation of the main language modeling experiments, the LRA results serve as a robustness check that BST's hybrid architecture transfers to non-language modalities and classification tasks. BST:SH:S4 achieves the highest average score (86.96) among methods that chunk input sequences, outperforming MEGA-CHUNK (85.66) and BRECT:FIXED:SKIP (60.80). BST wins on four of six LRA tasks (ListOps: 61.49, Image: 91.07, Pathfinder: 95.75, Path-X: 95.28) and is competitive on Text (87.63 vs. MEGA-CHUNK's 90.19) and Retrieval (90.51 vs. MEGA-CHUNK's 90.97). The average advantage over MEGA-CHUNK is 1.5%, which is meaningful given LRA is a standard benchmark. However, BST still lags MEGA without chunks (88.21 average), suggesting that chunking itself imposes a penalty on some LRA tasks.

Negative result with ReST<sup>EM</sup>-style training (mentioned in Section 5 of the reference example, but not present in this paper): The BST paper does not report negative results from alternative training recipes (unlike the Block-State Transformer paper's Appendix K, which is in the reference example but not in this paper). The closest to a negative result is the implicit acknowledgment that GSS-HYBRID required per-layer learning rate tuning while BST did not, suggesting that cross-attention integration is more stable than layer interleaving — but this is a stability property, not a failed experiment.

Critical Assessment

The experiments in this paper support its central claims, but with specific boundary conditions that the paper itself identifies. A critical reading reveals several dimensions where the evidence is strong, and several where it is more qualified than a first reading might suggest.

Claim 1: BST matches or improves upon BRECT's perplexity at lower compute cost.

This claim is well-supported for the ~200M parameter scale on PG19 and GitHub, where BST and BRECT are trained at the same or similar compute budgets and BST achieves comparable perplexity (Table 1). On PG19, BST:SH:UNSTRUCT (11.52 ± 1.1) and BRECT (11.55 ± 1.1) are statistically indistinguishable — the difference is smaller than the reported standard deviation. On GitHub, BST:MF:UNSTRUCT (2.03) and BRECT (2.04) are essentially tied. On arXiv, BRECT (2.36) holds a modest advantage over BST's best (2.44), so the claim is "matches or modestly underperforms" on arXiv rather than "matches or improves."

The "lower compute cost" part of the claim is nuanced. On PG19, BST trains at 0.5K TPUv4 hours vs. BRECT's 0.8K — a genuine 37.5% reduction. However, comparing total parameters, the BST variants are slightly larger: BST:SH:S4 at 202M vs. BRECT at 196M, BST:MF:S4 at 217M vs. 196M. The additional parameters in the SSM sublayer (the S4 state matrices and convolution machinery) contribute to perplexity independently of the architectural innovation, and the paper does not fully disentangle the effect of more parameters from the effect of the SSM-Transformer integration. A parameter-matched BRECT (e.g., by adding layers to reach 202M or 217M) would be a stronger test of whether the BST architecture is genuinely more parameter-efficient, or simply benefits from having more total capacity. This is a missing ablation.

Additionally, the GSS-HYBRID-L comparison at the ~400M scale shows BST-L outperforming GSS-HYBRID-L on all three datasets (Table 1), but the parameter counts are not precisely matched: GSS-HYBRID-L has 373M parameters, while BST:SH:S4-L has 366M (slightly fewer), BST:MF:S4-L has 383M (slightly more), and BST:MF:UNSTRUCT-L has 388M (15M more). The perplexity gaps (e.g., 10.37 vs. 10.52 on PG19, 2.41 vs. 2.51 on arXiv) are small enough that parameter count differences could account for some fraction of the improvement. A parameter-matched GSS-HYBRID variant or a parameter-count-normalized efficiency metric would clarify whether the architectural difference itself drives the gain.

Claim 2: BST achieves more than a tenfold increase in layer-level speed compared to BRECT.

This claim is the strongest in the paper, with direct empirical support from Figure 4 (left). The 6–11× speedup for BST:SH and 3–4× for BST:MH at the layer level on GPU is unambiguous, and the statement that BST reduces the BRECT-to-SLIDE speed gap "from ~15× to less than 2×" is a clean, interpretable framing. The speedup comes from removing the sequential for-loop over blocks, which is a direct consequence of the architectural innovation (SSM convolution pre-computes all context states in parallel).

However, there are important boundary conditions and caveats:

Layer-level speed ≠ end-to-end training speed. The benchmarks measure a single layer's forward pass, not the full training loop (which includes backward pass, optimizer steps, data loading, and multi-layer interactions). The paper acknowledges this implicitly by reporting total training TPUv4 hours in Table 1 rather than per-step time, and the training time comparisons show a more modest advantage: BST trains at 0.5K TPUv4 hours vs. BRECT at 0.8K on PG19, which is a 1.6× end-to-end speedup, not a 10× speedup. The backward pass through the SSM (which requires FFT derivatives), the fact that only 3 of 12 layers are BST layers (the other 9 are standard Block Transformer layers), and TPU-specific FFT bottlenecks all compress the theoretical layer-level speedup into a smaller end-to-end gain. The "tenfold" claim should be understood as a layer-level parallelism property under idealized hardware conditions, not a realized training speedup.
The speedup is measured on GPU, but training is on TPU. The paper states in Appendix A that "JAX FFT was 4× faster on GPUs" than on TPUs, and that the FFT operation is "an important speed bottleneck on TPUs." Since the training budgets in Table 1 are reported in TPUv4 hours, the realized training speedup on TPU is likely smaller than the layer-level GPU benchmarks in Figure 4 would suggest. The paper does not provide layer-level speed benchmarks on TPU, which would be more directly relevant to the reported training times.
Speedup depends on SSM variant. BST:MH is only 3–4× faster than BRECT (not 10×), because running H separate output channels in the SSM increases the FFT dimension. The MF variant (with S separate filters) is not benchmarked on the speed plot, but would likely fall between SH and MH depending on S. The "tenfold" claim applies to BST:SH specifically.
Speedup depends on hardware support for parallelism. The paper notes that the speed advantage is realized "provided there is hardware support to fully take advantage of parallelism." On hardware with limited parallel processing capability (e.g., CPUs with few cores, or edge devices), the gap between BST and BRECT would narrow because BRECT's sequential bottleneck matters less when parallelism is limited anyway. The results should be interpreted as applying to modern GPU/TPU accelerator hardware.

Claim 3: BST generalizes to sequences longer than training length, but only with structured SSM kernels.

This claim is strongly supported by Figure 3. The diagnostic separation between structured (S4) and unstructured (UNSTRUCT) variants at 65K evaluation length is clear and consistent across all three datasets. BST:SH:S4-L outperforms BST:MF:UNSTRUCT-L at 65K on PG19 (~9.5 vs. ~13–14 perplexity), arXiv (~3.0–3.5 vs. ~5.0), and GitHub (~1.8 vs. ~2.0). The paper's interpretation — that the structured kernel's length-independent parameterization enables generalization while the unstructured kernel's length-dependent positional encoding does not — is theoretically grounded and empirically validated.

However, the evaluation exposes several boundary conditions:

BRECT also generalizes, and sometimes better. On arXiv at 65K, BRECT:FIXED:SKIP-L actually outperforms BST:SH:S4-L (Figure 3, middle), which the authors attribute to BRECT's non-differentiable cross-sequence cache providing access to information beyond the 4K training chunk. This means BST's length generalization is not uniformly the best across datasets, and for arXiv-like documents with very long-range cross-references, BRECT's cache may provide a complementary mechanism that BST lacks.
Generalization is evaluated at fixed evaluation lengths {512, 16K, 65K}, not on tasks specifically designed to require long-range retrieval. Lower perplexity at 65K on PG19 does not necessarily mean the model is using the extra context effectively — it could be benefiting from a longer local smoothing window without actually retrieving information from 65K tokens ago. The LRA results (Appendix D) partially address this by testing on tasks designed to require long-range dependencies, but the language modeling perplexity improvement is an indirect measure of long-range utilization. A controlled experiment (e.g., synthetic tasks where the correct answer depends on a token exactly 50K steps back) would be a stronger test.
Only structured vs. unstructured is tested; other SSM variants (S5, H3) are not evaluated for length generalization. The paper focuses on S4 as the structured SSM and the Hyena-inspired design as the unstructured alternative, but does not test whether other recent SSMs (S5, which avoids FFT entirely; H3, which uses multiplicative gating) would show different generalization behavior in the BST framework.

Claim 4: Integrating SSM states into attention via cross-attention provides larger benefits than interleaving SSM and attention layers (GSS-HYBRID).

This is supported by the direct comparison in Table 1, where BST-L variants outperform GSS-HYBRID-L across all three datasets at matched compute. The paper also notes that this comparison is somewhat favorable to GSS-HYBRID (which benefited from learning rate grid search) while BST used a fixed learning rate, strengthening the conclusion. However, the comparison is between specific instantiations of two approaches, not a controlled ablation of the integration mechanism alone. Several architectural differences are confounded:

GSS-HYBRID uses Gated State Space (GSS) layers, which incorporate a gating mechanism not present in BST's S4 SSM. BST uses S4 or unstructured kernels. The SSM design differs in addition to the integration method.
GSS-HYBRID's Transformer layers use standard attention, while BST uses Block Transformer layers with cross-attention. The Transformer design differs.
GSS-HYBRID has SSM layers at every position with Transformers interleaved; BST has BST layers at specific positions (1, 7, 9 or 1, 5, 7, 9) with the rest being standard Block Transformer layers. The ratio of SSM to Transformer computation differs.

A cleaner ablation would fix the SSM type, fix the Transformer type, and vary only the integration mechanism (interleaving vs. cross-attention), keeping total SSM and attention FLOPs matched. This ideal ablation is not performed, so the claim that cross-attention is the cause of the improvement — rather than, say, the specific SSM or Transformer implementation details — is an interpretation supported by the results but not rigorously isolated.

Missing experiments that would strengthen the paper:

Parameter-matched comparisons with BRECT and GSS-HYBRID. The BST variants are consistently slightly larger than BRECT (202M–221M vs. 196M). A BRECT variant with matching parameter count or a BST variant scaled down to match BRECT's 196M would isolate the architectural contribution.
Ablation of down-sampling with quantitative perplexity impact. The claim that SSM input dimension reduction has "negligible impact" is stated without a supporting table.
Ablation of context IDs for MF. The necessity of learned context IDs for the MF variant is asserted based on design reasoning but not empirically validated by training MF without context IDs.
Layer-level speed benchmarks on TPU. The speed claims are GPU-based, but training is on TPU where FFT is reported to be 4× slower. TPU benchmarks would provide a more realistic picture of the training speed advantage.
Scaling beyond 400M parameters with matched training FLOPs. The scaling plot in Appendix C (Figure 5) goes to 1.3B parameters but only on PG19 with a single configuration (BST:SH:UNSTRUCT), and the training budgets at larger scales are not reported. A systematic scaling study with FLOPs-matched comparisons at multiple scales would strengthen the claim that BST scales better than BRECT.
Controlled evaluation of long-range utilization in language modeling. An experiment that probes whether the model actually retrieves information from 65K tokens ago — for example, by measuring perplexity improvement when context beyond the training window is added, broken down by dependency distance — would distinguish between benefiting from longer context (any model might) and actively using the SSM's compressed representation for long-range retrieval.

Overall assessment. The experiments genuinely support the paper's architectural claims: BST is faster than BRECT while matching its accuracy, generalizes to longer sequences when using structured SSMs, and outperforms GSS-HYBRID at matched compute. The evidence quality varies: the speed claims are the strongest (Figure 4, left) but are layer-level and GPU-specific; the perplexity comparisons are comprehensive (three datasets, two scales, multiple variants) but show small differences that could be partially explained by parameter count mismatches; the length generalization results (Figure 3) provide a clear diagnostic of the structured/unstructured kernel boundary but also show BRECT generalizing competitively on some datasets. The paper's most robust finding is the task-dependent tradeoff between SH and MF context states — this pattern reverses cleanly across datasets in Table 1 and represents a genuine architectural insight that future practitioners can use. The paper's most qualified claim is the "tenfold speedup" — this is a layer-level property under specific hardware conditions, not an end-to-end training speedup.

6. Limitations and Trade-offs

The FFT Bottleneck on TPUs Partially Undermines the "Tenfold Speedup" in Practice

The assumption or constraint. The paper's headline speed claim — "more than tenfold increase in speed at the layer level" — is measured on GPU (Figure 4 left), yet all model training is conducted on TPUs. Appendix A explicitly acknowledges that FFT operations, which are central to the SSM convolution, are "an important speed bottleneck on TPUs" and that the authors found "JAX FFT was 4× faster on GPUs." The paper does not provide TPU layer-level speed benchmarks, and the training-time comparisons in Table 1 tell a more modest story: on PG19, BST trains at 0.5K TPUv4 hours versus BRECT at 0.8K — an end-to-end speedup of roughly 1.6×, not 10×.

The consequence. A practitioner deploying BST on TPU hardware — the platform used for the paper's own training — will not realize the 6–11× layer-level speedup reported in Figure 4. The backward pass through the SSM (which requires FFT derivatives), the fact that only 3 of 12 layers are BST layers, and the 4× TPU-specific FFT slowdown all compress the theoretical layer-level speedup into a much smaller end-to-end training gain. The "tenfold" figure captures what is architecturally possible given sufficient hardware support for parallel FFT, but it does not reflect realized training speed on the authors' own experimental infrastructure. This distinction matters for anyone making hardware provisioning decisions based on the paper's claims.

What evidence exists in the paper. Appendix A contains the key admission: "We have found that the FFT operation is an important speed bottleneck on TPUs that needs to be resolved to better scale BST to many layers and larger models. While we are still investigating the reasons, we found that JAX FFT was 4× faster on GPUs." The training compute budgets in Table 1 show end-to-end speedups of 1.6× (0.5K vs. 0.8K TPUv4 hours on PG19) rather than the 6–11× layer-level figures. Figure 4 (left) is GPU-only and is explicitly a layer-level benchmark, not a full training measurement.

Mitigation status. The paper partially addresses this by suggesting that future SSM variants that bypass FFT entirely (such as S5, which uses a binary associative operator instead of FFT convolution) could be "simply plugged in" to BST. It also notes that "recent and faster hardware-specific I/O-aware implementations" could improve FFT performance. However, no TPU benchmarks or S5-based BST experiments are reported, so these remain forward-looking suggestions. The limitation is acknowledged transparently but not resolved.

Length Generalization Depends on Structured Kernels, but Structured Kernels Underperform at Training Length

The assumption or constraint. The paper demonstrates a clean diagnostic boundary: structured SSM kernels (S4) generalize to 65K tokens, while unstructured kernels (Hyena-inspired) degrade at unseen lengths (Figure 3). However, Table 1 reveals the inverse pattern at the training length of 4K: unstructured kernels consistently outperform structured kernels across all three datasets at both model scales. For example, on PG19 at ~200M parameters, BST:SH:UNSTRUCT achieves 11.52 (±1.1) versus BST:SH:S4 at 11.57; at ~400M, BST:SH:UNSTRUCT-L achieves 10.37 versus BST:SH:S4-L at 10.47. The same pattern holds on arXiv (2.44 vs. 2.48) and GitHub (2.03 vs. 2.07). The paper attributes this to unstructured kernels having "more free parameters" because they are "no longer restricted by A, B matrices."

The consequence. A practitioner faces a non-trivial deployment decision with no unified solution: train with unstructured kernels for the best perplexity at the training length (but sacrifice length generalization), or train with structured kernels to support arbitrary-length evaluation (but accept worse perplexity per unit of training compute). The paper does not provide guidance on when each choice is preferable — for example, whether the perplexity penalty of structured kernels decreases with scale (the gap appears to narrow from ~0.05 at 200M to ~0.10 at 400M on PG19, but with only two scale points, no trend is established). A deployment that requires both optimal training-length performance and occasional long-sequence evaluation would need to train two separate models or accept suboptimal performance in one regime.

What evidence exists in the paper. The pattern is visible across all relevant rows in Table 1: every UNSTRUCT variant outperforms its S4 counterpart at the training length. Figure 3 shows the crossover: UNSTRUCT degrades at 16K and 65K while S4 maintains or improves. The paper explicitly notes this tradeoff in Section 4.2: "BST:MF:UNSTRUCT-L performs best on the training sequence of 4K and is on-par for 512 with perplexity increasing for unseen 16K and 65K sequence lengths."

Mitigation status. The paper does not attempt to resolve this tradeoff. It does not explore whether structured kernels can be improved to match unstructured performance at the training length (e.g., by increasing the SSM state size N beyond 16, or by using more expressive A/B parameterizations). It does not explore whether unstructured kernels can be made extensible through interpolation of positional encodings. The tradeoff is presented as an intrinsic property of the two kernel families, and the paper does not suggest future work addressing it. This is a genuine unresolved design tension.

Cross-Sequence Information Loss Limits Performance on Very Long Documents

The assumption or constraint. BST's SSM processes sequences of a fixed training length L (4K tokens in the paper's experiments). When a document is longer than L, it must be split into consecutive sequences, and the SSM state at the end of one sequence is not passed to the next — the convolution computes context states only within each 4K chunk, with no mechanism to carry information across chunk boundaries. The paper contrasts this directly with BRECT, which "uses a non-differentiable cache that is carried from one sequence of size L to the next" — the recurrent state vectors from the end of one training sequence are stored and used to initialize the next sequence's recurrence, providing (weak) access to the entire document history beyond the current 4K window.

The consequence. For documents substantially longer than 4K tokens — which describes all three datasets evaluated (PG19 books are 50K–100K tokens, arXiv papers are comparably long, and GitHub repository concatenations are even larger) — BST's SSM context states represent only the last 4K tokens of history, not the entire document. Information from chapter 1 of a book is invisible to a BST model processing chapter 20, unless that information happens to be recapitulated within the last 4K tokens. This is a practical ceiling on the model's effective context length regardless of the SSM's theoretical ability to capture arbitrarily long dependencies. The length generalization experiments in Figure 3 evaluate single sequences of 65K tokens, but these are contiguous sequences — in deployment on a 100K-token book, the model would process the book as multiple 4K chunks, and the 65K evaluation does not measure cross-chunk information retention.

What evidence exists in the paper. Appendix B.1 explicitly acknowledges this gap: "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." The paper further suggests that "BST can further be improved by adding non-differentiable sequence cache for very long documents." The arXiv length generalization results in Figure 3 show BRECT:FIXED:SKIP-L outperforming BST:SH:S4-L at 65K on arXiv, which the authors hypothesize is due to BRECT's cross-sequence cache providing access to "past symbols, definitions, theorems or equations beyond the 4k training sequence length."

Mitigation status. The paper proposes a concrete direction (adding a non-differentiable sequence cache to BST) but does not implement or evaluate it. This is a recognized, but unresolved, limitation. The fix would likely involve storing the SSM's final state vector at the end of each chunk and using it to initialize or condition the SSM of the next chunk — conceptually straightforward but requiring non-trivial engineering to integrate with the convolution-based SSM training pipeline, which currently assumes complete sequences.

The MF vs. SH Tradeoff Is Dataset-Dependent and the Selection Criterion Is Post-Hoc

The assumption or constraint. The paper's redundancy-retrievability taxonomy (Section 3.3) frames the choice between Single-Head and Multi-Filter context states as a tradeoff: SH provides high redundancy/high retrievability (suitable for datasets where local context dominates), while MF provides low redundancy/high utilization (suitable for datasets with long-range structural dependencies). The empirical pattern in Table 1 supports this: SH outperforms MF on PG19 (books with strong local narrative structure), while MF outperforms SH on arXiv and GitHub (scientific articles with cross-references and code with cross-file dependencies). However, the paper provides this explanation after observing the results — there is no a priori prediction or held-out test of the taxonomy.

The consequence. A practitioner approaching a new dataset or task has no operationalized method for determining whether SH or MF will perform better before running experiments. The taxonomy is descriptive (explaining why the results came out as they did) rather than predictive (telling you what to choose for a new domain). Is the relevant criterion the average distance of token-token dependencies? The entropy of the dependency distribution? The ratio of local to long-range linguistic phenomena? None of these are quantified or thresholded. Without a measurable predictor, the choice between SH and MF requires running both variants — which doubles the experimental budget and, for large models, may be prohibitively expensive.

What evidence exists in the paper. Table 1 shows the task-dependent reversal clearly: on PG19, BST:SH:S4 (11.57) < BST:MF:S4 (11.63), meaning SH is better; on arXiv, BST:MF:S4 (2.48) < BST:SH:S4 (2.51), meaning MF is better; on GitHub, BST:MF:UNSTRUCT (2.03) < BST:SH:UNSTRUCT (2.09), meaning MF is better. The pattern is consistent across structured and unstructured kernel variants. However, the gap on PG19 (0.06 perplexity points) is small, and the standard deviation on BST:SH:S4 is ±1.1 — the entire difference falls within the noise. On arXiv and GitHub, the gaps are similarly small (~0.03–0.06). The statistical significance of the reversal is not established (only PG19 has error bars from multiple seeds; arXiv and GitHub results are single runs).

Mitigation status. The paper does not address the operationalization problem. It does not propose a measurable data statistic that predicts whether SH or MF will be superior, nor does it provide a decision rule beyond the post-hoc characterization. The taxonomy is a conceptual contribution that helps interpret results but does not yet serve as an engineering guide. Future work on characterizing datasets by their dependency distance distributions and correlating those with SH/MF performance would close this gap.

No Evaluation on Tasks Requiring Verifiable Long-Range Retrieval in Language Modeling

The assumption or constraint. The paper's primary evaluation metric is language modeling perplexity on PG19, arXiv, and GitHub test/validation sets. While lower perplexity at longer evaluation lengths (Figure 3) is evidence that the model benefits from additional context, perplexity improvements can arise from multiple mechanisms — including simply having a longer local smoothing window — that do not necessarily involve retrieving specific information from 65K tokens ago. The paper does not include controlled experiments that isolate whether BST is actually using the SSM's compressed long-range context for token prediction versus benefiting from more local statistics.

The consequence. A practitioner who needs guaranteed long-range retrieval — for example, a model that must correctly resolve a pronoun reference to a character introduced 50K tokens earlier, or a code model that must correctly predict a function call based on a definition in a distant file — cannot infer from perplexity improvements alone whether BST reliably performs such retrieval. It is possible that BST's perplexity gains come entirely from better local modeling (the SSM's temporal smoothing improves within-window predictions) while the model still fails at genuine long-range dependency resolution. The paper's LRA classification results (Appendix D) partially address this for non-language tasks, but LRA's Retrieval task uses synthetic key-value lookup, not natural language dependency resolution.

What evidence exists in the paper. The LRA results (Table 2) show BST:SH:S4 achieving 90.51 on the Retrieval task (where the correct answer depends on finding a matching key in a long sequence of key-value pairs), slightly below MEGA-CHUNK's 90.97. This is evidence of some long-range retrieval capability, but the LRA tasks use synthetic data and short sequence lengths relative to language modeling. For the main language modeling results, the paper does not include any probe or controlled evaluation that measures the model's ability to use information at specific long distances — for example, by measuring perplexity as a function of dependency distance, or by constructing synthetic language modeling tasks where the correct next token depends on a token exactly K steps back.

Mitigation status. The paper does not address this gap. Future work is not suggested for this specific limitation. The LRA results provide some evidence of retrieval capability in a classification setting, but the gap between LRA's synthetic retrieval and natural language long-range dependency resolution is substantial and unbridged in this paper.

Training Instability of ReST^EM-Style Revision Models — Not Applicable to BST

Note: This limitation — present in the reference example paper as "approximately 38% of correct answers get converted back to incorrect ones" and "ReST^EM experiment degraded" — does not apply to the BST paper. The BST paper does not use revision models, RL-based fine-tuning, or self-improvement loops. There is no equivalent finding in the BST paper. This limitation is omitted.

7. Implications and Future Directions

How This Work Changes the Landscape

The Block-State Transformer does not propose a fundamentally new class of sequence models — both SSMs and Transformers existed before, and BRECT already combined local attention with inter-block memory. Rather, BST makes a diagnostic and architectural contribution: it identifies the sequential bottleneck as the specific reason hybrid recurrent-attention architectures are slow, and it demonstrates that this bottleneck is eliminable by replacing the recurrent cell with a parallelizable SSM convolution whose outputs feed into a cross-attention interface. This is not a paradigm shift in the sense of replacing Transformers, but it is a frontier shift in the design space of efficient long-context architectures: BST occupies a previously empty region of the speed-accuracy tradeoff where inter-block memory and near-SLIDE speed coexist.

The paper's most important conceptual reframing is the integration-via-cross-attention principle. Prior hybrid work (GSS-HYBRID, MEGA) either kept SSM and attention in separate layers or mixed them through gating. BST treats the SSM's output as a queryable memory bank, making attention the interface through which local tokens access long-range context. This reframing matters because it unifies two previously separate concerns — local token-token comparison and long-range retrieval — under a single mechanism (attention queries), and because it is modular: any SSM can produce the memory, and any attention variant can query it. The paper's three context-state variants (SH, MH, MF) demonstrate this modularity, but the deeper point is that the interface is the contribution, not any single variant.

The paper also provides a diagnostic tool for length generalization in hybrid architectures. By comparing structured (S4) and unstructured (Hyena-inspired) kernels inside the same BST framework, the paper cleanly isolates the SSM kernel's parameterization as the determining factor for whether the full hybrid model generalizes to unseen lengths (Figure 3). This resolves a previously muddy question: is length generalization in hybrid models limited by the attention mechanism, the SSM design, or their interaction? The answer — it is dominated by the SSM kernel's extendability, and cross-attention is transparent to this property — is actionable for any future hybrid architecture.

The paper shifts the burden of improvement from search over integration strategies (which prior work did by grid-searching over layer interleaving patterns and learning rates) to improvement of the SSM backend itself. If cross-attention is a stable, modular interface, then progress on SSMs — better parameterizations, faster FFT implementations, FFT-free formulations — directly translates to improved BST performance without architectural re-engineering. This makes SSM research more immediately relevant to practical language modeling than it was when SSMs and Transformers were studied in isolation or only loosely coupled.

Finally, the redundancy-retrievability taxonomy — while currently descriptive rather than predictive — offers a language for discussing memory design in hybrid architectures that was previously missing. Instead of asking "which memory mechanism is best?" (recurrence vs. convolution vs. kNN), researchers can ask "what redundancy-retrievability tradeoff does my task require, and which memory construction provides it?" This reframes memory design from a mechanism-choice problem to a representation-quality problem, and the paper's three variants provide concrete reference points on that spectrum.

Follow-Up Research This Work Enables

Quantitative characterization of the redundancy-retrievability tradeoff in natural language data. The paper's post-hoc explanation — SH wins on PG19 (local narrative structure) and MF wins on arXiv/GitHub (long-range cross-references) — is intuitively plausible but not operationalized. A follow-up study could compute, for each dataset, the distribution of token-token dependency distances (using attention patterns from a trained Transformer, or by measuring how often the next token is predictable from tokens at various lags), and correlate those statistics with the SH-vs-MF performance gap. If average dependency distance or entropy of the dependency distribution predicts the optimal variant, the taxonomy becomes a practical decision rule. This would require training both SH and MF on a broader set of datasets (news, dialogue, legal documents, multi-lingual text) and measuring the correlation between data statistics and the performance delta. A strong result would be a scatter plot with a clear trend: datasets with X property consistently favor SH, datasets with Y property consistently favor MF.

Adding a non-differentiable cross-sequence cache to BST and measuring the gain on documents longer than the training chunk size. The paper explicitly identifies the absence of cross-sequence memory as a weakness relative to BRECT (Appendix B.1), hypothesizing that BRECT's length generalization on arXiv at 65K (Figure 3) is due to its non-differentiable cache that carries recurrent states from one 4K training sequence to the next. Implementing this for BST is non-trivial — the SSM convolution assumes a complete sequence, so the cache would need to initialize or condition the SSM state at the start of each new chunk — but is not architecturally impossible. One approach: store the SSM's final state vector $x_L$ (the $N$ -dimensional compressed representation at the end of chunk $k$ ) and use it to bias the convolution kernel or initialize the state for chunk $k+1$ . A successful implementation would close or reverse the gap with BRECT on arXiv at 65K. This is a direct stress test of whether SSMs can replace recurrence for all aspects of long-range modeling, not just within-chunk dependencies. A negative result — BST with cross-sequence cache still underperforms BRECT on arXiv — would be equally informative, suggesting that explicit attention-based recurrent states retain information differently than SSM state vectors.

Measuring whether BST actually uses its long-range context for token prediction, beyond local statistics. Language modeling perplexity improvements at longer evaluation lengths (Figure 3) could arise from better local smoothing rather than genuine long-range retrieval. A controlled experiment would construct synthetic sequences where the correct next token depends deterministically on a token exactly $K$ positions back (for $K \in \{1K, 2K, 4K, 8K, 16K, 32K, 64K\}$ ), and measure token-level accuracy as a function of $K$ . This would produce a "retrieval accuracy vs. distance" curve that directly measures whether BST:SH:S4 actually retrieves information from the compressed SSM context at long ranges, and how that curve compares to BRECT and to a pure sliding-window Transformer (which should fail beyond its window size). The LRA Retrieval task (90.51 in Table 2) provides partial evidence, but LRA uses synthetic key-value pairs, not natural language. A language-grounded version would use templates like "The key is [X]. [padding tokens]. The value is [Y]. [padding tokens]. The answer is [Y]" with controlled padding lengths, directly measuring retrieval at specific distances.

Scaling BST to the 1B+ parameter range with FLOPs-matched comparisons against BRECT and pure Transformers. The scaling plot in Appendix C (Figure 5) goes to 1.3B parameters on PG19 only, in a single configuration (BST:SH:UNSTRUCT), and does not report training FLOPs for the larger models. A systematic study would train BST, BRECT, and a standard Transformer at matched total FLOPs across at least four scales (200M, 400M, 800M, 1.6B parameters) on all three datasets, measuring both perplexity and wall-clock training time. This would answer whether the BST-vs-BRECT gap genuinely widens with scale (the paper's Figure 5 suggests it might, from 2.5% at 200M to 4.0% at 1.3B, but only two data points and no error bars) and whether BST's speed advantage compresses or expands at larger model sizes (larger models make FFT an increasing fraction of total compute, potentially eroding the parallelism gain). It would also reveal whether the structured-vs-unstructured kernel gap narrows with scale — if structured kernels catch up at 1B+ parameters, the length generalization tradeoff becomes less painful for large models.

Plugging FFT-free SSM variants (S5, H3) into BST and measuring the speed-accuracy-length generalization tradeoff. The paper identifies FFT as the TPU bottleneck (Appendix A) and suggests S5 (which uses a binary associative operator instead of FFT) as a drop-in replacement. A concrete experiment would replace the S4 SSM in BST:SH:S4 with an S5 layer, keeping all other architectural choices identical (same layer placement, same context state construction, same cross-attention interface), and measure: (a) training throughput on TPU vs. GPU (does S5 eliminate the 4× TPU slowdown?), (b) perplexity at 4K training length vs. BST:SH:S4 and BST:SH:UNSTRUCT, (c) length generalization to 65K (does S5's recurrence-based formulation generalize as well as S4's structured convolution?). This experiment directly tests whether the "structured kernel" property is necessary for length generalization or whether it is specifically the convolution formulation that matters — S5 has a structured state space but computes it via parallel scan rather than FFT convolution. A positive result (S5 generalizes as well as S4 while being faster on TPU) would make BST immediately more practical for TPU-based training. A negative result (S5 does not generalize despite its structured parameterization) would reveal that the convolution formulation specifically — with its explicitly extendable kernel — is the key mechanism for length generalization, not the structured state space alone.

Training a difficulty estimator to predict whether a given input chunk benefits more from SH or MF context states. Rather than choosing SH or MF globally for a dataset, a dynamic system could route different sequence chunks to different context-state variants based on the chunk's estimated local-vs-long-range dependency profile. A lightweight classifier (a small MLP or even a linear probe on the SSM's output features) could be trained to predict, for a given chunk, whether SH or MF would produce lower perplexity, using the SSM's own representations as input. This would be a direct application of the paper's insight that SH and MF serve different dependency structures, operationalizing the taxonomy as a per-chunk adaptive mechanism rather than a per-dataset architectural choice. The experiment would measure whether a model that dynamically switches between SH and MF per chunk outperforms either variant alone, and whether the routing decisions correlate with interpretable properties of the input (e.g., chunks with many pronouns route to SH, chunks with many import statements route to MF).

Practical Applications and Downstream Use Cases

On-device language models that need to process long documents under latency constraints. BST's 6–11× layer-level speedup over BRECT (Figure 4, left) and its near-SLIDE speed mean that a deployable model can process entire books or long articles on consumer hardware without the sequential bottleneck that made BRECT impractical. A reading-comprehension or summarization application that needs to ingest a 50K-token document could use BST:SH:S4 (trained at 4K, evaluated at 65K with length generalization) to process the document in large chunks rather than sliding a small window, improving coherence while maintaining interactive latency. The 1.6× end-to-end training speedup over BRECT on PG19 (Table 1, 0.5K vs. 0.8K TPUv4 hours) also makes BST cheaper to fine-tune for domain-specific long-document tasks.

Code understanding and generation over entire repositories. The GitHub results (Table 1) show BST:MF:UNSTRUCT achieving the best perplexity (2.03) among ~200M models, and BST:MF:UNSTRUCT-L achieving 1.83 at ~400M. A code assistant that needs to reason about cross-file dependencies — understanding a function call in one file based on its definition in another — can use the MF variant, whose independent filters can specialize in tracking different types of long-range structural dependencies (import graphs, class hierarchies, variable scopes). The ability to train at 4K tokens and evaluate at 65K (Figure 3, GitHub panel, where BST:SH:S4-L improves from ~2.0 to ~1.8 at 65K) means the model can ingest larger portions of a repository at inference time than it was trained on, without architecture changes or retraining.

Batch processing of scientific literature for information extraction. The arXiv results (Table 1) show BST:MF:UNSTRUCT-L achieving 2.41 perplexity, ~4% better than GSS-HYBRID-L (2.51), on a corpus where cross-references to theorems, equations, and citations span entire papers. An organization processing millions of scientific papers for automated extraction of methodology, results, or citation graphs could use BST:MF variants to model each paper as a single long sequence (rather than chunking it with loss of cross-reference context), improving extraction accuracy while keeping training costs manageable (BST trains at the same or lower TPUv4 hours budget as competing architectures). The modular SSM interface means the SSM backend can be upgraded (e.g., from S4 to a future FFT-free SSM) to improve throughput without changing the extraction pipeline.

Long-form text generation with maintained coherence. When generating long outputs — stories, reports, documentation — a language model must remain coherent over tens of thousands of tokens. BST's cross-attention interface allows the model to explicitly query its compressed history of what it has already written when generating each new token, rather than relying solely on the KV cache of recent tokens. An authoring tool using BST could maintain a running SSM state that compresses the entire generated text so far, and the cross-attention at each generation step could retrieve relevant earlier content (character descriptions, plot points, defined terms) without the quadratic cost of attending to the full generated sequence. The length generalization results (Figure 3) suggest that a model trained at 4K can generate coherently far beyond its training length, making long-form generation feasible without training on prohibitively long sequences.

When to Prefer This Method

The paper explicitly positions BST against BRECT, GSS-HYBRID, SLIDE:12L, and TRSF-XL:2048, and its results support clear decision rules anchored to specific conditions:

Prefer BST over BRECT when you need inter-block memory but training or inference speed matters. BST matches BRECT's perplexity (Table 1) while being 6–11× faster at the layer level on GPU (Figure 4, left). The main caveat: if your documents are much longer than the training sequence length and you cannot afford to train on longer sequences, BRECT's non-differentiable cross-sequence cache may provide better length generalization on datasets with very long-range cross-references (as on arXiv at 65K, Figure 3 middle), though the paper does not quantify how much this matters in practice.
Prefer BST over GSS-HYBRID when you want a stable, modular hybrid architecture that does not require per-layer learning rate tuning. BST matches or outperforms GSS-HYBRID-L on all three datasets at matched compute (Table 1) while using the same learning rate for all layers. If you anticipate needing to swap SSM backends (e.g., upgrading from S4 to a future SSM), BST's cross-attention interface is cleaner and more modular than GSS-HYBRID's layer interleaving.
Prefer BST with structured kernels (S4, DSS) over unstructured kernels when you need length generalization — evaluation on sequences longer than training length. BST:SH:S4-L achieves the best perplexity at 65K across all datasets (Figure 3). The cost is slightly worse perplexity at the training length (Table 1: BST:SH:UNSTRUCT outperforms BST:SH:S4 on all datasets at 4K).
Prefer BST with unstructured kernels (Hyena-inspired) over structured kernels when evaluation will always occur at the training sequence length. Unstructured kernels consistently achieve better perplexity at 4K (Table 1) and avoid the kernel recompilation step, providing a small engineering simplicity advantage. The paper's 1.3B-parameter scaling experiment (Appendix C, Figure 5) uses unstructured kernels, suggesting this choice becomes more attractive at larger scales, though the scaling data is limited.
Prefer SH over MF when your data has dominant local dependency structure (like narrative prose in PG19). SH's redundant adjacent context states make recent tokens highly retrievable. Prefer MF over SH when your data has distributed long-range structural dependencies (like cross-references in arXiv or import graphs in GitHub). MF's independent filters provide more efficient use of the context state budget. MH, which performs on par with SH but costs 8% more parameters, is not recommended by the paper's own evidence unless there is a specific reason to believe that per-head SSM readouts are beneficial for your task.