Linear Attention, RWKV, and RetNet

"For years they told me I was an attention mechanism, all softmax and quadratic glory, comparing every token to every other token like a gossip who never forgets a face. Then someone dropped the softmax, and I discovered the truth: under the costume I was a recurrence the whole time, carrying one little matrix of state forward, step by step, the way an RNN does. I had been a state-space model in a fancier coat."

A Linear Attention That Is Secretly a Recurrence

In Section 13.1 we built structured state-space models whose core was a linear recurrence $\mathbf{s}_t = \mathbf{A}\mathbf{s}_{t-1} + \mathbf{B}\mathbf{x}_t$ that could be unrolled as a convolution for parallel training and run as a recurrence for streaming inference. That duality, one model with a parallel face and a recurrent face, is the protagonist of this section too, but arriving from the opposite direction. There we started from a continuous linear system and discretized it. Here we start from the Transformer attention of Section 12.5, the workhorse of Part III, and ask what happens to its cost when we drop the one nonlinearity that makes it quadratic. The answer lands us back in the same place: a linear recurrence with a matrix-valued state. By the end of the chapter, Section 13.6 will close the loop and show that SSMs, linear attention, RWKV, and RetNet are literally the same computational object under four different parameterizations. This section builds the linear-attention half of that bridge. We use the unified notation of Appendix A: $\mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t$ for the query, key, and value at position $t$, $\mathbf{o}_t$ for the attention output, $d$ for the head dimension.

The stakes are practical, not merely aesthetic. Standard attention costs $O(T^2 d)$ time and, at inference, an $O(T)$ key-value cache that grows without bound as a generation lengthens. For the long sequences that motivate this entire chapter, a year of hourly readings, a high-frequency tape, a book-length context, that quadratic wall is exactly what we are trying to climb. Linear attention and its descendants trade a small amount of expressive power for a recurrence that trains in parallel like a Transformer yet generates in constant time and constant memory per token like an RNN. That is a genuinely new point on the design space, and understanding why it is possible is understanding the duality this section is built around.

The four competencies this section installs are these: to rewrite softmax-free attention as a running-sum recurrence and read off both its parallel-train and recurrent-infer forms; to recognize RWKV's WKV time-mixing and channel mixing as a decayed linear recurrence with an RNN-like inference mode; to state RetNet's retention mechanism in its parallel, recurrent, and chunkwise forms and explain the role of its decay-based positional structure; and to see all of these, together with the SSMs of Sections 13.1 and 13.2, as one linear-recurrence family, the unifying picture that Section 13.6 formalizes.

1. From Softmax Attention to a Linear Recurrence Intermediate

Recall the attention output from Section 12.5. With a causal mask, the output at position $t$ is a softmax-weighted average of the value vectors at positions $i \le t$, where the weights come from query-key similarities:

The reason this is quadratic is printed in the formula: the scalar weight $\exp(\mathbf{q}_t^\top \mathbf{k}_i)$ binds the query at $t$ and the key at $i$ together inside a single nonlinearity, the exponential, that does not factor into a part depending only on $\mathbf{q}_t$ times a part depending only on $\mathbf{k}_i$. Because it does not factor, we are forced to materialize the similarity for every one of the $T^2$ query-key pairs. The softmax is the entire source of the quadratic cost.

The linear-attention idea, introduced by Katharopoulos and colleagues in 2020, is to replace the exponential similarity with a factored one. Choose a feature map $\phi: \mathbb{R}^d \to \mathbb{R}^{m}$ and approximate $\exp(\mathbf{q}_t^\top \mathbf{k}_i) \approx \phi(\mathbf{q}_t)^\top \phi(\mathbf{k}_i)$. The point of $\phi$ is that the similarity now splits: it is an inner product of one vector built from the query alone and another built from the key alone. A common choice is the elementwise $\phi(\mathbf{u}) = \mathrm{elu}(\mathbf{u}) + 1$, which keeps the features positive so the denominator stays well behaved. Substituting and using the associativity of the sums:

The decisive step is pulling $\phi(\mathbf{q}_t)$ out of the sum over $i$, which the factorization permits and the softmax forbade. What remains inside the sums no longer depends on $t$ except through the running upper limit. Define a matrix-valued state and a vector-valued normalizer that accumulate over positions:

Both are cumulative sums, so each obeys a one-line recurrence: add the current term to the previous total. That is the entire content of linear attention. The model is a linear recurrence with a matrix-valued state $\mathbf{S}_t$, updated by a rank-one outer product per step, and read out by a contraction with the current query feature:

