Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Gu & Dao, 2024, Mamba (First Conference on Language Modeling (COLM 2024))

Subscriber deck — free during early access

2 cards in “2 State Space Models” · whole deck

2 State Space Models

  1. How does a structured state space model (S4) map an input sequence to an output?
    A continuous system \(h'(t) = A h(t) + B x(t)\), \(y(t) = C h(t)\) is discretized with a step size \(\Delta\) (e.g. zero-order hold, \(\bar A = \exp(\Delta A)\)), giving the recurrence \(h_t = \bar A h_{t-1} + \bar B x_t\), \(y_t = C h_t\), applied independently to each channel with an \(N\)-dimensional hidden state.
    • Gu & Dao, 2024, Mamba (First Conference on Language Modeling (COLM 2024)), §2
  2. Why could earlier SSMs be trained efficiently, and what did that require?
    Their parameters \((\Delta, A, B, C)\) were constant over time (linear time invariance), so the recurrence unrolls into a global convolution computable in parallel without materializing the large hidden state; recurrent mode was used for fast autoregressive inference.
    • Gu & Dao, 2024, Mamba (First Conference on Language Modeling (COLM 2024)), §2

Source

  • Albert Gu and Tri Dao. 2024. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. In First Conference on Language Modeling (COLM 2024). https://arxiv.org/abs/2312.00752