{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Module 22: Linear & Recurrent Attention" ], "id": "9c76154a-fd0d-4315-a931-f16b84964885" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [], "id": "d5c0aa48-4311-4da3-ad5c-70ddae3fbc83" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "0166abf7-97d6-4645-b07a-219fd5e23867" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction\n", "\n", "Two modules ago (m19) you built the leading *alternative* to attention —\n", "the state-space model, a linear recurrence over a fixed-size state. This\n", "module closes the loop by showing that the “alternative” was hiding\n", "inside attention all along.\n", "\n", "**Linear attention** is self-attention with one thing removed: the\n", "softmax. That single deletion turns attention from an $O(L^2)$ all-pairs\n", "comparison into a **linear RNN** — a recurrence over a fixed-size\n", "*matrix* state, exactly the shape of the SSM. The slogan from the paper\n", "that started this line of work says it plainly: *transformers are RNNs*.\n", "\n", "From there, one small addition — a **decay** on the state — gives\n", "**retention** (the RetNet architecture), which comes with three\n", "provably-identical faces: a parallel form for training, a recurrent form\n", "for cheap decoding, and a **chunkwise** form that gets both at once. And\n", "at the end, the payoff: linear attention, retention, and the SSMs of m19\n", "are literally the *same recurrence*, differing only in one gate. Knowing\n", "that one recurrence is how you read the whole zoo of sub-quadratic\n", "sequence models.\n", "\n", "Why it matters for LLMs:\n", "\n", "- **Constant-memory decoding.** Softmax attention’s KV cache (m09) grows\n", " with every token; linear attention keeps a fixed-size state, so each\n", " generated token costs the same no matter how long the context.\n", "- **Linear-time training and inference.** No $L \\times L$ score matrix —\n", " compute scales with $L$, not $L^2$.\n", "- **One mental model for many architectures.** RetNet, RWKV, gated\n", " linear attention, and Mamba are one idea with different gates. This\n", " module gives you that idea.\n", "\n", "### What You’ll Learn\n", "\n", "After this module, you can:\n", "\n", "- Explain why the softmax is the *only* thing forcing attention to be\n", " $O(L^2)$, and how a **kernel feature map** $\\varphi$ removes it.\n", "- Run **linear attention** as a masked matmul (parallel) *and* as a\n", " running-state recurrence, and see they are the same function.\n", "- Build **retention** (RetNet): a decayed state with a $\\gamma^{n-m}$\n", " decay mask, in its **parallel, recurrent, and chunkwise** forms — all\n", " provably equal.\n", "- Write the **one unifying recurrence**\n", " $S_t = A_t S_{t-1} + k_t^\\top v_t$ that specializes to linear\n", " attention ($A_t=1$), retention ($A_t=\\gamma$), and a diagonal SSM /\n", " gated linear attention ($A_t$ input-dependent).\n", "- Quantify the win: linear vs quadratic compute, constant vs growing\n", " decode memory.\n", "\n", "### Prerequisites\n", "\n", "This module requires familiarity with:\n", "\n", "- [Module 05: Attention](../m05_attention/lesson.qmd) — softmax\n", " self-attention, the mechanism we strip the softmax from.\n", "- [Module 09: Efficient\n", " Attention](../m09_efficient_attention/lesson.qmd) — the KV cache whose\n", " growth linear attention removes.\n", "- [Module 19: State-Space Models](../m19_state_space/lesson.qmd) — the\n", " fixed-state recurrence this module reveals inside attention.\n", "\n", "## Intuition: Attention Without the Softmax\n", "\n", "Recall causal self-attention for query position $i$:\n", "\n", "$$y_i = \\frac{\\sum_{j \\le i} \\exp(q_i \\cdot k_j)\\, v_j}{\\sum_{j \\le i} \\exp(q_i \\cdot k_j)} .$$\n", "\n", "The trouble is the $\\exp(q_i \\cdot k_j)$: it entangles $q_i$ and $k_j$\n", "inside one nonlinearity, so you cannot separate them. You are forced to\n", "compute a score for *every* pair $(i, j)$ — the $L \\times L$ matrix —\n", "before you can sum. That matrix is the whole $O(L^2)$ cost.\n", "\n", "Now suppose the similarity **factored** — suppose we could write\n", "$\\text{sim}(q_i, k_j) = \\varphi(q_i) \\cdot \\varphi(k_j)$ for some\n", "feature map $\\varphi$. Then the numerator regroups by associativity:\n", "\n", "$$\\sum_{j \\le i} \\big(\\varphi(q_i) \\cdot \\varphi(k_j)\\big)\\, v_j\n", "= \\varphi(q_i) \\underbrace{\\sum_{j \\le i} \\varphi(k_j)\\, v_j^\\top}_{S_i} .$$\n", "\n", "The sum $S_i$ no longer depends on $i$ except through its upper limit —\n", "it is a **running total** you can carry forward one token at a time. No\n", "$L \\times L$ matrix ever forms. The same regrouping is the difference\n", "between multiplying $(QK^\\top)V$ (build the big matrix first) and\n", "$Q(K^\\top V)$ (build a small state first). Step through both association\n", "orders and watch the intermediate shape:" ], "id": "638ce740-e430-42ec-b1e1-3b5989c57254" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "7e073bd1-77c2-4642-8ed6-6f748068a386" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "ea3c321d-94cf-433b-b626-b179bcbe643b" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> The softmax is the *only* thing that forces the $L \\times L$ matrix.\n", "> Remove it and replace the pairwise score with a factored kernel\n", "> $\\varphi(q)\\cdot\\varphi(k)$, and the associativity of matrix\n", "> multiplication lets you keep a fixed-size state instead — the same\n", "> move that makes an SSM linear-time.\n", "\n", "## The Math: Linear Attention\n", "\n", "We need a $\\varphi$ that keeps similarities **non-negative** (so the\n", "normalizer never turns negative) and is cheap. Katharopoulos et\n", "al. (2020) use\n", "\n", "$$\\varphi(x) = \\operatorname{elu}(x) + 1 ,$$\n", "\n", "which is $x+1$ for $x>0$ and $\\exp(x)\\in(0,1]$ for $x\\le 0$ — positive\n", "for any realistic activation. With\n", "$\\text{sim}(q,k)=\\varphi(q)\\cdot\\varphi(k)$, causal attention becomes a\n", "pair of running states:\n", "\n", "$$S_i = \\sum_{j\\le i} \\varphi(k_j)\\, v_j^\\top, \\qquad\n", "Z_i = \\sum_{j\\le i} \\varphi(k_j), \\qquad\n", "y_i = \\frac{\\varphi(q_i)^\\top S_i}{\\varphi(q_i)^\\top Z_i} .$$\n", "\n", "$S_i$ is the $(m \\times d_v)$ **state matrix** — a running, weighted\n", "memory of every value seen — and $Z_i$ is the running normalizer. Both\n", "are fixed-size and updated by one addition per token.\n", "`linear_attention.py` implements exactly this:" ], "id": "c99831be-ead8-46a5-8e71-1e8a9fcfefb5" }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "phi positive: True\n", "output shape: (6, 8)" ] } ], "source": [ "import torch\n", "from linear_attention import feature_map, linear_attention_recurrent\n", "\n", "torch.manual_seed(0)\n", "Q = torch.randn(6, 8) # (length, d)\n", "K = torch.randn(6, 8)\n", "V = torch.randn(6, 8) # (length, d_v)\n", "\n", "y = linear_attention_recurrent(Q, K, V)\n", "print(\"phi positive:\", bool((feature_map(K) > 0).all()))\n", "print(\"output shape:\", tuple(y.shape)) # (length, d_v), no L×L matrix built" ], "id": "7336a451" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Code: The Two Faces of Linear Attention\n", "\n", "Just like the SSM in m19 had a recurrent face and a convolutional face,\n", "linear attention has a **recurrent** face (the running state above,\n", "$O(L)$, cheap to decode) and a **parallel** face that keeps the familiar\n", "attention layout — build the masked score matrix\n", "$A_{ij}=\\varphi(q_i)\\cdot\\varphi(k_j)$ for $j\\le i$, normalize each row,\n", "and read out $V$:\n", "\n", "$$y_i = \\frac{\\sum_{j\\le i} A_{ij}\\, v_j}{\\sum_{j\\le i} A_{ij}} .$$\n", "\n", "The parallel face is one masked matmul (great for training on a GPU);\n", "the recurrent face never forms a matrix (great for decoding). For the\n", "same inputs they compute **bit-for-bit the same function** — the\n", "linear-attention analogue of m19’s recurrent≡convolutional identity:" ], "id": "79b2440b-6e01-4914-bd58-acb088646b0f" }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "parallel == recurrent: True" ] } ], "source": [ "from linear_attention import linear_attention_parallel\n", "\n", "y_par = linear_attention_parallel(Q, K, V)\n", "y_rec = linear_attention_recurrent(Q, K, V)\n", "print(\"parallel == recurrent:\", torch.allclose(y_par, y_rec, atol=1e-5))" ], "id": "02fdd430" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Bridge both to a plot: the recurrent dots land exactly on the parallel\n", "line, because they are the same function computed two ways." ], "id": "e0870172-0dc4-43c0-a1bc-0534540a1532" }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [], "id": "27ca26ae" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "7aacc95f-4710-4d28-8292-bdb7f78395f0" }, { "cell_type": "markdown", "metadata": {}, "source": [ "For **decoding**, wrap the running state in a small object and fold in\n", "one token at a time — the constant-work-per-token generation that a\n", "growing KV cache cannot match. `LinearAttentionState` does this:" ], "id": "75605059-2b11-44a4-bdef-d6cbce04d0e6" }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "incremental == full recompute: True" ] } ], "source": [ "from linear_attention import LinearAttentionState\n", "\n", "state = LinearAttentionState(m=8, d_v=8)\n", "ys = torch.stack([state.step(Q[i], K[i], V[i]) for i in range(6)])\n", "print(\"incremental == full recompute:\", torch.allclose(ys, y_rec, atol=1e-5))" ], "id": "460c9234" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> Train like attention, decode like an RNN. The parallel face gives\n", "> fast, parallelizable training; the recurrent face gives\n", "> constant-memory generation. The price is a matrix state of size\n", "> $m \\times d_v$ (typically $d \\times d$) instead of a growing cache —\n", "> constant in $L$, but quadratic in the model width.\n", "\n", "## Retention: Add a Decay to the State\n", "\n", "Linear attention keeps *every* past token with equal weight — its state\n", "$S_i$ is an undamped running sum, which can saturate and blur.\n", "**Retention** (the RetNet architecture, Sun et al., 2023) adds one\n", "ingredient: a scalar **decay** $\\gamma \\in (0,1]$ that shrinks the old\n", "state before each update:\n", "\n", "$$S_n = \\gamma\\, S_{n-1} + k_n^\\top v_n, \\qquad o_n = q_n\\, S_n .$$\n", "\n", "Two things are different from linear attention. First, older tokens fade\n", "geometrically — the state prefers recent context. Second, retention\n", "drops the normalizer entirely (no softmax, no division); RetNet\n", "stabilizes the output with a GroupNorm and a swish gate instead, but the\n", "core operation is this bare decayed state. Unrolling the recurrence,\n", "$S_n = \\sum_{m\\le n} \\gamma^{\\,n-m} k_m^\\top v_m$, so the output is\n", "\n", "$$o_n = \\sum_{m\\le n} \\gamma^{\\,n-m} (q_n\\cdot k_m)\\, v_m .$$\n", "\n", "That is ordinary $QK^\\top$ scores, causally masked, **weighted by a\n", "decay $\\gamma^{n-m}$** — a fixed lower-triangular **decay matrix** $D$.\n", "Drive $\\gamma$ and watch $D$ interpolate between two familiar limits:" ], "id": "20a6437c-a9da-49b3-86f8-ec2c31e9197b" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "94685d44-04de-4858-bf42-c7ca2a898de4" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "866e9255-7493-4db3-91c7-9bf1c23f4f4a" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try It!**\n", ">\n", "> 1. **Slide γ → 1.** Every causal entry becomes 1: the decay vanishes\n", "> and retention is *uniform causal attention* (unnormalized). This\n", "> is exactly linear attention’s state with $\\varphi=\\text{identity}$\n", "> and no forgetting.\n", "> 2. **Slide γ → 0.** Only the diagonal survives ($0^0=1$): each token\n", "> sees *only itself*. The decay is the dial between “remember\n", "> everything” and “remember only now”.\n", "> 3. **Mid-range γ.** A soft, fixed locality window — and note there is\n", "> no positional encoding anywhere. The decay *is* the position\n", "> information.\n", "\n", "## The Three Faces of Retention\n", "\n", "Retention’s real selling point is that the **same function** has three\n", "computational forms, and you pick whichever the situation wants:\n", "\n", "- **Parallel** — $\\text{Retention}(X) = (QK^\\top \\odot D)\\,V$. One\n", " masked matmul, fully parallel. Best for **training short-to-medium\n", " sequences**. $O(L^2)$.\n", "- **Recurrent** — $S_n = \\gamma S_{n-1} + k_n^\\top v_n$,\n", " $o_n = q_n S_n$. Constant state, one step at a time. Best for\n", " **decoding**. $O(L)$, $O(1)$ memory per token.\n", "- **Chunkwise** — split into chunks of size $B$; compute each chunk\n", " *internally* by the parallel form, and carry a fixed-size state $R$\n", " *between* chunks by the recurrent form. Best for **training on long\n", " sequences**: parallelism inside a chunk, linear scaling across chunks.\n", " $O(LB)$.\n", "\n", "The chunkwise form is the one production systems use. For a token at\n", "chunk-local index $t$ it adds an intra-chunk term (parallel) to a\n", "cross-chunk term that reads the carried state $R$, decayed by\n", "$\\gamma^{t+1}$:\n", "\n", "$$o_t = \\underbrace{\\sum_{s\\le t}\\gamma^{t-s}(q_t\\cdot k_s)v_s}_{\\text{intra-chunk (parallel)}}\n", "\\; + \\; \\underbrace{\\gamma^{\\,t+1}\\,(q_t R)}_{\\text{cross-chunk (recurrent)}},\n", "\\qquad\n", "R \\leftarrow \\gamma^{B} R + \\sum_s \\gamma^{\\,B-1-s} k_s^\\top v_s .$$\n", "\n", "All three are the same function. `linear_attention.py` proves it\n", "numerically:" ], "id": "7484f9b9-b83f-4f37-8811-b6731d137b23" }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "parallel == recurrent: True\n", "parallel == chunkwise: True" ] } ], "source": [ "from linear_attention import (retention_parallel, retention_recurrent,\n", " retention_chunkwise)\n", "\n", "torch.manual_seed(0)\n", "Qr = torch.randn(12, 6); Kr = torch.randn(12, 6); Vr = torch.randn(12, 6)\n", "gamma = 0.9\n", "\n", "y_par = retention_parallel(Qr, Kr, Vr, gamma)\n", "y_rec = retention_recurrent(Qr, Kr, Vr, gamma)\n", "y_chk = retention_chunkwise(Qr, Kr, Vr, gamma, chunk_size=4)\n", "\n", "print(\"parallel == recurrent:\", torch.allclose(y_par, y_rec, atol=1e-5))\n", "print(\"parallel == chunkwise:\", torch.allclose(y_par, y_chk, atol=1e-5))" ], "id": "a6d81a6d" }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [], "id": "879a401e" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "4bb96979-824c-4030-86f3-b0c07407f84a" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Real RetNet uses **multi-scale retention**: each head gets its own\n", "$\\gamma_h =\n", "1 - 2^{-5-h}$, so some heads keep long memory and others stay local —\n", "the decay analogue of multi-head attention’s different subspaces." ], "id": "9e947f93-06b7-4029-833e-f3c1faa8ba2b" }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "per-head γ: [0.9688, 0.9844, 0.9922, 0.9961]" ] } ], "source": [ "from linear_attention import multiscale_decays\n", "print(\"per-head γ:\", [round(v, 4) for v in multiscale_decays(4).tolist()])" ], "id": "c78be6f4" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Interactive Exploration: The One Recurrence Behind Them All\n", "\n", "Here is the idea this whole module builds toward. Linear attention,\n", "retention, and the diagonal SSM of m19 are the **same recurrence**\n", "\n", "$$S_t = A_t\\, S_{t-1} + k_t^\\top v_t, \\qquad y_t = q_t\\, S_t ,$$\n", "\n", "differing only in the **gate** $A_t$ that scales the old state:\n", "\n", "| Model | Gate $A_t$ | Behavior |\n", "|-------------------|----------------------------|--------------------------|\n", "| Linear attention | $1$ (identity) | Keep everything, equal weight |\n", "| Retention (RetNet) | $\\gamma$ (constant) | Forget geometrically |\n", "| SSM / gated linear attn / Mamba | $A_t = f(x_t)$ (input-dependent) | **Choose** what to keep, per token |\n", "\n", "The `unified_recurrence` function is literally one loop with a swappable\n", "$A_t$. Below, you *are* the gate: check which tokens are “important”.\n", "Important tokens get a high gate ($A_t \\approx 0.95$, the state holds);\n", "the rest get a low gate ($A_t \\approx 0.2$, the state forgets). Watch\n", "the memory integrate only what you select — this is **selectivity**, the\n", "input-dependent gate that lifts a fixed decay into Mamba’s content-based\n", "reasoning." ], "id": "1a387689-4abb-44cd-9c31-515c05414594" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "692f59db-cca5-4c70-a901-1105f667d811" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "55f0976a-163d-4951-bbe3-48fb3a9d2f84" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "f5900b3d-c8c4-4683-976e-86e84b7310c2" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now compare the three regimes on one fixed input stream — linear\n", "attention’s state grows without bound (nothing forgets), retention’s\n", "saturates (a constant decay), and the gated state rises and falls as the\n", "gates open and close:" ], "id": "3e08c770-07ef-46df-ba86-b26c67d67183" }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [], "id": "ec3bf31c" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "bcad37b6-72e9-4466-9f61-2954b9662446" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "5d8608cd-f7c6-44ca-9f9b-b2f18f662619" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> This is the m19↔attention duality (formalized by Mamba-2) made\n", "> runnable: the gate $A_t$ here is exactly the SSM’s discrete decay\n", "> $\\bar{A}$. The difference between a “transformer” and a “state-space\n", "> model” collapses to a single question — *is the gate constant or\n", "> input-dependent?*\n", "\n", "## Why It Matters: Cost\n", "\n", "The payoff, made quantitative. Softmax attention does two $O(L^2 d)$\n", "matmuls and stores a KV cache that grows linearly with the context.\n", "Linear attention does an $O(d^2)$ state update per token over $L$ tokens\n", "— $O(L d^2)$, **linear in $L$** — and carries a fixed $d \\times d$ state\n", "that never grows. Drive the context length and watch the compute curves\n", "diverge:" ], "id": "d41251b5-20f5-4649-b769-c435f1691d44" }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [], "id": "bf379206" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "68ad7b72-a48b-4cb8-9a6e-71d4d7982ac7" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "a11f7a99-1879-4f06-9bd0-bf5b4224367a" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "a39ed19f-f10d-442b-a7b0-54d5f84aa5bf" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "90565558-02a0-425a-876c-9f0dbef65123" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> Linear attention trades a *growing* KV cache for a *fixed* matrix\n", "> state. That state is larger than an SSM’s ($d \\times d$ vs\n", "> $d \\times N$ with a small $N$), which is why pure linear attention has\n", "> historically traded a little quality for that memory constant — and\n", "> why the frontier (gated linear attention, Mamba, hybrids) works to\n", "> recover the quality while keeping the constant-memory decode.\n", "\n", "## Common Pitfalls\n", "\n", "When building linear/recurrent attention, watch out for:\n", "\n", "1. **A non-positive feature map.** If $\\varphi$ can go negative, the\n", " normalizer $\\varphi(q_i)^\\top Z_i$ can hit zero or flip sign and the\n", " output explodes. Keep $\\varphi$ non-negative ($\\text{elu}+1$, or\n", " $\\text{ReLU}$, or a softmax-feature map).\n", "2. **Forgetting the normalizer in linear attention.** Linear attention\n", " divides by $\\varphi(q_i)^\\top Z_i$; retention does **not** (it uses\n", " a GroupNorm downstream instead). Mixing them up gives the wrong\n", " scale.\n", "3. **A decay $\\gamma \\ge 1$.** Retention needs $\\gamma \\in (0,1]$;\n", " $\\gamma>1$ makes $\\gamma^{n-m}$ grow with distance and the state\n", " diverges. (Same stability rule as the SSM’s $\\bar{A}\\in(0,1)$ in\n", " m19.)\n", "4. **Chunk-boundary bookkeeping.** The cross-chunk term must be scaled\n", " by $\\gamma^{t+1}$ and the carried state updated as\n", " $R \\leftarrow \\gamma^B R + \\sum_s\n", " \\gamma^{B-1-s}k_s^\\top v_s$. Off-by-one in these exponents silently\n", " breaks the equivalence — test chunkwise against parallel.\n", "5. **Expecting linear attention to match softmax exactly.** It is a\n", " *different* function, not an approximation of softmax. It trades\n", " some expressivity (the softmax’s sharp, data-dependent focus) for\n", " linear cost; the research since 2020 is largely about closing that\n", " quality gap.\n", "6. **Confusing the state size with the SSM’s.** Linear attention’s\n", " state is $d \\times d$; an SSM’s is $d \\times N$ with $N$ small (16).\n", " Both are constant in $L$, but linear attention’s is bigger — the\n", " price of using full key/value vectors as the “state coordinates”.\n", "\n", "## Exercises\n", "\n", "### Exercise 1: The two faces agree" ], "id": "1373733b-f386-4aa1-83dc-912ae9c0c0c9" }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "from linear_attention import linear_attention_parallel, linear_attention_recurrent\n", "import torch\n", "\n", "# Generate random Q, K, V of your choice. Confirm the parallel and recurrent\n", "# forms match, then explain in one sentence WHY they must: what property of the\n", "# sum over j lets you regroup it into a running state?\n", "\n", "# Your implementation here:" ], "id": "d4133f47" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 2: Decay recovers linear attention" ], "id": "6107fb3d-bd23-47ba-89f4-c5edc47263cd" }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "from linear_attention import retention_parallel\n", "import torch\n", "\n", "# Retention with gamma=1 (no decay) and the identity feature map is exactly\n", "# unnormalized causal attention: (Q Kᵀ ⊙ tril) V. Build that expression by hand\n", "# with torch.tril and confirm it equals retention_parallel(Q, K, V, gamma=1.0).\n", "# Then set gamma=0.0 and describe, in words, what the output becomes.\n", "\n", "# Your implementation here:" ], "id": "e7dd66ae" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 3: You are the gate" ], "id": "090a66d8-70da-4799-9781-6b10382ad614" }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [], "source": [ "from linear_attention import unified_recurrence\n", "import torch\n", "\n", "# Build Q, K, V (length 8). Using mode=\"gated\", design a `gate` vector that keeps\n", "# a HIGH gate (~0.95) only at positions 2 and 5 and a LOW gate (~0.1) elsewhere.\n", "# Show the output differs from mode=\"retention\" with a single fixed gamma — i.e.\n", "# input-dependent selectivity buys something a constant decay cannot.\n", "\n", "# Your implementation here:" ], "id": "af977e18" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Summary\n", "\n", "Key takeaways:\n", "\n", "1. **Linear attention is softmax attention minus the softmax.** Replace\n", " $\\exp(q\\cdot k)$ with a factored kernel $\\varphi(q)\\cdot\\varphi(k)$\n", " and associativity turns the $O(L^2)$ all-pairs sum into a running\n", " state $S_i = \\sum_{j\\le i}\\varphi(k_j)v_j^\\top$ — a linear RNN\n", " hiding inside attention.\n", "2. **It has two equal faces.** A parallel masked-matmul form for\n", " training and a constant-memory recurrent form for decoding compute\n", " the *same function*.\n", "3. **Retention adds a decay.** $S_n=\\gamma S_{n-1}+k_n^\\top v_n$ gives\n", " a $\\gamma^{n-m}$ decay mask (RetNet), which *is* the positional\n", " signal — no separate encoding — and comes in **parallel, recurrent,\n", " and chunkwise** forms, all provably identical.\n", "4. **Chunkwise is the practical form.** Parallel inside chunks,\n", " recurrent between them: GPU-parallel training that still scales\n", " linearly in sequence length.\n", "5. **It is all one recurrence.** $S_t=A_tS_{t-1}+k_t^\\top v_t$\n", " specializes to linear attention ($A_t=1$), retention ($A_t=\\gamma$),\n", " and a diagonal SSM / Mamba ($A_t=f(x_t)$). Transformer vs SSM\n", " collapses to *is the gate input-dependent?*\n", "6. **The cost win is long-context.** $O(Ld^2)$ vs $O(L^2d)$ and a fixed\n", " $d\\times d$ state vs a growing KV cache; the crossover is at\n", " $L\\approx d$, so linear attention pays off exactly where softmax’s\n", " quadratic wall bites.\n", "\n", "## What’s Next\n", "\n", "You have now seen both sides of the sub-quadratic coin — m19 approached\n", "it from control theory (SSMs), this module from attention (drop the\n", "softmax) — and watched them meet in one recurrence. The live frontier\n", "from here is **gated linear attention** and **DeltaNet**\n", "(input-dependent gates and a delta-rule state update that recover more\n", "of softmax’s expressivity) and **hybrid stacks** (Jamba, Griffin) that\n", "interleave a few full-attention layers among many linear/SSM layers to\n", "get global recall *and* linear cost. The throughline is the one this\n", "module makes concrete: the sequence-mixing layer is a design space, and\n", "the axis is *how the state forgets*.\n", "\n", "### Going Deeper\n", "\n", "**Core Papers:**\n", "\n", "- [Transformers are RNNs: Fast Autoregressive Transformers with Linear\n", " Attention](https://arxiv.org/abs/2006.16236) — Katharopoulos et\n", " al. (2020), the feature-map view and the causal running-state\n", " recurrence built here.\n", "- [Retentive Network: A Successor to Transformer for Large Language\n", " Models](https://arxiv.org/abs/2307.08621) — Sun et al. (2023),\n", " retention with decay and its parallel / recurrent / chunkwise forms.\n", "- [RWKV: Reinventing RNNs for the Transformer\n", " Era](https://arxiv.org/abs/2305.13048) — Peng et al. (2023), a\n", " per-channel-decay linear-attention RNN with token-shift, trained at\n", " scale.\n", "- [Gated Linear Attention Transformers with Hardware-Efficient\n", " Training](https://arxiv.org/abs/2312.06635) — Yang et al. (2023), the\n", " input-dependent gate that closes much of the quality gap.\n", "- [Transformers are SSMs: Generalized Models and Efficient Algorithms\n", " (Mamba-2)](https://arxiv.org/abs/2405.21060) — Dao & Gu (2024), the\n", " duality that makes “linear attention” and “state-space model” two\n", " views of one recurrence.\n", "\n", "**Practical Resources:**\n", "\n", "- [Linear Attention and Beyond (Songlin\n", " Yang)](https://sustcsonglin.github.io/blog/2024/deltanet-1/) — an\n", " annotated tour of linear-attention variants, chunkwise algorithms, and\n", " the delta rule.\n", "- [flash-linear-attention](https://github.com/fla-org/flash-linear-attention)\n", " — hardware-efficient reference kernels for linear attention, GLA,\n", " RetNet, and RWKV." ], "id": "8edb218a-dbf4-47af-b58d-0408a429e8f5" } ], "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "name": "python3", "display_name": "Python 3 (ipykernel)", "language": "python", "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3" }, "language_info": { "name": "python", "codemirror_mode": { "name": "ipython", "version": "3" }, "file_extension": ".py", "mimetype": "text/x-python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.14" } } }