This is the recurrent (inference) form: carry a fixed-size state $\mathbf{S}_t$ of shape $m \times d$, update it with one outer product, read out with one matrix-vector product, all independent of how many tokens have already gone by. Generation costs $O(md)$ time and $O(md)$ memory per token, constant in the sequence length, in place of the growing key-value cache of softmax attention. The same equations, read across the whole sequence at once rather than step by step, give the parallel (training) form: compute all $\phi(\mathbf{k}_i)\mathbf{v}_i^\top$ terms and take a cumulative sum, which a framework parallelizes across positions exactly as it parallelizes the SSM scan of Section 13.2. One model, two faces.

It is worth stating the cost ledger explicitly, because the whole motivation lives in it. The recurrent form's state $\mathbf{S}_t$ has shape $m \times d$, fixed once the head dimensions are chosen, so generating a new token costs exactly one outer product ($O(md)$), one state addition ($O(md)$), and one contraction ($O(md)$), with not a single term that depends on $t$. Contrast softmax attention, where token $t$ must compute a similarity against all $t$ stored keys, an $O(td)$ operation that grows as the generation lengthens, on top of a key-value cache whose memory is $O(td)$. Over a full length-$T$ generation, softmax attention pays $\sum_t O(td) = O(T^2 d)$, while linear attention pays $\sum_t O(md) = O(Tmd)$, linear in $T$. The asymptotic gap is the entire reason this family exists, and it is the same $O(T)$-versus-$O(T^2)$ improvement that the convolutional and state-space views of this chapter chase from their own directions.

One caveat sets honest expectations. The feature map $\phi$ is an approximation to the exponential kernel, not an equality, so linear attention is not a drop-in numerical match for softmax attention; it is a different, cheaper attention whose quality depends on the feature map. Early feature maps lost accuracy on language tasks, which is precisely the gap that RWKV and RetNet close with better-designed recurrences and decay structure. What is exact and unconditional is the duality: for any fixed $\phi$, the parallel and recurrent forms compute the same function, bit for bit up to floating-point order of operations. That equivalence is the load-bearing fact, and it is what we prove in code.

2. RWKV: An RNN-Transformer Hybrid Advanced

RWKV (Peng and colleagues, 2023, the name stands for the four core symbols R, W, K, V) takes the linear-recurrence template of subsection one and adds the one ingredient plain linear attention lacks: a learned decay that lets the state forget the distant past at a controllable rate. Where linear attention accumulates an undamped sum $\mathbf{S}_t = \mathbf{S}_{t-1} + \phi(\mathbf{k}_t)\mathbf{v}_t^\top$, RWKV weights each past contribution by an exponentially decaying factor, so a position $s$ steps back is discounted by roughly $e^{-sw}$ for a learned per-channel decay $w$. This is the same exponential-decay positional structure we will meet again in RetNet, and it is what gives the recurrence a finite effective memory rather than an ever-growing flat sum.

RWKV is built from two alternating sublayers, mirroring the attention-then-feedforward structure of a Transformer block but with both pieces made recurrent. The first is time mixing, which carries information across positions through the WKV recurrence; the second is channel mixing, a position-wise gated feedforward that mixes across feature channels with no cross-time interaction. The time-mixing output for a channel is a decayed, key-weighted average of values, which in its numerically-stabilized recurrent form reads

where $w$ is the learned channel-wise decay (the W), $u$ is a learned bonus that gives the current token extra weight, $k_t$ and $v_t$ are the key and value (the K and V), and $\mathbf{a}_t, \mathbf{b}_t$ are the running numerator and denominator states, the direct analogues of $\mathbf{S}_t$ and $\mathbf{z}_t$ from subsection one with a decay $e^{-w}$ multiplying the carried state. The receptance $\mathbf{r}_t = \sigma(\cdot)$ (the R) is a sigmoid gate that decides how much of the mixed signal passes through, the role attention's softmax weights play in a standard block. The output is $\mathbf{o}_t = \mathbf{r}_t \odot \mathrm{wkv}_t$ projected out.

The payoff is precisely the duality of subsection one, now in a competitive language model. Because the WKV recurrence is linear in its state, the whole sequence can be computed in a parallelizable form during training, so RWKV trains with Transformer-like throughput on a GPU. Because that same recurrence carries only the fixed-size pair $(\mathbf{a}_t, \mathbf{b}_t)$ forward, inference is pure RNN: each new token costs $O(d)$ time and $O(d)$ memory regardless of context length, with no key-value cache to grow. RWKV was the first architecture to demonstrate this combination at multi-billion-parameter scale on language modeling, reaching quality comparable to similarly-sized Transformers while generating in constant memory per token.

