Proofs for affine selective scan

The Mamba/S4 scan theorem is an algebra theorem about affine maps. A sequential recurrent update and a parallel prefix scan are equivalent because affine transition composition is associative:

$$ (a_2,b_2)\circ(a_1,b_1)=(a_2a_1,a_2b_1+b_2). $$

The tensor/CUDA implementation is allowed to choose an efficient scan schedule, but the mathematical contract is this file: prefix summaries denote the same state as the left-to-right recurrence.

@[simp]

theorem NN.MLTheory.StateSpace.ScalarAffineTransition.compose_apply {α : Type} [Semiring α] (t₂ t₁ : Spec.ScalarAffineTransition α) (h : α) : (t₂.compose t₁).apply h = t₂.apply (t₁.apply h)

Composing scalar affine transitions agrees with function composition.

@[simp]

theorem NN.MLTheory.StateSpace.ScalarAffineTransition.compose_id_left {α : Type} [Semiring α] (t : Spec.ScalarAffineTransition α) : Spec.ScalarAffineTransition.id.compose t = t

The identity transition is a left identity for composition.

@[simp]

theorem NN.MLTheory.StateSpace.ScalarAffineTransition.compose_id_right {α : Type} [Semiring α] (t : Spec.ScalarAffineTransition α) : t.compose Spec.ScalarAffineTransition.id = t

The identity transition is a right identity for composition.

@[simp]

theorem NN.MLTheory.StateSpace.ScalarAffineTransition.compose_assoc {α : Type} [Semiring α] (t₃ t₂ t₁ : Spec.ScalarAffineTransition α) : (t₃.compose t₂).compose t₁ = t₃.compose (t₂.compose t₁)

Scalar affine transition composition is associative.

@[simp]

theorem NN.MLTheory.StateSpace.ScalarAffineTransition.homogeneous_zero_fixed {α : Type} [Semiring α] (a : α) : { a := a, b := 0 }.apply 0 = 0

Zero is a fixed point of a homogeneous scalar transition.

@[simp]

theorem NN.MLTheory.StateSpace.DiagonalTransition.apply_getScalar {α : Type} [Add α] [Mul α] {stateDim : ℕ} (tr : Spec.DiagonalTransition α stateDim) (h : TorchLean.Tensor α [stateDim]) (i : Fin stateDim) : (tr.apply h).getScalar i = tr.a.getScalar i * h.getScalar i + tr.b.getScalar i

Applying a diagonal transition is exactly the scalar affine update in each channel.

@[simp]

theorem NN.MLTheory.StateSpace.DiagonalTransition.compose_apply_getScalar {α : Type} [Semiring α] {stateDim : ℕ} (t₂ t₁ : Spec.DiagonalTransition α stateDim) (h : TorchLean.Tensor α [stateDim]) (i : Fin stateDim) : ((t₂.compose t₁).apply h).getScalar i = (t₂.apply (t₁.apply h)).getScalar i

Composing diagonal transitions agrees channelwise with composing the corresponding scalar affine maps. This is the exact algebraic invariant used by the variable-coefficient selective-scan kernel: each flattened state lane is an independent affine scan.

theorem NN.MLTheory.StateSpace.scanArrayFrom_eq {State Input Output : Type} (step : State → Input → State × Output) (initial : State) (initialOutputs : Array Output) (xs : Array Input) : Spec.scanArrayFrom step initial initialOutputs xs = have result := Spec.scanArray step initial xs; (result.1, initialOutputs ++ result.2)

Seeding a scan with existing output preserves its state evolution and prepends that output.

theorem NN.MLTheory.StateSpace.scanArrayFrom_state_eq_foldl {State Input Output : Type} (step : State → Input → State × Output) (initial : State) (initialOutputs : Array Output) (xs : Array Input) : (Spec.scanArrayFrom step initial initialOutputs xs).1 = Array.foldl (fun (state : State) (input : Input) => (step state input).1) initial xs

The output buffer carried by a scan does not affect its final state.

theorem NN.MLTheory.StateSpace.scanArray_state_eq_foldl {State Input Output : Type} (step : State → Input → State × Output) (initial : State) (xs : Array Input) : (Spec.scanArray step initial xs).1 = Array.foldl (fun (state : State) (input : Input) => (step state input).1) initial xs

The state component of scanArray is the ordinary state-only left fold.

theorem NN.MLTheory.StateSpace.scanArray_append {State Input Output : Type} (step : State → Input → State × Output) (initial : State) (xs ys : Array Input) : Spec.scanArray step initial (xs ++ ys) = have first := Spec.scanArray step initial xs; have second := Spec.scanArray step first.1 ys; (second.1, first.2 ++ second.2)

A scan over appended inputs is the prefix scan followed by the state-dependent suffix scan.

