3. Selective-SSM Zero-Order-Hold Discretization

Problem

SSMs like S4 and Mamba are defined as continuous-time linear systems, x'(t) = A x(t) + B u(t), with a diagonal state matrix A (one scalar per (channel, state-dim) pair). To run on discrete audio/token timesteps, the continuous system must be discretized with a per-timestep step size delta_t — and because Mamba's delta_t is itself computed from the input (the "selection mechanism," covered in a later problem), this discretization must be recomputed at every timestep rather than once, offline.

The standard zero-order-hold (ZOH) discretization gives A_bar_t = exp(delta_t * A) and B_bar_t = (delta_t * A)^-1 * (exp(delta_t * A) - I) * delta_t * B_t. For diagonal A this collapses to elementwise ops (no matrix inverse needed). Writing B_bar's second factor as (exp(x) - 1) / x with x = delta_t * A, the naive (torch.exp(x) - 1) / x catastrophically cancels for small x (which happens often, since delta_t is a learned, data-dependent, often-small positive value) — the fix is torch.expm1(x) / x.

Source: src/3_ssm_zoh_discretization.py

def discretize_zoh(A: torch.Tensor, B: torch.Tensor, delta: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    ...

Examples:
>>> A = torch.tensor([[-1.0]])
>>> delta = torch.tensor([[[0.1]]])
>>> B = torch.tensor([[[2.0]]])
>>> A_bar, B_bar = discretize_zoh(A, B, delta)
>>> A_bar
tensor([[[[0.9048]]]])

Step-by-Step Approach

1. Check shapes: A is (dim, state), shared across time; delta is (batch, length, dim); B is (batch, length, state), broadcast across channels. The output shapes are both (batch, length, dim, state).
2. Broadcast delta against A to form x = delta.unsqueeze(-1) * A, shape (batch, length, dim, state) — this is the exponent that appears in both A_bar and B_bar.
3. Compute A_bar = torch.exp(x) directly — no numerical concern here, exp is well-behaved at any x in the realistic range (negative, since a stable continuous-time SSM needs A < 0).
4. For B_bar's (exp(x)-1)/x factor, use torch.expm1(x) / x instead of (torch.exp(x) - 1) / x — expm1 is implemented to stay accurate as x -> 0, where exp(x)-1 subtracts two nearly-equal floats and loses precision.
5. Guard the x == 0 edge case explicitly with torch.where (the true limit of (exp(x)-1)/x as x -> 0 is exactly 1, so substitute 1 rather than dividing by zero) — even though expm1(x)/x is far more stable than the naive form, literal x=0 still divides by zero.
6. Assemble B_bar = delta.unsqueeze(-1) * ratio * B.unsqueeze(2), broadcasting the per-timestep-per-state B_t across the channel axis.
7. Sanity-check the scalar case by hand: A=-1, delta=0.1, B=2.0 gives x=-0.1, A_bar=exp(-0.1)≈0.9048, and B_bar = 0.1 * (expm1(-0.1)/-0.1) * 2.0 ≈ 0.1903.

The key insight is that a numerically "obviously correct" formula ((exp(x)-1)/x) can silently degrade in float32 exactly in the regime this function is called most often — small delta_t — so recognizing when to reach for expm1 (or the analogous log1p, or a log-sum-exp trick elsewhere) is as much a part of "implementing the discretization" as getting the math right in the first place.

Reference solution

def discretize_zoh(A: torch.Tensor, B: torch.Tensor, delta: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    # x = delta_t * A, broadcast (batch,length,dim,1) * (dim,state) -> (batch,length,dim,state)
    x = delta.unsqueeze(-1) * A
    A_bar = torch.exp(x)  # exp(delta_t * A)
    # (exp(x) - 1) / x, computed via expm1 to avoid catastrophic cancellation as x -> 0
    ratio = torch.where(x.abs() < 1e-8, torch.ones_like(x), torch.expm1(x) / x)
    # B_bar = delta_t * ratio * B_t, broadcasting B_t over the channel/dim axis
    B_bar = delta.unsqueeze(-1) * ratio * B.unsqueeze(2)
    return A_bar, B_bar

Key Functions & Tricks

• torch.expm1(x) — numerically stable exp(x) - 1, accurate for small x where naive subtraction loses precision
• torch.where(cond, a, b) — branchless elementwise selection, used here to guard the x == 0 division
• tensor.unsqueeze(-1) / tensor.unsqueeze(2) — insert broadcast axes to align (batch,length,dim) and (dim,state)/(batch,length,state) shapes
• Zero-order-hold discretization — the standard continuous-to-discrete-time conversion for linear systems, elementwise for diagonal A
• Broadcasting instead of explicit loops/tiling — letting PyTorch's broadcast rules expand (dim,state) against (batch,length,dim,1) rather than manually repeating tensors

How to Recognize This Pattern