Selective scan specs

This file contains the small proof layer 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 kept 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.

def Spec.scanArrayFrom {State Input Output : Type} (step : State → Input → State × Output) (initial : State) (initialOutputs : Array Output) (xs : Array Input) : State × Array Output

Run a stateful step over an array, appending each emitted value to initialOutputs.

The initial output buffer makes this suitable for chunked execution without changing the state transition being specified.

def Spec.scanArray {State Input Output : Type} (step : State → Input → State × Output) (initial : State) (xs : Array Input) : State × Array Output

Run a stateful step over an array and return the final state and emitted values.

@[simp]

theorem Spec.scanArray_empty {State Input Output : Type} (step : State → Input → State × Output) (initial : State) : scanArray step initial #[] = (initial, #[])

Scanning an empty array returns the initial state and no outputs.

@[simp]

theorem Spec.scanArray_outputs_size {State Input Output : Type} (step : State → Input → State × Output) (initial : State) (xs : Array Input) : (scanArray step initial xs).2.size = xs.size

A stateful scan emits exactly one value for each input.

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

@[instance_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 : α) (transitions : Array (ScalarAffineTransition α)) : α

Sequentially run scalar affine transitions from an initial state.

def Spec.summarizeScalarAffine {α : Type} [Semiring α] (transitions : Array (ScalarAffineTransition α)) : ScalarAffineTransition α

Summarize an array of transitions 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 : α) : Array (ScalarAffineTransition α) → Array α

Return every recurrent state after each scalar affine transition.

@[simp]

theorem Spec.scalarAffineScan_size {α : Type} [Mul α] [Add α] (h0 : α) (transitions : Array (ScalarAffineTransition α)) : (scalarAffineScan h0 transitions).size = transitions.size

The scalar affine scan has one state per transition.

structure Spec.DiagonalTransition (α : Type) [TorchLean.Storage α] (stateDim : ℕ) : Type

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

• a : TorchLean.Tensor α [stateDim]

  Elementwise recurrent multiplier.

• b : TorchLean.Tensor α [stateDim]

  Elementwise additive token contribution.

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

Apply one diagonal affine state update.

def Spec.DiagonalTransition.compose {α : Type} [TorchLean.Storage α] [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} [TorchLean.Storage α] [Add α] [Mul α] {stateDim : ℕ} (h0 : TorchLean.Tensor α [stateDim]) (transitions : Array (DiagonalTransition α stateDim)) : TorchLean.Tensor α [stateDim]

Sequentially run diagonal transitions and return the final state.

def Spec.diagonalSelectiveScan {α : Type} [TorchLean.Storage α] [Add α] [Mul α] {stateDim : ℕ} (h0 : TorchLean.Tensor α [stateDim]) : Array (DiagonalTransition α stateDim) → Array (TorchLean.Tensor α [stateDim])

Return every hidden state from a diagonal selective scan.

@[simp]

theorem Spec.diagonalSelectiveScan_size {α : Type} [TorchLean.Storage α] [Add α] [Mul α] {stateDim : ℕ} (h0 : TorchLean.Tensor α [stateDim]) (transitions : Array (DiagonalTransition α stateDim)) : (diagonalSelectiveScan h0 transitions).size = transitions.size

The diagonal selective scan has one state per transition.