---

Mamba

Introduced in Linear-Time Sequence Modeling with Selective State Spaces.

Mamba is an efficient sequential model, that combines the strengths of [[Transformer]]s and [[Recurrent Neural Network]]s. It can also be seen as a general, learnable non-LTI filter. It builds on top of the State Space Model, adding a selection mechanism which improves the compression of the continuous memory state.

Overview

The classic [[Transformer]] architecture is limited in two ways: 1. Anything outside the context window isn't considered at all. 2. Compute scales quadratically with respect to the window length.

Thus the authors propose to add a selection mechanism to the State Space Model to combine the continuous memory with the strengths of attention in transformers. The normal SSM must be an LTI System, being invariant to time and input, so the recurrent state could be computed via convolution. Mamba circumvents this short-coming through the use of an optimized scan which computes the model recurrently in linear time. This allows Mamba to selectively attend to the current input without having to remember (cache) the previous inputs. Also, this attention is based on the continuous memory state and not on just a limited time-window.

The figure shows the SSM with the newly proposed selection mechanism (blue). The SSM assumes the learned matrices A, B, C and \(\Delta t\) to be constant across time. Mamba allows the matrices to change depending on the input by efficiently expanding the continuous memory \(h_t\) in SRAM.

SSM

Mamba is an SSM and uses a state-space system defined as

\[ \begin{align} h^{\prime}(t) &= Ah(t) + Bx(t)\\ y &= Ch(t) \end{align} \]

It uses [[Zero-Order Hold]] for discretization, so the recurrence defined as

\[ \begin{align} h_{t} &= \bar{A}h_{t-1}+\bar{B}x_{t}\\ y_{t} &= Ch(t) \end{align} \]

is given by the discretized matrices:

\[ \bar{A}=e^{\Delta A}\quad \bar{B}=(\Delta A)^{-1}(e^{\Delta A}-I)\Delta B \]

Mamba is a structured SSMs ([[S4]]), which used a structured transition matrix for efficiently computing the linear recurrence. Similar to [[S4]], a diagonal transition matrix is chosen and applied in SISO-style:

\[ A\in \mathbb{R}^{N\times N}\quad B\in \mathbb{R}^{N\times 1}\quad C\in \mathbb{R}^{1\times N} \]

Using reason for using real instead of complex numbers is that Mamba is designed for discrete tasks, where real numbers have empirically been sufficient. However, the authors hypothesize potential added value for continuous data tasks.

For stable system dynamics, the initialization of A is based on HIPPO theory, using S4D-Lin for complex and S4D-Real for real mamba. The authors expect other initialization to work as well.

The Mamba SISO SSM is embedded in a neural network architecture. Inspired by [[H3]] and Gated RNNs, the input \(x\) is expanded via a linear layer, followed by a convolution and non-linearity. This expansion allows a large bank of SISO Mamba SSMs to run on each channel in parallel. A gating mechanism mixes in the residual.

Selection Mechanism

Time-invariant SSMs are unable to effectively solve tasks with time-varying dynamics, like selective copying. Transformers solve the issue by remembering a large context window and attending to previous tokens. SSMs only carry a inner state, so to be able to model time-varying dynamics, they must themselves be time-varying and update the inner state selectively, dependent on the input. In practice, this makes the context carried in the state more compressed, and the SSM itself more efficient.

The matrices \(B\) and \(C\) are just linear projections of the channels into the same channel dimension. The time-step \(\Delta\) is chosen to be a one-dimensional projection, that is broadcast to all channels. This way, all channels have the same forgetting/remembering dynamics.

Given the discretization above, the selective parameters have the following interpretation:

• A large \(\Delta\) represents the model focusing on the current input for longer, so the it has a big impact on the state. A small \(\Delta\) conversely represents a very small impact of the input on the state. When \(\Delta=0\), an input is just ignored.
• A selective \(A\) is not tested, as the authors hypothesize that due to the interaction with the selective time-step (\(\bar{A}=e^{\Delta A}\)), selectivity is already baked into the discretized transition matrix.
• The selectivity of \(B\) and \(C\) allows the model to modulate the recurrent dynamics based on input and context respectively.

Optimization

