Linear Recurrence and State Space Models

Linear Recurrence and State Space Models

Gated RNNs fixed the memory problem and kept a compute problem: \mathbf{H}_t needs \mathbf{H}_{t-1}, so training is sequential in t while GPUs sit idle. CNNs (ch. 6) compute all positions at once.

Goal: parallel training, O(1)-state inference. Two passes at it:

  • Strip the GRU until its state update is linear → evaluate the whole sequence with a parallel scan (minGRU).
  • Rebuild properly from continuous time → state space models: principled gates (\Delta), stability by construction, recurrence \equiv convolution, HiPPO memory, S4D.

Why recurrence resists parallelism

\mathbf{H}_t = \phi(\mathbf{X}_t \mathbf{W}_{xh} + \mathbf{H}_{t-1} \mathbf{W}_{hh} + \mathbf{b})

Substitute it into itself: \mathbf{H}_{t-2} is buried under two nested \phi’s. Nonlinear compositions don’t simplify: the nonlinearity on the state path is the exact culprit.

Everything else is innocent: gates from inputs, nonlinear readouts, depth between layers. So: remove \mathbf{H}_{t-1} from inside nonlinearities, keep the rest.

minGRU: two cuts

From the GRU: gates and candidate from the input only (the reset gate then has nothing to reset, so it is deleted):

\mathbf{Z}_t = \sigma(\mathbf{X}_t \mathbf{W}_{xz} + \mathbf{b}_z),\qquad \tilde{\mathbf{H}}_t = \tanh(\mathbf{X}_t \mathbf{W}_{xh} + \mathbf{b}_h),

\mathbf{H}_t = (1 - \mathbf{Z}_t) \odot \mathbf{H}_{t-1} + \mathbf{Z}_t \odot \tilde{\mathbf{H}}_t.

Now affine in the state: \mathbf{H}_t = \mathbf{a}_t \odot \mathbf{H}_{t-1} + \mathbf{b}_t with \mathbf{a}_t = 1 - \mathbf{Z}_t, \mathbf{b}_t = \mathbf{Z}_t \odot \tilde{\mathbf{H}}_t with all coefficients computable for all t at once.

  • Memory is readable: input from k steps back is scaled by \prod_j \mathbf{a}_j \in (0,1): eigenvalues in plain sight (8.6).

The parallel scan

Represent step t by the affine map (\mathbf{a}_t, \mathbf{b}_t). Composition is again affine, and associative:

(\mathbf{a}_2, \mathbf{b}_2) \circ (\mathbf{a}_1, \mathbf{b}_1) = (\mathbf{a}_1 \odot \mathbf{a}_2,\ \mathbf{a}_2 \odot \mathbf{b}_1 + \mathbf{b}_2).

Any associative “running total” parallelizes (Blelloch, 1990):

Doubling strides 1, 2, 4: after \log_2 T rounds every position holds its prefix.

Implementation

One round per stride; each round is a batched multiply-add:

def scan_combine(prev, cur):
    a_prev, b_prev = prev
    a_cur, b_cur = cur
    return a_prev * a_cur, a_cur * b_prev + b_cur

def associative_scan(a, b):
    """Parallel prefix scan for h_t = a_t * h_{t-1} + b_t with h_0 = 0."""
    return jax.lax.associative_scan(scan_combine, (a, b))[1]
def sequential_scan(a, b):
    step = lambda h, ab: (ab[0] * h + ab[1],) * 2
    return jax.lax.scan(step, jnp.zeros_like(b[0]), (a, b))[1]

key1, key2 = jax.random.split(d2l.get_key())
a = jax.random.uniform(key1, (100, 4, 8))
b = jax.random.normal(key2, (100, 4, 8))
err = jnp.abs(associative_scan(a, b) - sequential_scan(a, b)).max()
print(f'maximum deviation: {float(err):.2e}')
maximum deviation: 4.77e-07

Trains like a CNN

Scan vs. sequential loop, batch 32, 128 units:

parallel_fn = jax.jit(associative_scan)
sequential_fn = jax.jit(sequential_scan)

def wall_clock(f, *args, reps=3):
    f(*args).block_until_ready()  # Warm up (and compile)
    start = time.time()
    for _ in range(reps):
        f(*args).block_until_ready()
    return (time.time() - start) / reps

lengths, times = [256, 1024, 4096, 16384], [[], []]
for T in lengths:
    key1, key2 = jax.random.split(jax.random.key(T))
    a = jax.random.uniform(key1, (T, 32, 128))
    b = jax.random.normal(key2, (T, 32, 128))
    times[0].append(wall_clock(parallel_fn, a, b) * 1e3)
    times[1].append(wall_clock(sequential_fn, a, b) * 1e3)
