Mathematics Behind Mamba Transformers: A Complete Guide

Continuous State Space Models

State Space Models originate from control theory and signal processing. In continuous time, they are defined by:

\[ \begin{align} h'(t) &= Ah(t) + Bx(t) \quad \text{(state equation)} \\ y(t) &= Ch(t) + Dx(t) \quad \text{(output equation)} \end{align} \tag{1}\]

Where:

• \(h(t) \in \mathbb{R}^N\) is the state vector at time t
• \(x(t) \in \mathbb{R}\) is the input signal
  
• \(y(t) \in \mathbb{R}\) is the output signal
• \(A \in \mathbb{R}^{N \times N}\) is the state transition matrix
• \(B \in \mathbb{R}^N\) is the input matrix
• \(C \in \mathbb{R}^{1 \times N}\) is the output matrix
• \(D \in \mathbb{R}\) is the feedthrough term (often set to 0)

The HiPPO Framework

For the Legendre polynomials case (LegS):

• The A matrix has entries: \(A_{nk} = (2n+1)^{1/2}(2k+1)^{1/2}\) if \(n > k\), and \(A_{nk} = n+1\) if \(n = k\)
• This choice ensures that the state vector maintains an optimal polynomial approximation of the input history

Discretization Process

To apply SSMs to discrete sequences, we discretize using a step size \(\Delta\):

The Zero-Order Hold (ZOH) discretization gives us:

\[ \begin{align} h_k &= \bar{A}h_{k-1} + \bar{B}x_k \\ y_k &= Ch_k \end{align} \tag{2}\]

Where:

\[ \begin{align} \bar{A} &= \exp(\Delta A) \\ \bar{B} &= (\Delta A)^{-1}(\exp(\Delta A) - I)\Delta B \end{align} \tag{3}\]

Computational Forms

The Core Innovation

\[ \begin{align} B &= s_B(x) \\ C &= s_C(x) \\ \Delta &= \tau(s_\Delta(x)) \end{align} \tag{4}\]

Where:

• \(s_B\), \(s_C\), \(s_\Delta\) are learnable projection functions
• \(\tau\) is typically the softplus function: \(\tau(x) = \log(1 + \exp(x))\)

Selection Functions

The selection functions are implemented as linear projections:

\[ \begin{align} s_B(x) &= \text{Linear}_B(x) \quad \in \mathbb{R}^{B \times N} \\ s_C(x) &= \text{Linear}_C(x) \quad \in \mathbb{R}^{B \times N} \\ s_\Delta(x) &= \text{Broadcast}(\text{Linear}_\Delta(x)) \quad \in \mathbb{R}^{B \times N} \end{align} \tag{5}\]

Where \(B\) is the batch size and \(N\) is the state dimension.

Mathematical Justification

The selection mechanism allows the model to:

1. Filter irrelevant information: By modulating \(B\), the model controls what information enters the state
2. Focus on specific aspects: By modulating \(C\), the model controls what information is output
   
3. Control information flow: By modulating \(\Delta\), the model controls the rate of state updates

Complete Block Definition

A Mamba block processes input \(x \in \mathbb{R}^{B \times L \times D}\) as follows:

# Pseudocode for Mamba block processing

def mamba_block(x):
    # 1. Input Projections
    x_prime = Linear_in(x)  # ∈ R^(B×L×2E) 
    x1, x2 = split(x_prime)  # each ∈ R^(B×L×E)
    
    # 2. Selection Parameters  
    B = s_B(x1)  # ∈ R^(B×L×N)
    C = s_C(x1)  # ∈ R^(B×L×N)
    Delta = softplus(s_Delta(x1))  # ∈ R^(B×L×N)
    
    # 3. Discretization
    A_bar = exp(Delta * A)  # ∈ R^(B×L×N×N) 
    B_bar = Delta * B       # ∈ R^(B×L×N)
    
    # 4. SSM Computation
    y1 = SSM(A_bar, B_bar, C)(x1)  # ∈ R^(B×L×E)
    
    # 5. Gating and Output
    y = y1 * SiLU(x2)
    output = Linear_out(y)
    
    return output

Selective Scan Algorithm

The core SSM computation for a sequence of length \(L\):

\[ \begin{align} h_0 &= 0 \\ \text{for } k &= 1 \text{ to } L: \\ h_k &= \bar{A}_k \odot h_{k-1} + \bar{B}_k \odot x_k \\ y_k &= C_k \odot h_k \end{align} \tag{6}\]

Where \(\odot\) denotes element-wise multiplication.

Parallel Scan Formulation

For parallel computation, we can express the recurrence as:

\[ h_k = \left(\prod_{i=1}^k \bar{A}_i\right) \odot h_0 + \sum_{j=1}^k \left(\prod_{i=j+1}^k \bar{A}_i\right) \odot (\bar{B}_j \odot x_j) \tag{7}\]

This can be computed efficiently using parallel prefix sum algorithms.

Matrix Form

The complete transformation can be written as:

\[ Y = \text{SSM}(X; A, B, C, \Delta) \tag{8}\]

Where each element is:

\[ Y[b,l,d] = \sum_{k=1}^l \sum_{n=1}^N C[b,l,n] \cdot \left(\prod_{j=k+1}^l \bar{A}[b,j,n]\right) \cdot \bar{B}[b,k,n] \cdot X[b,k,d] \tag{9}\]

Attention vs Selection

Information Flow

Initialization Strategies

1. A Matrix: Initialize using HiPPO-LegS or similar structured initialization
2. B, C Projections: Standard Gaussian initialization scaled by dimension
3. \(\Delta\) Projection: Initialize to encourage slow dynamics initially

Numerical Stability

Several techniques ensure stable computation:

Training Considerations

• Gradient Flow: The recurrent nature requires careful handling of gradients
• Truncated BPTT: May use truncated backpropagation for very long sequences
• Regularization: Apply dropout to projections rather than the state itself

Multi-Head Mamba

Similar to multi-head attention, Mamba can use multiple independent SSM heads:

\[ \text{MultiHead\_Mamba}(x) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O \]

where \(\text{head}_i = \text{Mamba}_i(x)\)

Bidirectional Processing

For non-causal applications, bidirectional Mamba processes sequences in both directions:

\[ y = \text{Mamba}_{\text{forward}}(x) + \text{Mamba}_{\text{backward}}(\text{reverse}(x)) \tag{12}\]

Integration with Other Mechanisms

Mamba blocks can be combined with:

• MLP blocks: Following similar patterns to transformer architectures
• Convolution: For local pattern recognition
  
• Attention: For hybrid architectures

Mathematical Notation Summary