RWKV Math

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.

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.

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.

Cost accounting against a KV cache

Put the two decode profiles side by side for a model with L layers, width d, and context N:

For L = 24, d = 2048, fp16: the KV cache costs 2 × 24 × 2048 × 2 = 196,608 bytes per token, about 0.19 MiB — roughly 6 GB at 32k context, all of which every generated token must stream. RWKV’s state is ~3 × 24 × 2048 floats, under 0.5 MB, identical at token 32,000 and token 1. The O(d^2) weight term is the same for both; the entire difference is the cache.

On a CPU that difference is the throughput, since decoding is memory-bandwidth-bound: tokens/s ≈ bandwidth / bytes_per_token. A 1.5B model at 4 bits is ~0.9 GB of weights; at ~30 GB/s that caps you near 33 tokens/s, and streaming a 6 GB cache on top drops it by an order of magnitude. RWKV reads its half-megabyte state instead, so throughput is flat in context length and RAM is predictable at load time.

RWKV-5, -6, -7: from vector state to matrix state

The version above is RWKV-4, and its weakness is visible in the algebra: a_t holds only d numbers, and w is a fixed learned constant that ignores what the token said. Later versions attack both.

RWKV-5 (Eagle) promotes the state from a vector to a per-head matrix S_t : [d_k, d_v] updated by a rank-1 outer product S_t = diag(decay) · S_(t-1) + k_t^T v_t, and drops the denominator entirely — the same linear-attention shape RetNet uses. With 64-wide heads that multiplies state capacity by roughly 64×. RWKV-6 (Finch) makes the decay data-dependent: w_t is computed per token from x_t via a low-rank (LoRA-style) projection, and token shift becomes dynamic too, so the model can choose to forget fast or hold on. RWKV-7 (Goose) generalizes the update to a delta-rule form with a non-diagonal transition, letting the state be selectively overwritten rather than only faded — which is what unlocks in-context state tracking that diagonal-decay models provably cannot express.

What a fixed state still costs you

No engineering escapes the information-theoretic fact: a fixed-size state is lossy compression of an unbounded history. Attention keeps every token verbatim and can retrieve any one exactly; RWKV keeps a summary and cannot. This shows up where you would predict — exact long-range recall, verbatim copying, needle-in-a-haystack retrieval — and it is why hybrid stacks that interleave a few full-attention layers among many linear ones remain popular.

The pitfalls follow the math. Initialize w across a spread of timescales — some channels near-zero decay for long memory, some fast for local detail; collapsing them to one rate wastes the channel dimension. Keep the (p, a, b) stabilization even if short-context tests pass without it. And note that w = exp(ω) is parameterized in log space because gradients through a strictly-positive decay behave far better there.