The previous section ended on a confession: the S4D is LTI. Its kernel weights the past by position, never by content: the model decides what to remember before it sees the input.
The LSTM had the opposite profile: gates that read the data, but a nonlinear recurrence that trains sequentially.
This section: make the dynamics a function of the input (Gu & Dao, 2023):
the forget gate falls out a third time (engineering → calculus → selectivity),
the scan survives; the kernel view does not,
packaged as the Mamba block (2023): selective recurrence competitive at language-model scaling — ground S4’s long-range wins had not reached.
A task that defeats time invariance
Selective copying: symbols scattered among filler at random positions; reproduce them, in order, at the query slots.
Store a token because it is a symbol, not because of where it sits.
LM-scale counterpart: associative recall: the induction head we built in ch. 10; a synthetic recall benchmark predicts most of the attention-vs-SSM gap (Arora et al., 2024).
Generating the task
Filler is token 0, queries are token 1, symbols are ids 2-9:
def selective_copy(num_seqs, num_steps, num_marked, num_symbols, seed=42):"""Sequences of filler (0) with num_marked symbols at random positions, followed by num_marked query slots (1); targets are the symbols.""" rng = np.random.default_rng(seed) X = np.zeros((num_seqs, num_steps + num_marked), dtype=np.int64) X[:, num_steps:] =1# Query slots Y = rng.integers(2, 2+ num_symbols, (num_seqs, num_marked)) pos = np.argsort(rng.random((num_seqs, num_steps)), axis=1) pos = np.sort(pos[:, :num_marked], axis=1) # Marked positions, in order np.put_along_axis(X, pos, Y, axis=1)return X, Y -2# Classes 0..num_symbols-1class SelectiveCopy(d2l.DataModule):def__init__(self, num_train=8192, num_val=1024, batch_size=128, num_steps=256, num_marked=4, num_symbols=8):super().__init__()self.save_hyperparameters() X, Y = selective_copy(num_train + num_val, num_steps, num_marked, num_symbols)self.X, self.Y = d2l.tensor(X), d2l.tensor(Y)def get_dataloader(self, train): idx =slice(0, self.num_train) if train elseslice(self.num_train, None)returnself.get_tensorloader([self.X, self.Y], train, idx)copy_data = SelectiveCopy()
Two witnesses from earlier sections
The S4D stack of the previous section (imported from d2l) vs. the LSTM, both ~30k parameters, same harness: embed, encode, classify the query slots.
class CopyModel(d2l.Classifier):"""Read a token sequence; predict the symbols at the query slots."""def__init__(self, encoder, num_hiddens, vocab_size=10, num_marked=4, num_symbols=8, lr=3e-3):super().__init__()self.save_hyperparameters()self.emb = nn.Embedding(vocab_size, num_hiddens)self.head = nn.LazyLinear(num_symbols)def forward(self, X): Y =self.encoder(self.emb(X.T)) # (T, batch, hiddens) Y = Y[0] ifisinstance(Y, tuple) else Y # RNNs return a statereturnself.head(Y[-self.num_marked:]).movedim(0, 1)def configure_optimizers(self):return torch.optim.Adam(self.parameters(), lr=self.lr)
copy_curves = {}def train_copy(name, model, data, epochs=32): trainer = d2l.Trainer(max_epochs=epochs, gradient_clip_val=1, num_gpus=1) trainer.fit(model, data) pts = model.board.data['val_acc'] copy_curves[name] = ([p.x for p in pts], [float(p.y) for p in pts])print(f'{name}: final validation accuracy {copy_curves[name][1][-1]:.3f}')s4d = CopyModel(nn.Sequential(*[S4DBlock(48, 4) for _ inrange(2)]), num_hiddens=48)train_copy('S4D', s4d, copy_data)
S4D: final validation accuracy 0.611
The LTI model stalls; the gated one solves it
names =list(copy_curves)d2l.plot([copy_curves[n][0] for n in names], [copy_curves[n][1] for n in names], 'epoch','validation accuracy', legend=names, figsize=(5, 3))
LSTM: climbs to solved. Gates read the data; “store this, it is a symbol” is learnable.
S4D: plateaus. Pointwise nonlinearities can suppress filler locally, but a fixed kernel cannot align “third symbol seen” with “third query slot” under variable spacing.
The task is built to stress content-independent dynamics — the mechanism the S4D lacks is content-dependent gating.
Selective SSMs: let the input choose the dynamics
One change to the S4D recipe: step size, input matrix, read-out become functions of the input,
Recall the step-size box: \Delta \to 0 freezes the state, \Delta large overwrites it. Input-dependent \Delta_t = forget and input gate in one scalar, on a linear state:
Price: time-varying coefficients kill the convolution view. No fixed \bar{\mathbf{K}}, no FFT. Of the three views only the recurrence survives.
Save: the recurrence is still an affine map of the state; the combine (\mathbf{a}_2, \mathbf{b}_2) \circ (\mathbf{a}_1, \mathbf{b}_1) never assumed constant coefficients. Same scan, per-step tensors:
num_steps, num_states =100, 4a_t = torch.rand(num_steps, num_states) # Per-step decays in (0, 1)b_t = torch.randn(num_steps, num_states) # Per-step inputsh, ys = torch.zeros(num_states), []for t inrange(num_steps): h = a_t[t] * h + b_t[t] ys.append(h)err = (associative_scan(a_t, b_t) - torch.stack(ys)).abs().max()print(f'time-varying scan vs loop: {float(err):.2e}')
time-varying scan vs loop: 4.77e-07
The selective SSM layer
Three changes vs. S4D: \Delta, \mathbf{B}, \mathbf{C} from linear heads; low-rank \Delta head (dt_rank); bias init gives the untrained layer S4D’s multi-timescale decays (\mathbf{B}_t,
\mathbf{C}_t are input-dependent from the start):
class SelectiveSSM(nn.Module):"""A diagonal SSM whose step size, input matrix, and read-out are functions of the input (Gu & Dao, 2023)."""def__init__(self, num_hiddens, num_states=4, dt_min=0.001, dt_max=0.1):super().__init__() H, N, R = num_hiddens, num_states, max(2, num_hiddens //16)self.log_a = nn.Parameter( torch.log(torch.arange(1., N +1)).repeat(H, 1))self.W_dt = nn.Sequential(nn.Linear(H, R), nn.Linear(R, H, bias=False)) dt = torch.exp(torch.rand(H) * math.log(dt_max / dt_min)+ math.log(dt_min))self.b_dt = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))self.W_B = nn.Linear(H, N, bias=False)self.W_C = nn.Linear(H, N, bias=False)self.D = nn.Parameter(torch.ones(H))def forward(self, u): # (num_steps, batch, num_hiddens) a =-torch.exp(self.log_a) # (H, N), Re(a) < 0 dt = F.softplus(self.W_dt(u) +self.b_dt) # (T, batch, H) B, C =self.W_B(u), self.W_C(u) # (T, batch, N) a_bar = torch.exp(dt.unsqueeze(-1) * a) # (T, batch, H, N) b_bar = (dt * u).unsqueeze(-1) * B.unsqueeze(-2) x = associative_scan(a_bar, b_bar)return (x * C.unsqueeze(-2)).sum(-1) +self.D * u
Why 2023 and not 2020: hardware-awareness
Our implementation materializes (T, batch, H, N) coefficient tensors: fine at textbook scale, ruinous at model scale (the compute is trivial; memory traffic sets the speed).
The Mamba kernel:
fuses discretization + scan + read-out in one pass, intermediates in on-chip SRAM,
recomputes them in the backward pass instead of storing (store-vs-recompute, cf. 8.6),
changes nothing about what is computed, only whether it was worth computing.
The Mamba block
Block and language model
A dozen lines around SelectiveSSM; the stack keeps the (inputs, state) interface, so d2l.RNNLM wraps it unchanged:
class MambaBlock(nn.Module):"""Conv + SiLU + selective SSM, gated, inside a pre-norm residual."""def__init__(self, num_hiddens, num_states=4, expand=2, conv_width=4, dropout=0):super().__init__() d = expand * num_hiddensself.ln = nn.LayerNorm(num_hiddens)self.W_in = nn.Linear(num_hiddens, 2* d)self.conv = nn.Conv1d(d, d, conv_width, groups=d, padding=conv_width -1)self.ssm = SelectiveSSM(d, num_states)self.W_out = nn.Linear(d, num_hiddens)self.drop = nn.Dropout(dropout)def forward(self, X): # (num_steps, batch, num_hiddens) u, gate =self.W_in(self.ln(X)).chunk(2, -1) u =self.conv(u.permute(1, 2, 0))[..., :X.shape[0]] # Causal: trim y =self.ssm(F.silu(u.permute(2, 0, 1)))return X +self.drop(self.W_out(y * F.silu(gate)))class Mamba(d2l.Module):"""A stack of Mamba blocks with the recurrent-cell interface."""def__init__(self, num_inputs, num_blocks=2, num_states=4, dropout=0):super().__init__()self.save_hyperparameters()self.num_hiddens = num_inputs # Output width, for headsself.blocks = nn.Sequential(*[ MambaBlock(num_inputs, num_states, dropout=dropout)for _ inrange(num_blocks)])self.ln = nn.LayerNorm(num_inputs)def forward(self, X, state=None):returnself.ln(self.blocks(X)), None
Capstone: the chapter’s three answers on one task
Gate it (LSTM) vs. linearize it (minGRU) vs. select it (Mamba); Time Machine BPE, 50k windows of 32, ten epochs, clip 1:
print(f'{"model":>7}{"val ppl":>8}{"bpb":>6}{"params":>9}{"s/epoch":>8}')for name, (ppl, bpb, params, secs) in results.items():print(f'{name:>7}{ppl:>8.1f}{bpb:>6.2f}{params:>9,}{secs:>8.1f}')
Mamba: best perplexity of the chapter in the PyTorch run; the minGRU edges it out in the JAX run.
At this scale, recipe and initialization move the board by amounts comparable to the architectural gap.
Caveats: Adam vs. the baselines’ SGD; more parameters; a small corpus rewards memorization.
Stepping the selective model
Selectivity does not break stepping: \Delta_t, \mathbf{B}_t,
\mathbf{C}_t come from the current token, so each step builds its own coefficients. One new piece of state: the causal conv needs its last 4 inputs → a rolling buffer: (conv buffer, SSM state) is exactly what production Mamba caches.
@d2l.add_to_class(SelectiveSSM)def step(self, u, x=None):"""Advance one token: u is (batch, H); x is the (batch, H, N) state.""" a =-torch.exp(self.log_a) dt = F.softplus(self.W_dt(u) +self.b_dt) # (batch, H) B, C =self.W_B(u), self.W_C(u) # (batch, N) a_bar = torch.exp(dt.unsqueeze(-1) * a) # (batch, H, N) b_bar = (dt * u).unsqueeze(-1) * B.unsqueeze(-2) x = b_bar if x isNoneelse a_bar * x + b_barreturn (x * C.unsqueeze(-2)).sum(-1) +self.D * u, x@d2l.add_to_class(MambaBlock)def step(self, X, state=None):"""Advance one token, carrying (conv buffer, SSM state).""" u, gate =self.W_in(self.ln(X)).chunk(2, -1)if state isNone: state = (u.new_zeros(*u.shape, self.conv.kernel_size[0]), None) buf, x = state buf = torch.cat([buf[..., 1:], u.unsqueeze(-1)], -1) # Roll the window u = (buf *self.conv.weight[:, 0]).sum(-1) +self.conv.bias y, x =self.ssm.step(F.silu(u), x)return X +self.drop(self.W_out(y * F.silu(gate))), (buf, x)@d2l.add_to_class(Mamba)def step(self, X, state=None):"""Advance the stack one token: X is (batch, d); one state per block.""" state = [None] *len(self.blocks) if state isNoneelselist(state)for i, blk inenumerate(self.blocks): X, state[i] = blk.step(X, state[i])returnself.ln(X), state
Stepped == scanned, on the trained LM
Push validation windows through the trained capstone twice: parallel scan vs. token-by-token step; the logits must agree (dropout off):
mamba_lm.eval() # Dropout offX = data.X[data.num_train:data.num_train +8].to(d2l.try_gpu())with torch.no_grad(): logits_scan = mamba_lm(X) # (batch, T, vocab) state, cols =None, []for t inrange(X.shape[1]): # One token at a time emb = mamba_lm.emb(X[:, t].unsqueeze(0)) # (1, batch, d) y, state = mamba_lm.rnn.step(emb[0], state) cols.append(mamba_lm.linear(y)) logits_step = torch.stack(cols, 1)err =float((logits_step - logits_scan).abs().max())scale =float(logits_scan.abs().max())print(f'stepped vs scanned logits: deviation {err:.2e}, 'f'relative {err / scale:.2e}')assert err <1e-3* scale
stepped vs scanned logits: deviation 6.20e-06, relative 5.47e-07
Sampling all three (8.7’s toolkit)
Baselines re-run the prefix per token; Mamba decodes through step: prefill the prompt once, then one state update per token: O(1), like a KV cache that never grows.
def step_fn(model):def step(ids): # Token ids in, numpy logits for the next token outwith torch.no_grad(): logits = model(d2l.tensor([ids], device=d2l.try_gpu()))return d2l.numpy(logits)[0, -1]return stepdef stepped_fn(model): state, seen =None, 0def step(ids): # Consume only the tokens the state has not absorbednonlocal state, seenwith torch.no_grad():for i in ids[seen:]: emb = model.emb(d2l.tensor([[i]], device=d2l.try_gpu())) y, state = model.rnn.step(emb[0], state) seen =len(ids)return d2l.numpy(model.linear(y))[0]return stepprefix = data.tokenizer.encode('the time traveller')for name, model, fn in [('LSTM', lstm_lm, step_fn), ('minGRU', mingru_lm, step_fn), ('Mamba', mamba_lm, stepped_fn)]: out = d2l.generate(fn(model), prefix, 25, strategy='sample', temperature=1.0, min_p=0.1, rng=np.random.default_rng(0))print(f'{name:>7}: {data.tokenizer.decode(out)!r}')
LSTM: 'the time travelleren the Time Traveller and diling with cry. Weve it will was the\nside and mon'
minGRU: 'the time travellered, and the unlination. And\nthe mean, and it seemed to the faint, and a'
Mamba: 'the time traveller, I was not believent the lever, and\nfolden in the sunt of the d'
Selective copying, revisited
names =list(copy_curves)d2l.plot([copy_curves[n][0] for n in names], [copy_curves[n][1] for n in names], 'epoch','validation accuracy', legend=names, figsize=(5, 3))
Mamba solves what the S4D could not, faster than the LSTM. The content-dependent gating we deleted to linearize is restored, without giving up the scan.
The mechanism, measured
Designed: filler collapses \Delta_t, symbols open it. Measured, on the trained model (one forward pass, no retraining):
The narrated gate is what training actually found.
Summary: one question, three answers (so far)
Selectivity from the state-space side: continuous dynamics, discretize, let the input set the step size. Trains as a scan, decodes as a constant-size state.
Next section, the other road: start from the linear-attention recurrence of ch. 10 and arrive at the same recurrence.
