{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Module 19: State-Space Models (Mamba)" ], "id": "e20ca71f-4683-4764-a572-e90743333e0c" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [], "id": "7310be7b-0c1e-49e8-9bc0-0f41cb3335b8" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "389c6a08-ba9d-4cb0-9910-f969d7cd88ef" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction\n", "\n", "Every module so far is built on **attention**. Attention is powerful,\n", "but it pays a tax: comparing every token to every other token costs\n", "$O(L^2)$ compute, and generating with a growing **KV cache** (m08) costs\n", "memory that climbs with the context. Double the context and attention\n", "quadruples its work.\n", "\n", "A **state-space model (SSM)** is the leading *alternative* — the\n", "architecture, made practical by **Mamba**, that many believe can rival\n", "the Transformer. An SSM reads a sequence the way an RNN does: it carries\n", "a **fixed-size hidden state** forward, one token at a time, folding each\n", "new token into that state. Because the state never grows, an SSM runs in\n", "**linear time** ($O(L)$) and **constant memory per generated token** —\n", "no cache that swells with the context.\n", "\n", "The catch that sank classic RNNs was training speed and long-range\n", "memory. SSMs fix both: a *time-invariant* SSM can be run as a **parallel\n", "convolution** for training, and a carefully structured state matrix\n", "gives it long memory. Mamba’s final ingredient — **selectivity** — lets\n", "the state decide *what to remember based on content*, closing the last\n", "quality gap with attention.\n", "\n", "Why it matters for LLMs:\n", "\n", "- **Long context, cheaply.** Linear scaling makes million-token\n", " sequences tractable where attention’s $L^2$ wall does not.\n", "- **Fast generation.** A constant-size state means each new token costs\n", " the same, no matter how long the context — Mamba reports ~5× the\n", " throughput of a Transformer.\n", "- **A different lens.** Understanding SSMs shows you which parts of “the\n", " Transformer” are essential and which are just one choice among\n", " several.\n", "\n", "### What You’ll Learn\n", "\n", "After this module, you can:\n", "\n", "- Explain a state-space model as a linear recurrence over a fixed-size\n", " state.\n", "- Discretize a continuous SSM with **zero-order hold** and run the\n", " discrete recurrence.\n", "- Run the *same* SSM as a **recurrence** and as a **convolution**, and\n", " see they are identical — the trick that trains in parallel and infers\n", " cheaply.\n", "- Explain why that recurrence is an **associative scan** that\n", " parallelizes in $O(\\log L)$ depth.\n", "- Build a **selective** SSM (Mamba’s S6) whose $\\Delta$, $B$, $C$ depend\n", " on the input.\n", "- Show *why* SSMs challenge attention: linear vs quadratic compute,\n", " constant vs growing decode memory.\n", "\n", "### Prerequisites\n", "\n", "This module requires familiarity with:\n", "\n", "- [Module 05: Attention](../m05_attention/lesson.qmd) — the\n", " quadratic-cost mechanism SSMs are the alternative to.\n", "- [Module 06: Transformer](../m06_transformer/lesson.qmd) — the block an\n", " SSM layer replaces.\n", "- [Module 08: Generation](../m08_generation/lesson.qmd) — the KV cache\n", " whose growth SSMs sidestep.\n", "\n", "## Intuition: A Sequence as a Running State\n", "\n", "Forget attention for a moment. Imagine reading a sentence with a small\n", "notebook: after each word you update a fixed set of numbers — your\n", "running summary — and you never allow the notebook to get bigger. That\n", "notebook is the **hidden state** $h$, and the rule for updating it is a\n", "**state-space model**:\n", "\n", "$$h_t = \\bar{A}\\, h_{t-1} + \\bar{B}\\, x_t, \\qquad y_t = C\\, h_t.$$\n", "\n", "$\\bar{A}$ says how much of the old state to *keep* (a decay), $\\bar{B}$\n", "says how much of the new token to *write in*, and $C$ reads an output\n", "off the state. That is the whole model. Feed it a single spike at $t=0$\n", "and watch the state light up, then decay — step through its **impulse\n", "response**:" ], "id": "fb60620d-be91-4b0d-a7c0-5d99c80a2d46" }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [], "id": "9a8d9fee" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "7b4c10d3-e650-4940-a886-749880f72f43" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "de15281d-b6e4-4c9b-a464-c83bb16326e0" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> The state is a **fixed-size** summary of everything seen so far.\n", "> Attention keeps *every* past token around (the KV cache) and re-reads\n", "> them; an SSM keeps one running state and updates it. That single\n", "> design choice is where the linear cost comes from.\n", "\n", "## The Math: The State-Space Recurrence\n", "\n", "SSMs start in **continuous** time, borrowed from control theory. A state\n", "vector $h(t) \\in \\mathbb{R}^N$ evolves under a linear system driven by\n", "the scalar input $x(t)$:\n", "\n", "$$h'(t) = A\\, h(t) + B\\, x(t), \\qquad y(t) = C\\, h(t).$$\n", "\n", "To use it on a discrete token stream we **discretize** it with a step\n", "size $\\Delta$, using the standard **zero-order hold (ZOH)** rule (it\n", "assumes the input is held constant across each step). For a **diagonal**\n", "$A$ — a vector of scalars $a$, which is exactly what Mamba uses — this\n", "is element-wise:\n", "\n", "$$\\bar{A} = \\exp(\\Delta A), \\qquad\n", "\\bar{B} = A^{-1}\\!\\left(\\exp(\\Delta A) - I\\right) B .$$\n", "\n", "Two things to notice. First, if every $a < 0$ then\n", "$\\bar{A} = \\exp(\\Delta a) \\in\n", "(0, 1)$ — a genuine **decay factor**, so the state forgets old inputs\n", "gracefully and never blows up. Second, as $a \\to 0$ the $\\bar{B}$ factor\n", "tends to $\\Delta$, so $\\bar{B} \\approx \\Delta B$ — the simple\n", "**forward-Euler** limit that Mamba’s kernel actually uses. Substituting\n", "gives the discrete recurrence from the intuition:\n", "\n", "$$h_t = \\bar{A}\\, h_{t-1} + \\bar{B}\\, x_t, \\qquad y_t = C\\, h_t .$$\n", "\n", "`discretize` in `ssm.py` implements the exact ZOH formula (with a stable\n", "small-$a$ fallback so the Euler limit is numerically safe):" ], "id": "46b37105-6085-4dcc-91c9-947400edfd3b" }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Ā (decay, in (0,1)): [0.779, 0.607, 0.223]\n", "B̄ (input gain): [0.442, 0.393, 0.259]" ] } ], "source": [ "import torch\n", "from ssm import discretize\n", "\n", "A = torch.tensor([-0.5, -1.0, -3.0]) # diagonal state matrix (must be < 0 to decay)\n", "B = torch.tensor([1.0, 1.0, 1.0])\n", "A_bar, B_bar = discretize(A, B, delta=0.5)\n", "print(\"Ā (decay, in (0,1)):\", [round(v, 3) for v in A_bar.tolist()])\n", "print(\"B̄ (input gain): \", [round(v, 3) for v in B_bar.tolist()])" ], "id": "65876f48" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Code: The Recurrent Scan\n", "\n", "Given the discrete coefficients, running the SSM is a short loop — the\n", "*inference* face, $O(L)$ time and $O(N)$ memory. `ssm_recurrent` returns\n", "both the outputs and the hidden state after each step:" ], "id": "bc8cc87d-a074-4764-acd4-4a9b0010db24" }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "impulse response: [1.0, 0.5, 0.25, 0.125, 0.0625]" ] } ], "source": [ "from ssm import ssm_recurrent\n", "\n", "A_bar = torch.tensor([0.5]) # a single-state SSM, decay 0.5 per step\n", "B_bar = torch.tensor([1.0])\n", "C = torch.tensor([1.0])\n", "\n", "x = torch.tensor([1.0, 0.0, 0.0, 0.0, 0.0]) # a unit impulse at t=0\n", "y, states = ssm_recurrent(A_bar, B_bar, C, x)\n", "print(\"impulse response:\", [round(v, 4) for v in y.tolist()]) # 1, .5, .25, ..." ], "id": "676ff600" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The output halves every step: the state remembers the spike, fading by\n", "$\\bar{A}=0.5$ each time. That decaying sequence is the SSM’s **impulse\n", "response** — hold onto it, because it is about to reappear as a\n", "convolution kernel.\n", "\n", "## Two Faces: Recurrent = Convolutional\n", "\n", "Here is the trick that makes SSMs trainable. Unroll the recurrence from\n", "$h_0 = 0$:\n", "\n", "$$y_t = \\sum_{i=0}^{t} \\big(C\\, \\bar{A}^{\\,i}\\, \\bar{B}\\big)\\, x_{t-i} .$$\n", "\n", "That is a **causal convolution** of the input with a fixed kernel\n", "\n", "$$\\bar{K} = \\big(C\\bar{B},\\; C\\bar{A}\\bar{B},\\; C\\bar{A}^2\\bar{B},\\; \\dots\\big) ,$$\n", "\n", "which is precisely the impulse response you just saw. So one SSM has\n", "**two faces**:\n", "\n", "- **Recurrent** — one step at a time, $O(L)$, cheap at inference (no\n", " cache).\n", "- **Convolutional** — the whole sequence at once with a single kernel,\n", " parallel on the GPU, ideal for **training**.\n", "\n", "Crucially, for time-invariant parameters they compute *exactly the same\n", "function*. `ssm_convolution` builds the kernel and convolves; it matches\n", "`ssm_recurrent` to floating-point tolerance:" ], "id": "043cf595-3ea8-41b6-a64a-c58395ff72e6" }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "kernel (first 4 taps): [0.942, 0.113, 0.074, 0.091]\n", "recurrent == convolutional: True" ] } ], "source": [ "from ssm import ssm_convolution, ssm_kernel\n", "\n", "torch.manual_seed(0)\n", "A_bar = torch.rand(4) * 0.9 + 0.05 # 4-dim state, each channel decays in (0.05, 0.95)\n", "B_bar = torch.randn(4)\n", "C = torch.randn(4)\n", "x = torch.randn(16)\n", "\n", "y_rec, _ = ssm_recurrent(A_bar, B_bar, C, x)\n", "y_conv = ssm_convolution(A_bar, B_bar, C, x)\n", "print(\"kernel (first 4 taps):\", [round(v, 3) for v in ssm_kernel(A_bar, B_bar, C, 4).tolist()])\n", "print(\"recurrent == convolutional:\", torch.allclose(y_rec, y_conv, atol=1e-5))" ], "id": "7a499227" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Bridge the two outputs to a plot: the convolutional dots land exactly on\n", "the recurrent line, because they are the same function computed two\n", "ways." ], "id": "b0d32624-da83-474b-8310-beff1d444c4d" }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "pts = [{\"t\": t, \"rec\": float(y_rec[t]), \"conv\": float(y_conv[t])} for t in range(len(x))]\n", "ojs_define(ssm_faces = pts)" ], "id": "dc9b014c" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "fa3be61d-ccdd-4b4c-b660-269613cebb98" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> Train like a CNN, infer like an RNN. The convolutional face gives\n", "> fast, parallel training on whole sequences; the recurrent face gives\n", "> cheap, cache-free generation. Both are the *same* SSM — this duality\n", "> is why SSMs are practical where plain RNNs were not.\n", "\n", "## The Parallel Scan\n", "\n", "The convolution needs a *fixed* kernel, which only exists when the\n", "parameters do not change with position. There is a second, more general\n", "route to parallelism that does not need that assumption. Look at the\n", "recurrence as\n", "\n", "$$h_t = a_t\\, h_{t-1} + b_t ,$$\n", "\n", "with $a_t = \\bar{A}$ and $b_t = \\bar{B}\\, x_t$. Composing two\n", "consecutive steps is **associative**:\n", "\n", "$$(a_2, b_2) \\circ (a_1, b_1) = (a_2 a_1,\\; a_2 b_1 + b_2) .$$\n", "\n", "Associativity is exactly the property a **parallel prefix scan** needs,\n", "so all of $h_1, \\dots, h_L$ can be computed in $O(\\log L)$ sequential\n", "depth instead of $L$ serial steps. `parallel_scan` implements this and\n", "reproduces the sequential loop exactly:" ], "id": "05cd6826-07e0-443e-8cef-c31c9934d324" }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "parallel scan h: [1.0, 1.1, -0.17, 0.181, 1.0362]\n", "sequential loop h: [1.0, 1.1, -0.17, 0.181, 1.0362]" ] } ], "source": [ "from ssm import parallel_scan\n", "\n", "a = torch.tensor([0.5, 0.9, 0.3, 0.7, 0.2]) # per-step decays\n", "b = torch.tensor([1.0, 0.2, -0.5, 0.3, 1.0]) # per-step inputs\n", "print(\"parallel scan h:\", [round(v, 4) for v in parallel_scan(a, b).tolist()])\n", "\n", "# Same numbers a plain left-to-right loop would give:\n", "h, seq = 0.0, []\n", "for i in range(5):\n", " h = a[i].item() * h + b[i].item()\n", " seq.append(round(h, 4))\n", "print(\"sequential loop h:\", seq)" ], "id": "8bbc1f27" }, { "cell_type": "markdown", "metadata": {}, "source": [ "This matters because it survives the next section. When Mamba makes the\n", "parameters *input-dependent*, the fixed-kernel convolution vanishes —\n", "but the scan does not (it never assumed constant $a_t$), so training\n", "stays parallel.\n", "\n", "## Selectivity: The Mamba Idea\n", "\n", "Everything so far is **time-invariant**: $\\bar{A}$, $\\bar{B}$, $C$ are\n", "the same at every step. That makes an SSM a fixed linear filter — and a\n", "fixed filter cannot do **content-based reasoning**. It cannot say “this\n", "token matters, remember it; that one is filler, skip it,” because its\n", "response to a token does not depend on the token.\n", "\n", "Mamba’s insight (the **S6** layer) is to make the parameters **functions\n", "of the input**: at each position, $\\Delta_t$, $B_t$, and $C_t$ are\n", "projected from $x_t$. Now $\\Delta_t$ acts as an input-driven **gate** —\n", "a large $\\Delta_t$ writes the token firmly into the state and resets old\n", "memory; a small $\\Delta_t$ (via $\\bar{A} = \\exp(\\Delta_t a) \\approx 1$,\n", "$\\bar{B} \\approx 0$) lets the token pass while the state holds its\n", "contents. The model **selects** what to keep.\n", "\n", "Drive it: below, a stream of tokens flows into a one-dimensional state.\n", "Choose which tokens are “important” (high $\\Delta$); the rest are\n", "near-skipped. Watch the state integrate only what you selected." ], "id": "2b138364-dfcb-47c3-864b-42f7fe50cb48" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "72f63d70-44e1-4ada-94d1-3cbc9a3726aa" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "c2e42c4e-0f1e-4a45-8d78-c03f135a18b9" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "4ed300d2-6775-4594-b0ec-1222c14ae492" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try It!**\n", ">\n", "> 1. **Select nothing.** With every token skipped ($\\Delta \\approx 0$),\n", "> the state barely moves — the model ignores the whole stream.\n", "> Selectivity off = no memory writes.\n", "> 2. **Select one token.** The state jumps at that token and then\n", "> *holds* (the skips have $\\bar{A}\\approx 1$), carrying that value\n", "> forward — the SSM is *copying* the selected token across time.\n", "> 3. **Select all.** Now every token writes and the state churns like\n", "> the time-invariant SSM. Selectivity is the dial between “ignore”\n", "> and “integrate everything.”\n", "\n", "## Code: A Selective SSM (S6)\n", "\n", "`SelectiveSSM` in `ssm.py` is a minimal, faithful S6 layer. The state\n", "matrix is diagonal and parameterized as $A = -\\exp(A_{\\log})$ so it\n", "stays negative (stable) for *any* weights; $\\Delta$, $B$, and $C$ are\n", "linear projections of the input, with $\\Delta = \\text{softplus}(\\cdot)$\n", "to keep it positive. The scan is the readable sequential\n", "`selective_scan` — Mamba’s engineering contribution is running that same\n", "scan fast on the GPU without ever materializing the state." ], "id": "0f1e7564-67e9-44f0-b843-c65f9724b045" }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "input shape: (2, 32, 16)\n", "output shape: (2, 32, 16)\n", "A all negative (stable): True\n", "input-dependent response: True" ] } ], "source": [ "from ssm import SelectiveSSM\n", "\n", "layer = SelectiveSSM(d_model=16, d_state=8)\n", "x = torch.randn(2, 32, 16) # (batch, length, channels)\n", "y = layer(x)\n", "print(\"input shape:\", tuple(x.shape))\n", "print(\"output shape:\", tuple(y.shape))\n", "\n", "# Selectivity check: A is guaranteed negative (stable) whatever the parameters are.\n", "A = -torch.exp(layer.A_log)\n", "print(\"A all negative (stable):\", bool((A < 0).all()))\n", "\n", "# Because B, C, Δ are read from the input, two different inputs get different\n", "# effective kernels — a fixed convolution could never do this.\n", "y2 = layer(torch.randn(2, 32, 16))\n", "print(\"input-dependent response:\", not torch.allclose(y, y2))" ], "id": "1c2d9647" }, { "cell_type": "markdown", "metadata": {}, "source": [ "A full **Mamba block** wraps this SSM with an input projection, a short\n", "causal depthwise convolution, and a SiLU gate (the same gated-MLP idea\n", "as m06’s SwiGLU) — but the *state-space* content is exactly what you\n", "built here.\n", "\n", "## Why They Challenge Attention\n", "\n", "Now the payoff, made quantitative. A self-attention layer does two\n", "matmuls that scale with the sequence: $QK^\\top$ and the weighted sum\n", "over $V$, each $O(L^2 d)$ — **quadratic**. An SSM does an $O(d\\,N)$\n", "update per token over $L$ tokens — $O(L\\,d\\,N)$, **linear**. At\n", "inference the difference is just as stark: attention’s KV cache holds\n", "$O(L)$ vectors and grows with every token; an SSM carries a $d\\times N$\n", "state that never grows. Drive the context length and watch the curves\n", "diverge:" ], "id": "170525f2-af9a-466d-9d54-7a1511712dc6" }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "from ssm import demonstrate_complexity\n", "lengths = [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536]\n", "rows = demonstrate_complexity(lengths, d_model=1024, d_state=16)\n", "ojs_define(ssm_cost = rows)" ], "id": "27c73373" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "d62f835e-c638-41e5-868f-47bb07b86625" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "30d5617c-f43c-40f1-8889-340f850122a4" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "8c61041c-de02-4090-acb7-e4dd48ac11b2" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> The gap is not a constant factor — it *widens* with context. That is\n", "> why SSMs are so attractive for long sequences: at 64K tokens the SSM\n", "> does a tiny fraction of the compute and keeps a state thousands of\n", "> times smaller than the KV cache. Mamba reports roughly **5× the\n", "> generation throughput** of a same-size Transformer.\n", "\n", "## Common Pitfalls\n", "\n", "When building SSMs, watch out for:\n", "\n", "1. **Forgetting $A$ must be negative.** Stability lives in\n", " $\\bar{A} = \\exp(\\Delta a)\n", " \\in (0, 1)$, which needs $a < 0$. Parameterize $A = -\\exp(A_{\\log})$\n", " so it can never drift non-negative and make the state explode.\n", "2. **Discretization sign errors.** ZOH is $\\bar{A}=\\exp(\\Delta A)$ (not\n", " $\\exp(-\\Delta\n", " A)$) and $\\bar{B}=A^{-1}(\\exp(\\Delta A)-I)B$. The near-$a=0$ case\n", " needs the Euler limit $\\bar{B}\\approx\\Delta B$ or you divide by\n", " zero.\n", "3. **Expecting a convolution after adding selectivity.** The moment\n", " $B$, $C$, or $\\Delta$ depend on the input, the kernel is no longer\n", " fixed and the convolutional face is gone. Use the **scan** — it does\n", " not assume constant coefficients.\n", "4. **Δ without softplus.** $\\Delta$ is a step size and must be\n", " positive; feeding a raw linear projection lets it go negative and\n", " flips decay into growth. Mamba uses $\\Delta=\\text{softplus}(\\cdot)$.\n", "5. **Confusing $N$ with $d$.** The state dimension $N$ (Mamba’s default\n", " is 16) is small and per-channel; the model width $d$ is large. SSM\n", " cost is $O(L\\,d\\,N)$ — linear in $L$ *because* $N$ is a small\n", " constant, not a second sequence axis.\n", "\n", "## Exercises\n", "\n", "### Exercise 1: The impulse response is the kernel" ], "id": "4d9d85c4-6c02-495a-a29c-e3034223b99e" }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "from ssm import ssm_recurrent, ssm_kernel\n", "import torch\n", "\n", "# Pick any diagonal SSM. Feed it a unit impulse (x = [1, 0, 0, ...]) through\n", "# ssm_recurrent, and separately build ssm_kernel of the same length. Show the two\n", "# are identical — the impulse response IS the convolution kernel. Then explain, in\n", "# one sentence, why that means recurrent and convolutional forms must agree.\n", "\n", "# Your implementation here:" ], "id": "a34f0072" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 2: ZOH vs forward Euler" ], "id": "42d59699-396a-4487-b169-ae17995c188e" }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "from ssm import discretize\n", "import torch\n", "\n", "# discretize uses exact ZOH. Write a forward-Euler discretization by hand\n", "# (A_bar = 1 + delta*A, B_bar = delta*B) and compare to ZOH for a = -1 across\n", "# delta in {0.1, 0.5, 1.0, 2.0}. Where do they agree, and where does Euler break\n", "# (hint: what happens to 1 + delta*A when delta*A < -1)?\n", "\n", "# Your implementation here:" ], "id": "106d8fcb" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 3: Selective copy" ], "id": "e9675846-853b-41a2-9168-422895f473bc" }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "from ssm import selective_scan\n", "import torch\n", "\n", "# Build a batch of 1, length 8, d_model 1 input. Set delta HIGH (e.g. 1.5) at one\n", "# chosen position and LOW (e.g. 0.02) everywhere else, with B = C = 1 and A = -1.\n", "# Show the output \"copies\" the chosen token's value and holds it afterwards — the\n", "# selective-copy task Mamba is designed to solve and a time-invariant SSM cannot.\n", "\n", "# Your implementation here:" ], "id": "bc665e88" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Summary\n", "\n", "Key takeaways:\n", "\n", "1. **An SSM is a linear recurrence over a fixed-size state.**\n", " $h_t = \\bar{A}h_{t-1}\n", " + \\bar{B}x_t$, $y_t = Ch_t$ — a running summary that never grows,\n", " giving $O(L)$ time and $O(1)$ decode memory.\n", "2. **Discretization connects continuous to discrete.** Zero-order hold\n", " turns $(A, B)$ into $\\bar{A}=\\exp(\\Delta A)$ and\n", " $\\bar{B}=A^{-1}(\\exp(\\Delta A)-I)B$; $A<0$ makes $\\bar{A}$ a stable\n", " decay in $(0,1)$.\n", "3. **One SSM has two faces.** Time-invariant, it is both a cheap\n", " **recurrence** (inference) and a parallel **convolution** (training)\n", " with kernel $\\bar{K} =\n", " (C\\bar{B}, C\\bar{A}\\bar{B}, \\dots)$ — provably the same function.\n", "4. **The recurrence is an associative scan.** $h_t = a_t h_{t-1} + b_t$\n", " composes associatively, so training parallelizes in $O(\\log L)$\n", " depth even without a fixed kernel.\n", "5. **Selectivity (Mamba’s S6) makes it content-aware.** Letting\n", " $\\Delta, B, C$ depend on the input turns a fixed filter into a model\n", " that gates what to remember — at the cost of the convolution,\n", " recovered by the selective scan.\n", "6. **SSMs challenge attention on cost.** Linear vs quadratic compute\n", " and constant vs growing decode memory; the advantage widens with\n", " context, which is why SSMs are a serious contender for long-sequence\n", " modeling.\n", "\n", "## What’s Next\n", "\n", "You now have a from-scratch alternative to attention itself — the last\n", "major architectural axis this book opens. The frontier from here is\n", "**hybrid stacks** that interleave SSM and attention layers (each\n", "covering the other’s weakness), and other sub-quadratic sequence mixers\n", "(linear and recurrent attention, RWKV, RetNet). The throughline is the\n", "one this module makes concrete: *the Transformer is one point in a\n", "design space*, and knowing the space is how you read — and build — what\n", "comes next.\n", "\n", "### Going Deeper\n", "\n", "**Core Papers:**\n", "\n", "- [Mamba: Linear-Time Sequence Modeling with Selective State\n", " Spaces](https://arxiv.org/abs/2312.00752) — Gu & Dao (2023), the\n", " selective SSM (S6), hardware-aware scan, and ~5× throughput built\n", " toward here.\n", "- [Efficiently Modeling Long Sequences with Structured State Spaces\n", " (S4)](https://arxiv.org/abs/2111.00396) — Gu, Goel & Ré (2021), the\n", " structured-SSM predecessor that first made SSMs competitive on\n", " long-range tasks.\n", "- [HiPPO: Recurrent Memory with Optimal Polynomial\n", " Projections](https://arxiv.org/abs/2008.07669) — Gu et al. (2020), the\n", " theory of how to initialize $A$ for long memory.\n", "- [Transformers are SSMs: Generalized Models and Efficient Algorithms\n", " (Mamba-2)](https://arxiv.org/abs/2405.21060) — Dao & Gu (2024), the\n", " duality tying state-space models and attention together.\n", "\n", "**Practical Resources:**\n", "\n", "- [The Annotated S4](https://srush.github.io/annotated-s4/) — Rush &\n", " Karamcheti, a from-scratch, runnable walkthrough of structured state\n", " spaces.\n", "- [state-spaces/mamba](https://github.com/state-spaces/mamba) — the\n", " official reference implementation of the selective scan." ], "id": "4c625fd1-0c31-407c-b33e-758e9d8d67df" } ], "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" } } }