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-1
class 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 else slice(self.num_train, None)
return self.get_tensorloader([self.X, self.Y], train, idx)
copy_data = SelectiveCopy()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).
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, num_features=None, lr=3e-3, rngs=None):
super().__init__()
self.save_hyperparameters(ignore=['encoder', 'rngs'])
rngs = nnx.Rngs(1) if rngs is None else rngs
num_features = num_hiddens if num_features is None else num_features
self.encoder = encoder
# Embed by one-hot matmul: on a table this tiny (ten rows), the
# scatter-add in an embedding lookup's gradient makes 33,000 colliding
# updates to ten rows and serializes the whole training step on GPU;
# as a dense product, forward and backward are ordinary matmuls
self.emb = nnx.Linear(vocab_size, num_hiddens, use_bias=False,
rngs=rngs)
self.head = nnx.Linear(num_features, num_symbols, rngs=rngs)
def forward(self, X):
E = self.emb(jax.nn.one_hot(X.T, self.vocab_size))
Y = self.encoder(E) # (T, batch, hiddens)
Y = Y[0] if isinstance(Y, tuple) else Y # RNNs return a state
return self.head(Y[-self.num_marked:]).swapaxes(0, 1)
def configure_optimizers(self):
return optax.adam(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(nnx.Sequential(*[S4DBlock(48, 4) for _ in range(2)]),
num_hiddens=48)
train_copy('S4D', s4d, copy_data)
S4D: final validation accuracy 0.658The 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.
The Mamba block
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.