Block-Biased Mamba for Long-Range Sequence Processing

TL;DR

This work analyzes why Mamba underperforms on long-range sequence tasks through expressiveness, inductive bias, and training stability perspectives, establishing that shared channel weights and input-dependent dynamics limit effective width and memory retention. It then introduces Block-Biased-S6 (B2S6), which combines a block-structured, multihead-like setup with channel-specific bias to restore expressiveness and provide a gentler inductive bias, while stabilizing training. Theoretical results show B2S6 regains universal approximation properties under block or bias components and exhibits a milder inductive bias (robust to large input magnitudes) with improved stability, complemented by training strategies. Empirically, B2S6 achieves state-of-the-art performance on Long-Range Arena benchmarks and matches Mamba’s language-modeling perplexity on SlimPajama, indicating strong cross-domain applicability and improved long-range processing. Overall, the paper advances long-range sequence modeling by introducing a principled extension to Mamba that enhances expressiveness and stability without sacrificing versatility for language tasks.

Abstract

Mamba extends earlier state space models (SSMs) by introducing input-dependent dynamics, and has demonstrated strong empirical performance across a range of domains, including language modeling, computer vision, and foundation models. However, a surprising weakness remains: despite being built on architectures designed for long-range dependencies, Mamba performs poorly on long-range sequential tasks. Understanding and addressing this gap is important for improving Mamba's universality and versatility. In this work, we analyze Mamba's limitations through three perspectives: expressiveness, inductive bias, and training stability. Our theoretical results show how Mamba falls short in each of these aspects compared to earlier SSMs such as S4D. To address these issues, we propose $\text{B}_2\text{S}_6$, a simple extension of Mamba's S6 unit that combines block-wise selective dynamics with a channel-specific bias. We prove that these changes equip the model with a better-suited inductive bias and improve its expressiveness and stability. Empirically, $\text{B}_2\text{S}_6$ outperforms S4 and S4D on Long-Range Arena (LRA) tasks while maintaining Mamba's performance on language modeling benchmarks.

Figures (7)

• Figure 1: Comparison of S4D, S6, and $\text{B}_2\text{S}_6$ units. S4D uses independent linear SSM units for each channel, giving it high capacity (or width) but no input-dependent selectivity. S6 introduces a selective mechanism that modulates its internal dynamics based on the input, but shares parameters across channels, limiting its effective width and expressiveness. Our proposed $\text{B}_2\text{S}_6$ unit partitions the input into smaller blocks, enabling selective behavior across multiple subspaces (see also dao2024transformers). Additionally, it includes a channel-specific, input-independent bias term that further increases model capacity. These design choices improve the performance on long-range sequence tasks.
• Figure 2: The architecture of the neural network \ref{['eq.onelayer']}. In this picture, a horizontal operator is applied channel-wise to every sequence in a channel, and a vertical operation is applied element-wise to every position in a sequence. A green color indicates a linear operator while an orange color indicates a nonlinear one.
• Figure 3: The mean loss $\|\tilde{\boldsymbol{g}} - \boldsymbol{g}\|_2$ between the true coefficient $\boldsymbol{g}$ and the model prediction $\tilde{\boldsymbol{g}}$. Every model has a single layer and is trained for $10$ epochs. Here, $d$ is the number of channels in a model. For the $\text{B}_2\text{S}_6$ model, $h$ is the number of blocks and $p$ is the number of channels in each block.
• Figure 4: The mean relative loss $|{G}(\mathbf{u}) - \tilde{G}(\mathbf{u})| / \sigma_1$ for different choices of $\sigma_1$ and $\sigma_2$. The S6 model cannot make useful predictions when $\sigma_1$ is small and $\sigma_2$ is large; $\text{B}_2\text{S}_6$ fixes this. In all experiments, we fix $d = 32$. For $\text{B}_2\text{S}_6$, we set $h = 8$ and $p = 4$.
• Figure 5: Numerical experiments to verify \ref{['thm.stability']}, where we compute the ratio between the gradients with respect to the S6 parameters and S4D parameters. For the first two figures, we fix $L = 100$; for the last figure, we fix $c = 1$. The gradients are computed using closed algebraic formulas. The red reference lines in the three log-log plots have slopes of $3$, $2$, and $1/2$, respectively.

Theorems & Definitions (14)

• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• Theorem 5
• Lemma 1
• proof
• proof : Proof of \ref{['thm.UAT']}
• proof : Proof of \ref{['thm.B2S6UAT']}
• Lemma 2