SlowGuess/NeurIPS_2023_OCR huggingface dataset:
BST:SH:S4 is composed of four BST layers (no BRT layers are interleaved) and two S4 layers on top. We use the same standard block length of 512 for BST and BRT. However, we train BST and BRT on the full sequences (up to 16K for Path-X). We use AdamW as our optimizer [24] with a warmup for the learning rate, where we start from a value of  $1e^{-7}$  and increase the learning rate linearly up a specified value  $\in \{1e^{-3}, 2e^{-3}, 4e^{-3}\}$  for the first  $10\%$  of training. This is followed by cosine annealing for the rest of training down to a value of  $1e^{-7}$ . All layers are bidirectional, including the S4 layer in BST:SH:S4 as described in [13]. Our weight decay is chosen from  $\{0, 0.05, 0.1, 0.15\}$  and our dropout is chosen from  $\{0, 0.1\}$ . Except for Path-X experiments, we use weight decays  $\in \{0.03, 0.05, 0.07\}$  for all parameters except S4D matrices A and B. Also, for Path-X, the initialization range of our discretization time step  $\Delta$  for PathX is decreased from  $(\Delta_{\min}, \Delta_{\max}) = (0.001, 0.1)$  to  $(\Delta_{\min}, \Delta_{\max}) = (0.0001, 0.01)$ .
Our results on LRA are very promising and show that, compared to other state-of the art methods that chunk sequences into blocks, BST is able to model long range dependencies. For example, BST outperforms MEGA-CHUNK [27] on four out of six LRA tasks and by  $1.5\%$  on the average score. However, BST still needs to improve (perhaps by extending the block size) to catch up to MEGA (without chunks).
E Ablation Studies
In the following section, we perform ablations to investigate (1) the placement of a single SSM layer in Table 3 in the overall architecture, (2) the effects of the number of SSM layers added in Table 4, and (3) the size  $D$  of the SSM state in Table 5. For the ablations, we use the  $\sim 200\mathrm{M}$  parameter BST:SH:S4, since it is the fastest model, and assess various configurations on PG19.
Table 3: A single BST at various layer index.
Layer index	Perplexity
3	12.41
7	11.92
9	11.88
12	12.03
Table 4: Multiple BST layers at various locations.
Num layers	Perplexity
2	11.69
3	11.57
4	11.21
5	11.20
Table 5: Increasing BST's S4 model state size  $D$
State Size	Perplexity	Step Time
8	11.95	×0.7
16	11.57	×1.0
32	11.55	×1.8
64	11.54	×3.2
In Table 3, we experiment adding a single BST layer at layer indices 3, 6, 9, 12. We notice that a single BST layer with state size  $D = 16$  located closer to the middle of the whole Block Transformer stack, at index  $= 9$ , has the greatest effect on perplexity. This finding is inline with findings in prior work [42, 21].
In Table 4, we test if adding multiple BST layers yields improvements on performance. We start with BST layers with state size  $D = 16$  at indices 0, 9. We follow by adding another BST layer at index 7 for a total of three BST layers and then another at index 5, followed by another at index 12. Adding more BST layers lowers perplexity. However, the results seem to plateau at 5 BST layers. We note also that there is a  $3.5\%$  training step time increase for each added layer.
In Table 5, we train our models with different state sizes  $D$ . For the state size ablation, we use three BST layers at indices 0, 7, 9. We find that increasing  $D$  improves perplexity to the detriment of training speed (step time). For this reason, we chose  $D = 16$  for Table 1 BST results.
F JAX Implementation of BST
Pseudocode 1 contains a function that implements convolution of multiple filters over the same input sequence using FFT and inverse FFT operations. Pseudocodes 2, 3 and 4 respectively implement context state collection of BST variants: Single-Head (SH), Multi-Head (MH) and Multi-Filter (MF). Finally, Pseudocode 5 runs the Block Transformer sublayer in parallel by feeding the context states to their corresponding block.
Pseudocode 1: Unstructured filters and convolutions.
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
```
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``"
``}
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
`'
```c
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`'
`"
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
```c
' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' */
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
    ' "
