Causality of S4/Mamba-style recurrent blocks

This file proves the sequence-causality property expected of state-space sequence models: appending future tokens cannot change outputs already emitted for a prefix.

We state the theorem at the array-runner level rather than for a particular CUDA kernel. Runtime implementations may use chunked or parallel selective scan, but they must refine these spec runners. Combined with NN.MLTheory.StateSpace.diagonalSelectiveScan_append, this gives the proof layer contract for Mamba/S4-style causal sequence processing.

References:

• Albert Gu and Tri Dao, "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", COLM 2024.
• Tri Dao and Albert Gu, "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality", ICML 2024.

theorem NN.MLTheory.StateSpace.diagonalS4_runArray_append_outputs_prefix {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim : ℕ} (m : Models.DiagonalS4Spec α inputDim stateDim outputDim) (h0 : TorchLean.Tensor α [stateDim]) (xs ys : Array (TorchLean.Tensor α [inputDim])) : (m.runArray h0 (xs ++ ys)).2.take xs.size = (m.runArray h0 xs).2

Diagonal S4 prefix causality.

Appending future tokens ys cannot change the outputs already produced for prefix xs.

theorem NN.MLTheory.StateSpace.compactMamba_runArray_append_outputs_prefix {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim : ℕ} (m : Models.MambaBlockSpec α inputDim stateDim outputDim) (h0 : TorchLean.Tensor α [stateDim]) (xs ys : Array (TorchLean.Tensor α [inputDim])) : (m.runArray h0 (xs ++ ys)).2.take xs.size = (m.runArray h0 xs).2

Compact Mamba prefix causality.

If a sequence xs has already been processed, appending future tokens ys cannot change the outputs for xs. This is the recurrent-model analogue of causal attention non-anticipation.

theorem NN.MLTheory.StateSpace.selectiveMamba_runArrayWithHistory_append_outputs_prefix {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim innerDim convWidth : ℕ} (m : Models.SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h0 : TorchLean.Tensor α [innerDim, stateDim]) (history : Array (TorchLean.Tensor α [innerDim])) (xs ys : Array (TorchLean.Tensor α [inputDim])) : (m.runArrayWithHistory h0 history (xs ++ ys)).2.take xs.size = (m.runArrayWithHistory h0 history xs).2

Full selective Mamba prefix causality for the internal runner.

The internal runner carries a newest-first causal convolution history. Even with that extra state, future input tokens only affect future outputs.

theorem NN.MLTheory.StateSpace.selectiveMamba_runArray_append_outputs_prefix {α : Type} [TorchLean.Storage α] [Context α] {inputDim stateDim outputDim innerDim convWidth : ℕ} (m : Models.SelectiveMambaBlockSpec α inputDim innerDim stateDim outputDim convWidth) (h0 : TorchLean.Tensor α [innerDim, stateDim]) (xs ys : Array (TorchLean.Tensor α [inputDim])) : (m.runArray h0 (xs ++ ys)).2.take xs.size = (m.runArray h0 xs).2

Full selective Mamba prefix causality for the public runner.

This is the user-facing theorem: extending the input stream preserves all previously produced outputs.