3. RetNet: The Retention Mechanism Advanced

RetNet (Sun and colleagues, 2023) reaches the same linear-recurrence template from a more explicitly attention-flavored direction, and contributes the cleanest statement of why these models have three computational forms rather than two. Its core operation, retention, replaces the softmax with a fixed exponential decay applied along the sequence. The parallel form looks almost exactly like attention with the softmax removed and a decay mask inserted:

where $\gamma \in (0,1)$ is a per-head decay constant and $\mathbf{D}$ is a causal mask whose entries decay geometrically with the distance $t - i$ between query and key. The decay $\gamma^{t-i}$ is RetNet's positional structure: instead of adding learned position embeddings, the model encodes "how long ago" directly into the strength of each interaction, so recent positions count more than distant ones by a fixed multiplicative schedule. This is the same exponential discounting that RWKV applies through $e^{-w}$ and that an SSM applies through its state matrix $\mathbf{A}$; RetNet just writes it as an explicit mask on the attention matrix.

The decay mask is also RetNet's answer to a question every attention model must settle: how does the model know when a token occurred? Softmax Transformers bolt on a separate positional encoding (sinusoids or rotary embeddings) because the bare dot product is permutation-invariant. Retention bakes position into the mechanism itself: the factor $\gamma^{t-i}$ is monotone in the gap $t - i$, so a token's influence fades smoothly with age and the model reads distance directly off the interaction strength. This is a more parsimonious design, no extra parameters, no separate embedding table, and it generalizes more gracefully to sequences longer than those seen in training, because the decay is defined for any gap rather than learned per absolute position.

Because the decay factorizes as $\gamma^{t-i} = \gamma^{t}/\gamma^{i}$, the same retention admits a recurrent form. Carry a matrix-valued state $\mathbf{S}_t$ and update it with a decay-then-add step, then read out with the current query:

This is exactly the linear-attention recurrence of subsection one with a decay $\gamma$ on the carried state and the feature map taken as the identity (no normalizer, since retention drops the denominator). The parallel form trains fast; the recurrent form generates in $O(d^2)$ time and memory per token, constant in context length. RetNet's distinctive third contribution is the chunkwise form, a hybrid that splits the sequence into chunks, runs the parallel form within each chunk, and the recurrent form across chunks by carrying the state $\mathbf{S}$ from one chunk to the next. This gives the best of both: parallel throughput inside a chunk and a fixed-size state between chunks, so a very long sequence is trained in $O(T)$ total work with bounded memory, neither the $O(T^2)$ of full parallel attention nor the strict sequential dependency of pure recurrence. The chunkwise form is the practical default for training RetNet on long contexts, and it is the same idea that the chunked scans of Section 13.2 use for SSMs.

4. The Unifying Picture: One Recurrence, Four Parameterizations Advanced

Stand back from the three mechanisms and a single shape is visible through all of them, the shape that also describes the structured state-space models of Sections 13.1 and 13.2. Every one of these models carries a state forward through a linear recurrence of the form "decay the old state, add a contribution from the new input, read out by a projection". Written generically:

and the four architectures are four choices of the decay $\mathbf{G}_t$ and of how $\mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t$ are produced. Plain linear attention sets the decay to one (no forgetting) and applies a feature map $\phi$ to queries and keys. RWKV sets the decay to a learned per-channel constant $e^{-w}$ and adds a receptance gate. RetNet sets the decay to a per-head scalar $\gamma$ and uses the identity feature map. A diagonal structured SSM (S4, S5, Mamba) sets the decay to its discretized state matrix $\bar{\mathbf{A}}$, which in the selective case is itself input-dependent, exactly the gated decay $\mathbf{G}_t$. The differences that look so large at the architecture level, attention versus state space versus RNN, are differences in three slots of one template.

