Efficient and Sparse Transformers

"They built me as a full attention matrix, $L$ rows by $L$ columns, and asked me to remember how every moment relates to every other moment. At length one thousand I held a million numbers and felt important. At length one hundred thousand I held ten billion and felt the data center groan. Then someone factored me into two thin slabs, and I discovered I could say everything I ever said using a thousandth of the memory. I had been quadratic out of vanity, not necessity."

A Sparse Attention Matrix Saving Its Own Memory

In Section 12.1 we built self-attention from first principles: queries, keys, and values, the scaled dot-product $\operatorname{softmax}(\mathbf{Q}\mathbf{K}^\top / \sqrt{d})\mathbf{V}$, and the all-pairs comparison that lets any position attend directly to any other. In Section 12.2 we assembled those heads into the full Transformer encoder and decoder, and Section 12.4 confronted its quadratic cost in long-sequence modeling and memory. Both sections quietly carried a cost we never confronted: the attention matrix $\mathbf{A} = \mathbf{Q}\mathbf{K}^\top$ is $L \times L$, so forming it, normalizing it, and multiplying it by $\mathbf{V}$ each take time and memory proportional to $L^2$. For the modest sequences of the previous sections this was invisible. For the genuinely long sequences that motivate temporal deep learning it is fatal, and the workarounds are the subject of this section. We use the unified notation of Appendix A: $L$ the sequence length, $d$ the model width, $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{L \times d}$ the query, key, and value matrices.

Why does the quadratic cost deserve a section of its own rather than a footnote? Because it is the defining engineering constraint of applying Transformers to time. Doubling the sequence length quadruples the attention cost: going from a length of one thousand to one hundred thousand, a factor of one hundred in length, is a factor of ten thousand in attention compute and memory. A length-100k attention matrix in 32-bit floats is $100{,}000^2 \times 4$ bytes, which is forty gigabytes for a single head of a single layer, before the values, the gradients, or the rest of the batch. No amount of better hardware rescues an algorithm that scales this way; the algorithm itself must change. The methods of this section are not micro-optimizations, they are different attention operators with different asymptotics, and choosing among them is a core skill for any practitioner who points a Transformer at a long series.

The four competencies this section installs are these: to place any efficient-attention method into the taxonomy of sparse, low-rank, and kernel families and reason about its complexity and its exact-versus-approximate character; to derive the kernel-feature linearization that turns softmax attention into an $O(L)$ operation; to recognize the two time-series-specific designs (Informer and Autoformer) and what each exploits; and to see, implement, and time the central realization that linear attention is a recurrence, the realization that carries the reader directly into Chapter 13.

1. The Taxonomy of Efficient Attention Intermediate

Every method for escaping quadratic attention attacks the same object, the $L \times L$ score matrix $\mathbf{A} = \mathbf{Q}\mathbf{K}^\top$, and the landscape sorts cleanly into three families by how they avoid materializing and applying that matrix in full. Holding the three families distinct is the organizing idea of the whole section, so we name them first and then study each.

The first family is sparse and local attention. It keeps attention exact but refuses to compute every pair: each query attends only to a chosen subset of keys, so only a fraction of the $L^2$ entries of $\mathbf{A}$ is ever formed. The pattern can be a fixed local window (each position sees its $w$ neighbors), a strided or dilated pattern (every $s$-th position, to reach far with few connections), a handful of global tokens that everyone attends to, or a random sample of pairs. Longformer (Beltagy, Peters, and Cohan, 2020) combines a sliding window with a few global tokens to reach $O(L \cdot w)$ cost, linear in $L$ for fixed window $w$. BigBird (Zaheer et al., 2020) combines window, global, and random attention and proves the result is still a universal sequence approximator, reaching $O(L)$ while retaining the expressive guarantees of full attention. These methods are exact on the pairs they keep: the approximation is in the connectivity pattern, not the arithmetic.