Introducing selection makes the model time-varying and efficient computation using convolutions is not possible, as \(C,B,\Delta\) vary from step to step. To solve the issue, Mamba makes recurrent computation feasible using three hardware optimizations:

• Hardware Specialization: Usually, neural networks are processed in the High Bandwidth Memory (HBM) of the GPU because of its capacity. The authors of Mamba propose to load \(\Delta\), A, B and C from the HBM into the SRAM, which is the cache directly located within the GPU die. While having a much lower capacity, the cache is directly available to the GPU during processing with minimal latency. Mamba uses a fused kernel to compute the discretization, scan and multiplication with CUDA in a single kernel directly in SRAM. Because the inner state is never actually materialized in HBM but only t

• Scan Algorithm: The parallel scan algorithm (Baloch 1990) is an algorithm that converts any sequentially applied associative operation into a much more efficient parallel operation. It does so by iteratively computing and aggregating partial results. This can algorithm can be used to efficiently compute the matrix multiplications of the recurrence.

• Recomputation: For [[Backpropagation]], the intermediate states of the continuous memory are required. However, storing and reading them requires writing and reading \(BLDN\) elements from the slow HBM. In contrast, the costs of loading the inputs and outputs to the SRAM and recomputing the intermediate states only costs \(2\cdot (BLN+DN)\) reads, which is a significant speed up.

Mamba 2

Mamba 2 aims to optimize mambas inference speed even further. They do so by introducing a framework called structured state space duality (SSD), which allows to reframe the computation of SSMs as a matrix multiplication algorithm.

[!Sequential Semiseperable (SSS) Matrix]- This is based on the representation of an SSM as a transformation by a sequential semiseperable (SSS) matrix.

First we represent the each output \(y_{t}\) as the read-out of the sum of previous states:

\[ \begin{align} y_{t} &= C_{t}(A_{t}\ldots A_{1}B_{0}x_{0} + \ldots + A_{t}A_{t-1}B_{t-1}x_{t-1}+B_{t}x_{t}) \\ &=\sum\limits_{s=0}^{t}C_{t}^{\mathsf{T}}A^{\times}_{t:s}B_{s}x_{s} \end{align} \]

Each component of the sum can be written as:

\[ M_{ji}:=C^{\mathsf{T}}_{j}A_{j}\ldots A_{i+1}B_{i} \]

which is a N-SSS for \(B_{k},C_{k}\in \mathbb{R}^{N}\) and \(A_{k}\in \mathbb{R}^{N}\). In the 1-dimensional case, \(B_{k}\) and \(C_{k}\) can be ignored as they are just a scalar multiplication, and represent the SSM as:

Applying \(y=Mx\) results in the linear recurrence of an SSM:

\[ y_{t}=a_{t:0}x_{0}+\ldots+a_{t:t}x_{t}=a_{t}(a_{t-1:0}x_{0}+\ldots+a_{t-1:t-1}x_{t})=a_{t}y_{t-1}+x_{t} \]

This relation holds for higher dimensions as well

This representation has requires only \(O(\mathsf{T})\) parameters and is efficiently computed.

The SSS representation allows to derive a quadratic form, for which an equivalence to transformer-like masked-attention methods can be shown. Attention is just an unfolded quadratic representation of a linear recurrence. In turn, the class of 1-semiseperable masks is equivalent to a scalar-identity (\(A=aI\)) transition matrix. This scalar-identity matrix allows for formulating of the recurrence as

\[ \begin{align} M_{ji} &= A_{j:i}\cdot(C^\mathsf{T}_{j}B_{j})\\ Y &= MX \end{align} \]

which can be computed very efficient through a vectorized implementation that leverages both the linear (for the diagonal blocks) and quadratic (for the off-diagonal blocks) form. Effectively, small chunks are computed in the quadratic form, and then the linear recurrence is applied to the processed chunks. The authors also transform several advancement in attention (multi-head, multi-query, multi-value, normalizations) into their state-space framework and optimize the tensor-operations for multi-GPU training. The resulting SSM performs on par or even outperforms Mamba 1 and Transformer++ while being significantly more efficient in computation.

Mamba 3

Mamba 3 builds on top of Mamba 2, introducing three major changes: - Exponential-Trapezoidal Discretization - Complex-valued State Space - Moving from per-channel SISO to MIMO

These performance, expressivity and especially inference latency.