11. Streaming Step vs. Full-Sequence Forward Equivalence

Problem

A recurrent or SSM-style layer has two faces. During training you usually run it over a whole sequence at once (forward), because that's what a dataset of complete utterances gives you. In a real-time voice product, inference instead has to consume one audio frame at a time as it arrives over the wire and emit a hidden state immediately, without waiting for the rest of the utterance — a step that advances the recurrence by exactly one timestep.

These two code paths must be mathematically identical: whatever forward computes for timestep t must equal what you get from calling step t+1 times in a row. A layer whose step-wise and batched paths silently disagree is a correctness bug that a non-streaming unit test would never catch.

Source: src/11_streaming_step_vs_forward.py

class GatedLinearRecurrentLayer:
    def __init__(self, W_a: torch.Tensor, W_b: torch.Tensor): ...
    def step(self, x_t: torch.Tensor, h_prev: torch.Tensor) -> torch.Tensor: ...
    def forward(self, x: torch.Tensor, h0: torch.Tensor) -> torch.Tensor: ...

>>> layer = GatedLinearRecurrentLayer(torch.randn(3, 4), torch.randn(3, 4))
>>> h0 = torch.zeros(2, 4)
>>> x = torch.randn(2, 5, 3)
>>> layer.forward(x, h0).shape
torch.Size([2, 5, 4])

Step-by-Step Approach

1. Define the recurrence once: a_t = sigmoid(x_t @ W_a) (a data-dependent decay gate in (0, 1)), b_t = x_t @ W_b, and h_t = a_t * h_{t-1} + b_t.
2. Implement step(x_t, h_prev) as exactly that single-timestep update — no loop, no dependence on anything but the arguments it's given.
3. Implement forward(x, h0) as nothing more than calling step once per timestep in a Python loop over seq_len, carrying the returned hidden state into the next call, and stacking the results.
4. Resist the temptation to give forward a separate, "more vectorized" implementation of the same math — if forward and step are two independent code paths, they can drift apart even when both look individually correct.
5. Verify equivalence directly: call forward once over a full sequence, then call step in a loop over the same sequence, and assert the two hidden-state trajectories match exactly.

The key insight is that forward should be defined in terms of step, not as a parallel reimplementation of the same recurrence — that's what guarantees a real streaming server (which only ever calls step) behaves identically to whatever was validated offline with forward.

Reference solution

import torch


class GatedLinearRecurrentLayer:
    def __init__(self, W_a: torch.Tensor, W_b: torch.Tensor):
        self.W_a = W_a
        self.W_b = W_b

    def step(self, x_t: torch.Tensor, h_prev: torch.Tensor) -> torch.Tensor:
        # (batch, input_size) @ (input_size, hidden_size) -> (batch, hidden_size)
        a_t = torch.sigmoid(x_t @ self.W_a)
        b_t = x_t @ self.W_b
        return a_t * h_prev + b_t

    def forward(self, x: torch.Tensor, h0: torch.Tensor) -> torch.Tensor:
        batch, seq_len, _ = x.shape
        h = h0
        outputs = []
        # forward is defined as nothing more than repeated step() calls --
        # this loop IS the streaming path, just run to completion up front
        for t in range(seq_len):
            h = self.step(x[:, t, :], h)
            outputs.append(h)
        return torch.stack(outputs, dim=1)

Key Functions & Tricks

• torch.sigmoid — keeps the decay gate in (0, 1) so the recurrence is a genuine convex-ish blend of old state and new input, not an exploding one.
• x[:, t, :] slicing — pulls out a single timestep's batch of inputs, shape (batch, input_size), to feed into step.
• torch.stack(outputs, dim=1) — reassembles the list of per-timestep hidden states back into a single (batch, seq_len, hidden) tensor.
• Defining forward in terms of step rather than duplicating the math — a design pattern, not a tensor op, but the actual "trick" this problem is testing.

How to Recognize This Pattern