S4 / HiPPO: Continuous to Discrete

The continuous SSM

The starting point is the state-space representation — a standard way to describe dynamical systems in control theory and signal processing. Any system of linear ordinary differential equations (ODEs) can be written in this form: The SSM deep learning literature adopted this formulation from the HiPPO paper Gu et al. (2020) .

where $\statevec(t) \in \R^\statedim$ is the hidden state (a vector that summarizes the system’s memory — not to be confused with “state” in the ML training sense), $u(t)$ is the scalar input, $y(t)$ is the scalar output, and the system is defined by four matrices: $\statemat \in \R^{\statedim \times \statedim}$ (dynamics — how the state evolves), $\inputmat \in \R^{\statedim \times 1}$ (how input enters the state), $\outputmat \in \R^{1 \times \statedim}$ (how state maps to output), and $\feedmat \in \R$ (direct input-to-output skip connection). These matrices are time-invariant: the same values apply at every time step. This is the Linear Time-Invariant (LTI) assumption that makes the convolutional view possible (later section). Mamba (Week 7) breaks LTI by making $\inputmat$, $\outputmat$, $\stepsize$ input-dependent, gaining content-based selection but losing the convolutional view.

The connection to Week 1 is direct: the damped oscillator $\ddot{q} + \gamma \dot{q} + \omega^2 q = 0$ is exactly the continuous SSM with $\statemat = \begin{pmatrix} 0 & 1 \\ -\omega^2 & -\gamma \end{pmatrix}$ and $\inputmat = 0$. Now $\statemat$ will be learned, not hand-crafted — but the stability requirements are identical: eigenvalues in the left half-plane.

HiPPO: principled initialization

The central question is: how do we initialize $\statemat$? Random initialization leads to vanishing or exploding dynamics — the same pathology that makes recurrent neural networks difficult to train.

HiPPO (High-order Polynomial Projection Operators) Gu et al. (2020) provides a principled answer: initialize $\statemat$ so that the state $\statevec(t)$ encodes the best $\statedim$-term Legendre polynomial approximation of the input history $u(\tau)$ for $\tau \in [0, t]$. Legendre polynomials are an orthogonal basis for representing functions on an interval — the optimal choice for approximating smooth signals with a finite number of coefficients. The state vector’s $k$-th entry holds the $k$-th Legendre coefficient of the input history, giving the SSM a mathematically principled “memory” of the past.

Discretization: continuous to discrete

Discretization here means converting a continuous-time ODE ($\dot{\statevec} = \statemat\statevec + \inputmat u$, defined for all $t \in \R$) into a discrete-time recurrence ($\statevec_{t+1} = \discA\statevec_t + \discB u_t$, defined at integer time steps). This is the same process as converting a differential equation into a difference equation for numerical simulation — but here the step size $\stepsize$ is a learnable parameter, not a fixed choice. Note: “discretization” in NLP often means tokenization; here it strictly means ODE-to-recurrence conversion.

Given $\stepsize > 0$, we seek discrete matrices $\discA, \discB$ such that:

Zero-Order Hold (ZOH)

Zero-Order Hold (ZOH) assumes the input $u(t)$ is piecewise constant over each discretization interval $[t, t{+}\stepsize)$ — that is, the input doesn’t change between time steps. This is a reasonable assumption when the step size $\stepsize$ is small relative to the input’s rate of change, which is the typical regime for SSMs. ZOH is the default discretization across S4, S4D, S5, and Mamba reference implementations. [Convergence] The resulting formulas are:

We avoid computing $\statemat^{-1}$ by using the stacked matrix exponential: build the augmented matrix $\begin{pmatrix} \statemat\stepsize & \inputmat\stepsize \\ 0 & 0 \end{pmatrix}$, compute its expm, and extract $\discA$, $\discB$ from the result. Implemented in s4_hippo.py:discretize_zoh .

Bilinear (Tustin) transform

The bilinear transform maps the imaginary axis to the unit circle:

This is second-order accurate, vs. first-order for forward Euler.

Three views of SSM computation

An LTI SSM can be computed in two equivalent ways, enabling the train-fast / infer-cheap trick that makes S4 practical.

Recurrent view

The discrete state update is a direct recurrence computed via lax.scan with $\bigO{\seqlen}$ sequential steps and $\bigO{\statedim}$ memory per token. This is the same scan primitive from Week 1, with $\text{carry} = \statevec_t$ and $\text{input} = u_t$.

Convolutional view

The term “kernel” here means a convolution filter — a fixed sequence of coefficients that defines a weighted combination of nearby inputs. This is the signal processing sense, unrelated to kernel methods in classical ML (support vector machines, Gaussian processes) or operating system kernels.

Expanding the recurrence with $\statevec_0 = 0$ gives:

