6. Selective SSM Forward (Data-Dependent Gating)

Problem

Classic S4 uses a fixed B, C, and step size delta shared across the whole sequence, which makes it a linear time-invariant (LTI) system — efficient, but unable to selectively attend to or ignore specific inputs the way attention can. Mamba's central idea is to make B, C, and delta functions of the input at each timestep: B_t, C_t, and delta_t are all computed by cheap linear projections of the current input u_t. This "selection mechanism" is what lets the model gate information in and out of its hidden state based on content, closing much of the gap to attention while keeping the O(length) recurrent (or O(length log length) scan) cost.

This problem assembles the full per-timestep forward: project the input to get data-dependent delta_t (via a linear layer + softplus, which keeps delta_t positive), B_t, and C_t; discretize A (fixed, diagonal, continuous-time) into A_bar_t using the simplified Euler-style B_bar_t = delta_t * B_t; run the elementwise recurrence to update the (dim, state) hidden state; and read out y_t = C_t^T h_t via a contraction over the state dimension.

Source: src/6_selective_gating_mamba.py

def selective_ssm_forward(
    u: torch.Tensor, A: torch.Tensor,
    W_B: torch.Tensor, W_C: torch.Tensor, W_delta: torch.Tensor, delta_bias: torch.Tensor,
) -> torch.Tensor:
    ...

Examples:
>>> u = torch.tensor([[[1.0], [2.0]]])
>>> A = torch.tensor([[-1.0]])
>>> W_B = W_C = torch.tensor([[1.0]])
>>> W_delta = torch.tensor([[0.0]])
>>> delta_bias = torch.tensor([0.0])
>>> selective_ssm_forward(u, A, W_B, W_C, W_delta, delta_bias)
tensor([[[0.6931], [6.2383]]])

Step-by-Step Approach

1. Compute delta = F.softplus(u @ W_delta + delta_bias), shape (batch, length, dim) — softplus guarantees delta > 0, which the discretization requires (a negative or zero step size is meaningless).
2. Compute B_t = u @ W_B and C_t = u @ W_C, each shape (batch, length, state) — these are shared across channels (broadcast over dim later), a simplification of Mamba's actual per-channel-selective projections.
3. Discretize: A_bar = torch.exp(delta.unsqueeze(-1) * A), shape (batch, length, dim, state), and the simplified Euler B_bar = delta.unsqueeze(-1) * B_t.unsqueeze(2) (also broadcast to the same shape).
4. Initialize the hidden state h to zeros of shape (batch, dim, state) — note this is one dimension richer than the earlier scan problems' (batch, dim), since each channel now carries its own state-sized memory.
5. Loop over timesteps: update h = A_bar[:, t] * h + B_bar[:, t] * u[:, t].unsqueeze(-1) — note u_t (shape (batch, dim)) must be unsqueezed to broadcast against the state axis.
6. Read out y_t = torch.einsum("bdn,bn->bd", h, C_t[:, t]) — a contraction over the state axis n, per batch and channel, turning the (dim, state) hidden state into a (dim,) output.
7. Stack the per-timestep outputs to get y of shape (batch, length, dim), matching the input's shape.

The key insight is that "selective" doesn't change the shape of the recurrence at all — it's still h_t = A_bar_t * h_{t-1} + B_bar_t * u_t elementwise per (dim, state) — it only changes where A_bar_t, B_bar_t, and the readout C_t come from: fixed parameters in S4, versus linear projections of the current input in Mamba. Getting comfortable tracing which projections feed which discretization step is most of what this problem tests.

Reference solution

def selective_ssm_forward(
    u: torch.Tensor, A: torch.Tensor,
    W_B: torch.Tensor, W_C: torch.Tensor, W_delta: torch.Tensor, delta_bias: torch.Tensor,
) -> torch.Tensor:
    batch, length, dim = u.shape
    state = A.shape[1]

    # data-dependent parameters: all functions of u, not fixed weights (the "selection")
    delta = F.softplus(u @ W_delta + delta_bias)  # (batch, length, dim), softplus keeps delta > 0
    B_t = u @ W_B  # (batch, length, state)
    C_t = u @ W_C  # (batch, length, state)

    # discretize: exact exp for A_bar, simplified Euler approx (delta * B) for B_bar
    A_bar = torch.exp(delta.unsqueeze(-1) * A)          # (batch, length, dim, state)
    B_bar = delta.unsqueeze(-1) * B_t.unsqueeze(2)       # (batch, length, dim, state)

    h = torch.zeros(batch, dim, state, dtype=u.dtype, device=u.device)
    ys = []
    for t in range(length):
        # h_t = A_bar_t * h_{t-1} + B_bar_t * u_t  (u_t broadcasts over the state axis)
        h = A_bar[:, t] * h + B_bar[:, t] * u[:, t].unsqueeze(-1)
        # y_t = C_t^T h_t, contracting over the state dimension per channel
        y_t = torch.einsum("bdn,bn->bd", h, C_t[:, t])
        ys.append(y_t)
    return torch.stack(ys, dim=1)  # (batch, length, dim)

Key Functions & Tricks

• F.softplus(x) — smooth, always-positive alternative to ReLU; used to constrain delta > 0
• torch.einsum("bdn,bn->bd", h, C_t) — explicit contraction over the state axis, clearer than manually broadcasting + summing
• u @ W_B — matmul-as-projection, the standard way to turn raw input into per-timestep parameters
• tensor.unsqueeze(-1) / unsqueeze(2) — align (batch,length,dim), (batch,length,state), and (dim,state) shapes for broadcasting
• Data-dependent (selective) SSM parameters — replacing fixed B, C, delta with per-timestep projections of the input
• Simplified Euler discretization for B_bar — delta * B instead of the exact ZOH form, a documented approximation in the Mamba paper

How to Recognize This Pattern