Mamba-style selective state-space blocks

Mamba replaces quadratic attention with a linear-time selective state-space recurrence. In full models, the token controls discretization and input/output state parameters.

This file exposes two layers:

• MambaBlockSpec: a compact theorem-friendly diagonal SSM block for scan laws and kernel validation.
• SelectiveMambaBlockSpec: a fuller Mamba-style block with input/gate projections, causal depthwise convolution, SiLU, token-dependent Delta/B/C, diagonal selective scan, gated output, and output projection.

The compact block is intentionally retained: it is the smallest reusable core for proving scan algebra and for validating CUDA kernels. The full block builds the paper-style Mamba dataflow on top of the same affine-scan idea.

• recurrent selective scan (h ← A ⊙ h + B ⊙ x_state),
• a gated state readout,
• tokenwise input/output projections.

Implementation status

nn.mamba uses the trainable selective block in NN/Runtime/Autograd/Model/Layers/Mamba.lean. The recurrence in NN/Runtime/Autograd/Model/Mamba.lean follows the SelectiveMambaBlockSpec dataflow through generic differentiable operations: causal depthwise convolution and SiLU produce the feature used for softplus time steps and token-dependent B/C, then the diagonal state update feeds the skip connection, SiLU gate, and output projection. The runtime stores logA, so this Spec's rate tensor corresponds to exp(logA). It uses the dense time-step projection described below and does not call a fused variable-coefficient scan.

The layer has eleven trainable parameter tensors. Its expanded channel count is innerWidth = expansion * hiddenWidth, where hiddenWidth is the output feature width. Expansion, state width, and convolution width determine the saved tensor shapes. Checkpoints from the former gated recurrence use a different parameter layout and cannot be loaded into this layer unchanged. Each layer call starts with zero recurrent state and empty convolution history; Runtime.Autograd.Model.Mamba.runArray accepts and returns both a state tensor of shape [innerWidth, stateWidth] and newest-first projected-token history for continuation across chunks.

The causality (prefix-preservation) theorems in NN/MLTheory/Proofs/StateSpace/MambaCausality.lean concern MambaBlockSpec.runArray, SelectiveMambaBlockSpec.runArray, and SelectiveMambaBlockSpec.runArrayWithHistory, built on the scan algebra in NN/MLTheory/Proofs/StateSpace/Scan.lean. They do not establish equivalence between the generic differentiable runtime and these Spec runners or prove its gradients.

PyTorch import caveats

• convKernel is indexed (tap, channel) with taps newest-first; PyTorch's depthwise conv1d weight is (channel, 1, tap) with taps oldest-first, so import needs a transpose and a reversal.
• dtProj is one square innerDim × innerDim map; the reference model's dt_rank factorization (x_proj then dt_proj) is not represented and must be multiplied out before loading.
• bProj and cProj are stored as separate innerDim × stateDim matrices; the reference model packs B, C, and the dt input into one x_proj output.

References:

• Gu, Dao. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", COLM 2024.
• Dao, Gu. "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality" (Mamba-2), ICML 2024.

structure Models.MambaBlockSpec (α : Type) [TorchLean.Storage α] (inputDim stateDim outputDim : ℕ) : Type

Parameters for a compact diagonal Mamba-style block.

• inProj : TorchLean.Tensor α [inputDim, stateDim]

  Input projection into SSM state channels.

• gateProj : TorchLean.Tensor α [inputDim, stateDim]

  Gate projection. The gate is sigmoid(x @ gateProj).

• outProj : TorchLean.Tensor α [stateDim, outputDim]

  Output projection from gated state channels.

• ssm : Spec.Dynamics.DiagonalSSM α stateDim

  Diagonal state-space core.

def Models.MambaBlockSpec.projectInput {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim : ℕ} (m : MambaBlockSpec α inputDim stateDim outputDim) (x : TorchLean.Tensor α [inputDim]) : TorchLean.Tensor α [stateDim]

Input-to-state projection.

def Models.MambaBlockSpec.gate {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim : ℕ} (m : MambaBlockSpec α inputDim stateDim outputDim) (x : TorchLean.Tensor α [inputDim]) : TorchLean.Tensor α [stateDim]

Token-dependent sigmoid gate.

def Models.MambaBlockSpec.step {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim : ℕ} (m : MambaBlockSpec α inputDim stateDim outputDim) (h : TorchLean.Tensor α [stateDim]) (x : TorchLean.Tensor α [inputDim]) : TorchLean.Tensor α [stateDim] × TorchLean.Tensor α [outputDim]

One Mamba-style token step, returning (new_state, output).