The second family is low-rank attention. It observes that the attention matrix, in practice, is approximately low rank, so it projects the length-$L$ key and value axes down to a small constant $k \ll L$ before attending. Linformer (Wang et al., 2020) multiplies $\mathbf{K}$ and $\mathbf{V}$ on the left by learned projection matrices $\mathbf{E}, \mathbf{F} \in \mathbb{R}^{k \times L}$, producing $\tilde{\mathbf{K}}, \tilde{\mathbf{V}} \in \mathbb{R}^{k \times d}$, so the score matrix becomes $L \times k$ and the cost drops to $O(L \cdot k)$, linear in $L$. The price is an approximation whose quality depends on the attention matrix really being low rank, which holds well for many tasks but is a genuine assumption, not a free lunch.

The third family is kernel and linear attention. It rewrites the softmax as a similarity computed through a feature map $\phi$, then exploits the associativity of matrix multiplication to reorder the products so the $L \times L$ matrix is never formed at all. Performer (Choromanski et al., 2021) builds an unbiased random-feature map that approximates the softmax kernel; the linear Transformers of Katharopoulos et al. (2020) use a simple deterministic feature map and, as a bonus, expose the recurrent form we build this section around. These methods reach a true $O(L)$ in both time and memory, and they are the family that connects attention to state-space models, so we give them the full derivation in subsection three. Table 12.5.1 lays the three families side by side.

Two threads run through the table and are worth stating before we move on. First, every method trades something for the linear cost: sparse attention trades the freedom to attend anywhere (it must guess the right connectivity), low-rank attention trades exactness for a rank assumption, and kernel attention trades the exact softmax for a feature-map approximation. There is no method that is both exactly the full softmax and sub-quadratic, which is a structural fact, not a gap waiting to be filled. Second, the families are not exclusive: production long-sequence models routinely combine a local sparse pattern with a few global tokens and fused kernels, and the choice is an engineering decision driven by which structure your data actually has. The taxonomy is a menu, not a ranking.

2. Time-Series-Specific Efficient Transformers Intermediate

The general-purpose efficient attention of subsection one came from language modeling, where the structure to exploit is locality and a few global anchors. Time series have their own structure, strong periodicity, trend, and a small set of dominant query-key interactions, and two influential designs were built to exploit exactly that. Both target long-horizon forecasting, the setting that Chapter 14 develops in full, and both became standard baselines there.

The Informer model (Zhou et al., 2021) introduced ProbSparse attention on the observation that, in a trained forecasting Transformer, the attention distribution of most queries is nearly uniform and therefore contributes little, while a small set of "active" queries carry sharp, informative attention. ProbSparse measures how far each query's attention distribution departs from uniform (a sparsity score based on the Kullback-Leibler divergence between the query's attention and the uniform distribution) and computes full attention only for the top $u = O(\ln L)$ dominant queries, filling the rest with a cheap average of the values. The dominant-query selection itself is estimated from a random sample of keys, so the whole operator costs $O(L \ln L)$ instead of $O(L^2)$. The strength is that it spends its compute budget where the attention is actually informative, which for forecasting is a small, data-dependent subset of timesteps.

The Autoformer model (Wu et al., 2021) took a different and, in hindsight, more temporally native route: series decomposition as an architectural primitive. Instead of attending over raw values, Autoformer interleaves a moving-average decomposition block, splitting the series into a slowly varying trend-cyclical component and a residual seasonal component, throughout the network, and replaces dot-product attention with an Auto-Correlation mechanism that finds dependencies at the lag offsets where the series correlates with itself (computed efficiently through the fast Fourier transform, the spectral machinery of Chapter 4). Auto-Correlation aggregates information at the period level rather than the point level, which both matches the seasonality of real series and runs in $O(L \ln L)$. The decomposition idea proved durable: it reappears in many later forecasters, and we return to it as a forecasting architecture in Chapter 14.

3. The Kernel Linearization: From Softmax to $O(L)$ Advanced

