{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-1 fw=all -->\n",
    "\n",
    "# The Matrix State: From Linear Attention to Mamba-2\n",
    "\n",
    "Two roads have been leading to the same place. At the end of\n",
    "that section, dropping the softmax turned attention into\n",
    "a recurrence: a matrix-valued state updated by an outer-product write at\n",
    "every token, trained in parallel, stepped in constant memory at generation\n",
    "time. That section closed with a promise, that Mamba-2's *state space\n",
    "duality* makes the correspondence between such recurrences and the selective\n",
    "state space models of that section exact. The last two sections built\n",
    "the other road: linear recurrences discretized from continuous time, made\n",
    "selective, packaged into Mamba. This section is where the roads meet. We\n",
    "first take stock of what the matrix state actually is as a memory, and\n",
    "measure its capacity; a family of models then appears as one-line variations\n",
    "on a single recurrence, differing only in how they *forget*. Unrolling that\n",
    "recurrence proves the promised duality: a gated linear recurrence and masked\n",
    "attention are the same matrix, computed in two different orders, and a third\n",
    "order, chunkwise, is how every model in this family trains at scale. We\n",
    "close with the family table that organizes these models, from\n",
    "RetNet to Mamba-2 to xLSTM.\n",
    "\n",
    "*Prerequisites: the linear-attention recurrence and its normalizer from\n",
    "that section; the selective SSM of\n",
    "that section and the parallel scan of\n",
    "that section; the KV-cache accounting of\n",
    "that section.*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-state-the-matrix-state-from-linear-attention-to-mamba-2"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import numpy as np\n",
    "import time"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-2 fw=all -->\n",
    "\n",
    "## Two Roads to One Recurrence\n",
    "\n",
    "Recall the object that fell out of kernelizing attention,\n",
    "the equation: a state\n",
    "$\\mathbf{S}_t \\in \\mathbb{R}^{d_k \\times d_v}$, with keys indexing rows,\n",
    "written by outer products and read by the query. Every model in this\n",
    "section is an instance of the template\n",
    "\n",
    "$$\n",
    "\\mathbf{S}_t = \\mathbf{D}_t\\, \\mathbf{S}_{t-1} + \\mathbf{k}_t \\mathbf{v}_t^\\top,\n",
    "\\qquad\n",
    "\\mathbf{o}_t = \\mathbf{S}_t^\\top \\mathbf{q}_t,\n",
    "$$\n",
    "\n",
    "where the *transition* $\\mathbf{D}_t$ decides what survives of the past and\n",
    "the *write* $\\mathbf{k}_t \\mathbf{v}_t^\\top$ files the current token under\n",
    "its key. Linear attention is the case $\\mathbf{D}_t = \\mathbf{I}$: never\n",
    "forget. The minGRU and the selective SSM obeyed the same shape one level\n",
    "down, as the elementwise affine recurrence the equation\n",
    "evaluated by the parallel scan of that section; here the\n",
    "state is a matrix rather than a vector, but the algebra, affine maps\n",
    "composing into affine maps, is identical.\n",
    "\n",
    "Before building on this template, we should record what it keeps\n",
    "from that section and what it drops. What you verified\n",
    "there was a *pair* of states: $\\mathbf{S}_t$ together with a normalizer\n",
    "$\\mathbf{z}_t = \\sum_{s \\le t} \\phi(\\mathbf{k}_s)$, read as\n",
    "$\\phi(\\mathbf{q}_t)^\\top \\mathbf{S}_t / \\phi(\\mathbf{q}_t)^\\top \\mathbf{z}_t$,\n",
    "with the feature map $\\phi = \\mathrm{elu} + 1$ keeping everything positive\n",
    "so that the denominator could not vanish. The modern family drops both.\n",
    "Queries and keys come out of learned projections whose scale training can\n",
    "set freely, so the feature map goes; and rather than tracking a normalizer\n",
    "state, production layers normalize the *output*, passing $\\mathbf{o}_t$\n",
    "through an RMSNorm or GroupNorm before the residual add\n",
    "[@Yang.Wang.Shen.ea.2024; @Dao.Gu.2024]. One member of the family, the\n",
    "mLSTM at the end of this section, keeps an explicit normalizer state; when\n",
    "we meet it, the pair $(\\mathbf{S}, \\mathbf{z})$ of that section will be waiting.\n",
    "\n",
    "### What the Memory Costs\n",
    "\n",
    "A matrix as associative memory is an old idea, and its failure mode is just\n",
    "as old. Suppose we store $n$ pairs with $\\mathbf{D}_t = \\mathbf{I}$, so\n",
    "$\\mathbf{S} = \\sum_{i=1}^{n} \\mathbf{k}_i \\mathbf{v}_i^\\top$, and read back\n",
    "key $j$. If the keys have unit norm (an assumption we make throughout this\n",
    "subsection; normalized keys are also the shipped default in this model\n",
    "family), the read-out is\n",
    "\n",
    "$$\n",
    "\\mathbf{S}^\\top \\mathbf{k}_j\n",
    "= \\mathbf{v}_j + \\sum_{i \\neq j} (\\mathbf{k}_i^\\top \\mathbf{k}_j)\\, \\mathbf{v}_i .\n",
    "$$\n",
    "\n",
    "The stored value comes back exactly, plus a sum of every *other* value\n",
    "weighted by how much its key overlaps ours. Only mutually orthogonal keys\n",
    "retrieve without error, and $\\mathbb{R}^{d_k}$ holds at most $d_k$ of\n",
    "those: the capacity of the memory is capped by its width, not by time. As\n",
    "Songlin Yang puts it, the enemy of such a memory is not time, it is other\n",
    "memories [@Yang.Wang.Zhang.ea.2024].\n",
    "\n",
    "How fast does interference bite? The answer is a proposition whose\n",
    "assumptions matter as much as its formula.\n",
    "\n",
    "**Proposition.** Let the keys $\\mathbf{k}_1, \\ldots, \\mathbf{k}_n$ be\n",
    "independent and isotropic on the unit sphere of $\\mathbb{R}^{d_k}$\n",
    "(normalized Gaussians, for instance), and let the stored values be any\n",
    "unit-norm vectors chosen independently of the keys. Then the expected\n",
    "squared read error at every stored key $j$ is\n",
    "\n",
    "$$\n",
    "\\mathbb{E}\\,\\big\\| \\mathbf{S}^\\top \\mathbf{k}_j - \\mathbf{v}_j \\big\\|^2\n",
    "= \\frac{n-1}{d_k},\n",
    "$$\n",
    "\n",
    "where the expectation is over the keys and the squared norm sums over\n",
    "all $d_v$ coordinates of the read-out; coordinate $c$ alone contributes\n",
    "$\\sum_{i \\neq j} v_{i,c}^2 / d_k$ in expectation.\n",
    "\n",
    "**Proof.** The error is $\\sum_{i \\neq j} (\\mathbf{k}_i^\\top\n",
    "\\mathbf{k}_j)\\, \\mathbf{v}_i$. For $i \\neq l$, the cross term\n",
    "$\\mathbb{E}[(\\mathbf{k}_i^\\top \\mathbf{k}_j)(\\mathbf{k}_l^\\top\n",
    "\\mathbf{k}_j)]\\, \\mathbf{v}_i^\\top \\mathbf{v}_l$ vanishes: conditioned\n",
    "on $\\mathbf{k}_j$, the factors are independent with mean zero. What\n",
    "remains is $\\sum_{i \\neq j} \\mathbb{E}[(\\mathbf{k}_i^\\top\n",
    "\\mathbf{k}_j)^2]\\, \\|\\mathbf{v}_i\\|^2$, and by isotropy each\n",
    "$\\mathbb{E}[(\\mathbf{k}_i^\\top \\mathbf{k}_j)^2] = 1/d_k$, with $n-1$\n",
    "terms. $\\blacksquare$\n",
    "\n",
    "Error linear in how much you store, inverse in the width you store it\n",
    "into — *for keys that are independent and spread evenly*. Both\n",
    "assumptions are load-bearing: correlated keys revive the cancelled\n",
    "cross terms, and learned keys are rarely isotropic. The cell below\n",
    "measures the law under its own assumptions and then breaks them on\n",
    "purpose. It fills a memory with $n$ random pairs whose values come from\n",
    "a codebook of 256 unit vectors, reads every key back, and decodes by\n",
    "nearest codebook entry, sweeping $n$ across three widths; a final curve\n",
    "mixes every key with one shared direction ($\\rho = 0.5$) to correlate\n",
    "them."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-state-what-the-memory-costs"
   },
   "source": [
    "rng = np.random.default_rng(0)\n",
    "\n",
    "def hebbian_memory(d, nums, num_values=256, trials=20, rho=0):\n",
    "    \"\"\"Read error and recall of S = sum_i k_i v_i^T over random unit keys;\n",
    "    rho > 0 mixes every key with one shared direction (correlated keys).\"\"\"\n",
    "    codebook = rng.standard_normal((num_values, d))\n",
    "    codebook /= np.linalg.norm(codebook, axis=1, keepdims=True)\n",
    "    errs, accs = [], []\n",
    "    for n in nums:\n",
    "        err = acc = 0\n",
    "        for _ in range(trials):\n",
    "            K = rng.standard_normal((n, d))\n",
    "            K /= np.linalg.norm(K, axis=1, keepdims=True)  # Unit-norm keys\n",
    "            if rho:                    # Break independence: shared component\n",
    "                c = rng.standard_normal(d)\n",
    "                K = np.sqrt(1 - rho ** 2) * K + rho * c / np.linalg.norm(c)\n",
    "                K /= np.linalg.norm(K, axis=1, keepdims=True)\n",
    "            ids = rng.integers(0, num_values, n)\n",
    "            S = K.T @ codebook[ids]                        # Write: sum k v^T\n",
    "            read = K @ S                                   # Rows are S^T k_j\n",
    "            err += ((read - codebook[ids]) ** 2).sum(1).mean()\n",
    "            acc += ((read @ codebook.T).argmax(1) == ids).mean()\n",
    "        errs.append(err / trials)\n",
    "        accs.append(acc / trials)\n",
    "    return errs, accs\n",
    "\n",
    "fig, axes = d2l.plt.subplots(1, 2, figsize=(9, 3.2))\n",
    "for d in [32, 64, 128]:\n",
    "    nums = [d // 4, d // 2, d, 2 * d, 4 * d, 8 * d]\n",
    "    errs, accs = hebbian_memory(d, nums)\n",
    "    axes[0].loglog(nums, errs, marker='o', label=f'$d_k$ = {d}')\n",
    "    axes[0].loglog(nums, [(n - 1) / d for n in nums], 'k--', lw=1)\n",
    "    axes[1].semilogx(nums, accs, marker='o', label=f'$d_k$ = {d}')\n",
    "nums = [16, 32, 64, 128, 256, 512]                 # Width 64, correlated keys\n",
    "errs, accs = hebbian_memory(64, nums, rho=0.5)\n",
    "axes[0].loglog(nums, errs, marker='s', ls=':', color='C1',\n",
    "               label=r'$d_k$ = 64, $\\rho$ = 0.5')\n",
    "axes[1].semilogx(nums, accs, marker='s', ls=':', color='C1')\n",
    "for ax, ylabel in zip(axes, ['squared read error', 'recall accuracy']):\n",
    "    ax.set_xlabel('stored pairs n')\n",
    "    ax.set_ylabel(ylabel)\n",
    "    ax.grid(linestyle='--', alpha=0.4)\n",
    "axes[0].legend();"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 76268 characters]"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-3 fw=all -->\n",
    "\n",
    "The left panel is the capacity law (both panels: independent isotropic\n",
    "unit keys, except the dotted curve): the measured error sits on the\n",
    "dashed $(n-1)/d_k$ prediction so closely that the curves are hard to\n",
    "tell apart, across a 256-fold range of $n$ and three widths. The dotted\n",
    "curve breaks the proposition's assumptions as promised, and the\n",
    "proposition's cancellation is what it loses: with a common component in\n",
    "every key, the cross terms no longer vanish in expectation, the read\n",
    "error at $d_k = 64$ sits about five-fold above the independent curve\n",
    "already at small $n$ and grows faster than linearly, and recall\n",
    "degrades well before $n$ reaches $d_k$.\n",
    "Width stops helping when the keys all resemble one another; this is why\n",
    "the family normalizes its keys and why key *geometry*, not just key\n",
    "count, decides what a matrix memory holds. The right panel is the\n",
    "law's consequence. Recall is essentially perfect while the interference\n",
    "stays small against the unit-norm signal, then collapses as $n$ grows\n",
    "past $d_k$; at a fixed error level, doubling the width doubles how many\n",
    "pairs fit. Nothing in either panel depends on *when* a pair was stored.\n",
    "This is the number that matters for everything ahead: a fixed-size\n",
    "state does not lose memories to time, it loses them to crowding, and\n",
    "that section will price exactly this quantity when deciding\n",
    "how much full attention a production model must keep.\n",
    "\n",
    "One caution about the word \"capacity\", which this chapter uses for\n",
    "four distinct quantities. *Algebraic rank*: $\\mathbf{S}$ has rank at\n",
    "most $d_k$, so no write rule stores more than $d_k$ pairs with exact\n",
    "recall. *Interference capacity*: the proposition above, an expected\n",
    "read error under random independent keys. *Information capacity*: at\n",
    "finite precision, $d_k d_v$ numbers hold at most a fixed number of\n",
    "bits, the counting behind the copying lower bound of\n",
    "that section. *Task-usable capacity*: what a trained model\n",
    "retrieves in practice, which that section measures with its\n",
    "recall sweep and which none of the first three numbers guarantees.\n",
    "\n",
    "### The Decay Ladder\n",
    "\n",
    "If crowding is the failure, forgetting is the first fix: shrink the past\n",
    "before each write and the interference sum stops growing. The family did\n",
    "this in three steps, each a one-line change to\n",
    "the equation:\n",
    "\n",
    "$$\n",
    "\\mathbf{S}_t = \\gamma\\, \\mathbf{S}_{t-1} + \\mathbf{k}_t \\mathbf{v}_t^\\top,\n",
    "\\qquad\n",
    "\\mathbf{S}_t = a_t\\, \\mathbf{S}_{t-1} + \\mathbf{k}_t \\mathbf{v}_t^\\top,\n",
    "\\qquad\n",
    "\\mathbf{S}_t = \\mathrm{diag}(\\boldsymbol{\\alpha}_t)\\, \\mathbf{S}_{t-1} + \\mathbf{k}_t \\mathbf{v}_t^\\top .\n",
    "$$\n",
    "\n",
    "A *fixed scalar* $\\gamma \\in (0,1)$ is RetNet's *retention*\n",
    "[@Sun.Dong.Huang.ea.2023]: each head gets its own constant decay, so\n",
    "different heads keep different horizons. (If you worked the closing\n",
    "exercise of that section, you have already built this\n",
    "rung and verified that its parallel form weights past values by\n",
    "$\\gamma^{t-s}$.) An *input-dependent scalar* $a_t$, one number per head\n",
    "per token, is the transition of Mamba-2 [@Dao.Gu.2024]: the model\n",
    "reads the token and decides how much of the whole memory survives it, the\n",
    "forget gate of that section acting on a matrix state. And an\n",
    "input-dependent *diagonal* $\\mathrm{diag}(\\boldsymbol{\\alpha}_t)$ decays\n",
    "each key coordinate at its own rate, the design of gated linear attention\n",
    "(GLA) [@Yang.Wang.Shen.ea.2024] and RWKV-6\n",
    "[@Peng.Goldstein.Anthony.ea.2024]. the figure\n",
    "draws the ladder. The selective SSM of that section lives on the\n",
    "diagonal rung too, run per channel: $e^{\\Delta_t \\mathbf{a}}$ in\n",
    "the equation is an input-dependent diagonal decay in\n",
    "everything but notation.\n",
    "\n",
    "Each rung also has a statistical reading, which we note here and\n",
    "develop in that section. Stack the keys and\n",
    "values seen so far as rows of $\\mathbf{K}$ and $\\mathbf{V}$, and let\n",
    "the diagonal matrix $\\mathbf{W}$ hold each row's accumulated decay. The\n",
    "linear read-out minimizing the decay-weighted squared error\n",
    "$\\sum_i w_i \\|\\mathbf{S}^\\top \\mathbf{k}_i - \\mathbf{v}_i\\|^2$ is the\n",
    "weighted least-squares solution\n",
    "\n",
    "$$\n",
    "\\mathbf{S}^\\star\n",
    "= \\big(\\mathbf{K}^\\top \\mathbf{W} \\mathbf{K} + \\lambda \\mathbf{I}\\big)^{-1}\n",
    "  \\mathbf{K}^\\top \\mathbf{W} \\mathbf{V},\n",
    "$$\n",
    "\n",
    "with a small ridge $\\lambda$ (or a pseudoinverse at $\\lambda = 0$)\n",
    "covering rank-deficient keys. The decayed state on any rung of the\n",
    "ladder stores only the *weighted cross-moment*\n",
    "$\\mathbf{K}^\\top \\mathbf{W} \\mathbf{V}$: the shortcut that deletes the\n",
    "covariance correction\n",
    "$(\\mathbf{K}^\\top \\mathbf{W} \\mathbf{K} + \\lambda \\mathbf{I})^{-1}$.\n",
    "For orthonormal keys the deletion costs nothing beyond the decay's own\n",
    "discounting of each stored value; for overlapping keys it also pays the\n",
    "interference sum of the equation. The decay sets how\n",
    "fast old evidence expires, and which architectures restore the deleted\n",
    "correction is part of that section's subject.\n",
    "\n",
    "![The decay ladder. Left to right, the transition $\\mathbf{D}_t$ of the equation gains structure: identity (linear attention), a fixed scalar per head (RetNet), an input-dependent scalar (Mamba-2), an input-dependent diagonal (GLA, RWKV-6, and per channel the selective SSM). Every rung keeps the same additive write; the write rule that can edit the state is the next section's subject.](../img/mdl-modernrnn-decay-ladder.svg)\n",
    "\n",
    "## The State-Space Duality\n",
    "\n",
    "We now make good on the promise of that section. Take the scalar-decay rung of the\n",
    "ladder, the simplest one that forgets, and *unroll* it, exactly as we\n",
    "unrolled the LTI recurrence into a convolution in\n",
    "that section. Substituting the update into itself, the state\n",
    "at time $t$ is a decayed sum of every write so far, and the read-out is\n",
    "\n",
    "$$\n",
    "\\mathbf{o}_t\n",
    "= \\sum_{s \\le t} \\big(a_t\\, a_{t-1} \\cdots a_{s+1}\\big)\\,\n",
    "  (\\mathbf{q}_t^\\top \\mathbf{k}_s)\\, \\mathbf{v}_s,\n",
    "\\qquad \\textrm{i.e.} \\qquad\n",
    "\\mathbf{Y} = \\big(\\mathbf{L} \\circ \\mathbf{Q}\\mathbf{K}^\\top\\big)\\, \\mathbf{V},\n",
    "\\quad\n",
    "L_{ts} = \\!\\!\\prod_{r=s+1}^{t}\\!\\! a_r,\n",
    "$$\n",
    "\n",
    "where $\\circ$ is the elementwise product and $\\mathbf{L}$ is lower\n",
    "triangular with ones on the diagonal. Now stare at $\\mathbf{L}$. If every\n",
    "$a_t = 1$, then $L_{ts} = 1$ for $t \\ge s$ and zero otherwise: $\\mathbf{L}$\n",
    "*is the causal mask*, and the equation is masked linear\n",
    "attention, scores $\\mathbf{Q}\\mathbf{K}^\\top$, mask, value mixing, exactly\n",
    "as in that section. A gated linear recurrence is\n",
    "attention whose 0/1 causal mask has been replaced by a *learned, decaying*\n",
    "mask. One object, two readings; @Dao.Gu.2024 named the\n",
    "correspondence *state space duality* (SSD), and matrices of $\\mathbf{L}$'s\n",
    "form, every submatrix below the diagonal having rank one, are called\n",
    "1-semiseparable. Keep the equivalence's scope in view: it is an identity\n",
    "for this structured family, a scalar-gated (1-semiseparable) mask on\n",
    "*linear* attention, not a statement about softmax attention, whose\n",
    "exponential kernel is no masked linear read.\n",
    "\n",
    "The duality is also a statement about computation. The right-hand side of\n",
    "the equation is a triple product, and the two ways of\n",
    "parenthesizing it are the two modes we already know. Materialize\n",
    "$\\mathbf{L} \\circ \\mathbf{Q}\\mathbf{K}^\\top$ first and you compute\n",
    "attention: quadratic in $T$, all positions in parallel, no state in sight.\n",
    "Sweep the sum over $s$ from the left instead, one $t$ at a time, and the\n",
    "partial sums *are* the state $\\mathbf{S}_t$: linear in $T$, constant\n",
    "memory, one token at a time. Same tensor contraction, two orders. Let's\n",
    "verify that on real numbers, recurrent form first:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-state-the-state-space-duality-1"
   },
   "source": [
    "d_k, d_v, T = 64, 64, 512\n",
    "\n",
    "def matrix_state_recurrent(Q, K, V, a):\n",
    "    \"\"\"S_t = a_t S_{t-1} + k_t v_t^T, read o_t = S_t^T q_t.\"\"\"\n",
    "    def step(S, qkva):\n",
    "        q, k, v, a_t = qkva\n",
    "        S = a_t * S + k[:, None] * v[None, :]           # Decay, then write\n",
    "        return S, q @ S                                 # Read S^T q\n",
    "    S0 = jnp.zeros((Q.shape[-1], V.shape[-1]))\n",
    "    return jax.lax.scan(step, S0, (Q, K, V, a))[1]\n",
    "\n",
    "keys = jax.random.split(jax.random.key(0), 4)\n",
    "Q, K, V = (jax.random.normal(k, (T, d))\n",
    "           for k, d in zip(keys, (d_k, d_k, d_v)))\n",
    "a = 0.65 + 0.3 * jax.random.uniform(keys[3], (T,))      # Decays in (0.65, 0.95)\n",
    "y_rec = matrix_state_recurrent(Q, K, V, a)"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-4 fw=all -->\n",
    "\n",
    "The quadratic dual builds $\\mathbf{L}$ in log space: the product of decays\n",
    "from $s$ to $t$ is a difference of cumulative log-decays, and masking with\n",
    "$-\\infty$ *before* exponentiating keeps the entries above the diagonal\n",
    "from overflowing (they would be products of *inverse* decays). We assert\n",
    "agreement rather than eyeball it; this check is the duality."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-state-the-state-space-duality-2"
   },
   "source": [
    "def matrix_state_dual(Q, K, V, a):\n",
    "    \"\"\"Y = (L o Q K^T) V with L_ts = a_t ... a_{s+1} for t >= s, else 0.\"\"\"\n",
    "    cum = jnp.cumsum(jnp.log(a))\n",
    "    logL = cum[:, None] - cum[None, :]                  # Sum of logs over (s, t]\n",
    "    causal = jnp.tril(jnp.ones((len(a), len(a)), bool))\n",
    "    L = jnp.exp(jnp.where(causal, logL, -jnp.inf))\n",
    "    return (L * (Q @ K.T)) @ V\n",
    "\n",
    "with jax.default_matmul_precision('highest'):\n",
    "    err = (jnp.abs(matrix_state_dual(Q, K, V, a) - y_rec).max()\n",
    "           / jnp.abs(y_rec).max())\n",
    "    print(f'relative deviation, dual vs recurrence: {float(err):.2e}')\n",
    "    assert err < 1e-3                      # Two orders, one computation\n",
    "\n",
    "    ones = jnp.ones(T)                                  # All a_t = 1\n",
    "    err = jnp.abs(matrix_state_dual(Q, K, V, ones)\n",
    "                  - jnp.tril(Q @ K.T) @ V).max()\n",
    "    print(f'a_t = 1, deviation from masked linear attention: {float(err):.2e}')\n",
    "    assert err < 1e-6                      # L is literally the causal mask"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "relative deviation, dual vs recurrence: 4.14e-06\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "a_t = 1, deviation from masked linear attention: 0.00e+00\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-5 fw=all -->\n",
    "\n",
    "Ten lines of quadratic code, a scan-shaped loop, and they agree to\n",
    "floating-point rounding; with the decays pinned to one, the dual reduces\n",
    "*exactly* to the masked linear attention of that section. The correspondence\n",
    "that that section promised is now a computation you\n",
    "have run, not an analogy. It also locates Mamba-2 precisely: take\n",
    "that section's selective model, restrict its per-coordinate decay\n",
    "to a single input-dependent scalar per head, and the layer *is*\n",
    "the equation, trainable in whichever mode is cheaper.\n",
    "Which raises the practical question: if the recurrence costs\n",
    "$\\mathcal{O}(T)$ but crawls token by token, and the dual is one big matmul\n",
    "but costs $\\mathcal{O}(T^2)$, is there something in between?\n",
    "\n",
    "## Chunked Computation: Mostly Matmul, a Little Scan\n",
    "\n",
    "There is, and it is how GLA, Mamba-2, and the DeltaNet family of the next\n",
    "section all train. Partition the sequence into chunks of $C$ tokens and\n",
    "partition $\\mathbf{L} \\circ \\mathbf{Q}\\mathbf{K}^\\top$ accordingly, as in\n",
    "the figure. The blocks on the diagonal are the\n",
    "interactions *within* a chunk: each is a $C \\times C$ masked-attention\n",
    "product, quadratic only in the chunk length, and all of them run in\n",
    "parallel. Everything below the diagonal is the influence of earlier chunks\n",
    "on later ones, and here the semiseparable structure pays off: those blocks\n",
    "have rank at most $d_k$, because all a later chunk needs from the entire\n",
    "past is the $d_k \\times d_v$ state at its boundary. So the algorithm is:\n",
    "compute the diagonal blocks as matmuls, pass the state chunk to chunk with\n",
    "a short recurrence of $T/C$ steps, and add each chunk's read of its\n",
    "incoming state. Attention inside, recurrence across; mostly matmul, a\n",
    "little scan [@Dao.Gu.2024].\n",
    "\n",
    "![One matrix, three schedules. The mask $\\mathbf{L}$ of the equation, computed for random decays and partitioned into chunks; the shading fades as the decays compound. Diagonal blocks (blue) are $C \\times C$ masked-attention matmuls computed in parallel; every off-diagonal block (orange) has rank at most $d_k$, so its contribution flows through the $d_k \\times d_v$ state passed from chunk to chunk (arrows).](../img/mdl-modernrnn-semiseparable.svg)\n",
    "\n",
    "The implementation needs one helper. `segsum` builds, for one chunk, the\n",
    "matrix of *segment sums* of log-decays, $\\sum_{s < r \\le t} \\log a_r$,\n",
    "whose exponential is the within-chunk $\\mathbf{L}$; as in the dual, the\n",
    "invalid upper triangle is masked to $-\\infty$ in log space, before any\n",
    "exponential can overflow. The chunked form is then a dozen lines: the\n",
    "decayed attention matmul inside each chunk, and a loop over chunk\n",
    "boundaries that carries the state, decays each key's write to its chunk's\n",
    "end, and lets each chunk read the state it inherited."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-state-chunked-computation-mostly-matmul-a-little-scan-1"
   },
   "source": [
    "def segsum(log_a):\n",
    "    \"\"\"Segment sums: (..., C) -> (..., C, C) with entries sum over (s, t].\"\"\"\n",
    "    C = log_a.shape[-1]\n",
    "    cum = jnp.cumsum(log_a, -1)\n",
    "    diff = cum[..., :, None] - cum[..., None, :]\n",
    "    keep = jnp.tril(jnp.ones((C, C), bool))\n",
    "    return jnp.where(keep, diff, -jnp.inf)              # Mask in log space\n",
    "\n",
    "def matrix_state_chunked(Q, K, V, a, chunk_size=64):\n",
    "    \"\"\"The same Y, chunk by chunk: attention inside, recurrence across.\"\"\"\n",
    "    T, C = Q.shape[0], chunk_size\n",
    "    Qc, Kc = Q.reshape(-1, C, Q.shape[-1]), K.reshape(-1, C, K.shape[-1])\n",
    "    Vc, log_a = V.reshape(-1, C, V.shape[-1]), jnp.log(a).reshape(-1, C)\n",
    "    L = jnp.exp(segsum(log_a))                          # (T/C, C, C)\n",
    "    b = jnp.exp(jnp.cumsum(log_a, -1))                  # Decay since chunk entry\n",
    "    Y = (L * (Qc @ Kc.transpose(0, 2, 1))) @ Vc         # Diagonal blocks\n",
    "    def boundary(S, chunk):                             # T/C boundary steps\n",
    "        Qb, Kb, Vb, bb, Yb = chunk\n",
    "        Yb = Yb + (Qb * bb[:, None]) @ S                # Read the carried state\n",
    "        S = bb[-1] * S + (Kb * (bb[-1] / bb)[:, None]).T @ Vb\n",
    "        return S, Yb\n",
    "    S0 = jnp.zeros((Q.shape[-1], V.shape[-1]))\n",
    "    _, Y = jax.lax.scan(boundary, S0, (Qc, Kc, Vc, b, Y))\n",
    "    return Y.reshape(T, -1)\n",
    "\n",
    "with jax.default_matmul_precision('highest'):\n",
    "    for chunk_size in [16, 64, 256]:\n",
    "        err = (jnp.abs(matrix_state_chunked(Q, K, V, a, chunk_size)\n",
    "                       - y_rec).max() / jnp.abs(y_rec).max())\n",
    "        print(f'chunk size {chunk_size:4d}: relative deviation {float(err):.2e}')\n",
    "        assert err < 1e-3                  # Chunked == recurrence"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "chunk size   16: relative deviation 2.13e-07\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "chunk size   64: relative deviation 3.84e-07\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "chunk size  256: relative deviation 1.54e-06\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-6 fw=all -->\n",
    "\n",
    "One numerical note before we time these forms. Within a chunk, `segsum`\n",
    "works in log space and masks with $-\\infty$ *before* exponentiating, so\n",
    "the diagonal blocks cannot overflow. The boundary step is less careful:\n",
    "it exponentiates the prefix products $b$ and then divides, and $b$\n",
    "shrinks geometrically with position in the chunk. For our decay range\n",
    "the smallest prefix product at $C = 256$ is around $10^{-28}$, still\n",
    "inside float32's range; by $C = 512$ it underflows to exactly zero and\n",
    "the late-chunk ratios $b_C / b_t$ degenerate to $0/0$. The $C = 1024$\n",
    "row of the timing table below therefore measures the right arithmetic\n",
    "shape — the work and memory traffic are unchanged — on values that have\n",
    "already turned to NaN; correctness holds for the chunk sizes asserted\n",
    "above. Production kernels carry the boundary decays in log space too,\n",
    "or rescale chunk by chunk, for exactly this reason; we keep the plain\n",
    "ratio because it leaves the algebra visible.\n",
    "\n",
    "The chunk size interpolates between the two modes we proved dual: at\n",
    "$C = 1$ each diagonal block shrinks to the $1 \\times 1$ self-write\n",
    "$(\\mathbf{q}_t^\\top \\mathbf{k}_t)\\,\\mathbf{v}_t$ — within-chunk\n",
    "interaction between *distinct* tokens is gone, and everything earlier\n",
    "flows through the state recurrence; at $C = T$ there is one block and\n",
    "no recurrence. Where between the endpoints should we sit? Time and\n",
    "memory answer differently, so we measure both, on the same footing as\n",
    "the attention measurements of that section."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-state-chunked-computation-mostly-matmul-a-little-scan-2"
   },
   "source": [
    "T = 4096\n",
    "keys = jax.random.split(jax.random.key(1), 4)\n",
    "Q, K, V = (jax.random.normal(k, (T, d))\n",
    "           for k, d in zip(keys, (d_k, d_k, d_v)))\n",
    "a = 0.65 + 0.3 * jax.random.uniform(keys[3], (T,))\n",
    "\n",
    "def wall_clock(f, *args, reps=5):\n",
    "    f(*args).block_until_ready()  # Warm up (and compile)\n",
    "    start = time.time()\n",
    "    for _ in range(reps):\n",
    "        f(*args).block_until_ready()\n",
    "    return (time.time() - start) / reps\n",
    "\n",
    "forms = {'recurrence': jax.jit(matrix_state_recurrent)}\n",
    "for C in [16, 64, 256, 1024]:\n",
    "    forms[f'chunked ({C})'] = jax.jit(\n",
    "        lambda q, k, v, a, C=C: matrix_state_chunked(q, k, v, a, C))\n",
    "forms['quadratic dual'] = jax.jit(matrix_state_dual)\n",
    "\n",
    "print(f'{\"form\":>15} {\"time (ms)\":>10} {\"temp (MiB)\":>11}')\n",
    "for name, f in forms.items():\n",
    "    stats = f.lower(Q, K, V, a).compile().memory_analysis()\n",
    "    print(f'{name:>15} {wall_clock(f, Q, K, V, a) * 1e3:10.2f} '\n",
    "          f'{stats.temp_size_in_bytes / 2**20:11.1f}')\n",
    "print(f'({jax.devices()[0].device_kind}, float32, jax {jax.__version__}; '\n",
    "      f'teaching implementations, all forms compiled)')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "           form  time (ms)  temp (MiB)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "     recurrence      49.63         0.0\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "   chunked (16)       2.80         1.1\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "   chunked (64)       1.62         1.1\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "  chunked (256)       0.46         5.0\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " chunked (1024)       0.37        21.0\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " quadratic dual       0.45        68.0\n(NVIDIA GeForce RTX 4090, float32, jax 0.10.2; teaching implementations, all forms compiled)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-7 fw=all -->\n",
    "\n",
    "Provenance first: these are our teaching implementations, timed in\n",
    "float32 on the printed device after a warm-up pass, and the two\n",
    "frameworks do not measure quite the same thing — the PyTorch recurrence\n",
    "is an eager Python loop paying per-step dispatch, while JAX compiles\n",
    "all three forms, which flatters the recurrence. Read the orders of\n",
    "magnitude, not the exact ratios. The pattern in our runs is the one the\n",
    "derivation predicts. The token\n",
    "recurrence is by far the slowest, a couple of orders of magnitude behind\n",
    "the best schedule: its $T$ steps are serialized, whether as eager kernel\n",
    "launches or as a compiled scan, so the accelerator's parallelism sits\n",
    "unused, though its memory footprint is trivial. The quadratic dual is the\n",
    "mirror image: matmul-fast, but its $T \\times T$ matrices dominate the\n",
    "memory column and grow fourfold with every doubling of $T$. The chunked\n",
    "schedule takes both prizes at once: with chunks of a few hundred tokens it\n",
    "runs within a small factor of the dual's speed (in some runs ahead of it)\n",
    "at an order of magnitude less memory, and its constituent operations are\n",
    "exactly the dense matmuls the hardware wants to run. The FLOP-optimal chunk size, which an exercise asks you to derive,\n",
    "sits near $C \\approx d_k$; the measured optimum drifts higher because\n",
    "kernel-launch overhead penalizes many small chunks.\n",
    "\n",
    "### What the Hardware Bought\n",
    "\n",
    "This schedule, not any change in the model class, is the main story of\n",
    "Mamba-2 relative to Mamba. Recall from\n",
    "that section that an accelerator moves bytes far more\n",
    "slowly than it multiplies them; its full arithmetic rate is reserved for\n",
    "dense matrix products on tensor cores. Mamba-1's parallel scan is\n",
    "elementwise work, so it runs at memory bandwidth and leaves the matmul\n",
    "units idle. Restricting the transition to a scalar per head buys the\n",
    "chunked form above, whose inner loop is matmuls, and the practical\n",
    "consequence — the Mamba-2 paper's empirical scaling observation, not a\n",
    "theorem — was startling: training got several times faster *and* the\n",
    "state dimension could jump from $N = 16$ to $64$--$256$ at little extra\n",
    "wall clock, because a bigger state just makes the matmuls better shaped\n",
    "for the hardware [@Dao.Gu.2024]. The capacity law of\n",
    "that section says what that buys: an order of magnitude\n",
    "more pairs held at the same read error. The trade is not one-sided, and\n",
    "it runs in both directions: Mamba-1's per-coordinate\n",
    "decay is a more expressive transition per unit of state, and for pure\n",
    "recurrent inference, where there are no chunks and every step is\n",
    "elementwise either way, nothing forces the scalar restriction. Production\n",
    "systems pick per workload. The frontier keeps moving on the same axis:\n",
    "Mamba-3 upgrades the discretization itself, replacing the Euler-style\n",
    "update on the input path with a trapezoidal rule (and adding complex-valued\n",
    "states and multi-input, multi-output reads that feed the decode-time\n",
    "matmuls) [@Lahoti.Li.Chen.ea.2026].\n",
    "\n",
    "One accounting cell closes the systems thread that\n",
    "that section opened. A transformer's per-layer cache grows\n",
    "linearly with context; every matrix state in this section is a constant.\n",
    "Priced per layer and per sequence, at $b$ bytes per element, the two\n",
    "sides of the ledger are\n",
    "\n",
    "$$\n",
    "\\underbrace{\\big(d_{\\textrm{inner}}\\, d_{\\textrm{state}}\n",
    "+ d_{\\textrm{inner}}\\, (w - 1)\\big)\\, b}_{\\textrm{recurrent: SSM state + conv buffer, constant}}\n",
    "\\qquad \\textrm{vs.} \\qquad\n",
    "\\underbrace{2\\, n\\, h_{\\textrm{kv}}\\, d_{\\textrm{head}}\\, b}_{\\textrm{KV cache, grows with the } n \\textrm{ tokens kept}},\n",
    "$$\n",
    "\n",
    "where $w$ is the causal-convolution width whose rolling buffer decoding\n",
    "must also carry (that section) and $h_{\\textrm{kv}}$ the\n",
    "number of KV heads; a batch multiplies both sides by its size, and a\n",
    "whole model multiplies both by its layer count. The cell prices the\n",
    "$64 \\times 64$ head state of our cells and one fully specified\n",
    "production instance — a Mamba-2 2.7B layer: $d_{\\textrm{model}} = 2560$,\n",
    "expansion $2$, so $d_{\\textrm{inner}} = 5120$, $d_{\\textrm{state}} =\n",
    "128$, $w = 4$, hence $5120 \\times 128 = 655{,}360$ state elements per\n",
    "layer — against the grouped-query KV settings of that section.\n",
    "Headline numbers use fp16 ($b = 2$) on *both* sides, the common serving\n",
    "precision; the formula re-prices any row for other dtypes."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-state-what-the-hardware-bought"
   },
   "source": [
    "def mib(elements, bytes_per=2):             # fp16 headline; re-price at will\n",
    "    return elements * bytes_per / 2**20\n",
    "\n",
    "d_inner, d_state, w = 5120, 128, 4          # Mamba-2 2.7B layer\n",
    "ssm, conv = d_inner * d_state, d_inner * (w - 1)\n",
    "print('per layer, per sequence, fp16:')\n",
    "print(f'matrix state, one 64x64 head:  {mib(64 * 64):8.2f} MiB, constant')\n",
    "print(f'Mamba-2 2.7B, SSM {ssm:,} + conv {conv:,} elements: '\n",
    "      f'{mib(ssm + conv):.2f} MiB, constant')\n",
    "for n in [4096, 131072, 1048576]:\n",
    "    kv = 2 * n * 8 * 128                    # 2 x tokens x kv-heads x head-dim\n",
    "    print(f'GQA KV cache at n = {n:>9,}: {mib(kv):8.2f} MiB and growing')"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "per layer, per sequence, fp16:\nmatrix state, one 64x64 head:      0.01 MiB, constant\nMamba-2 2.7B, SSM 655,360 + conv 15,360 elements: 1.28 MiB, constant\nGQA KV cache at n =     4,096:    16.00 MiB and growing\nGQA KV cache at n =   131,072:   512.00 MiB and growing\nGQA KV cache at n = 1,048,576:  4096.00 MiB and growing\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-8 fw=all -->\n",
    "\n",
    "The recurrent side of the ledger is constant in context length; the\n",
    "cache side crosses it within a few thousand tokens and keeps growing\n",
    "with every token kept. Both sides compose with quantization, and the\n",
    "cache side with sharing and eviction schemes, so the constants move;\n",
    "the shapes do not. What the flat line costs, we measured at\n",
    "the start of this section: a capacity ceiling set by the state's width.\n",
    "Holding both numbers in mind at once is the whole hybrid-design problem of\n",
    "that section, which prices whole models with exactly\n",
    "the equation.\n",
    "\n",
    "## The Family, So Far\n",
    "\n",
    "We can now lay out the family that this section and its two predecessors\n",
    "have assembled, one row per model, every row an instance of\n",
    "the equation. The table is worth reading column by column:\n",
    "the *transition* column is the decay ladder, and the *write* column is,\n",
    "for now, monotonous by design.\n",
    "\n",
    "> **Reading the papers.** The literature does not agree on which way the\n",
    "> state matrix points. @Katharopoulos.Vyas.Pappas.ea.2020 write\n",
    "> $\\mathbf{S} = \\sum_s \\phi(\\mathbf{k}_s)\\mathbf{v}_s^\\top$ read by\n",
    "> $\\phi(\\mathbf{q})^\\top \\mathbf{S}$, as we do; RetNet and GLA carry the\n",
    "> same state under transposed notation; the DeltaNet line stores the\n",
    "> transpose $\\mathbf{S} \\in \\mathbb{R}^{d_v \\times d_k}$, writes\n",
    "> $\\mathbf{v}_t \\mathbf{k}_t^\\top$, and applies transitions from the\n",
    "> *right*. All of these are one recurrence up to transposition. This\n",
    "> chapter fixes a single convention, keys index rows:\n",
    "> $\\mathbf{S}_t = \\mathbf{D}_t \\mathbf{S}_{t-1} + \\mathbf{k}_t \\mathbf{v}_t^\\top$\n",
    "> with read $\\mathbf{o}_t = \\mathbf{S}_t^\\top \\mathbf{q}_t$, and\n",
    "> translates every paper into it, as in the table.\n",
    "\n",
    "<!-- tbl-caption:The matrix-state family. Every model is the equation with a choice of transition $\\mathbf{D}_t$ and an additive write; the columns for state and transition vary, the write column does not. -->\n",
    "\n",
    "| model | state | transition $\\mathbf{D}_t$ | write | origin, in one line |\n",
    "|:--|:--|:--|:--|:--|\n",
    "| linear attention (2020) | $\\mathbf{S}$ and normalizer $\\mathbf{z}$ | $\\mathbf{I}$ | add $\\phi(\\mathbf{k}_t)\\mathbf{v}_t^\\top$ | attention minus softmax |\n",
    "| RetNet (2023) | $\\mathbf{S}$ per head | $\\gamma\\,\\mathbf{I}$, fixed per head | add $\\mathbf{k}_t\\mathbf{v}_t^\\top$ | linear attention plus decay |\n",
    "| Mamba (2023) | vector $\\mathbf{x}_t \\in \\mathbb{R}^N$ per channel | input-dep. diagonal $e^{\\Delta_t \\mathbf{a}}$ | add $\\Delta_t u_t \\mathbf{B}_t$ | selective SSM (that section) |\n",
    "| GLA (2024) | $\\mathbf{S}$ per head | input-dep. diagonal $\\mathrm{diag}(\\boldsymbol{\\alpha}_t)$ | add $\\mathbf{k}_t\\mathbf{v}_t^\\top$ | per-coordinate forgetting |\n",
    "| Mamba-2 (2024) | $\\mathbf{S}$ per head | input-dep. scalar $a_t\\,\\mathbf{I}$ | add $\\mathbf{k}_t\\mathbf{v}_t^\\top$ | the duality of this section |\n",
    "| RWKV-6 (2024) | $\\mathbf{S}$ per head | input-dep. diagonal $\\mathrm{diag}(\\mathbf{w}_t)$ | add $\\mathbf{k}_t\\mathbf{v}_t^\\top$ | the RNN lineage arrives |\n",
    "| mLSTM (2024) | $\\mathbf{S}$ and normalizer $\\mathbf{z}$ | scalar gate $f_t\\,\\mathbf{I}$ | add $i_t\\,\\mathbf{k}_t\\mathbf{v}_t^\\top$ | LSTM with a matrix cell |\n",
    "\n",
    "Three research lineages fill the rows: attention seeking linearity\n",
    "(linear attention, RetNet, GLA), state space models seeking selectivity\n",
    "(Mamba, Mamba-2), and recurrent networks seeking parallelism (RWKV-6,\n",
    "mLSTM). That they converge on one recurrence is the strongest evidence in\n",
    "this book that the recurrence, rather than any single architecture, is the\n",
    "right object of study. Mamba's row is the one non-matrix entry: its\n",
    "per-channel vector state is the same mathematics run single-input,\n",
    "single-output, and Mamba-2's move to the matrix state is what the duality\n",
    "section derived.\n",
    "\n",
    "### The One That Kept the Normalizer\n",
    "\n",
    "The last row deserves its promised closer look, because it answers a\n",
    "question that section left open: whatever happened to the normalizer\n",
    "$\\mathbf{z}$? The mLSTM, the matrix-memory cell of xLSTM\n",
    "[@Beck.Poppel.Spanring.ea.2024], is the LSTM rebuilt on\n",
    "the equation: a forget gate $f_t$ and input gate $i_t$,\n",
    "computed from the token, gating a matrix cell state and, alongside it, the\n",
    "running normalizer that the rest of the family traded for output\n",
    "normalization:\n",
    "\n",
    "$$\n",
    "\\mathbf{S}_t = f_t\\, \\mathbf{S}_{t-1} + i_t\\, \\mathbf{k}_t \\mathbf{v}_t^\\top,\n",
    "\\qquad\n",
    "\\mathbf{z}_t = f_t\\, \\mathbf{z}_{t-1} + i_t\\, \\mathbf{k}_t,\n",
    "\\qquad\n",
    "\\mathbf{o}_t = \\frac{\\mathbf{S}_t^\\top \\mathbf{q}_t}\n",
    "                    {\\max\\big(|\\mathbf{z}_t^\\top \\mathbf{q}_t|,\\, 1\\big)},\n",
    "$$\n",
    "\n",
    "with $i_t = \\exp(\\tilde{i}_t)$ and $f_t = \\sigma(\\tilde{f}_t)$ or\n",
    "$\\exp(\\tilde{f}_t)$ for learned pre-activations $\\tilde{i}_t, \\tilde{f}_t$\n",
    "(the paper's $\\mathbf{C}_t$ and $\\mathbf{n}_t$ are our $\\mathbf{S}_t$ and\n",
    "$\\mathbf{z}_t$ transposed). The exponential gates are the novelty and the\n",
    "hazard: they let a single strong token outweigh the entire accumulated\n",
    "past, and they overflow float32 after a few dozen steps of ordinary gate\n",
    "values. The fix is one we have already met. Exactly as online softmax\n",
    "the equation carries a running maximum and rescales the\n",
    "past, the mLSTM carries a stabilizer state $m_t$\n",
    "[@Milakov.Gimelshein.2018]:\n",
    "\n",
    "$$\n",
    "m_t = \\max\\big(\\tilde{f}_t + m_{t-1},\\, \\tilde{i}_t\\big),\n",
    "\\qquad\n",
    "f'_t = \\exp\\big(\\tilde{f}_t + m_{t-1} - m_t\\big),\n",
    "\\qquad\n",
    "i'_t = \\exp\\big(\\tilde{i}_t - m_t\\big),\n",
    "$$\n",
    "\n",
    "for the exponential-forget case. Running the recurrence with\n",
    "$(f'_t, i'_t)$ scales state and normalizer by the same factor\n",
    "$e^{-m_t}$, so the read-out is unchanged provided the constant $1$ in the\n",
    "denominator is scaled too, to $e^{-m_t}$. We implement the recurrence\n",
    "directly from these equations (the reference xLSTM code base is under a\n",
    "restrictive license, so reimplementing from the paper is also the legally\n",
    "clean route) and verify both claims at once: the unstabilized float32\n",
    "recurrence dies, and the stabilized one matches an unstabilized float64\n",
    "reference to float precision."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-state-the-one-that-kept-the-normalizer"
   },
   "source": [
    "d_k, d_v, T = 8, 8, 64\n",
    "kk = jax.random.split(jax.random.key(2), 5)\n",
    "q, k, v = (jax.random.normal(key, (T, d))\n",
    "           for key, d in zip(kk, (d_k, d_k, d_v)))\n",
    "i_pre = jax.random.normal(kk[3], (T,))               # Gate pre-activations\n",
    "f_pre = 2.0 + jax.random.normal(kk[4], (T,))\n",
    "\n",
    "def mlstm_naive(q, k, v, i_pre, f_pre, dtype):\n",
    "    \"\"\"The unstabilized recurrence (NumPy), in a given precision.\"\"\"\n",
    "    q, k, v = (np.asarray(x, dtype) for x in (q, k, v))\n",
    "    i_g, f_g = np.exp(np.asarray(i_pre, dtype)), np.exp(np.asarray(f_pre, dtype))\n",
    "    S, z = np.zeros((d_k, d_v), dtype), np.zeros(d_k, dtype)\n",
    "    outputs = []\n",
    "    for t in range(T):\n",
    "        S = f_g[t] * S + i_g[t] * np.outer(k[t], v[t])\n",
    "        z = f_g[t] * z + i_g[t] * k[t]\n",
    "        outputs.append(q[t] @ S / max(abs(q[t] @ z), 1.0))\n",
    "    return np.stack(outputs)\n",
    "\n",
    "def mlstm_stabilized(q, k, v, i_pre, f_pre):\n",
    "    \"\"\"Carry m_t and rescale the past, as in online softmax.\"\"\"\n",
    "    def step(carry, x):\n",
    "        S, z, m = carry\n",
    "        q_t, k_t, v_t, i_t, f_t = x\n",
    "        m_new = jnp.maximum(f_t + m, i_t)\n",
    "        f_g, i_g = jnp.exp(f_t + m - m_new), jnp.exp(i_t - m_new)\n",
    "        S = f_g * S + i_g * (k_t[:, None] * v_t[None, :])\n",
    "        z = f_g * z + i_g * k_t\n",
    "        o = q_t @ S / jnp.maximum(jnp.abs(q_t @ z), jnp.exp(-m_new))\n",
    "        return (S, z, m_new), o\n",
    "    init = (jnp.zeros((d_k, d_v)), jnp.zeros(d_k), -jnp.inf)\n",
    "    return jax.lax.scan(step, init, (q, k, v, i_pre, f_pre))[1]\n",
    "\n",
    "exact = mlstm_naive(q, k, v, i_pre, f_pre, np.float64)\n",
    "with np.errstate(over='ignore', invalid='ignore'):   # The overflow is the point\n",
    "    naive = mlstm_naive(q, k, v, i_pre, f_pre, np.float32)\n",
    "bad = ~np.isfinite(naive).all(-1)           # argmax(bool) would report 0\n",
    "first_bad = int(bad.argmax()) if bad.any() else 'none'\n",
    "stab = mlstm_stabilized(q, k, v, i_pre, f_pre)\n",
    "err = jnp.abs(stab - exact).max() / jnp.abs(exact).max()\n",
    "print(f'float32 unstabilized: first non-finite output at step {first_bad}')\n",
    "print(f'float32 stabilized vs float64: relative deviation {float(err):.2e}')\n",
    "assert bool(jnp.isfinite(stab).all()) and err < 1e-3"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "float32 unstabilized: first non-finite output at step 47\nfloat32 stabilized vs float64: relative deviation 1.38e-06\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=matrix-state-md-9 fw=all -->\n",
    "\n",
    "The mLSTM is one half of the xLSTM architecture; its scalar sibling sLSTM\n",
    "keeps the LSTM's recurrent memory mixing (and therefore steps\n",
    "sequentially, gaining state-tracking power the parallel family lacks), a\n",
    "trade we noted in the gate table of that section.\n",
    "\n",
    "### The Column Left Open\n",
    "\n",
    "Look down the write column of the table once more. Seven\n",
    "models, three lineages, and every single write is an *addition*: not one\n",
    "model in the table can change what it wrote about key $\\mathbf{k}$ without\n",
    "also decaying everything else it knows. The decay ladder manages crowding,\n",
    "it does not cure it, because forgetting is indiscriminate where\n",
    "interference the equation is specific. The fast weight\n",
    "programmers saw this in the 1990s and proposed the repair\n",
    "[@Schlag.Irie.Schmidhuber.2021]: before writing to key\n",
    "$\\mathbf{k}_t$, first *read* what the memory already holds there, and\n",
    "write only the correction. That single change to the write rule, a memory\n",
    "that edits rather than accumulates, is the subject of\n",
    "that section.\n",
    "\n",
    "## Summary\n",
    "\n",
    "Linear attention's matrix state and the selective SSM's recurrence are one\n",
    "template, $\\mathbf{S}_t = \\mathbf{D}_t \\mathbf{S}_{t-1} + \\mathbf{k}_t\n",
    "\\mathbf{v}_t^\\top$ read by the query. As a memory it has a measured\n",
    "capacity law: after $n$ writes with independent isotropic unit keys\n",
    "into width $d_k$, the expected squared read error is $(n-1)/d_k$ — and\n",
    "measurably worse once keys correlate — so capacity is set by width, not\n",
    "by time. The\n",
    "decay ladder, fixed scalar (RetNet), input-dependent scalar (Mamba-2),\n",
    "input-dependent diagonal (GLA, RWKV-6, and per channel the selective SSM),\n",
    "adds forgetting one line at a time. Unrolling the scalar-decay recurrence\n",
    "gives $\\mathbf{Y} = (\\mathbf{L} \\circ \\mathbf{Q}\\mathbf{K}^\\top)\\mathbf{V}$\n",
    "with a 1-semiseparable $\\mathbf{L}$; at $a_t = 1$ the mask $\\mathbf{L}$ is\n",
    "literally the causal mask, so masked linear attention and the gated\n",
    "recurrence are the same matrix, and the quadratic and recurrent modes are\n",
    "two orders of one contraction. That is Mamba-2's state space duality,\n",
    "verified here by assertion, and its chunked middle order, attention within\n",
    "chunks, a short state recurrence across them, is how the family trains:\n",
    "matmul-heavy, linear in sequence length, and the schedule the Mamba-2\n",
    "paper credits for growing state sizes an order of magnitude at little\n",
    "wall-clock cost. The family table collects the\n",
    "members, including the mLSTM, which keeps the normalizer state of that section and\n",
    "stabilizes its exponential gates with the online-softmax rescaling. Every\n",
    "write rule in the table only adds; editing the memory is next.\n",
    "\n",
    "What this section's evidence does and does not show: the capacity law\n",
    "is a *proposition* under independent isotropic random keys, confirmed\n",
    "by measurement under exactly those assumptions and shown failing under\n",
    "correlated keys; it bounds nothing about trained or whitened key\n",
    "geometries. The duality and chunking checks are *identities* for the\n",
    "scalar-gated family, verified to float tolerance. The timing and\n",
    "memory tables are *illustrations* from one device, dtype, and teaching\n",
    "implementation: constants that move across hardware and kernels around\n",
    "asymptotic shapes that do not.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. [conceptual] Count FLOPs for the three schedules of the equation at\n",
    "   sequence length $T$, chunk size $C$, and $d_k = d_v = d$: the token\n",
    "   recurrence, the quadratic dual, and the chunked form (diagonal blocks\n",
    "   plus state passing). Show that the chunked total is minimized near\n",
    "   $C \\approx d$ and find the ratio of chunked to dual FLOPs at\n",
    "   $T = 4096$, $d = 64$, $C = 256$. Then reconcile your formula with the\n",
    "   measured table: why does the measured optimum sit at a larger $C$ than\n",
    "   the FLOP-optimal one?\n",
    "1. [short-code] Extend the capacity sweep to the decayed write\n",
    "   $\\mathbf{S}_t = \\gamma \\mathbf{S}_{t-1} + \\mathbf{k}_t \\mathbf{v}_t^\\top$.\n",
    "   Store $n$ pairs sequentially, then query all of them, plotting recall\n",
    "   separately for the oldest and newest quarter of the pairs as $\\gamma$\n",
    "   varies over $\\{1, 0.99, 0.95, 0.9\\}$. What does decay buy the newest\n",
    "   items, what does it cost the oldest, and why does no value of $\\gamma$\n",
    "   raise the total number of pairs recallable at once above the\n",
    "   $\\mathcal{O}(d_k)$ ceiling?\n",
    "1. [short-code] Restore the construction of that section inside the dual form: apply\n",
    "   $\\phi = \\mathrm{elu} + 1$ to queries and keys, add a normalizer row to\n",
    "   the computation, and verify that with $a_t = 1$ your result matches the\n",
    "   `linear_attention_parallel` of that section\n",
    "   to float tolerance. Which of the two changes (feature map, normalizer) alters the\n",
    "   attention pattern, and which only rescales it?\n",
    "1. [short-code] In the chunked form, replace the scalar decay $a_t$ by a per-coordinate\n",
    "   vector gate $\\boldsymbol{\\alpha}_t \\in (0,1)^{d_k}$ (the GLA rung of\n",
    "   the ladder). Which of `segsum`, the diagonal-block product, and the\n",
    "   boundary recurrence change, and to what shapes? Implement it and verify\n",
    "   against a per-coordinate token loop.\n",
    "1. [short-code] The mLSTM demonstration used forget pre-activations around $+2$.\n",
    "   Compute analytically at which step the unstabilized float32 recurrence\n",
    "   must overflow as a function of the mean of $\\tilde{f}_t$, and check\n",
    "   your prediction by varying the mean over $\\{1, 2, 4\\}$. What changes\n",
    "   when $f_t = \\sigma(\\tilde{f}_t)$ instead, and why does the stabilizer\n",
    "   still matter then?\n",
    "1. [conceptual] RetNet fixes per-head decays $\\gamma_h = 1 - 2^{-5-h}$ for heads\n",
    "   $h = 0, 1, \\ldots$ Using the effective-horizon definition of\n",
    "   that section's exercises (the lag at which influence drops below\n",
    "   $0.01$), compute each head's memory horizon for eight heads. Why might\n",
    "   fixed, staggered horizons be preferable to learned ones at small scale,\n",
    "   and what does the model lose relative to the input-dependent rungs of\n",
    "   the ladder?\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