def Models.MambaBlockSpec.runArray {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim : ℕ} (m : MambaBlockSpec α inputDim stateDim outputDim) (h0 : TorchLean.Tensor α [stateDim]) (xs : Array (TorchLean.Tensor α [inputDim])) : TorchLean.Tensor α [stateDim] × Array (TorchLean.Tensor α [outputDim])

Run an array of tokens through the recurrent block.

@[simp]

theorem Models.MambaBlockSpec.runArray_empty {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim : ℕ} (m : MambaBlockSpec α inputDim stateDim outputDim) (h0 : TorchLean.Tensor α [stateDim]) : m.runArray h0 #[] = (h0, #[])

An empty token sequence leaves the hidden state untouched and emits nothing.

@[simp]

theorem Models.MambaBlockSpec.runArray_outputs_size {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim : ℕ} (m : MambaBlockSpec α inputDim stateDim outputDim) (h0 : TorchLean.Tensor α [stateDim]) (xs : Array (TorchLean.Tensor α [inputDim])) : (m.runArray h0 xs).2.size = xs.size

A Mamba recurrent pass emits one output token per input token.

structure Models.SelectiveMambaBlockSpec (α : Type) [TorchLean.Storage α] (inputDim innerDim stateDim outputDim convWidth : ℕ) : Type

Parameters for a fuller Mamba-style selective SSM block.

Shape conventions:

• inputDim: token/input feature width,
• innerDim: expanded channel width used by Mamba's convolution and SSM path,
• stateDim: per-channel diagonal SSM state size,
• outputDim: output feature width,
• convWidth: causal depthwise-convolution width.

The recurrence state has shape [innerDim, stateDim]. This mirrors the common implementation view of Mamba where each expanded channel carries a small diagonal state vector.

• xProj : TorchLean.Tensor α [inputDim, innerDim]

  Content/input projection x -> x_path.

• zProj : TorchLean.Tensor α [inputDim, innerDim]

  Gate projection x -> z_path.

• convKernel : TorchLean.Tensor α [convWidth, innerDim]

  Causal depthwise-convolution kernel, indexed by (tap, channel) with tap 0 applied to the current token and tap t to the token t steps back. PyTorch's conv1d weight of shape (innerDim, 1, convWidth) stores taps oldest-first along its last axis, so importing a checkpoint requires transposing to (tap, channel) and reversing the tap axis.

• convBias : TorchLean.Tensor α [innerDim]

  Causal depthwise-convolution bias.

• dtProj : TorchLean.Tensor α [innerDim, innerDim]

  Projection from activated convolution features to per-channel time steps Delta.

  This is a single square map. The reference implementation factors it through a low-rank bottleneck as x_proj (innerDim -> dt_rank) followed by dt_proj (dt_rank -> innerDim); the product of those two matrices can be loaded here, but the factorization itself is not modelled.

• dtBias : TorchLean.Tensor α [innerDim]

  Bias before the softplus time-step nonlinearity.

• A : TorchLean.Tensor α [innerDim, stateDim]

  Positive diagonal state rates A[d,n] used as exp(-Delta[d] * A[d,n]).

• bProj : TorchLean.Tensor α [innerDim, stateDim]

  Token-dependent input-state projection B_t = u_t @ bProj.

• cProj : TorchLean.Tensor α [innerDim, stateDim]

  Token-dependent state-output projection C_t = u_t @ cProj.

• dSkip : TorchLean.Tensor α [innerDim]

  Per-channel residual/skip coefficient.

• outProj : TorchLean.Tensor α [innerDim, outputDim]

  Output projection from expanded channels to output features.