d2l.plot(lengths, times, 'sequence length', 'time per call (ms)',
         legend=['parallel scan', 'sequential loop'],
         xscale='log', yscale='log', figsize=(5, 3))

The loop grows linearly in T; the scan’s few rounds saturate the GPU. Inference is unchanged: still one token, one state update at a time.

A minGRU cell

Drop-in replacement for the gated cells of the previous section, with no time loop; a carried state folds into \mathbf{b}_1:

class MinGRU(nnx.Module):
    """The minimal GRU: input-only gates and a linear state path."""
    def __init__(self, num_inputs, num_hiddens, rngs=None):
        rngs = nnx.Rngs(0) if rngs is None else rngs
        self.num_inputs, self.num_hiddens = num_inputs, num_hiddens
        self.W_xz = nnx.Linear(num_inputs, num_hiddens, rngs=rngs)
        self.W_xh = nnx.Linear(num_inputs, num_hiddens, rngs=rngs)

    def __call__(self, inputs, H=None):
        Z = jax.nn.sigmoid(self.W_xz(inputs))    # (num_steps, batch, hiddens)
        H_tilde = jnp.tanh(self.W_xh(inputs))
        a, b = 1 - Z, Z * H_tilde
        if H is not None:  # Fold the carried-in state into the first step
            b = jnp.concatenate([b[:1] + a[:1] * H, b[1:]])
        outputs = associative_scan(a, b)
        return outputs, outputs[-1]

Language modeling, same recipe as the LSTM’s

Time Machine, 1,024-token BPE, 50k windows of 32, emb 64, hidden 128, 10 epochs, clip 1:

data = d2l.TimeMachine(batch_size=1024, num_steps=32,
                       num_train=50000, num_val=5000)
mingru = MinGRU(num_inputs=64, num_hiddens=128)
model = d2l.RNNLM(mingru, vocab_size=len(data.vocab), lr=4)
trainer = d2l.Trainer(max_epochs=10, gradient_clip_val=1, num_gpus=1)
model.board.yscale = 'log'
start = time.time()
trainer.fit(model, data)
fit_time = time.time() - start

minGRU result

def val_ppl(model):
    total_loss = num_tokens = 0
    for X_val, y_val in data.val_dataloader():
        losses = model.loss(model(X_val), y_val, averaged=False)
        total_loss += float(losses.sum())
        num_tokens += losses.size
    return math.exp(total_loss / num_tokens)

pred = model.predict('the time traveller', 30, data.tokenizer)
print(f'validation perplexity {val_ppl(model):.1f}, '
      f'{fit_time:.0f}s wall clock')
print(pred)
validation perplexity 85.3, 16s wall clock
the time traveller.

'I had
dism.

'I had
dism.

'I had
dism
  • A few points behind the GRU (low-to-mid eighties vs. high seventies to low eighties) at ~1/4 of its recurrent parameters: 2h(d{+}1) vs. 3h(d{+}h{+}1).
  • Modest claim: nothing essential lost, training became a scan.
  • Price paid: the decay \mathbf{a}_t is decided by the input alone (remember that).

State space models

The principled route to the same recurrence: continuous-time linear dynamics:

\dot{\mathbf{x}}(t) = \mathbf{A}\mathbf{x}(t) + \mathbf{B}u(t), \qquad y(t) = \mathbf{C}\mathbf{x}(t) + Du(t).

Why continuous? It parameterizes a law of motion, not a table of step weights: one extra scalar \Delta (step size) converts (\mathbf{A}, \mathbf{B}) into a recurrence at any sampling rate.

ZOH discretization: the step size is a gate

Hold u constant over \Delta, integrate exactly:

\bar{\mathbf{A}} = \exp(\Delta \mathbf{A}), \qquad \bar{\mathbf{B}} = \left(\textstyle\int_0^{\Delta} e^{\tau \mathbf{A}} d\tau\right) \mathbf{B} = \mathbf{A}^{-1}(\exp(\Delta \mathbf{A}) - \mathbf{I})\, \mathbf{B} \ \textrm{(invertible } \mathbf{A}\textrm{)}.

  • \Delta \to 0: \bar{\mathbf{A}} \to \mathbf{I}, \bar{\mathbf{B}} \to 0. Copy the state, ignore the input.
  • \Delta large: \bar{\mathbf{A}} \to 0. Overwrite from the input.
  • A learned \Delta is an update gate, not bolted on: it fell out of the calculus. (Next section: \Delta becomes input-dependent.)
  • Identity in the worked scalar mode a=-1: backward Euler gives x_t = (1-z)x_{t-1} + z u_t with z = \sigma(\log \Delta): the GRU gate is a discretization rule there (Gu, 2023).
  • Menu of rules: ZOH / bilinear (S4’s choice) / backward Euler all keep stability at any \Delta; forward Euler does not. In code: expm1.

