RetNet (the Retentive Network, Sun et al., Microsoft, 2023) asks a pointed question: can one mechanism train in parallel like a Transformer and generate tokens at constant cost per step like an RNN, without giving up quality? Its answer is retention — a decay-weighted operation that replaces the softmax in attention with a fixed exponential decay along positions. The trick is that retention has three mathematically identical forms: a parallel form for training, a recurrent form for cheap inference, and a chunkwise-recurrent form for long sequences. Same numbers out, three different compute schedules in. This piece derives all three, shows why they are equal, works through the decay factor γ numerically, and places retention next to its close cousins — linear attention and state-space models — so you can see exactly what it borrows and what it changes.
From softmax attention to retention
Standard attention computes Attention(X) = softmax(QK^T / √d_k) V. The softmax does two jobs: it makes the weights non-negative and it normalizes each row to sum to 1. It is also what blocks a cheap recurrent form — because exp does not factor across positions, you cannot accumulate a running state, so decoding needs the full growing KV cache and costs O(N) per token.
Retention throws the softmax out. In its place it uses a deterministic decay by relative position: a token m steps in the past is weighted by γ^m for a fixed scalar 0 < γ < 1. The core operation becomes a plain bilinear form Q_n K_m^T scaled by γ^(n-m) and summed over the past. Because that decay does factor — γ^(n-m) = γ^n / γ^m — the same computation can be rolled into a constant-size recurrent state. Positional information rides along via an xPos/rotary-style complex rotation on Q and K, so retention encodes relative position in both a phase (the rotation) and a magnitude (the γ decay).
The recurrent form: a constant-size state
The recurrent form is where retention earns its inference story. Instead of keeping every past key and value, it keeps a single state matrix S_n of fixed shape [d_k, d_v] that summarizes the entire history:
State: S_n = γ · S_(n-1) + K_n^T V_n S_n : [d_k, d_v], S_0 = 0
Output: o_n = Q_n · S_n o_n : [1, d_v]
Q_n, K_n, V_n : [1, d] (row vectors for position n)
K_n^T V_n : [d_k, d_v] rank-1 outer-product updateEach step does one outer product to update the state and one vector–matrix product to read it out. The old state is simply faded by γ before the new token is added. Crucially, S_n never grows: its size is d_k × d_v regardless of how many tokens have gone by. So generation costs O(1) time and O(1) memory per token — no KV cache that swells with context. This is the RNN-style inference profile that Transformers cannot match.
The parallel form: a decay mask instead of softmax
The recurrent form is sequential, which is death for training throughput. So retention has an equivalent parallel form that processes the whole sequence at once, just like attention — but with the softmax replaced by an elementwise decay mask D:
Retention(X) = (Q K^T ⊙ D) V
D_(nm) = γ^(n-m) if n ≥ m (causal + exponential decay)
= 0 if n < m (no peeking at the future)
Q, K, V = X W_Q, X W_K, X W_V Q,K,V : [N, d]
QK^T : [N, N] ⊙ D (elementwise) → × V → [N, d_v]D is a single lower-triangular matrix that folds two things together: causal masking (zeros above the diagonal) and the decay (γ^(n-m) below it). There is no row-wise softmax — the weights are fixed by position, not learned per query, though a GroupNorm on the output plus a swish gate restores stable scale. This form is fully parallel across positions and maps onto the same dense matmul hardware attention already uses, so training runs at Transformer speed.
Why the two forms are the same computation
The magic is that the parallel and recurrent forms are not approximations of each other — they are algebraically identical. Unroll the recurrence from S_0 = 0:
S_n = Σ_(m=1..n) γ^(n-m) K_m^T V_m
o_n = Q_n S_n = Σ_(m=1..n) γ^(n-m) (Q_n K_m^T) V_m
= Σ_(m=1..n) (Q_n K_m^T · D_(nm)) V_mThat last line is precisely row n of (QK^T ⊙ D)V. The factorization γ^(n-m) = γ^n · γ^(-m) is what lets the double sum collapse into a single running state: the γ^n pulls out as the fade applied to the whole accumulated state each step. This is the same identity that powers linear attention; retention just adds the decay. The upshot is a rare luxury — you train with the parallel form and, with the same learned weights, serve with the recurrent form. No distillation, no re-training, no mismatch. One model, two schedules.
The chunkwise-recurrent form: parallel within, recurrent across
The parallel form is O(N^2) — fine for training on moderate context, painful for very long sequences. The pure recurrent form is O(N) but sequential, wasting the parallel hardware. The chunkwise-recurrent form is the hybrid: split the sequence into chunks of size B, run the parallel form inside each chunk, and carry a recurrent state between chunks.
For chunk i (inner positions j = 1..B):
Inner (parallel, in-chunk): I_i = (Q_i K_i^T ⊙ D) V_i
Cross (recurrent carry): C_i = (Q_i R_(i-1)) ⊙ ξ, ξ_j = γ^j
State update: R_i = γ^B R_(i-1) + K_i^T (V_i ⊙ ζ), ζ_j = γ^(B-j)
Retention(X_i) = I_i + C_iInside a chunk you pay the quadratic cost, but only O(B^2); across chunks you pay a single state carry. Total cost is O(NBd + Nd^2) — linear in sequence length for a fixed chunk size, while still doing most work as big parallel matmuls. This is the form you reach for to train or prefill on 100K-token context without the N^2 memory wall.
The decay factor γ and multi-scale retention
γ is the whole personality of a retention head. It sets an effective memory horizon: the weight on a token k steps back is γ^k, so influence decays geometrically. Work it numerically for γ = 0.9:
distance k : 0 1 5 10 20 50
γ^k (0.9) : 1.00 0.90 0.59 0.35 0.12 0.005A token 50 back contributes essentially nothing. A useful rule of thumb: the effective window is about 1/(1-γ) — roughly 10 tokens for γ=0.9, 100 for γ=0.99. A single fixed decay would be a straitjacket, so RetNet uses multi-scale retention (MSR): each head h gets its own γ_h, spread geometrically (roughly γ_h = 1 - 2^(-5-h)). Short-γ heads capture local detail; long-γ heads (near 1) hold long-range context. The decays are fixed, not learned — a deliberate simplification that keeps the three forms clean while covering many timescales at once.
Complexity and the impossible triangle
Retention’s pitch is that it grabs all three corners of a triangle usually considered to allow only two: parallel training, cheap O(1) inference, and strong performance. Transformers own the first and third but pay O(N) per decode step; classic RNNs own the second but cannot parallelize training; older linear-attention models got the first two but lagged on quality.
| Form | Time | Per-step decode | Used for |
|---|---|---|---|
| Parallel | O(N^2 d) | — | Training |
| Recurrent | O(N d^2) | O(1) time & memory | Inference / generation |
| Chunkwise | O(NBd + N d^2) | — | Long-sequence train/prefill |
| Transformer | O(N^2 d) | O(N) (KV cache grows) | (for contrast) |
The constant-memory decode is the standout: whether you are 1K or 100K tokens deep, the retention state is the same d_k × d_v matrix, so throughput and memory stay flat with context length — the opposite of a KV cache that grows without bound.
Cousin one: linear attention
Retention is best understood as linear attention with a decay gate. Linear attention removes the softmax and applies a kernel feature map φ, exploiting associativity: φ(Q)(φ(K)^T V) can be computed as a running sum S_n = S_(n-1) + φ(K_n)^T V_n, then o_n = φ(Q_n) S_n. That is exactly retention’s recurrence — except linear attention sums with no decay (an implicit γ = 1).
Two changes turn one into the other. First, retention inserts the γ fade, giving a stable, bounded state with a built-in recency bias instead of an ever-growing sum that dilutes as the sequence lengthens. Second, retention drops the feature map (φ = identity) and leans on the xPos rotation plus GroupNorm for stability rather than a positive kernel. That single decay term is a big part of why retention reports Transformer-level quality where vanilla linear attention historically did not — the fade keeps the state well-conditioned and gives each head a defined receptive field.
Cousin two: state-space models and Mamba
Retention’s recurrence is also a state-space model (SSM) in disguise. The canonical linear SSM is h_n = A h_(n-1) + B x_n, y_n = C h_n. Set A = γ I (a scalar diagonal decay), let B inject K_n^T V_n and C read out via Q_n, and you have retention. SSMs like S4 instead learn a structured A (HiPPO-initialized) for principled long-range memory, and parallelize training as a long convolution (FFT) rather than a masked matmul.
The sharpest contrast is with Mamba. Retention’s decay is data-independent — γ is fixed per head, the same for every token. Mamba makes its transition selective: A, B, C become functions of the input, so the model can context-dependently choose what to remember or forget. That expressivity costs the clean, attention-shaped parallel form; Mamba relies on a hardware-aware parallel scan instead. All three are the same family: linear recurrences with a decay/transition and a parallelizable schedule, differing mainly in how rich, and how input-dependent, that transition is.
Practical implications and pitfalls
For CPU and small-model serving, the constant-memory decode is the headline. No growing KV cache means the memory budget for a conversation is fixed and predictable, and per-token latency does not creep up as context fills — attractive when RAM bandwidth, not FLOPs, is the ceiling. You also get one weight set that trains in parallel and serves recurrently, so long-context prefill (chunkwise) and streaming generation (recurrent) reuse the same parameters.
The pitfalls are real, though. The fixed γ decay imposes a structural recency bias: exact recall of a single token far in the past — the needle-in-a-haystack task attention aces — is harder when its contribution has faded by γ^k, which multi-scale heads mitigate but do not erase. The state S_n is a lossy d_k × d_v summary, not a perfect transcript, so it trades unbounded precise lookup for bounded cost. And the ecosystem is young: the kernels and scaling recipes that surround attention are still maturing for retention. Treat it as a compelling architecture with a clean cost story — and verify recall-heavy behavior on your own workload first.
γ^(n-m) along relative position, and that one substitution unlocks three equivalent forms of the same computation: a parallel form (QK^T ⊙ D)V for fast training, a recurrent form S_n = γS_(n-1) + K_n^T V_n, o_n = Q_n S_n for O(1)-per-step, constant-memory inference, and a chunkwise-recurrent form that makes long sequences linear in length. Because the forms are algebraically identical, you train one way and serve another with the very same weights. That is the impossible triangle — parallel training, cheap inference, and strong quality at once. Read retention as linear attention plus a decay gate, and as a scalar-diagonal cousin of state-space models like S4 and Mamba, which trade its fixed decay for richer, input-dependent transitions. The catch: a fixed decay builds in a recency bias, so verify long-range exact recall before you rely on it.