where $K = (\outputmat\discB, \outputmat\discA\discB, \outputmat\discA^2\discB, \ldots)$ is the SSM kernel. This is a causal convolution (output at time $t$ depends only on inputs at times $\leq t$), computed in $\bigO{\seqlen \log \seqlen}$ via FFT.

Our implementation computes $K$ via iterative matrix powers in a lax.scan — $\bigO{\statedim^2\seqlen}$, deliberately naive. Production S4 exploits DPLR (Diagonal Plus Low-Rank) structure to evaluate $K$ via Cauchy matrices (matrices with entries $1/(x_i - y_j)$, enabling fast structured products) in $\bigO{\statedim \log^2 \statedim}$.

Mode switching

S4’s practical trick: convolutional mode for training ($\bigO{\seqlen \log \seqlen}$ parallel) and recurrent mode for inference ($\bigO{1}$ per token). Both modes share the same parameters; only the computation changes.

The S4 layer in Flax NNX

The S4Layer module wraps the SSM as a Flax NNX module, with parameters split by type for dual-optimizer treatment:

• SSMParam: $\statemat$, $\inputmat$, $\log\stepsize$ — SSM optimizer (Adam, LR $\leq 10^{-3}$, no weight decay)
• nnx.Param: $\outputmat$, $\feedmat$, encoder, decoder — backbone optimizer (AdamW, LR $= 8 \times 10^{-4}$, WD $0.01$)

This is the same SSMParam split from Week 2. The week 2 harness works unchanged: the dual optimizer filters SSMParam variables with wrt=SSMParam.

class S4Layer(nnx.Module):
    def __init__(self, d_state=32, *, rngs):
        A, B = make_hippo_legs(d_state)
        self.A = SSMParam(A)
        self.B = SSMParam(B.squeeze())
        self.C = nnx.Param(random.normal(rngs.params(), (d_state,)))
        self.D = nnx.Param(jnp.ones(()))
        self.log_dt = SSMParam(...)  # log-uniform in [log(0.001), log(0.1)]

The S4SeqModel is a drop-in for SimpleSeqModel (Week 2), preserving the $(B, L, 1) \to (B, L, 1)$ contract. With $H = 32$ features and $\statedim = 32$, the S4SeqModel has ~35K parameters vs. ~2K for SimpleSeqModel.

Post-training eigenvalue analysis

After training, we extract the learned $\statemat$ matrices and verify stability:

• Continuous: $\operatorname{Re}(\lambda_k) < 0$ (left half-plane)
• Discrete: $\abs{\bar{\lambda}_k} < 1$ (inside unit disk)

The low learning rate on SSMParam keeps eigenvalue drift small. [Official] The two-LR recipe (Adam $\leq 10^{-3}$ on SSM params, AdamW on backbone) originates in the S4 paper Gu et al. (2020) and reproduces across S4D, S5, and Mamba reference code.

Week 18 performs full Lyapunov stability analysis of trained SSMs, using the QR-based infrastructure from Week 1.

Reflections

What transfers

• Discretization awareness: The choice of integrator (Euler, ZOH, bilinear) determines training stability — exactly as in numerical PDE solvers.
• Eigenvalue monitoring: Checking $\operatorname{Re}(\lambda) < 0$ after training is the SSM equivalent of checking CFL conditions.
• Dual computation modes: The recurrent/convolutional duality is unique to LTI systems and is lost in selective models (Week 7+).

What requires adaptation

1. $\statemat$ is learned, not given. In your PhD, the state matrix comes from physics. In S4, HiPPO provides a principled initialization, but training moves $\statemat$ — possibly off the stability manifold. This is why the ZOH stability proposition only bounds initialization; post-training stability is a Week 18 question.
2. Diagonalization is not free. Week 5’s S4D simplification abandons the HiPPO theoretical interpretation. The diagonal form reproduces S4’s empirical performance but loses the Legendre-polynomial memory interpretation — a theory-for-speed trade-off with no clean analogue in numerical PDE solvers.

Forward map

• Week 5: S4D makes $\statemat$ diagonal, collapsing the $\bigO{\statedim^2\seqlen}$ kernel to $\bigO{\statedim\seqlen}$.
• Week 6: S5 replaces lax.scan with associative scan for $\bigO{\log \seqlen}$ parallel depth.
• Week 7: Mamba makes $\inputmat, \outputmat, \stepsize$ input-dependent, breaking the convolutional view.
• Week 18: Lyapunov analysis of trained S4 layers using Week 1 QR infrastructure.
• Week 9 (Julia): ZOH, bilinear, and exponential-trapezoidal compared on a forced oscillator, with empirical order-of-accuracy verification in experiments/julia/week09/discretization.jl.
• Week 20: Stability regions plotted as contour heatmaps in the complex $z$ plane, experiments/julia/week20/stability_regions.jl.