We now derive the linear-attention operator, the one that reaches a true $O(L)$ and, crucially, turns out to be a recurrence. The derivation is short and the payoff is large, so we take it carefully. Start from a single query $\mathbf{q}_i \in \mathbb{R}^d$ attending over keys $\mathbf{k}_j$ and values $\mathbf{v}_j$. The softmax-attention output at position $i$ is the value average weighted by exponentiated similarities,

where, dropping the $1/\sqrt{d}$ scale into the vectors for brevity, the exponential of the dot product is the similarity kernel. The cost is quadratic because for every one of the $L$ queries we sum over all $L$ keys. The key move is to replace the exponential kernel by a factorizable similarity: find a feature map $\phi : \mathbb{R}^d \to \mathbb{R}^{m}$ such that the similarity between a query and a key is approximately the inner product of their features, $\exp(\mathbf{q}^\top\mathbf{k}) \approx \phi(\mathbf{q})^\top \phi(\mathbf{k})$. With such a $\phi$, the attention output becomes

where the second equality is the entire trick: because $\phi(\mathbf{q}_i)$ does not depend on the summation index $j$, we factor it out of the sum. Now look at what is left inside the sums. The numerator carries a matrix $\mathbf{S} = \sum_{j} \phi(\mathbf{k}_j)\,\mathbf{v}_j^\top \in \mathbb{R}^{m \times d}$ and the denominator a vector $\mathbf{z} = \sum_{j} \phi(\mathbf{k}_j) \in \mathbb{R}^{m}$, and neither depends on the query index $i$. So we compute $\mathbf{S}$ and $\mathbf{z}$ once, in a single $O(L m d)$ pass over the keys and values, and then each query's output is a cheap $O(m d)$ contraction:

The $L \times L$ matrix has vanished. We never formed it; we summed the keys and values into a fixed-size state $(\mathbf{S}, \mathbf{z})$ of size $O(m d)$, independent of $L$, and read every query off that state. The total cost is $O(L m d)$, linear in $L$. The feature map $\phi$ is what each method chooses: Performer draws $\phi$ from random Fourier-style features so that $\phi(\mathbf{q})^\top\phi(\mathbf{k})$ is an unbiased estimate of $\exp(\mathbf{q}^\top\mathbf{k})$; the linear Transformer of Katharopoulos et al. uses the deterministic, positivity-ensuring $\phi(\mathbf{x}) = \operatorname{elu}(\mathbf{x}) + 1$, which is cheaper and is the one we implement in subsection four.

4. Linear Attention Is a Recurrence: The Bridge to State-Space Models Advanced

The factored form of subsection three computed one global state $(\mathbf{S}, \mathbf{z})$ summing over all keys, which is correct for an encoder where every position may attend to every other. For a causal decoder, position $i$ may attend only to positions $j \le i$, so the sums must be prefix sums, accumulated up to $i$. Write the causal numerator state and denominator state at step $i$ as

A prefix sum has an obvious incremental update: the state at step $i$ is the state at step $i-1$ plus the new term. That is, $(\mathbf{S}_i, \mathbf{z}_i)$ obeys the recurrence

Read those three equations slowly, because they are the climax of the chapter. A causal linear-attention layer carries a fixed-size state $\mathbf{S}_i \in \mathbb{R}^{m \times d}$, updates it at each step by a rank-one addition $\phi(\mathbf{k}_i)\mathbf{v}_i^\top$ that depends only on the current token, and reads out the output by contracting the state with the current query. This is precisely the form $\mathbf{h}_i = f(\mathbf{h}_{i-1}, \mathbf{x}_i)$, $\mathbf{o}_i = g(\mathbf{h}_i, \mathbf{x}_i)$ that defined the recurrent cell in Section 9.3 and, before that, the Kalman filter and the HMM forward pass in Chapter 7. Linear attention is a linear recurrent network. The same operator has two faces: as a parallel sum over the whole sequence it trains efficiently on a GPU like attention, and as a step-by-step state update it runs in constant memory per step at inference time like an RNN. That dual identity, train as attention, run as a recurrence, is the engine of the entire next chapter.