Stability by construction: store a_n = -e^{\theta_n}, then |\bar{a}_n| = e^{\Delta\,\mathrm{Re}(a_n)} < 1 for every parameter value; diagonal (normal) \mathbf{A} ⇒ no transients either. No clipping, no knife-edge (8.6).

Recurrence is convolution

LTI: the same \bar{\mathbf{A}}, \bar{\mathbf{B}}, \mathbf{C} at every step. Unroll:

y_t = \sum_k \mathbf{C}\bar{\mathbf{A}}^k \bar{\mathbf{B}}\, u_{t-k}, \qquad \bar{\mathbf{K}} = (\mathbf{C}\bar{\mathbf{B}}, \mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}, \mathbf{C}\bar{\mathbf{A}}^2\bar{\mathbf{B}}, \ldots).

One model, three views: ODE, recurrence, convolution.

Three views, one answer

num_states, num_steps, delta = 8, 64, 0.1
a = -jnp.arange(1., num_states + 1)                 # Re(a) < 0
b, c = jnp.ones(num_states), jax.random.normal(d2l.get_key(), (num_states,))
a_bar = jnp.exp(delta * a)                          # Zero-order hold
b_bar = jnp.expm1(delta * a) / a * b                # expm1: accurate e^x - 1
u = jax.random.normal(d2l.get_key(), (num_steps,))

def step(x, u_t):                                   # (i) recurrence
    x = a_bar * x + b_bar * u_t
    return x, jnp.dot(c, x)
_, y_loop = jax.lax.scan(step, jnp.zeros(num_states), u)

xs = associative_scan(jnp.tile(a_bar, (num_steps, 1)),
                      b_bar * u[:, None])           # (ii) parallel scan
y_scan = xs @ c

k = jnp.arange(num_steps)[:, None]                  # (iii) convolution
kernel = (c * b_bar * a_bar ** k).sum(-1)
y_conv = jnp.array(np.convolve(u, kernel)[:num_steps])

print(f'scan vs loop: {float(jnp.abs(y_scan - y_loop).max()):.2e}, '
      f'conv vs loop: {float(jnp.abs(y_conv - y_loop).max()):.2e}')
scan vs loop: 8.94e-08, conv vs loop: 1.19e-07

S4 trains through the conv view with FFTs (O(T \log T)); we keep the scan, because it survives the input-dependent dynamics of the next section; the kernel view will not.

What should A be? HiPPO

The state is N numbers summarizing all history. Which N lose least? HiPPO: make \mathbf{x}(t) the coefficients of the optimal polynomial approximation of the past; maintaining it online is a linear, time-varying ODE with a fixed matrix pattern: \dot{\mathbf{x}} = (\mathbf{A}\mathbf{x} + \mathbf{B}u)/t (A_{nn} = -(n{+}1), B_n = \sqrt{2n+1}) — the 1/t tracks the uniform measure over a growing past.

Reconstructing 1,000 steps of history from the final state alone.

Three objects: online LegS projection (owns the theorem) → S4 = LTI model initialized at the HiPPO matrix (DPLR kernels, solved Long Range Arena) → S4D = diagonal approximation (matches S4 empirically; init a_n = -(n{+}1) is multiscale, not optimal). Plus S5: the scan. We implement S4D’s parameterization with S5’s scan.

The S4D layer

Per channel: N-state diagonal SSM; learnable \log \Delta per channel (log-uniform init: gates of many speeds); \mathbf{B} = \mathbf{1}; three lines of ZOH + the scan:

