S4: Structured State Spaces for Sequences.

Implements the S4 architecture from "Efficiently Modeling Long Sequences with Structured State Spaces" (Gu et al., ICLR 2022). S4 introduced the key idea of using HiPPO-initialized diagonal state matrices for stable long-range sequence modeling.

Key Innovation: HiPPO Initialization

The state matrix A is initialized using the HiPPO framework, which produces matrices that optimally compress continuous signals into finite-dimensional state. The diagonal parameterization enables efficient parallel computation:

Continuous SSM:
  x'(t) = A x(t) + B u(t)
  y(t)  = C x(t) + D u(t)

Discretization (ZOH):
  A_bar = exp(dt * A)
  B_bar = (A_bar - I) * A^{-1} * B   (simplified to dt * B for diagonal A)
  h[t]  = A_bar * h[t-1] + B_bar * u[t]
  y[t]  = C * h[t]

Architecture

Input [batch, seq_len, embed_dim]
      |
      v
+-----------------------+
| Input Projection      |
+-----------------------+
      |
      v
+-----------------------+
| S4 Block x N          |
|  LayerNorm            |
|  SSM (HiPPO A)        |
|  Dropout + Residual   |
|  FFN Block            |
+-----------------------+
      |
      v
+-----------------------+
| Final LayerNorm       |
+-----------------------+
      |
      v
[batch, hidden_size]    (last timestep)

Usage

model = S4.build(
  embed_dim: 287,
  hidden_size: 256,
  state_size: 64,
  num_layers: 4
)

Reference

• Paper: "Efficiently Modeling Long Sequences with Structured State Spaces"
• arXiv: https://arxiv.org/abs/2111.00396