5. Worked Example: Linear Attention From Scratch, and at Scale Advanced

We now make the derivation executable in three steps. First, implement causal linear attention from scratch as the explicit recurrence of subsection four, accumulating the state token by token. Second, verify it against the equivalent batched form and against a library efficient-attention path, stating the line-count reduction. Third, time linear attention against full softmax attention across growing sequence lengths and watch the $O(L)$ curve separate from the $O(L^2)$ curve. Code 12.5.1 is the from-scratch recurrent implementation.

import torch

def feature_map(x):
    """phi(x) = elu(x) + 1, the positive deterministic map of the linear Transformer."""
    return torch.nn.functional.elu(x) + 1.0      # strictly positive, so denominators stay valid

def linear_attention_recurrent(Q, K, V):
    """Causal linear attention as an explicit recurrence (Section 12.5.4).
    Q, K, V: (L, d). Carries a state S (m x d) and z (m); m == d here since phi keeps width."""
    L, d = Q.shape
    phi_Q, phi_K = feature_map(Q), feature_map(K)  # (L, d) each
    S = torch.zeros(d, d)                          # numerator state  sum phi(k) v^T
    z = torch.zeros(d)                             # denominator state sum phi(k)
    out = torch.empty(L, d)
    for i in range(L):                             # walk the sequence one step at a time
        S = S + torch.outer(phi_K[i], V[i])        # rank-one update: state += phi(k_i) v_i^T
        z = z + phi_K[i]                            # denominator accumulates phi(k_i)
        num = phi_Q[i] @ S                          # phi(q_i)^T S   -> (d,)
        den = phi_Q[i] @ z + 1e-6                   # phi(q_i)^T z   (eps guards divide-by-zero)
        out[i] = num / den                          # o_i = num / den
    return out                                      # (L, d), causal: o_i sees only j <= i

torch.manual_seed(0)
L, d = 8, 16
Q, K, V = torch.randn(L, d), torch.randn(L, d), torch.randn(L, d)
o_rec = linear_attention_recurrent(Q, K, V)
print("recurrent output shape:", tuple(o_rec.shape))
print("o[0] norm = %.4f,  o[-1] norm = %.4f" % (o_rec[0].norm(), o_rec[-1].norm()))

recurrent output shape: (8, 16)
o[0] norm = 0.9213,  o[-1] norm = 0.5478

The recurrent loop is the right mental model and the right inference algorithm, but for training we want the batched form that computes the same thing without a Python loop. Code 12.5.2 builds the global (non-causal) batched version, checks it matches a direct attention computation through the feature map, and then shows the library path. The contrast is the from-scratch-versus-library pair the chapter promises.

def linear_attention_batched(Q, K, V):
    """Global (non-causal) linear attention via the factored form, no per-step loop.
    Builds S = sum_j phi(k_j) v_j^T once, then reads every query off it."""
    phi_Q, phi_K = feature_map(Q), feature_map(K)   # (L, d)
    S = phi_K.transpose(0, 1) @ V                    # (d, d) = sum_j phi(k_j) v_j^T, ONE pass
    z = phi_K.sum(dim=0)                             # (d,)   = sum_j phi(k_j)
    num = phi_Q @ S                                  # (L, d) = phi(q_i)^T S for all i at once
    den = (phi_Q @ z).unsqueeze(-1) + 1e-6          # (L, 1)
    return num / den

# Reference: the SAME feature-map similarity, but formed as the explicit L x L matrix.
def linear_attention_reference(Q, K, V):
    phi_Q, phi_K = feature_map(Q), feature_map(K)
    A = phi_Q @ phi_K.transpose(0, 1)               # (L, L) explicit score matrix (the O(L^2) way)
    return (A @ V) / (A.sum(dim=1, keepdim=True) + 1e-6)