class S4D(nnx.Module):
    """A diagonal state space layer: one SSM per feature channel."""
    def __init__(self, num_hiddens, num_states=4, dt_min=0.001, dt_max=0.1,
                 rngs=None):
        rngs = nnx.Rngs(0) if rngs is None else rngs
        H, N = num_hiddens, num_states
        self.log_a = nnx.Param(jnp.tile(jnp.log(jnp.arange(1., N + 1)),
                                        (H, 1)))
        self.log_dt = nnx.Param(
            rngs.params.uniform((H, 1)) * math.log(dt_max / dt_min)
            + math.log(dt_min))
        self.C = nnx.Param(rngs.params.normal((H, N)) / math.sqrt(N))
        self.D = nnx.Param(jnp.ones(H))

    def __call__(self, u):                   # (num_steps, batch, num_hiddens)
        a = -jnp.exp(self.log_a[...])                 # (H, N), Re(a) < 0
        da = jnp.exp(self.log_dt[...]) * a
        a_bar = jnp.exp(da)
        b_bar = jnp.expm1(da) / a                     # ZOH with B = 1
        a_elems = jnp.broadcast_to(                   # Same at every step
            a_bar[None, None], (u.shape[0], 1, *a_bar.shape))
        b_elems = b_bar * u[..., None]                # (T, batch, H, N)
        x = associative_scan(a_elems, b_elems)
        return (x * self.C).sum(-1) + self.D * u

The S4D block

SSM mixes time, gated MLP mixes channels and restores nonlinearity:

class S4DBlock(nnx.Module):
    def __init__(self, num_hiddens, num_states, rngs=None):
        rngs = nnx.Rngs(0) if rngs is None else rngs
        self.ln1 = nnx.LayerNorm(num_hiddens, rngs=rngs)
        self.ssm = S4D(num_hiddens, num_states, rngs=rngs)
        self.ln2 = nnx.LayerNorm(num_hiddens, rngs=rngs)
        self.W_v = nnx.Linear(num_hiddens, 2 * num_hiddens, rngs=rngs)
        self.W_g = nnx.Linear(num_hiddens, 2 * num_hiddens, rngs=rngs)
        self.W_o = nnx.Linear(2 * num_hiddens, num_hiddens, rngs=rngs)

    def __call__(self, X):
        X = X + self.ssm(self.ln1(X))
        Y = self.ln2(X)
        return X + self.W_o(self.W_v(Y) * jax.nn.sigmoid(self.W_g(Y)))

Sequential Fashion-MNIST

784 pixels, one at a time. The readout decides the question: mean-pool all outputs = mixing benchmark; read the final step only = retention diagnostic. Same head, data, epochs, clip 1 (the LSTM destabilizes without it; the S4D can’t explode); S4D stack vs. parameter-matched LSTM (~30k params each):

data = d2l.FashionMNIST(batch_size=128)
results = {}

def train_and_report(name, model, epochs=10):
    trainer = d2l.Trainer(max_epochs=epochs, gradient_clip_val=1, num_gpus=1)
    start = time.time()
    trainer.fit(model, data)
    params = sum(p.size for p in
                 jax.tree.leaves(nnx.state(model, nnx.Param)))
    results[name] = (params, float(model.board.data['val_acc'][-1].y),
                     (time.time() - start) / epochs)

s4d = SeqClassifier(nnx.Sequential(*[S4DBlock(48, 4) for _ in range(2)]),
                    num_hiddens=48)
train_and_report('S4D', s4d)

s4d_final = SeqClassifier(
    nnx.Sequential(*[S4DBlock(48, 4) for _ in range(2)]),
    num_hiddens=48, pool='last')
train_and_report('S4D (final)', s4d_final)

lstm = SeqClassifier(d2l.LSTM(num_inputs=48, num_hiddens=64),
                     num_hiddens=48, num_features=64)
train_and_report('LSTM', lstm)

Results

print(f'{"model":>12} {"params":>8} {"val acc":>8} {"s/epoch":>8}')
for name, (params, acc, secs) in results.items():
    print(f'{name:>12} {params:>8,} {acc:>8.3f} {secs:>8.1f}')
       model   params  val acc  s/epoch
         S4D   30,058    0.827      6.1
 S4D (final)   30,058    0.823      5.8
        LSTM   29,674    0.805     48.6
  • S4D (mean-pool): low-to-mid 80s in both frameworks, every rerun. S4D (final state): within run-to-run noise of it — the memory really is in the state after 784 steps, not rescued by early outputs.
  • LSTM: 60-84% across reruns, tracking init defaults: plain uniform init leaves memory to chance, orthogonal weights improve the odds, and parity takes the full folklore (+ forget bias 1).
  • The \Delta-init hands the S4D multi-scale memory by design; the LSTM gets it only from initialization folklore.
  • LRA (16k steps) is where this family first beat everything.

Inference, one token at a time

Training used the scan; the bargain’s other half is recurrent inference. Nothing to derive: the discretized update is the algorithm. State: one (H, N) block per layer; the gated MLP is position-wise and carries none.

