RetNet Attention

Multi-Scale Retention mechanism (RetNet).

Implements the parallel representation of RetNet’s multi-scale retention, which replaces softmax attention with an exponential decay matrix D. Supports three computational paradigms: parallel (training), recurrent (O(1) inference), and chunkwise (hybrid). This implementation uses the parallel formulation with per-head decay rates (gamma) and a Swish-gated output pathway.

Reference:

Sun et al. (2023), “Retentive Network: A Successor to Transformer for Large Language Models”, arXiv:2307.08621.

class src.model.attention.retnet.MultiScaleRetention(*args: Any, **kwargs: Any)

Bases: Module

Multi-scale retention with exponential decay (Sun et al. 2023).

Projects the input into query, key, value, and gate tensors. Computes scaled dot-product scores, then multiplies element-wise by a per-head exponential decay matrix D where D_{nm} = gamma^{|n-m|} for encoder mode (or gamma^{n-m} for decoder mode with causal masking). The decay rates gamma are log-uniformly spaced across heads, providing multi-scale temporal receptive fields. A Swish-gated pathway modulates the output.

Complexity:

Training: O(n^2 * d) parallel. Inference: O(1) recurrent, no KV cache.

Reference:

Sun et al. (2023), “Retentive Network: A Successor to Transformer for Large Language Models”, arXiv:2307.08621.

Parameters:

config – Model configuration object with attributes hidden_size, retention_heads, dropout, use_bitnet, norm_type, and optionally mode ("encoder" or "decoder").

dim

Dimensionality of the input and output embeddings.

heads

Number of retention heads.

head_dim

Dimensionality of each retention head (dim // heads).

scale

Scaling factor 1 / sqrt(head_dim) applied to dot products.

q_proj

Linear (or BitLinear) projection for queries.

k_proj

Linear (or BitLinear) projection for keys.

v_proj

Linear (or BitLinear) projection for values.

g_proj

Linear (or BitLinear) projection for the gate pathway.

out_proj

Linear (or BitLinear) output projection.

swish

SiLU (Swish) activation for the gate pathway.

norm

Normalization layer applied after retention aggregation.

decay_mask

Precomputed per-head gamma decay rates (buffer).

mode

"encoder" for bidirectional, "decoder" for causal.

forward(x)

Compute multi-scale retention in parallel mode.

Parameters:

x – Input tensor of shape (batch_size, seq_len, hidden_size).

Returns:

Output tensor of shape (batch_size, seq_len, hidden_size).