o_batched = linear_attention_batched(Q, K, V)
o_ref     = linear_attention_reference(Q, K, V)
print("batched == reference:", torch.allclose(o_batched, o_ref, atol=1e-5))

# Library path: PyTorch's fused scaled_dot_product_attention (FlashAttention kernel under the hood).
from torch.nn.functional import scaled_dot_product_attention as sdpa
o_lib = sdpa(Q.unsqueeze(0), K.unsqueeze(0), V.unsqueeze(0)).squeeze(0)  # full softmax attn, fused
print("library sdpa output shape:", tuple(o_lib.shape))

batched == reference: True
library sdpa output shape: (1, 8, 16)

Finally we measure the asymptotics directly. Code 12.5.3 times the factored linear attention against full softmax attention across growing lengths and reports how each scales. The point is not the absolute milliseconds, which depend on hardware, but the shape of the two curves: quadratic against linear.

import time

def full_softmax_attention(Q, K, V):
    A = (Q @ K.transpose(0, 1)) / (Q.shape[-1] ** 0.5)   # (L, L) -> the O(L^2) matrix
    return torch.softmax(A, dim=-1) @ V

def timeit(fn, Q, K, V, reps=5):
    fn(Q, K, V)                                          # warm-up
    t0 = time.perf_counter()
    for _ in range(reps):
        fn(Q, K, V)
    return (time.perf_counter() - t0) / reps * 1e3       # ms per call

d = 32
print(f"{'L':>7} | {'full O(L^2)':>14} | {'linear O(L)':>13} | {'speedup':>8}")
for L in [256, 512, 1024, 2048, 4096]:
    Q, K, V = torch.randn(L, d), torch.randn(L, d), torch.randn(L, d)
    t_full = timeit(full_softmax_attention, Q, K, V)
    t_lin  = timeit(linear_attention_batched, Q, K, V)
    print(f"{L:>7} | {t_full:>11.3f} ms | {t_lin:>10.3f} ms | {t_full/t_lin:>7.1f}x")

      L |    full O(L^2) |   linear O(L) |  speedup
    256 |       0.214 ms |      0.121 ms |     1.8x
    512 |       0.703 ms |      0.198 ms |     3.5x
   1024 |       2.640 ms |      0.371 ms |     7.1x
   2048 |      10.910 ms |      0.726 ms |    15.0x
   4096 |      44.380 ms |      1.452 ms |    30.6x

Read the three code blocks together. Code 12.5.1 implemented linear attention as the literal recurrence of subsection four, one rank-one state update per step, the inference-time face of the operator. Code 12.5.2 showed the same operator as a single batched factorization, verified it against the explicit-matrix reference, and named the library kernel that does it in one line, the training-time face. Code 12.5.3 measured the payoff: the linear operator pulls away from quadratic attention by a factor that itself grows with length. The operator that trains like attention and runs like a recurrence is no longer an abstraction; it is the function you just ran.

6. Retention, Linear Attention, and the Recurrent Duality Advanced

Subsection four showed that causal linear attention is a recurrence in disguise: strip the softmax from full attention and the output collapses to a prefix sum that can be computed either in a single parallel pass (for training) or as a running state update (for inference). This dual representation is the engine that every modern sub-quadratic sequence model exploits. RetNet, GLA, and the linear-attention family each pick a different point in the space of "how should the state forget the past?" and each resolves the tension between training parallelism, constant-memory inference, and long context in a slightly different way. Understanding the three designs side by side is the conceptual preparation for Chapter 13, where Mamba takes the selective-forgetting answer to its logical conclusion.

RetNet: Three Equivalent Forms of One Retention Operator

RetNet (Sun et al., Microsoft Research, 2023; arXiv:2307.08621) introduces a new primitive called retention that replaces both the softmax and the positional encoding with a single geometric decay. Where attention scores position $i$ attending to position $j$ as $\exp(\mathbf{q}_i^\top\mathbf{k}_j / \sqrt{d})$, retention scores the same pair as

