RetNet

Self-attention made Transformers the dominant architecture for sequence modeling, but it carries a well-known inference penalty. During training, the attention mechanism can be computed in parallel across all positions, making efficient use of GPU throughput. During autoregressive generation, however, the model must cache historical keys and values—the KV Cache—for every previously generated token. As the sequence grows, both GPU memory consumption and per-step computation increase linearly, turning long-context inference into a serious bottleneck.

RetNet, short for Retentive Network, offers a different path. Rather than patching attention with approximations, it replaces the core interaction with a retention mechanism designed to satisfy three properties at once: parallelizable training, constant-size recurrent inference, and efficient long-sequence modeling. It is not merely linear attention under another name. Through a causally decaying retention operator, RetNet mathematically unifies parallel, recurrent, and chunked recurrent computations within a single framework.
Given an input sequence (X \in \mathbb{R}^{N \times d}), RetNet projects it into queries, keys, and values:

[
Q = XW_Q,\quad K = XW_K,\quad V = XW_V
]

The output of a single retention head is:

[
O = (QK^T \odot D)V
]

Here (D) is a causal decay matrix whose entries are defined by:

[
D_{nm} =
\begin{cases}
\gamma^{n-m}, & n \ge m \
0, & n < m
\end{cases}
]

The decay factor (\gamma), ranging between 0 and 1, determines how quickly older information fades. Distant positions receive exponentially smaller weights, while the causal mask prevents any token from attending to the future. Unlike standard self-attention, retention omits the softmax normalization. This structural difference is what allows the same operation to be expressed as a recurrence, opening the door to inference that does not scale with sequence length.

Parallel Representation: Training-Friendly Computation

The parallel form computes retention over the entire sequence at once, closely resembling transformer attention from an implementation perspective. It is well suited to GPU acceleration during training:

import torch

def retention_parallel(q, k, v, gamma):
    # q, k, v: [B, H, N, D]
    N = q.size(-2)
    idx = torch.arange(N, device=q.device)
    diff = idx[:, None] - idx[None, :]  # n - m
    D = torch.where(
        diff >= 0,
        gamma ** diff,
        torch.zeros_like(diff, dtype=q.dtype)
    )
    scores = torch.matmul(q, k.transpose(-1, -2)) * D
    out = torch.matmul(scores, v)
    return out

The naive parallel implementation costs (O(N^2d)) time, which is comparable to standard attention. The advantage appears when this form is combined with chunked processing: memory pressure can be capped without sacrificing the parallelism that modern hardware expects.

Recurrent Representation: Constant-Time Inference

Retention can also be written recursively. Define a hidden state (S_n) that summarizes everything seen up to position (n):

[
S_n = \gamma S_{n-1} + K_n^T V_n
]

The output at position (n) is then:

[
O_n = Q_n S_n
]

Because (S_n) has shape (d_k \times d_v), its size is independent of the sequence length. At inference time, the model performs a single state update per generated token and never needs to revisit earlier keys or values:

def retention_recurrent(q, k, v, gamma):
    # q, k, v: [B, H, N, D]
    B, H, N, D = q.shape
    S = torch.zeros(B, H, D, D, device=q.device)
    outs = []
    for n in range(N):
        S = gamma * S + k[:, :, n].unsqueeze(-1) * v[:, :, n].unsqueeze(-2)
        o = torch.einsum("bhd,bhdv->bhv", q[:, :, n], S)
        outs.append(o)
    return torch.stack(outs, dim=2)

Consequently, RetNet inference incurs constant memory and compute per step. There is no growing KV Cache, and the cost of generating the thousandth token is essentially the same as generating the first.

Chunked Recurrent Representation: A Practical Middle Ground

Pure parallel computation is fast on GPUs but memory-hungry for long inputs; pure recurrence is memory-light but underutilizes parallel hardware. RetNet resolves this tension with a chunked recurrent representation. The sequence is divided into blocks of size (B). Within each block, retention is computed in parallel. Between blocks, the recurrent state carries historical information forward.

This arrangement reduces complexity to roughly (O(NB)) while keeping memory usage manageable. Common block sizes include 128, 256, and 512, though the exact choice depends on available hardware and target sequence length. During training, the chunked form delivers high throughput; during inference, the model can switch to the pure recurrent form for minimal per-step overhead.

Multi-Scale Retention and Network Architecture

RetNet extends single-head retention to Multi-Scale Retention (MSR), where different heads use different decay factors. For example:

[
\gamma_h = 1 - 2^{-5-h}
]

Heads with fast decay specialize in short-range dependencies, while heads with slow decay retain broader context. The outputs of all heads are concatenated and passed through a linear projection.

A RetNet layer looks familiar to transformer practitioners. It typically contains:

• Multi-Scale Retention
• A position-wise feed-forward network (FFN)
• Pre-LayerNorm and residual connections
• Optionally, Group