This unification is not just tidy bookkeeping; it is the conceptual backbone of the rest of the chapter and a recurring beat of the whole book's temporal thread. Part II's Kalman filter and the RNN of Chapter 10 were already this same loop with a hand-set versus learned transition; here the loop returns once more, now with the transition designed so the recurrence is linear and therefore parallelizable. That linearity is the whole trick: a nonlinear recurrence like the RNN's $\tanh$ cell cannot be unrolled into a parallel scan, because the nonlinearity blocks the associativity that a parallel prefix-sum needs, which is exactly why Chapter 10's RNNs were stuck with sequential training. Linear attention, RWKV, RetNet, and the SSMs buy back parallel training precisely by keeping the cross-time recurrence linear and pushing all nonlinearity into the per-position feature maps and gates, where it does not obstruct the scan. The bridge of Section 13.6 will turn this informal "they are all linear recurrences" into the precise state-space-duality statement that Mamba-2 proved, and Section 13.4 will then leave the discrete linear recurrence behind for the continuous-time models, Neural ODEs and their relatives, that handle irregular sampling.

5. Worked Example: Parallel Equals Recurrent Advanced

We now make the central claim executable: linear attention computed in its parallel form and in its recurrent form produce the identical output sequence. The plan is the cleanest possible proof of correctness. Implement the parallel form, which materializes the position-by-position cumulative state exactly as a training kernel would; implement the recurrent form, which carries the running state one step at a time exactly as inference would; run both on the same random queries, keys, and values; then assert with np.allclose that every output matches to floating-point precision. If they agree, the duality of subsection one is verified by independent computation. Code 13.3.1 is the parallel form.

import numpy as np

rng = np.random.default_rng(0)
T, d, m = 6, 4, 5                       # sequence length, value dim, feature-map dim

# Random per-position queries, keys, values for one attention head.
Q = rng.normal(0, 1.0, size=(T, d))     # queries q_t
K = rng.normal(0, 1.0, size=(T, d))     # keys    k_t
V = rng.normal(0, 1.0, size=(T, d))     # values  v_t

def phi(u):
    """Feature map: elu(u) + 1, kept positive so the normalizer is well behaved."""
    return np.where(u > 0, u + 1.0, np.exp(u))   # elu(u)+1, here mapped d -> d (m == d path)

def linear_attention_parallel(Q, K, V):
    """Parallel form: cumulative sum of phi(k) v^T over positions, contracted by phi(q)."""
    PhiQ, PhiK = phi(Q), phi(K)                  # feature-mapped queries and keys, (T, d)
    O = np.zeros((T, d))
    for t in range(T):                           # one cumulative sum per query position
        S = np.zeros((d, d)); z = np.zeros(d)    # state and normalizer up to t, rebuilt
        for i in range(t + 1):                   # sum over all i <= t (the causal prefix)
            S += np.outer(PhiK[i], V[i])         # accumulate phi(k_i) v_i^T
            z += PhiK[i]                          # accumulate phi(k_i)
        num = PhiQ[t] @ S                         # phi(q_t)^T S_t
        den = PhiQ[t] @ z + 1e-6                  # phi(q_t)^T z_t (epsilon for stability)
        O[t] = num / den
    return O

O_par = linear_attention_parallel(Q, K, V)
print("parallel  O[5] =", np.round(O_par[5], 6))

parallel  O[5] = [ 0.273944 -0.105421  0.486318 -0.041907]

Now the recurrent form. Code 13.3.2 carries the running state $\mathbf{S}_t$ and normalizer $\mathbf{z}_t$ forward one step at a time, updating each by a single rank-one addition, exactly the inference recurrence boxed in subsection one. It never rebuilds the prefix sum; it accumulates it. We then compare the two outputs with np.allclose.

def linear_attention_recurrent(Q, K, V):
    """Recurrent form: carry running state S_t and normalizer z_t, one outer product per step."""
    PhiQ, PhiK = phi(Q), phi(K)
    S = np.zeros((d, d)); z = np.zeros(d)        # S_0 = 0, z_0 = 0 (empty history)
    O = np.zeros((T, d))
    for t in range(T):                           # single forward pass, no inner loop
        S = S + np.outer(PhiK[t], V[t])          # S_t = S_{t-1} + phi(k_t) v_t^T
        z = z + PhiK[t]                           # z_t = z_{t-1} + phi(k_t)
        O[t] = (PhiQ[t] @ S) / (PhiQ[t] @ z + 1e-6)
    return O

O_rec = linear_attention_recurrent(Q, K, V)
print("recurrent O[5] =", np.round(O_rec[5], 6))
print("parallel == recurrent :", np.allclose(O_par, O_rec))   # the duality, certified

recurrent O[5] = [ 0.273944 -0.105421  0.486318 -0.041907]
parallel == recurrent : True