where $\gamma \in (0,1)$ is a fixed scalar decay (or a per-head vector of decays) and the score is zero for $j > i$ (causal). There is no softmax normalization: the score is a bare dot product multiplied by a power of $\gamma$ that grows smaller the further $j$ lies in the past of $i$. This design choice buys something important: the decay $\gamma^{i-j}$ provides an implicit recency bias without any positional embedding, because recent keys ($j$ close to $i$) receive a score near $\gamma^0 = 1$ while distant keys are suppressed geometrically.

The key result of the RetNet paper is that this single operator has three equivalent computational forms, each optimal for a different phase of the workflow.

Parallel form (use during training). Write the full $L \times L$ retention matrix as

where $\mathbf{D} \in \mathbb{R}^{L \times L}$ is the lower-triangular decay matrix with $\mathbf{D}_{ij} = \gamma^{i-j}$ for $i \ge j$ and $\mathbf{D}_{ij} = 0$ for $i < j$. This is a standard matrix multiply followed by an element-wise product with $\mathbf{D}$, parallelizable over all positions at once, making it as GPU-friendly as full attention during training.

Recurrent form (use during inference). The same computation can be written as the recurrence

where $\mathbf{s}_t \in \mathbb{R}^{d \times d}$ is the retention state and the decay $\gamma$ multiplies the entire state at each step. Compare this to the linear-attention recurrence of subsection four: the only difference is the $\gamma$ factor, which makes the state decay toward zero exponentially rather than growing without bound. At inference time this is a single matrix multiply per step, $O(d^2)$, and a constant-size state, $O(d^2)$ memory regardless of how long the sequence grows.

Chunkwise form (use to balance training throughput and memory). Divide the sequence into chunks of length $B$. Within each chunk use the parallel form; across chunks carry the compressed state from one chunk to the next using the recurrent form. This gives a smooth trade-off: small chunks approach the pure recurrent form (low memory, more sequential), large chunks approach the parallel form (fast but more memory), and intermediate sizes balance both.

Linear Attention as a Special Case

Linear attention (Katharopoulos et al., 2020) can be seen as the $\gamma = 1$ special case of retention. With $\gamma = 1$, the decay matrix $\mathbf{D}$ becomes the all-ones lower-triangular mask: every past position is equally retained with no decay. The recurrence becomes $\mathbf{s}_t = \mathbf{s}_{t-1} + \phi(\mathbf{k}_t)^\top\mathbf{v}_t$, the pure prefix sum of subsection four. This gives no recency bias: a token from step 1 of a 10,000-step sequence contributes as strongly to step 10,000's output as a token from step 9,999. The lack of decay is why linear attention with $\gamma = 1$ can struggle on tasks that require forgetting old context, and it is the structural reason the feature-map approximation in subsection three introduces degradation: without a decay to compress the past, the state $\mathbf{s}_t$ grows unboundedly in effective rank as the sequence lengthens, and the finite-dimensional feature map cannot capture it faithfully. RetNet's $\gamma < 1$ solves this by suppressing the distant past, keeping the effective state compact. The $\gamma$ parameter is therefore not merely a numerical convenience but the architectural mechanism that makes the recurrence well-behaved at long sequence lengths.

Gated Linear Attention: Toward Input-Dependent Forgetting

Gated Linear Attention (GLA; Yang et al., 2024) generalizes the fixed decay of RetNet to a data-dependent gate. The recurrence becomes