theorem NN.MLTheory.StateSpace.scanArray_append_outputs_take {State Input Output : Type} (step : State → Input → State × Output) (initial : State) (xs ys : Array Input) : (Spec.scanArray step initial (xs ++ ys)).2.take xs.size = (Spec.scanArray step initial xs).2

Appending future inputs cannot change outputs already emitted by a stateful scan.

@[simp]

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

A stateful scan emits exactly one value for every input.

theorem NN.MLTheory.StateSpace.runScalarAffine_append {α : Type} [Mul α] [Add α] (h0 : α) (xs ys : Array (Spec.ScalarAffineTransition α)) : Spec.runScalarAffine h0 (xs ++ ys) = Spec.runScalarAffine (Spec.runScalarAffine h0 xs) ys

Running appended scalar transitions factors through the state reached after the prefix.

@[simp]

theorem NN.MLTheory.StateSpace.runScalarAffine_singleton {α : Type} [Mul α] [Add α] (h0 : α) (tr : Spec.ScalarAffineTransition α) : Spec.runScalarAffine h0 #[tr] = tr.apply h0

Running one scalar transition is the same as applying it.

theorem NN.MLTheory.StateSpace.summarizeScalarAffine_apply_eq_run {α : Type} [Semiring α] (h0 : α) (transitions : Array (Spec.ScalarAffineTransition α)) : (Spec.summarizeScalarAffine transitions).apply h0 = Spec.runScalarAffine h0 transitions

The affine summary denotes the same state as the sequential recurrence.

theorem NN.MLTheory.StateSpace.summarizeScalarAffine_append {α : Type} [Semiring α] (xs ys : Array (Spec.ScalarAffineTransition α)) : Spec.summarizeScalarAffine (xs ++ ys) = (Spec.summarizeScalarAffine ys).compose (Spec.summarizeScalarAffine xs)

Prefix summaries compose across array append in execution order.

theorem NN.MLTheory.StateSpace.summarizeScalarAffine_append_apply {α : Type} [Semiring α] (h0 : α) (xs ys : Array (Spec.ScalarAffineTransition α)) : (Spec.summarizeScalarAffine (xs ++ ys)).apply h0 = ((Spec.summarizeScalarAffine ys).compose (Spec.summarizeScalarAffine xs)).apply h0

Prefix summaries composed across append have the expected denotation.

@[simp]

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

The scalar affine scan has one state per transition.

theorem NN.MLTheory.StateSpace.scalarAffineScan_append {α : Type} [Mul α] [Add α] (h0 : α) (xs ys : Array (Spec.ScalarAffineTransition α)) : Spec.scalarAffineScan h0 (xs ++ ys) = Spec.scalarAffineScan h0 xs ++ Spec.scalarAffineScan (Spec.runScalarAffine h0 xs) ys

Scanning appended scalar transitions is the prefix scan followed by the suffix scan.

@[simp]

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

The diagonal tensor scan has one state per transition.

theorem NN.MLTheory.StateSpace.runDiagonalTransitions_append {α : Type} [Add α] [Mul α] {stateDim : ℕ} (h0 : TorchLean.Tensor α [stateDim]) (xs ys : Array (Spec.DiagonalTransition α stateDim)) : Spec.runDiagonalTransitions h0 (xs ++ ys) = Spec.runDiagonalTransitions (Spec.runDiagonalTransitions h0 xs) ys

Running appended diagonal transitions factors through the state after the prefix.

theorem NN.MLTheory.StateSpace.diagonalSelectiveScan_append {α : Type} [Add α] [Mul α] {stateDim : ℕ} (h0 : TorchLean.Tensor α [stateDim]) (xs ys : Array (Spec.DiagonalTransition α stateDim)) : Spec.diagonalSelectiveScan h0 (xs ++ ys) = Spec.diagonalSelectiveScan h0 xs ++ Spec.diagonalSelectiveScan (Spec.runDiagonalTransitions h0 xs) ys

The diagonal scan of an append is the prefix scan followed by the state-dependent suffix scan.

theorem NN.MLTheory.StateSpace.abs_homogeneous_apply_le (a ρ h : ℝ) (ha : |a| ≤ ρ) : |{ a := a, b := 0 }.apply h| ≤ ρ * |h|

A homogeneous affine transition over $\mathbb{R}$ is Lipschitz with factor $\rho$ whenever $|a|\leq\rho$.

This is the one-channel stability lemma used to lift diagonal SSMs into contraction proofs.

theorem NN.MLTheory.StateSpace.abs_homogeneous_apply_le_self (a h : ℝ) (ha : |a| ≤ 1) : |{ a := a, b := 0 }.apply h| ≤ |h|

A homogeneous scalar transition with $|a|\leq 1$ is non-expansive.