RWKV

Attention scores every pair. RWKV asks what happens if the score between two positions is a fixed function of their distance and the key alone — no query–key interaction at all.

The numerator is a weighted sum of values; the denominator normalises it. The weight on position $i$ is $e^{k_i}$ — how much that token wants to be attended — multiplied by $e^{-(t-1-i)w}$, a per-channel exponential decay in distance. $u$ is a separate bonus for the current token, present because otherwise the decay would treat “now” as merely the nearest thing to the past.

The receptance $r_t = \sigma(W_r x_t)$ then gates the output, which is the LSTM’s output gate under a different name. R, W, K, V.

Why it is both parallel and recurrent

Eq. 1 has no term coupling two positions through a learned matrix, so the sums telescope.

Training uses the parallel form and a scan; generation uses the recurrence, with a state of two vectors per layer regardless of how many tokens have gone by. Not an approximation of one by the other — the same arithmetic, regrouped.

What it gives up

The decay $w$ is learned per channel and, in the original formulation, fixed — independent of content. A transformer can decide that token 900 is exactly what token 4000 needs; RWKV can only decide that some channels forget slowly. Recall of a specific distant token is the measurable weakness, and it is the same weakness every fixed-decay model has.

The Eagle and Finch revision addresses it directly: Finch makes $w$ depend on the input — a data-dependent decay, which is precisely Mamba’s selection arriving by a different route — and moves the state from a vector to a matrix, giving the recurrence somewhere to put more than one thing at a time.

That convergence is the interesting part. Mamba started from a continuous linear system and added input dependence; RWKV started from an RNN and added it; gated linear attention started from attention and added it. All three landed on the same structure, and state space duality is the paper that says so explicitly.

import torch
from torch import Tensor


def wkv_step(
    state: tuple[Tensor, Tensor, Tensor], k: Tensor, v: Tensor, w: Tensor, u: Tensor
):
    """One decode step. state = (numerator, denominator, running max)."""
    a, b, p = state

    # Current token: bonus u instead of decay, and it is not yet in the state.
    q = torch.maximum(p, u + k)
    e1, e2 = torch.exp(p - q), torch.exp(u + k - q)
    out = (e1 * a + e2 * v) / (e1 * b + e2)

    # Now fold it in, decaying what was there.
    q = torch.maximum(p - w, k)
    e1, e2 = torch.exp(p - w - q), torch.exp(k - q)
    return out, (e1 * a + e2 * v, e1 * b + e2, q)

Every exponent in that function is non-positive by construction — the running maximum is subtracted before any exp. Written the obvious way instead, matching eq. 1 term for term, the model produces inf/inf and then NaN somewhere around the two-hundredth token, and it does so only at long context, which makes it a bug that survives testing.