"A fixed memory forgets on a schedule. I decided to forget only when the signal told me nothing."
A Selectively Attentive AI Agent
Prerequisites
This section builds directly on the structured state-space models of Section 16.2: the linear recurrence \(h_t = \bar{A}h_{t-1} + \bar{B}x_t\), the diagonal-plus-low-rank parameterization of S4, and the dual view where the same operator runs as a fast parallel convolution during training and a constant-memory recurrence at inference. It assumes the quadratic-cost argument for attention from Section 16.1 and comfort with the multivariate sensor tensor of shape \((\text{channels} \times \text{time})\) from Chapter 3. Familiarity with gated recurrences from Chapter 14 will make the selection mechanism feel natural.
The Big Picture
S4 gave us a linear-time sequence model, but it paid for that speed with a rigid memory: its state-update matrices are the same at every timestep, chosen once and applied blindly whether the current sample is a meaningful event or dead air between events. That is fine for smooth, stationary signals and clumsy for the bursty, event-sparse streams that dominate sensing, where a footstep, a heartbeat, or a fault transient hides in long stretches of nothing. Mamba fixes this by making the state-space parameters functions of the input. At each step the model reads the current sample and decides how much to write into memory and how fast to let old memory decay. This one change, called selection, lets a linear-time recurrence do attention's job, deciding what matters, without attention's quadratic bill. The price is that the fast convolution trick of S4 no longer applies, so Mamba is made practical by a hardware-aware parallel scan that keeps the whole recurrence on the GPU's fast memory.
The limitation selection removes: linear time-invariance
Every S4-style model of Section 16.2 is linear time-invariant (LTI). The discretized matrices \(\bar{A}\), \(\bar{B}\), and \(\bar{C}\) are learned parameters that do not depend on \(t\) or on the data flowing through them. That property is exactly what enables the convolutional view: because the recurrence weights are constant, unrolling it produces a single fixed kernel that can be applied by FFT in \(O(L \log L)\). The constancy that buys the speed is also the model's cage. An LTI system processes the tenth sample of a quiet baseline with the identical dynamics it applies to a sharp impact spike. It cannot say "this sample matters, hold it" or "this stretch is noise, let it fade," because the coefficients that would make that decision are frozen at training time.
For sensor streams this is a real handicap. Consider an inertial measurement unit on a wrist: minutes of near-constant gravity vector punctuated by a half-second gesture. A useful model must remember the gesture for many steps while aggressively discarding the flat baseline, and which is which depends on content, not position. LTI models solve a fixed convolution problem; they have no mechanism to route information conditionally. Attention has exactly that mechanism, comparing content to content, which is why Chapter 15 models excel at selective recall, and why they cost \(O(L^2)\) to do it. Selection is the attempt to get the routing without the quadratic bill.
The selection mechanism: input-dependent parameters
Mamba, introduced by Gu and Dao in 2023 (the architecture is also called S6, for "selective S4"), makes three of the state-space quantities functions of the input token \(x_t \in \mathbb{R}^{D}\). Instead of a single learned \(B\), \(C\), and step size \(\Delta\), each is produced by a small linear projection of the current sample:
$$B_t = \mathrm{Lin}_B(x_t), \qquad C_t = \mathrm{Lin}_C(x_t), \qquad \Delta_t = \mathrm{softplus}\big(\mathrm{Lin}_\Delta(x_t)\big)$$
The continuous \(A\) matrix stays fixed and diagonal, but the discretization now uses the per-step \(\Delta_t\), so the effective transition \(\bar{A}_t = \exp(\Delta_t A)\) and input map \(\bar{B}_t\) change at every timestep. The recurrence becomes time-varying:
$$h_t = \bar{A}_t \, h_{t-1} + \bar{B}_t \, x_t, \qquad y_t = C_t \, h_t$$
Read \(\Delta_t\) as a data-driven gate, the same idea as the forget gate of an LSTM from Chapter 14, now expressed in continuous-time language. A large \(\Delta_t\) means "take a big step": strongly incorporate the current input and let older state decay quickly, the right response to a salient event. A small \(\Delta_t\) means "barely move": ignore this sample and preserve the accumulated memory, the right response to a flat, uninformative stretch. Because \(B_t\) and \(C_t\) are also input-dependent, the model additionally chooses what of the current sample to write into state and what to read out, giving it content-based routing along the whole sequence.
Key Insight
Selection converts the step size \(\Delta_t\) into a learned, per-sample resolution knob. Where classical sampling theory (Chapter 3) fixes your temporal resolution once at acquisition time, Mamba lets the model re-choose an effective resolution at every timestep: fine when something is happening, coarse when nothing is. This is why a selective SSM can skim thousands of idle samples cheaply and still snap to full attention on the rare informative burst, exactly the behavior event-sparse sensor streams demand.
Common Misconception
It is tempting to read "input-dependent \(B_t\), \(C_t\), \(\Delta_t\)" as "Mamba gets attention's retrieval power for free." It does not. Attention keeps every past token available and looks any of them up by content, at \(O(L^2)\) cost. Selection only controls how the current sample is written into and read from a single fixed-size state \(h_t \in \mathbb{R}^{D \times N}\); once history is compressed into that state, the specific timestep that produced it is gone, and there is no way to later query "what happened at \(t=50\)" the way attention can. Selection decides how much to remember, not where to look back, which is why streams that need precise, long-range, content-addressed retrieval push toward the hybrid SSM-attention designs of Section 16.6 instead of a selective SSM alone.
Why selection breaks convolution, and how the scan restores speed
Making the parameters vary with the input destroys linear time-invariance, and with it the convolutional kernel: there is no longer a single fixed filter to FFT, because the "filter" is different at every position. Naively that pushes you back to a sequential recurrence, \(L\) steps that cannot be parallelized, which is death on accelerators that want thousands of lanes busy at once. Mamba's second contribution is what makes it usable: the time-varying recurrence is computed with a parallel associative scan rather than run step by step. Write each timestep as a pair \((a_t, b_t) = (\bar{A}_t,\ \bar{B}_t x_t)\) and define the combination \((a_2,b_2)\circ(a_1,b_1) = (a_2 a_1,\ a_2 b_1 + b_2)\); this operator is associative, so segments of the sequence can be merged in any order and still land on the result of merging strictly left to right. A plain recurrence lacks that property; a Blelloch-style scan exploits it by building a balanced binary tree over the \(L\) steps, combining adjacent segments pairwise on an upward pass, then distributing the accumulated prefixes back down, so every \(h_t\) is ready after \(O(\log L)\) sequential rounds instead of \(L\). Total work is still \(O(L)\), the same multiply-adds as the sequential loop, but spread over \(\log L\) parallel rounds rather than \(L\) serial ones, which is why the trick wins on a GPU with thousands of idle lanes and barely helps a single CPU core, where depth and work cost the same wall clock. It stays cheap because \(A_t\) is diagonal, so each merge is an elementwise product, not a matrix multiply.
The authors go further with a hardware-aware implementation: the expanded state (batch \(\times\) channels \(\times\) state dimension) is large enough that shuttling it to and from slow GPU high-bandwidth memory at every step would erase the speedup, so Mamba's kernel keeps the scan inside fast on-chip SRAM, materializing the big state only transiently and recomputing it during the backward pass instead of storing it, a fusion strategy in the spirit of FlashAttention. The result trains in parallel and at inference runs as a plain recurrence with \(O(1)\) memory per step, the constant-state property that makes SSMs attractive for streaming and on-device sensing in the first place.
Practical Example: continuous fall detection on a hearing-aid IMU
A hearing-aid maker wants fall detection to run forever on a tiny battery with a 6-axis IMU sampled at 100 Hz. A transformer over a 30-second context is 3000 tokens per channel with attention memory that grows with the window, so the team is forced into short windows that miss the slow pre-fall instability. They switch the encoder to a two-layer Mamba block. During ordinary walking and sitting, the learned \(\Delta_t\) stays small, so the recurrence coasts and the state barely changes, keeping average energy low; when the acceleration signature of a stumble arrives, \(\Delta_t\) spikes, the model writes the transient sharply into state, and \(C_t\) surfaces it to the classifier. Because inference is a fixed-size recurrence, memory footprint is constant whether the context is 30 seconds or 30 minutes. On their test set the selective model matches the short-window transformer's precision while cutting steady-state inference energy by roughly half, the difference between a feature that ships and one that drains the battery by noon.
A minimal selective recurrence
The listing below implements the core of one selective SSM channel as a readable sequential recurrence, the inference-time form. It is not the fused parallel-scan kernel you would train with, but it makes the time-varying update concrete: notice that \(\Delta\), \(B\), and \(C\) are recomputed from each input sample inside the loop, which is precisely what an LTI S4 layer does not do.
import torch
import torch.nn.functional as F
def selective_ssm_step(x, A, W_delta, W_B, W_C):
# x: (T, D) one sequence, D channels. A: (D, N) diagonal state dynamics (log-neg).
T, D = x.shape
N = A.shape[1]
h = torch.zeros(D, N) # constant-size hidden state
ys = []
for t in range(T): # sequential = the streaming form
xt = x[t] # (D,)
delta = F.softplus(x[t] @ W_delta) # (D,) input-dependent step size
B = x[t] @ W_B # (N,) input-dependent write map
C = x[t] @ W_C # (N,) input-dependent read map
A_bar = torch.exp(delta[:, None] * A) # (D, N) per-step decay, A<0
B_bar = delta[:, None] * B[None, :] # (D, N) discretized input gain
h = A_bar * h + B_bar * xt[:, None] # selective state update
ys.append(h @ C) # (D,) read out
return torch.stack(ys) # (T, D)delta and the maps B, C are recomputed from every sample, so the dynamics vary with content. Training uses a mathematically equivalent parallel scan instead of this Python loop.Library Shortcut
You never write the scan or the fused CUDA kernel yourself. The reference mamba-ssm package exposes the whole thing as from mamba_ssm import Mamba; layer = Mamba(d_model=128, d_state=16, d_conv=4, expand=2), a single line that replaces roughly 150 lines of hand-rolled discretization, associative-scan logic, and memory-aware backward pass, dispatching to the hardware-aware kernel automatically. Feed sensor tensors as \((\text{batch} \times \text{time} \times \text{channels})\), just like a transformer block. When the CUDA kernel is unavailable (CPU or a bare edge target), the same library offers a slower pure-PyTorch fallback so your code stays identical across environments.
Research Frontier
The current state of the art is Mamba-2 (Dao and Gu, 2024), whose State Space Duality framework reframes the selective scan as structured masked attention, unifying the SSM and attention views and unlocking a matmul-based kernel several times faster to train while allowing much larger state dimensions. Selection has since moved from research idea to production component: hybrid models such as AI21's Jamba and Zyphra's Zamba2 interleave Mamba-2 blocks with a handful of full-attention layers, keeping most of the compute linear while still getting exact retrieval where it counts. On the sensing side, selective SSMs are the backbone of a wave of 2024 to 2026 architectures covered next in Section 16.4, including bidirectional variants that let non-causal, offline sensor analysis see both past and future, an option streaming attention cannot cheaply match. Open questions remain around how well a scalar-gated linear recurrence captures very long-range, precisely-indexed dependencies compared with full attention, which motivates the hybrid designs of Section 16.6.
Exercise
Take the selective_ssm_step function above and construct a 500-step input that is zero everywhere except for a single unit impulse at \(t = 50\). Run it twice: once forcing delta to a small constant (say 0.01) and once forcing it large (say 2.0). Plot the resulting hidden-state magnitude over time for both. Explain, in terms of \(\bar{A}_t = \exp(\Delta_t A)\), why the small-\(\Delta\) run remembers the impulse far longer than the large-\(\Delta\) run, and connect this to when you would want each behavior on a real event-sparse stream. Then make \(\Delta\) input-dependent again and verify that the model preserves the impulse while decaying the surrounding zeros.
Self-Check
- Why does making \(B\), \(C\), and \(\Delta\) input-dependent prevent Mamba from being trained as a single fixed convolution, and what replaces the convolution?
- An LTI S4 layer and a selective S6 layer both run at inference as an \(O(1)\)-memory recurrence. What capability does the selective layer gain, and on what kind of sensor signal would you expect the gap between them to be largest?
- The parallel scan has \(O(L)\) work but \(O(\log L)\) depth. Why does that depth number, rather than the work number, explain why Mamba trains efficiently on a GPU?
What's Next
In Section 16.4, we take the selective state-space block from a general-purpose primitive to purpose-built sensor and time-series architectures: SiMBA's channel-mixing for multivariate streams, TimeMachine's multi-scale design, and the bidirectional Bi-Mamba variants that let offline analysis attend to both past and future. You will see how the one idea developed here, input-dependent dynamics inside a linear-time recurrence, gets wired into models that compete directly with the transformers of Chapter 15 on long-horizon sensor tasks.