Mamba -- State-Space Models¶

Mamba represents a fundamental departure from the transformer paradigm. Introduced by Gu and Dao (2023)¹, Mamba is a Selective State Space Model (SSM) that achieves linear-time sequence processing with constant memory per step. Where transformers compute pairwise interactions between all tokens (quadratic in sequence length), Mamba maintains a compressed hidden state that evolves as each token is processed -- analogous to a recurrent neural network but with carefully structured dynamics that enable efficient parallel training.

1. Architecture Overview¶

2. Key Innovations¶

2.1 Continuous-Time State-Space Model¶

The foundation of Mamba is a linear dynamical system in continuous time:

2.2 Discretization¶

To process discrete token sequences, the continuous system is discretized using a step size \( \Delta \):

2.3 Selective Scan (The Mamba Innovation)¶

In S4, the parameters \( \Delta, B, C \) are fixed across the sequence. Mamba makes them input-dependent:

2.4 Efficient Parallel Scan¶

During training, the sequential recurrence is computed efficiently using a parallel prefix scan (also called parallel associative scan):

3. Architecture Diagram¶

flowchart TD
    INPUT["Token IDs"] --> EMB["Token Embedding"]
    EMB --> BLOCKS

    subgraph BLOCKS["N x MambaBlock"]
        NORM["RMSNorm"]
        NORM --> PROJ["Input Projection\n(expand by factor E)"]
        PROJ --> SPLIT["Split"]
        SPLIT --> CONV["Conv1D\n(kernel=4)"]
        SPLIT --> GATE["Gate Branch\n(SiLU)"]
        CONV --> SILU["SiLU Activation"]
        SILU --> SSM_PARAMS["Compute delta, B, C\n(input-dependent)"]
        SSM_PARAMS --> SSM["Selective Scan\n(h_t = A_bar h_{t-1} + B_bar x_t)"]
        SSM --> MUL["Element-wise Multiply"]
        GATE --> MUL
        MUL --> OUT_PROJ["Output Projection\n(contract by factor E)"]
        OUT_PROJ --> RES["Residual Add"]
    end

    BLOCKS --> FINAL_NORM["Final RMSNorm"]
    FINAL_NORM --> LM_HEAD["LM Head"]

4. Configuration Parameters¶

5. Mathematical Formulation¶

5.1 Full MambaBlock Forward Pass¶

For input \( x \in \mathbb{R}^{s \times d} \):

Step 1: Input projection and gating

[ z = x W_{\text{in}}, \quad z \in \mathbb{R}^{s \times 2d_{\text{inner}}} ] [ u = z_{:, \, :d_{\text{inner}}}, \quad g = z_{:, \, d_{\text{inner}}:} ]

Step 2: Causal convolution

[ u' = \text{Conv1D}(u, \text{kernel_size}=4) ] [ u'' = \text{SiLU}(u') ]

Step 3: SSM parameter computation

[ \Delta = \text{softplus}(u'' W_{\Delta} + b_{\Delta}) ] [ B = u'' W_B, \quad C = u'' W_C ]

Step 4: Discretization and selective scan

[ \bar{A}t = \exp(\Delta_t \odot A), \quad \bar{B}_t = \Delta_t \odot B_t ] [ h_t = \bar{A}_t \odot h_t \odot u''_t ] [ y_t = (C_t \cdot h_t) + D \cdot u''_t ]} + \bar{B

Step 5: Gated output

5.2 Comparison: Attention vs. SSM¶

6. Zig Implementation¶

6.1 MambaConfig¶

pub const MambaConfig = struct {
    n_layers: u32,
    d_model: u32,
    d_state: u32 = 16,         // SSM state dimension N
    d_conv: u32 = 4,           // causal convolution kernel size
    expand_factor: u32 = 2,    // inner dimension = d_model * expand_factor
    vocab_size: u32 = 50280,
    dt_rank: ?u32 = null,      // null means auto = ceil(d_model / 16)
    dt_min: f32 = 0.001,
    dt_max: f32 = 0.1,
    norm_eps: f32 = 1e-5,

    pub fn dInner(self: MambaConfig) u32 {
        return self.d_model * self.expand_factor;
    }

    pub fn dtRank(self: MambaConfig) u32 {
        return self.dt_rank orelse
            @intFromFloat(@ceil(@as(f32, @floatFromInt(self.d_model)) / 16.0));
    }
};

6.2 MambaBlock¶

pub const MambaBlock = struct {
    norm: RMSNorm,
    in_proj: Linear,       // d_model -> 2 * d_inner
    conv1d: Conv1D,        // causal conv, kernel_size = d_conv
    dt_proj: Linear,       // d_inner -> dt_rank -> d_inner
    A_log: Tensor(f32),    // log of diagonal A matrix [d_inner, d_state]
    D: Tensor(f32),        // skip connection parameter [d_inner]
    x_proj: Linear,        // d_inner -> dt_rank + 2 * d_state
    out_proj: Linear,      // d_inner -> d_model

    ssm_state: Tensor(f32),  // [d_inner, d_state] -- persistent state

    pub fn forward(self: *MambaBlock, x: []f32) ![]f32 {
        const normed = self.norm.forward(x);

        // Project and split into main path + gate
        const projected = self.in_proj.forward(normed);
        const u = projected[0..self.config.dInner()];
        const gate = projected[self.config.dInner()..];

        // Causal convolution + activation
        const conv_out = self.conv1d.forward(u);
        const activated = silu(conv_out);

        // Compute input-dependent SSM parameters
        const x_db = self.x_proj.forward(activated);
        const dt = softplus(self.dt_proj.forward(x_db[0..dt_rank]));
        const B = x_db[dt_rank .. dt_rank + d_state];
        const C = x_db[dt_rank + d_state ..];

        // Selective scan
        const A = negate(exp(self.A_log.data));
        const y = self.selectiveScan(activated, dt, A, B, C);

        // Gated output
        const gated = elementMul(y, silu(gate));
        return self.out_proj.forward(gated);
    }

    fn selectiveScan(
        self: *MambaBlock,
        u: []f32,     // input
        dt: []f32,    // step sizes
        A: []f32,     // state matrix (diagonal)
        B: []f32,     // input matrix
        C: []f32,     // output matrix
    ) []f32 {
        // For each element in d_inner:
        //   A_bar = exp(dt * A)
        //   B_bar = dt * B
        //   h = A_bar * h + B_bar * u
        //   y = C^T * h + D * u
        // ... (see source for full implementation)
    }
};

7. Variants¶

7.1 Mamba-2 (Structured State Space Duality)¶

Mamba-2 (Dao and Gu, 2024)³ reformulates the selective SSM as a structured matrix operation, revealing a deep connection between SSMs and attention:

7.2 Model Variants¶

8. Educational Value¶

9. References¶