.
...
def multichannel_convolution(u, h, b):
    '''Multichannel convolution function. 
Args:
    u: input of shape (seq_length, dim)
    h: filters of shape (seq_length, channels, dim)
    b: bias of shape (channels, dim)
'''t = jnp.linspace(0.0, 1.0, seq_length)
h = jnp.exp(-alpha * t) * dense positional emb(t))
b = get.bias()
return h, b
def multichannel_convolution(u, h, b):
    '''Multichannel convolution function. 
Args:
    u: input of shape (seq_length, dim)
    h: filters of shape (seq_length, channels, dim)
    b: bias of shape (channels, dim)
'''h = rearrange(h, "l c d -> c d l")
fft_size = seq_length * 2
u_f = jnp fft.rfft(x, n=fft_size)
h_f = jnp fft.rfft(h, n=fft_size)
y = jnp fft.irqft(h_f * x_f, n=fft_size, norm="forward")[ ... , :seq_length] # (c, d, l)
y = y + x * b[... , None] # (c, d, l)
y = rearrange(y, "cd l -> ld c") 
return y
```
Context state collection for BST-SH variant."">
num_heads = 8 # (h)
num_states = 32 # (s)
# (SH): Single-Head
def SH_context_states(u):
    """Single-Head Context Collection."""
    h, b = get_filters_[unstruct/s4](channels=1)
Pseudocode 2: Context state collection for BST-SH variants.
y_1 = multichannel_convolution(u, h, b) # y_1: (l, d, 1)
# lift to multiple heads
y_h = dense(y_1) # y_h: (l, d, h)
context_states = jnp.split(   y_h, seq_length // win_length, axis=0)
return context_states # (l/w, w, d, h)
Pseudocode 3: Context state collection for BST-MH variants.
"Context state collection for BST-MH variant."
# (MH): Multi-Head
def MH_context_states(u):
    ""Multi-Head Context Collection."
h, b = get_filters_[unstruct/s4](channels=num_heads)
    y_h = multichannel_convolution(u, h, b) # y_h: (l, d, h)
    context_states = jnp.split(
        y_h, seq_length // win_length, axis=0)
    return context_states # (l/w, w, d, h)
Pseudocode 4: Context state collection for BST-MF variants.
"Context state collection for BST-MF variant."
# (MF): Multi-Filter
def MF_context_states(u):
    ""Multi-Filter Context Collection."
h, b = get_filters_[unstruct/s4](channels=num_states)
y_s = multichannel_convolution(u, h, b) # y_s: (l, d, s)
context_states = jnp.split(
    y_s, seq_length // win_length, axis=0)
# context_states: (l/w, w, d, s)
# collect the last context states
context_states = context_states[:, -1, ...] # (l/w, d, s)
context_states = rearrange(
    context_states, "lw d s -> lw s d")
# shift context states corresponding to windows
context_states = jnp-roll(context_states, 1, axis=1)
# replace the initial window with trainable weights
init_context = get_init_context(num_states) # (d, s)
context_states[0] = init_context
# lift to multiple heads
context_states = dense(context_states)
return context_states # (l/w, s, d, h)
""Block-State Transformer Layer."
# Block Transformers are non-recurrent and parallelizable. block_transformer =jax.vmap(BRecT.nonrecurrent_cell)
def BST(u):
    ""Block-State Transformer Layer."
"
global MF # True if Multi-Filter, False otherwise (SH/MH)  
# split inputs into windows (l/w, w, d)  
u = jnp.split(u, seq_length // win_length, axis=0)  
# collect context states from SSM outputs  
context_states = [SH/MH/MF]_context_states(u)  
# pass the contexts in place of recurrent states  
y = block_transformer(token_embeddingings=u,  
    recurrent_state=context_states,  
    use CROSS attn_causal_mask=not MF,  
    use CROSS_positional_emb=MF, # context IDs)  
return rearrange(y, "lw w d -> (lw w) d") # (l, d)
Pseudocode 5: Block-State Transformer Layer.