8. Data-Controlled Gated Linear Recurrence

Problem

Real-valued forget gates (as in an LSTM, or a_t in the earlier scan problems) can only scale a hidden state up or down. GateLoop generalizes the linear-recurrence family (S4, S5, LRU, RetNet) by making the per-timestep transition complex-valued and fully data-controlled: rather than a fixed or purely-decaying gate, each timestep's transition has a data-dependent magnitude (bounded in (0, 1) for stability, exactly like a forget gate) and a data-dependent phase (a rotation), so the state can be both scaled and rotated based on the input. GateLoop shows this same recurrence admits both an O(length) sequential recurrent form and an O(length log length) parallel-scan form — the same tradeoff explored in problems 1 and 2 of this set, but over a strictly more expressive complex-valued transition.

Implement the sequential form: h_t = g_t * h_{t-1} + k_t * v_t, where g_t = sigmoid(gate_logits_t) * exp(i * phase_t) is a complex number with magnitude in (0, 1) and a data-dependent phase, and k_t * v_t is treated as a real-valued input cast to complex. Note that when phase_t = 0 for all t, g_t is purely real and this recurrence reduces exactly to the h_t = a_t * h_{t-1} + b_t recurrence from problems 1–2 — the complex phase is what adds rotation on top of the familiar real-valued gating.

Source: src/8_gated_linear_recurrence_gateloop.py

def gated_linear_recurrence(
    v: torch.Tensor, k: torch.Tensor, gate_logits: torch.Tensor, phase: torch.Tensor,
    h0: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    ...

Examples:
>>> v = torch.tensor([[[2.0], [3.0]]])
>>> k = torch.tensor([[[1.0], [1.0]]])
>>> gate_logits = torch.zeros(1, 2, 1)
>>> phase = torch.zeros(1, 2, 1)
>>> h_real, h_imag = gated_linear_recurrence(v, k, gate_logits, phase)
>>> h_real
tensor([[[2.], [4.]]])

Step-by-Step Approach

1. Compute the gate magnitude: magnitude = torch.sigmoid(gate_logits), guaranteed in (0, 1) — this is what keeps the recurrence stable no matter what the phase does, since |g_t| < 1 always.
2. Build the complex gate with torch.polar(magnitude, phase), which constructs magnitude * exp(i * phase) directly from polar coordinates — avoid manually computing magnitude * (cos(phase) + i*sin(phase)), which is more error-prone and less readable.
3. Cast the real-valued k * v input term to complex with .to(torch.complex64), so it can be added to the complex hidden state.
4. Initialize h_prev to h0 (cast to complex, real-valued default) or zeros of dtype torch.complex64.
5. Loop over timesteps: h_prev = g[:, t] * h_prev + kv[:, t] — ordinary PyTorch elementwise ops work transparently on complex tensors, no special-casing needed.
6. Stack the per-timestep states and split into real and imaginary parts with .real / .imag for the return value — downstream code (e.g. an output projection) typically consumes the real part, but returning both lets a caller use the phase information if needed.
7. Validate with the zero-phase reduction: if phase is all zeros, h_imag should come out all zeros and h_real should exactly match the real-valued scan from problem 1 — a strong, cheap correctness check before trusting the complex-valued general case.

The key insight is that PyTorch's native complex dtype support (torch.complex64) means this generalization requires almost no new mechanics over the real-valued scan — the recurrence, the elementwise multiply-add, even the loop structure are identical; only the dtype and the gate's construction (torch.polar instead of a plain scalar) change, which is exactly why complex-valued state transitions are a tractable generalization rather than a totally different algorithm.

Reference solution

def gated_linear_recurrence(v, k, gate_logits, phase, h0=None):
    batch, length, dim = v.shape

    # magnitude bounded in (0,1) via sigmoid keeps the recurrence stable (no blow-up over time)
    magnitude = torch.sigmoid(gate_logits)
    # torch.polar(abs, angle) builds a complex tensor from magnitude + phase directly
    g = torch.polar(magnitude, phase)  # complex64, (batch, length, dim)

    kv = (k * v).to(torch.complex64)  # real input cast into the complex state space
    h_prev = torch.zeros(batch, dim, dtype=torch.complex64, device=v.device)
    if h0 is not None:
        h_prev = h0.to(torch.complex64)

    outs = []
    for t in range(length):
        h_prev = g[:, t] * h_prev + kv[:, t]  # complex elementwise: scales AND rotates h_prev
        outs.append(h_prev)
    h = torch.stack(outs, dim=1)
    return h.real, h.imag

Key Functions & Tricks

• torch.polar(abs, angle) — build a complex tensor from magnitude/phase (polar form) directly, avoiding manual trig
• torch.sigmoid(gate_logits) — bound a learned/data-dependent gate magnitude to (0, 1) for a stable recurrence
• tensor.to(torch.complex64) — cast a real tensor into the complex dtype so it can combine with complex state
• tensor.real / tensor.imag — extract components from a complex tensor
• Complex-valued elementwise recurrence — ordinary PyTorch ops (*, +) work unmodified on complex64/complex128 tensors
• Zero-phase reduction as a correctness check — verifying a generalized implementation collapses to the known-correct special case

How to Recognize This Pattern