The trade RWKV is trying to dodge

Softmax attention is expensive in two ways with different causes. Training costs O(N^2 · d) because every pair of positions is scored. Decoding costs O(N · d) per token because the model re-reads a KV cache that grows with context. A classical RNN inverts both: O(1) per token, but strictly sequential training.

The blocker on the attention side is exp(q_n · k_m). That term couples a query and a key inside an exponential, so it cannot split into a part depending only on n and a part depending only on m. No factorization means no running state, which means the cache. RWKV replaces it with something that does factor: a per-channel exponential decay in the distance n - m, times a content term exp(k_m) that depends on the past token alone.

Advertisement

Token shift: a two-tap causal convolution

Before any mixing happens, RWKV blends each token with its immediate predecessor. For a learned per-channel vector μ with entries in (0, 1):

x̄_t = μ ⊙ x_t + (1 - μ) ⊙ x_(t-1)        x_t : [d]

r_t = W_r · x̄_t^(r),   k_t = W_k · x̄_t^(k),   v_t = W_v · x̄_t^(v)
(separate μ_r, μ_k, μ_v — each projection gets its own blend)

This is a two-tap causal convolution: one shift and one elementwise lerp. What it buys is a per-channel dial between “now” and “one step ago.” Channels that learn μ ≈ 0 effectively read the previous token, giving the block a bigram-style comparison — a crude but useful stand-in for the induction behaviour attention gets for free. Ablating token shift measurably hurts RWKV.

Advertisement

The WKV operator: attention with a fixed decay

The core of the time-mixing block is wkv_t, a weighted average of all past values:

          Σ_(i=1..t-1) exp(-(t-1-i)·w + k_i) ⊙ v_i  +  exp(u + k_t) ⊙ v_t
wkv_t = ---------------------------------------------------------------
          Σ_(i=1..t-1) exp(-(t-1-i)·w + k_i)        +  exp(u + k_t)

w = exp(ω) > 0   per-channel decay rate   w, u : [d]
u                  bonus for the current token

Read the exponent: -(t-1-i)·w is position and nothing else — a token m steps back is damped by e^(-mw). k_i is content and nothing else — a large key makes that token loud regardless of who is asking. There is no query. The current token is pulled out of the decay schedule and given its own weight exp(u + k_t), since otherwise “now” would sit on the same curve as the recent past. Numerator and denominator share weights, so wkv_t is a convex combination of the v_i and can never blow up.

Why it collapses into a recurrence

The whole architecture rests on one line of algebra:

exp(-(t-1-i)·w + k_i) = e^(-w) · exp(-(t-2-i)·w + k_i)

Every term in the sum at step t is the corresponding term at step t-1 multiplied by the same scalar e^(-w). So the sums themselves are a recurrence. Keep two running accumulators — a numerator a_t and a denominator b_t, each of shape [d]:

wkv_t = (a_(t-1) + exp(u + k_t) ⊙ v_t) / (b_(t-1) + exp(u + k_t))

a_t   = e^(-w) ⊙ a_(t-1) + exp(k_t) ⊙ v_t          a_0 = b_0 = 0
b_t   = e^(-w) ⊙ b_(t-1) + exp(k_t)

Two elementwise multiply-adds per token, and the state is 2d floats whether t is 10 or 10 million. Training uses the same recurrence, but the expensive parts — the W_r, W_k, W_v projections — run for all positions at once as dense matmuls, leaving only a cheap elementwise scan over time that parallelizes across batch and channels: O(N · d), not O(N^2 · d).

Numerical stability: carry the exponent

Written naively the recurrence overflows: exp(k_t) with k_t = 40 is already past fp32 range, and a_t accumulates such terms. Kernels therefore store the state factored as (p_t, a_t, b_t), where the true accumulator is e^(p_t) · a_t and p_t tracks the running maximum exponent:

q  = max(p, u + k_t)
wkv_t = (e^(p-q)⊙a + e^(u+k_t-q)⊙v_t) / (e^(p-q)⊙b + e^(u+k_t-q))

p' = max(p - w, k_t)
a' = e^(p-w-p')⊙a + e^(k_t-p')⊙v_t
b' = e^(p-w-p')⊙b + e^(k_t-p')

Every exponent fed to exp is now ≤ 0, so the worst case is underflow to zero — harmless — rather than inf/inf = NaN. This is the log-sum-exp shift that keeps softmax safe, applied incrementally; kernels also clamp k as a second guard. If you reimplement WKV and see NaNs at long context, this is the missing piece.

A worked example, one channel

Take a single channel with decay w = 0.7 (so e^(-w) = 0.4966) and bonus u = 0.5. Feed three tokens with k = [1.0, 0.0, 0.5] and v = [2.0, 6.0, 4.0], then evaluate wkv_3 straight from the definition:

i=1:  exp(-(3-1-1)·0.7 + 1.0) = exp(0.3) = 1.3499   × v=2  →  2.6997
i=2:  exp(-(3-1-2)·0.7 + 0.0) = exp(0.0) = 1.0000   × v=6  →  6.0000
now:  exp(0.5 + 0.5)             = exp(1.0) = 2.7183   × v=4  → 10.8731

wkv_3 = (2.6997 + 6.0000 + 10.8731) / (1.3499 + 1.0000 + 2.7183)
      = 19.5728 / 5.0682 = 3.8619

Now the recurrence: a_1 = e^1·2 = 5.4366, b_1 = 2.7183; then a_2 = 0.4966·5.4366 + 6 = 8.6997 and b_2 = 0.4966·2.7183 + 1 = 2.3499. Plugging in: (8.6997 + 10.8731) / (2.3499 + 2.7183) = 3.8619. Identical — two numbers of state reproduced a sum over the entire history, and the result lands inside [2, 6], as a convex combination must.

Receptance: the gate on the way out

The averaged wkv_t is always something: even with nothing useful to contribute, the block emits a weighted mean of past values. Receptance is the valve that fixes this:

o_t = W_o · ( σ(r_t) ⊙ wkv_t )        σ = sigmoid, elementwise

Because σ(r_t) lies in (0, 1) per channel, the block decides independently how much of each channel of the mixed history gets through: → 0 mutes it for this token, → 1 passes it intact. This is a GRU-style gate, and it is where query-like selectivity partially returns: RWKV cannot choose which past token to look at, but r_t is computed from the current token, so it can choose which channels of the summary to act on.

Channel mixing: the feed-forward half

Each layer pairs the time-mixing block above with a channel-mixing block playing the role of the Transformer’s FFN. It too starts from a token shift and ends in a receptance gate:

k'_t = W_k' · x̄_t        [d] → [d_ff],  d_ff ≈ 3.5× to 4× d
v'_t = W_v' · max(k'_t, 0)^2                   squared ReLU
o'_t = σ(r'_t) ⊙ v'_t                       back to [d]

The nonlinearity is squared ReLU, not GELU or SwiGLU: exactly zero on the negative half-line, so activations are genuinely sparse, and quadratic on the positive side, sharpening the contrast between weak and strong units. The receptance gate reappears, so both halves of every layer are gated by a sigmoid of a token-shifted projection. Parameter-wise this block dominates the layer, as the FFN does in a Transformer — time mixing is the clever part, not the expensive one.