What RetNet Is

Retention — recurrence with exponential decay

RetNet (Sun et al., July 2023, arXiv:2307.08621, Microsoft Research Asia) introduced the retention mechanism as a unifying primitive that can be viewed as either recurrence or as a constrained kind of attention. The core update is delightfully simple: s_n = γ · s_{n-1} + K_n^T · V_n, with output o_n = Q_n · s_n. The state s is exactly like the SSM hidden state — fixed size, summarizing all past tokens.

The decay γ is the design choice. RetNet sets γ to be fixed per attention head, with different heads getting different γ values. Heads with γ ≈ 1 maintain long memory (slow decay). Heads with γ ≈ 0 maintain short memory (fast decay). This naturally creates a multi-scale memory hierarchy: some heads pay attention to the recent past, others to the deep past, all in parallel.

Positional encoding: xPos

RetNet uses xPos for positional encoding — a complex-exponential rotation closely related to RoPE but designed for the retention mechanism's recurrent form. xPos buys translation invariance (the model doesn't care about absolute positions, only relative offsets) while remaining compatible with the recurrent computation. This is the kind of detail that doesn't matter much at the conceptual level but matters a lot at the implementation level.

Why this is interesting

The mathematical neatness is the headline. Retention is recurrence (you can write it as a per-token state update) and it is also attention (the recurrent form unrolls into a particular kind of attention matrix with exponential decay). This unification — that recurrence and attention aren't fundamentally different, but points on a continuum — is RetNet's lasting intellectual contribution. Every architecture that came after on the linear-attention/SSM side absorbed some version of this insight.

Code

RetNet retention update·python

# Per timestep n, given input x_n:
# Q_n, K_n, V_n produced from x_n by linear projections
# gamma is FIXED per-head (no input dependence — the key constraint)

# Recurrent form (constant memory at inference)
s_n = gamma * s_prev + K_n.transpose(-1, -2) @ V_n
o_n = Q_n @ s_n

# Parallel form (training): equivalent to attention with exponential-decay mask
# attn_mask[i, j] = gamma ** (i - j) for j <= i, else 0

Exercise

Implement the retention operation in PyTorch in both its parallel and recurrent forms, and verify they produce identical outputs on a small (1, 64, 8, 16) Q/K/V/gamma input. The fact that the same parameters yield the same output under two different evaluation patterns is the core property that gives RetNet its production flexibility — and that property doesn't carry over to architectures that don't have this dual structure.