Read what the pair establishes. Code 13.3.1 computed linear attention the way a training kernel does, as a cumulative sum over positions. Code 13.3.2 computed it the way an inference loop does, as a one-step-at-a-time recurrence with a fixed-size state. They returned the same numbers. That is the parallel-equals-recurrent duality of subsection one, not asserted but demonstrated, and it is the same duality that makes RWKV and RetNet trainable in parallel yet generative in constant memory. From-scratch the two paths took about thirty lines between them; the library shortcut below collapses both to a single call.

import torch
from fla.layers import GatedLinearAttention   # pip install flash-linear-attention

gla = GatedLinearAttention(hidden_size=512, num_heads=8)  # data-dependent decay built in
x = torch.randn(2, 1024, 512)                              # (batch, length, features)
y, _ = gla(x)            # parallel chunked kernel for training; pass a state cache for O(1) inference
print(y.shape)           # torch.Size([2, 1024, 512])

import torch
import numpy as np

# ---- From-scratch SSD kernel: parallel and sequential forms ----
def ssd_parallel(Q, K, V, alpha):
    """SSD parallel form: causal masked linear attention with decay mask M[t,i] = alpha^(t-i).
    Q, K: (L, n) -- queries and keys (here B and C in SSM language)
    V:    (L,)   -- scalar values (the input sequence u)
    Returns output y of shape (L,).
    """
    L, n = Q.shape
    # Build the full L x L causal decay mask.
    t_idx = np.arange(L)[:, None]          # column of row indices
    i_idx = np.arange(L)[None, :]          # row of col indices
    M = np.where(i_idx <= t_idx, alpha ** (t_idx - i_idx), 0.0)  # (L, L)
    scores = (Q @ K.T) * M                 # (L, L): masked dot products
    y = scores @ V                         # (L,): weighted sum over positions
    return y

def ssd_sequential(B, C, u, alpha):
    """SSD sequential form: scalar-times-identity SSM recurrence.
    B, C: (n,) -- input and output projection vectors
    u:    (L,) -- scalar input sequence
    Returns output y of shape (L,).
    """
    n = len(B)
    s = np.zeros(n)
    y = []
    for u_t, c_t in zip(u, C[None, :].repeat(len(u), axis=0)):
        s = alpha * s + B * u_t            # s_t = alpha * s_{t-1} + B * u_t
        y.append(float(C @ s))             # y_t = C^T s_t
    return np.array(y)

# Verify duality on a small example.
rng = np.random.default_rng(42)
L, n_dim = 8, 4
B_vec = rng.normal(0, 1, n_dim)
C_vec = rng.normal(0, 1, n_dim)
u_seq = rng.normal(0, 1, L)
alpha = 0.85

# For the parallel form, Q = C (broadcast to all positions), K = B (broadcast).
Q_mat = np.tile(C_vec, (L, 1))            # (L, n): same C at every query position
K_mat = np.tile(B_vec, (L, 1))            # (L, n): same B at every key position

y_par = ssd_parallel(Q_mat, K_mat, u_seq, alpha)
y_seq = ssd_sequential(B_vec, C_vec, u_seq, alpha)
print("parallel  y[-1]  =", round(float(y_par[-1]), 6))
print("sequential y[-1] =", round(float(y_seq[-1]), 6))
print("SSD duality verified:", np.allclose(y_par, y_seq))

parallel  y[-1]  = -1.847312
sequential y[-1] = -1.847312
SSD duality verified: True

from mamba_ssm import Mamba2   # pip install mamba-ssm>=2.0
import torch

# Mamba-2 layer: d_model channels, d_state is the per-head state size (P in the paper),
# headdim is the head dimension P, ngroups controls multi-head grouping.
layer = Mamba2(d_model=512, d_state=64, headdim=64, ngroups=1)

# Training: parallel matrix form (flash-style chunked computation, O(L) memory).
x_train = torch.randn(4, 2048, 512)      # (batch, length, channels)
y_train = layer(x_train)                 # SSD parallel form internally
print("train output shape:", y_train.shape)   # torch.Size([4, 2048, 512])

# Inference (autoregressive): sequential recurrent form, O(1) per step, O(d_state) memory.
# Pass inference_params with a pre-allocated state cache to run in recurrent mode.
from mamba_ssm.utils.generation import InferenceParams
inference_params = InferenceParams(max_seqlen=2048, max_batch_size=1)
x_step = torch.randn(1, 1, 512)          # single new token
y_step = layer(x_step, inference_params=inference_params)  # recurrent step
print("step   output shape:", y_step.shape)   # torch.Size([1, 1, 512])