def Models.SelectiveMambaBlockSpec.projectX {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (x : TorchLean.Tensor α [inputDim]) : TorchLean.Tensor α [innerDim]

Projection feeding the content path before convolution and selective state updates.

def Models.SelectiveMambaBlockSpec.projectZ {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (x : TorchLean.Tensor α [inputDim]) : TorchLean.Tensor α [innerDim]

Projection feeding the multiplicative gate path in the selective state-space block.

def Models.SelectiveMambaBlockSpec.siluVec {α : Type} [TorchLean.Storage α] [Context α] {innerDim : ℕ} (x : TorchLean.Tensor α [innerDim]) : TorchLean.Tensor α [innerDim]

SiLU/Swish applied channelwise.

def Models.SelectiveMambaBlockSpec.causalDepthwiseConv {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (history : Array (TorchLean.Tensor α [innerDim])) : TorchLean.Tensor α [innerDim]

Causal depthwise convolution from a newest-first history of projected tokens.

history[0] is the current projected token, history[1] is the previous token, etc. Missing history entries are treated as zero padding.

def Models.SelectiveMambaBlockSpec.delta {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (u : TorchLean.Tensor α [innerDim]) : TorchLean.Tensor α [innerDim]

Token-dependent positive time steps Delta = softplus(u @ dtProj + dtBias).

def Models.SelectiveMambaBlockSpec.bToken {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (u : TorchLean.Tensor α [innerDim]) : TorchLean.Tensor α [stateDim]

Token-dependent input-state vector B_t.

def Models.SelectiveMambaBlockSpec.cToken {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (u : TorchLean.Tensor α [innerDim]) : TorchLean.Tensor α [stateDim]

Token-dependent state-output vector C_t.

def Models.SelectiveMambaBlockSpec.selectiveStateStep {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h : TorchLean.Tensor α [innerDim, stateDim]) (u : TorchLean.Tensor α [innerDim]) : TorchLean.Tensor α [innerDim, stateDim]

One selective diagonal SSM update:

h'[d,n] = exp(-Delta[d] * A[d,n]) * h[d,n] + (Delta[d] * B_t[n]) * u[d].

def Models.SelectiveMambaBlockSpec.stateReadout {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h : TorchLean.Tensor α [innerDim, stateDim]) (u : TorchLean.Tensor α [innerDim]) : TorchLean.Tensor α [innerDim]

Read out expanded channels from the updated state using C_t, plus the Mamba skip path.

def Models.SelectiveMambaBlockSpec.stepWithHistory {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h : TorchLean.Tensor α [innerDim, stateDim]) (history : Array (TorchLean.Tensor α [innerDim])) (z : TorchLean.Tensor α [innerDim]) : TorchLean.Tensor α [innerDim, stateDim] × TorchLean.Tensor α [outputDim]

One full Mamba token step from an already-updated convolution history.

The history argument is newest-first and must include the current projected content token.

def Models.SelectiveMambaBlockSpec.stepWithConvolutionHistory {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (state : TorchLean.Tensor α [innerDim, stateDim] × Array (TorchLean.Tensor α [innerDim])) (x : TorchLean.Tensor α [inputDim]) : (TorchLean.Tensor α [innerDim, stateDim] × Array (TorchLean.Tensor α [innerDim])) × TorchLean.Tensor α [outputDim]

One recurrent step while carrying the newest-first convolution history.

The carried history is truncated to the convWidth most recent projected tokens. Older entries can never be read by causalDepthwiseConv, so dropping them changes no output while keeping the state size bounded over arbitrarily long sequences.

def Models.SelectiveMambaBlockSpec.runArrayWithHistory {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h0 : TorchLean.Tensor α [innerDim, stateDim]) (history : Array (TorchLean.Tensor α [innerDim])) (xs : Array (TorchLean.Tensor α [inputDim])) : TorchLean.Tensor α [innerDim, stateDim] × Array (TorchLean.Tensor α [outputDim])

Recurrent runner from an existing state and newest-first convolution history.

def Models.SelectiveMambaBlockSpec.runArray {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h0 : TorchLean.Tensor α [innerDim, stateDim]) (xs : Array (TorchLean.Tensor α [inputDim])) : TorchLean.Tensor α [innerDim, stateDim] × Array (TorchLean.Tensor α [outputDim])

Run a sequence through the full selective Mamba block.

@[simp]

theorem Models.SelectiveMambaBlockSpec.runArrayWithHistory_empty {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h0 : TorchLean.Tensor α [innerDim, stateDim]) (history : Array (TorchLean.Tensor α [innerDim])) : m.runArrayWithHistory h0 history #[] = (h0, #[])

With no tokens the convolution history is never consulted and the state is returned as is.

@[simp]

theorem Models.SelectiveMambaBlockSpec.runArray_empty {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h0 : TorchLean.Tensor α [innerDim, stateDim]) : m.runArray h0 #[] = (h0, #[])

Same for the selective block: no tokens in, no tokens out.

@[simp]

theorem Models.SelectiveMambaBlockSpec.runArrayWithHistory_outputs_size {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h0 : TorchLean.Tensor α [innerDim, stateDim]) (history : Array (TorchLean.Tensor α [innerDim])) (xs : Array (TorchLean.Tensor α [inputDim])) : (m.runArrayWithHistory h0 history xs).2.size = xs.size

The full Mamba recurrent pass emits one output token per input token.

@[simp]

theorem Models.SelectiveMambaBlockSpec.runArray_outputs_size {α : Type} [TorchLean.Storage α] [Context α] {inputDim innerDim stateDim outputDim convWidth : ℕ} (m : SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h0 : TorchLean.Tensor α [innerDim, stateDim]) (xs : Array (TorchLean.Tensor α [inputDim])) : (m.runArray h0 xs).2.size = xs.size

The public full Mamba runner emits one output token per input token.