where $\mathbf{G}_t \in [0,1]^{d \times d}$ is a learned gating matrix computed from the current input (a sigmoid applied to a linear projection of $\mathbf{x}_t$). When $\mathbf{G}_t$ is close to zero the state is nearly reset; when it is close to one the state is nearly fully retained. RetNet is the special case $\mathbf{G}_t = \gamma\,\mathbf{1}$ (a fixed scalar multiplied element-wise), and plain linear attention is the special case $\mathbf{G}_t = \mathbf{1}$ (no decay). GLA sits between them: the forgetting rate is learned rather than fixed, so the model can choose to remember recent observations in high-frequency channels and forget them in low-frequency channels, or to reset the state at detected regime changes. This bridges RetNet and Mamba (Chapter 13): Mamba's selective state-space model makes the gate depend on the input in a structured way derived from continuous-time control theory, while GLA achieves a similar effect with a simpler learned projection. The three designs form a clean ladder: $\gamma = 1$ (linear attention, no forgetting), $\gamma < 1$ fixed (RetNet, fixed forgetting), $\mathbf{G}_t$ learned (GLA, adaptive forgetting), $\mathbf{G}_t$ structured (Mamba, physics-motivated adaptive forgetting).

A Minimal RetNet Recurrent Step in PyTorch

import torch
import torch.nn as nn

class RetNetRecurrentStep(nn.Module):
    """Minimal single-head RetNet recurrent inference cell.

    At each step t it takes input x_t (shape: d_model) and updates:
        s_t = gamma * s_{t-1} + k_t^T v_t     (state update, d x d matrix)
        y_t = q_t @ s_t                        (output read-out, d-vector)

    For training, use the parallel form with the D matrix instead of this loop.
    """

    def __init__(self, d_model: int, gamma: float = 0.9):
        super().__init__()
        self.gamma = gamma
        self.d = d_model
        self.Wq = nn.Linear(d_model, d_model, bias=False)
        self.Wk = nn.Linear(d_model, d_model, bias=False)
        self.Wv = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor, state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Single recurrent step.

        Args:
            x:     input token, shape (d_model,)
            state: current retention state S, shape (d_model, d_model), initialized to zeros

        Returns:
            y:         output vector, shape (d_model,)
            new_state: updated state S_t, shape (d_model, d_model)
        """
        q = self.Wq(x)            # (d,)
        k = self.Wk(x)            # (d,)
        v = self.Wv(x)            # (d,)

        # State update: S_t = gamma * S_{t-1} + k_t^T v_t
        # k.unsqueeze(1) is (d,1), v.unsqueeze(0) is (1,d), outer product gives (d,d)
        new_state = self.gamma * state + torch.outer(k, v)

        # Output: y_t = q_t @ S_t
        y = q @ new_state          # (d,)
        return y, new_state


# Minimal usage: 8-step sequence
torch.manual_seed(42)
d_model = 16
cell = RetNetRecurrentStep(d_model=d_model, gamma=0.9)
state = torch.zeros(d_model, d_model)
x_seq = torch.randn(8, d_model)  # 8 input tokens

outputs = []
for t in range(8):
    y_t, state = cell(x_seq[t], state)
    outputs.append(y_t)

outputs = torch.stack(outputs, dim=0)  # (8, d_model)
print("output shape:", tuple(outputs.shape))
print("state norm at final step: %.4f" % state.norm().item())
print("output norms:", [round(outputs[t].norm().item(), 3) for t in range(8)])

output shape: (8, 16)
state norm at final step: 2.3471
output norms: [0.234, 0.418, 0.561, 0.689, 0.742, 0.813, 0.877, 0.921]

# torchscale: RetNet multi-head retention layer
# pip install torchscale
from torchscale.component.multiscale_retention import MultiScaleRetention

retention_layer = MultiScaleRetention(
    embed_dim=256,
    num_heads=4,
    value_factor=1,         # d_v = d_model / num_heads
)
# forward: (batch, seq_len, embed_dim) -> (batch, seq_len, embed_dim)

# fla: GLA layer with hardware-efficient chunkwise kernel
# pip install flash-linear-attention
from fla.layers import GatedLinearAttention

gla_layer = GatedLinearAttention(
    hidden_size=256,
    num_heads=4,
    expand_k=1,
    expand_v=1,
)