State-space dynamical systems

State-space sequence models are recurrent dynamical systems:

$$ h_{t+1}=A_t h_t+B_t x_t,\qquad 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 reuse DynamicalSystem / DrivenSystem from NN/Spec/Dynamics/System.lean, so existing notions such as iterate, trajectory, IsFixedPoint, and isContractive apply immediately.

structure Spec.Dynamics.DiagonalSSM (α : Type) [TorchLean.Storage α] (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 : TorchLean.Tensor α [dim]

  Elementwise recurrent multiplier. Stability usually requires $\lvert A_i\rvert<1$.

• B : TorchLean.Tensor α [dim]

  Elementwise input-to-state gain.

• C : TorchLean.Tensor α [dim]

  Elementwise state-to-output gain.

• D : TorchLean.Tensor α [dim]

  Elementwise residual / skip gain.

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

One recurrent state update $h'=A\odot h+B\odot x$.

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

Readout $y=C\odot h+D\odot x$.

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

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

@[simp]

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

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

def Spec.Dynamics.DiagonalSSM.scan {α : Type} [TorchLean.Storage α] [Add α] [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (h0 : TorchLean.Tensor α [dim]) (xs : Array (TorchLean.Tensor α [dim])) : Array (TorchLean.Tensor α [dim])

Run a full input sequence and return all hidden states.

@[simp]

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

Scanning no input produces no hidden states.

@[simp]

theorem Spec.Dynamics.DiagonalSSM.scan_size {α : Type} [TorchLean.Storage α] [Add α] [Mul α] {dim : ℕ} (m : DiagonalSSM α dim) (h0 : TorchLean.Tensor α [dim]) (xs : Array (TorchLean.Tensor α [dim])) : (m.scan h0 xs).size = xs.size

A diagonal SSM scan returns one hidden state for each input token.

def Spec.Dynamics.DiagonalSSM.drivenSystemSpecScalar {dim : ℕ} (m : DiagonalSSM SpecScalar dim) : DrivenSystem (Shape.dim dim Shape.scalar) (Shape.dim dim 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 Spec.Dynamics.DiagonalSSM.drivenSystemSpecScalar_step {dim : ℕ} (m : DiagonalSSM SpecScalar dim) (h x : TorchLean.Tensor SpecScalar [dim]) : m.drivenSystemSpecScalar.step h x = m.step h x

Viewing a diagonal SSM as a driven dynamical system does not change its step.

The wrapper exists so the dynamics theorems apply to the layer; this is the rfl that says the wrapper carries no extra content, so a bound proved about the system holds of the layer.