@d2l.add_to_class(S4D)
def step(self, u, x=None):
    """Advance one token: u is (batch, H); x is the (batch, H, N) state."""
    a = -jnp.exp(self.log_a[...])
    da = jnp.exp(self.log_dt[...]) * a
    a_bar = jnp.exp(da)                               # Same ZOH as forward
    b_bar = jnp.expm1(da) / a
    if x is None:
        x = jnp.zeros((*u.shape, a_bar.shape[-1]))
    x = a_bar * x + b_bar * u[..., None]              # One recurrence step
    return (x * self.C).sum(-1) + self.D * u, x

@d2l.add_to_class(S4DBlock)
def step(self, X, x=None):
    """Advance the block one token; only the SSM carries state."""
    y, x = self.ssm.step(self.ln1(X), x)
    X = X + y
    Y = self.ln2(X)
    return X + self.W_o(self.W_v(Y) * jax.nn.sigmoid(self.W_g(Y))), x

Stepped == scanned, on the trained model

Replay validation images through the trained encoder both ways: all 784 pixels by scan, then one at a time by step:

X_val = jnp.asarray(next(iter(data.val_dataloader()))[0][:64])
seq = s4d.emb(X_val.reshape(X_val.shape[0], -1, 1).transpose(1, 0, 2))
# GPU matmuls default to TF32; run both schedules in full fp32 so the
# comparison isolates the reassociation error
with jax.default_matmul_precision('float32'):
    y_scan = s4d.encoder(seq)                    # All 784 steps at once
    states, ys = [None] * len(s4d.encoder.layers), []
    for t in range(seq.shape[0]):                # One pixel at a time
        X, new_states = seq[t], []
        for blk, x in zip(s4d.encoder.layers, states):
            X, x = blk.step(X, x)
            new_states.append(x)
        states, ys = new_states, ys + [X]
err = float(jnp.abs(jnp.stack(ys) - y_scan).max())
scale = float(jnp.abs(y_scan).max())
print(f'stepped vs scanned: deviation {err:.2e} '
      f'on activations of scale {scale:.0f}: relative {err / scale:.2e}')
assert err < 1e-3 * scale
stepped vs scanned: deviation 1.67e-03 on activations of scale 79: relative 2.12e-05

Same arithmetic, different association order: float-level agreement. The assert doubles as a regression guard.

The stopwatch and the memory bill

Cost of absorbing the next token after a prefix of P, batch of 32 streams:

def one_step(model, X, states):
    for blk, x in zip(model.layers, states):
        X, x = blk.step(X, x)
    return X

step_fn = nnx.jit(one_step)
rerun_fn = nnx.jit(lambda model, seq: model(seq))

lengths, times = [256, 1024, 4096, 16384], [[], []]
states = [jnp.zeros((32, 48, 4)) for _ in s4d.encoder.layers]
X1 = jnp.zeros((32, 48))
for P in lengths:
    seq = jax.random.normal(jax.random.key(P), (P, 32, 48))
    times[0].append(wall_clock(step_fn, s4d.encoder, X1, states) * 1e3)
    times[1].append(wall_clock(rerun_fn, s4d.encoder, seq) * 1e3)
d2l.plot(lengths, times, 'prefix length', 'ms per new token',
         legend=['recurrent step', 're-run the prefix'],
         xscale='log', yscale='log', figsize=(5, 3))
print(f'carried state: {sum(4 * x.size for x in states) // 32} '
      f'bytes per sequence, at every prefix length')

carried state: 1536 bytes per sequence, at every prefix length
  • Re-running the prefix grows with P; the step is flat.
  • Entire memory: ~1.5 KiB of state per sequence, at any length: vs. a KV cache that grows tens of KiB per token (ch. on transformers).

Recap

  • Linear state path → affine maps → associative scan: O(\log T) depth training, O(1)-state inference.
  • minGRU ≈ GRU quality on our LM recipe.
  • SSMs: \Delta is a gate; eigenvalues inside the unit circle by construction; recurrence \equiv convolution \equiv ODE.
  • HiPPO: a fixed-size state provably carries a long past — via a time-varying system; frozen, its matrix initializes S4/S4D.
  • Inference delivered: step == scan to float precision; flat cost per token; ~1 KiB of state vs. a KV cache that grows with context.
  • But everything is LTI: the kernel is fixed before the model sees the data — the state update is content-blind (the MLPs around it are not). Fixing that (and keeping the scan) is the next section.