State-space dynamical systems

State-space sequence models are recurrent dynamical systems:

h_{t+1} = A_t h_t + B_t x_t, y_t = C_t h_t + D_t x_t.

The full S4/Mamba family uses structured matrices and input-dependent ("selective") parameters. This file starts with the diagonal/channelwise version because it is:

• the primitive implemented by a compact CUDA selective-scan kernel,
• enough to express the formal affine-scan algebra,
• and a useful educational baseline for examples.

The definitions intentionally reuse DynamicalSystem / DrivenSystem from NN/Spec/Dynamics/System.lean, so existing notions such as iterate, trajectory, IsFixedPoint, and isContractive apply immediately.

structure NN.Spec.Dynamics.DiagonalSSM (α : Type) (dim : ℕ) : Type

A diagonal state-space model with elementwise input, state, and output channels.

This is the smallest useful SSM shape: it captures the recurrent memory path used by Mamba/S4 while remaining easy to verify. Richer models can wrap this core with linear projections and gates.

• A : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)

  Elementwise recurrent multiplier. Stability usually requires |Aᵢ| < 1.

• B : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)

  Elementwise input-to-state gain.

• C : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)

  Elementwise state-to-output gain.

• D : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)

  Elementwise residual / skip gain.

def NN.Spec.Dynamics.DiagonalSSM.step {α : Type} [Add α] [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (h x : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)

One recurrent state update h' = A ⊙ h + B ⊙ x.

def NN.Spec.Dynamics.DiagonalSSM.readout {α : Type} [Add α] [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (h x : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)

Readout y = C ⊙ h + D ⊙ x.

def NN.Spec.Dynamics.DiagonalSSM.transition {α : Type} [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (x : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)) : Spec.DiagonalTransition α dim

Convert a token vector into the corresponding affine transition for the hidden state.

@[simp]

theorem NN.Spec.Dynamics.DiagonalSSM.transition_apply_eq_step {α : Type} [Add α] [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (h x : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)) : (m.transition x).apply h = m.step h x

Applying the affine transition generated by x is exactly the SSM state update.

def NN.Spec.Dynamics.DiagonalSSM.scan {α : Type} [Add α] [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (h0 : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)) (xs : List (Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar))) : List (Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar))

Run a full input sequence and return all hidden states.

@[simp]

theorem NN.Spec.Dynamics.DiagonalSSM.scan_nil {α : Type} [Add α] [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (h0 : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)) : m.scan h0 [] = []

@[simp]

theorem NN.Spec.Dynamics.DiagonalSSM.scan_cons {α : Type} [Add α] [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (h0 x : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)) (xs : List (Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar))) : m.scan h0 (x :: xs) = m.step h0 x :: m.scan (m.step h0 x) xs

def NN.Spec.Dynamics.DiagonalSSM.drivenSystemSpecScalar {dim : ℕ} (m : DiagonalSSM Spec.SpecScalar dim) : DrivenSystem (Spec.Shape.dim dim Spec.Shape.scalar) (Spec.Shape.dim dim Spec.Shape.scalar)

View a scalar-specialized diagonal SSM recurrence as a DrivenSystem.

DrivenSystem uses TorchLean's default SpecScalar, while the reusable SSM record above is polymorphic over scalar backends. This adapter is the bridge to the existing dynamics API.

@[simp]

theorem NN.Spec.Dynamics.DiagonalSSM.drivenSystemSpecScalar_step {dim : ℕ} (m : DiagonalSSM Spec.SpecScalar dim) (h x : Spec.Tensor Spec.SpecScalar (Spec.Shape.dim dim Spec.Shape.scalar)) : m.drivenSystemSpecScalar.step h x = m.step h x