Selective scan specs

This file contains the small proof-facing core behind state-space sequence models such as S4 and Mamba.

The key observation, used by Mamba's hardware-aware parallel scan, is that each per-token recurrent update can be viewed as an affine map

h ↦ A_t h + b_t.

Affine maps compose associatively. A recurrent scan can therefore be implemented either by a left-to-right recurrence or by a parallel prefix scan over affine summaries. The scalar definitions below are intentionally compact so that NN/MLTheory/Proofs/StateSpace/Scan.lean can prove the algebra without depending on a particular runtime backend. The diagonal tensor definitions are the direct TorchLean spec analogue used by the model and CUDA contracts.

References:

• Gu, Goel, Ré. "Efficiently Modeling Long Sequences with Structured State Spaces" (S4), ICLR 2022.
• 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 Spec.ScalarAffineTransition (α : Type) : Type

A scalar affine transition h ↦ a*h + b.

• a : α

  Linear multiplier. In diagonal SSMs this is one channel of the discretized state matrix.

• b : α

  Additive input contribution for the current token.

def Spec.instReprScalarAffineTransition.repr {α✝ : Type} [Repr α✝] : ScalarAffineTransition α✝ → ℕ → Std.Format

@[implicit_reducible]

instance Spec.instReprScalarAffineTransition {α✝ : Type} [Repr α✝] : Repr (ScalarAffineTransition α✝)

def Spec.ScalarAffineTransition.apply {α : Type} [Mul α] [Add α] (tr : ScalarAffineTransition α) (h : α) : α

Apply a scalar affine transition.

def Spec.ScalarAffineTransition.id {α : Type} [One α] [Zero α] : ScalarAffineTransition α

Identity affine transition.

def Spec.ScalarAffineTransition.compose {α : Type} [Mul α] [Add α] (t₂ t₁ : ScalarAffineTransition α) : ScalarAffineTransition α

Compose two affine transitions.

compose t₂ t₁ means "first apply t₁, then apply t₂".

def Spec.runScalarAffine {α : Type} [Mul α] [Add α] (h0 : α) : List (ScalarAffineTransition α) → α

Sequentially run a list of scalar affine transitions from an initial state.

def Spec.summarizeScalarAffine {α : Type} [Semiring α] : List (ScalarAffineTransition α) → ScalarAffineTransition α

Summarize a transition list as one affine transition.

This is the algebraic payload used by parallel selective scan: prefix summaries can be produced by any associative scan algorithm, and applying the summary to h0 is equivalent to recurrence.

def Spec.scalarAffineScan {α : Type} [Mul α] [Add α] (h0 : α) : List (ScalarAffineTransition α) → List α

Return every recurrent state after each scalar affine transition.

structure Spec.DiagonalTransition (α : Type) (stateDim : ℕ) : Type

A diagonal vector affine transition h ↦ a ⊙ h + b.

• a : Tensor α (Shape.dim stateDim Shape.scalar)

  Elementwise recurrent multiplier.

• b : Tensor α (Shape.dim stateDim Shape.scalar)

  Elementwise additive token contribution.

def Spec.DiagonalTransition.apply {α : Type} [Add α] [Mul α] {stateDim : ℕ} (tr : DiagonalTransition α stateDim) (h : Tensor α (Shape.dim stateDim Shape.scalar)) : Tensor α (Shape.dim stateDim Shape.scalar)

Apply one diagonal affine state update.

def Spec.DiagonalTransition.compose {α : Type} [Add α] [Mul α] {stateDim : ℕ} (t₂ t₁ : DiagonalTransition α stateDim) : DiagonalTransition α stateDim

Compose diagonal affine transitions channelwise.

The order is the same as ScalarAffineTransition.compose: compose t₂ t₁ is first t₁, then t₂.

def Spec.runDiagonalTransitions {α : Type} [Add α] [Mul α] {stateDim : ℕ} (h0 : Tensor α (Shape.dim stateDim Shape.scalar)) : List (DiagonalTransition α stateDim) → Tensor α (Shape.dim stateDim Shape.scalar)

Sequentially run diagonal transitions and return the final state.

def Spec.diagonalSelectiveScan {α : Type} [Add α] [Mul α] {stateDim : ℕ} (h0 : Tensor α (Shape.dim stateDim Shape.scalar)) : List (DiagonalTransition α stateDim) → List (Tensor α (Shape.dim stateDim Shape.scalar))

Return every hidden state from a diagonal selective scan.