mamba, from scratch | khushi keswani

what i built

the mechanism, not the library — no mamba_ssm. the recurrence starts as a small numpy prototype, becomes input-dependent, and only then turns into a pytorch module.
a reference scan — a sequential selective scan written as the correctness reference, plus a parallel-style prefix formulation the prototype checks against it.
the block and the stack — MambaBlock — input projection, causal depthwise conv, silu, Δ/B/C projection, selective scan, gated output and residual — stacked four layers deep at d_model=256, d_state=16.
stability guards — A parameterized as -exp(A_log), Δ clamped before softplus, Ā capped at 0.999, gradient norm clipped to 1.0, with finite-output and gradient checks after every build.
a real task — trained as a binary anomaly classifier on 11,412 hdfs event traces and evaluated on a held-out split: 0.9504 test f1, 0.9825 roc-auc, 0.9621 pr-auc.

how it fits together

figure 1each trace becomes event ids, gets embedded into 256 dimensions, passes through four mamba blocks, and is classified from the final timestep.(scroll the figure sideways)

figure 2inside one block: the input projection splits into an ssm branch and a gate branch. Δ, B and C are generated per token, discretized, scanned, gated, projected back and added to the residual.(scroll the figure sideways)

from ssm to selective ssm

a classic state-space model keeps a hidden state and updates it with fixed matrices, so every token is treated the same way. mamba makes the step size Δ and the matrices B and C functions of the current input. after discretization, the model chooses per token what to write into the state and what to ignore.

h_t = Ā_t \cdot h_(t-1) + B̄_t \cdot x_t y_t = C_t \cdot h_t + D \cdot x_t

the input-dependent parameters come from learned projections, and Δ goes through a numerically stable softplus:

Δ_t = softplus(W_Δ \cdot x_t + b_Δ) B_t = W_B \cdot x_t C_t = W_C \cdot x_t

for the continuous-time transition to stay stable, A is parameterized as -exp(A_log) and discretized by exp(ΔA), which gives Ā and B̄:

A = -exp(A_log) Ā_t = exp(Δ_t \cdot A) B̄_t = Δ_t \cdot B_t

that single change is the whole idea. cost grows linearly with sequence length, and generation needs a fixed-size state instead of a growing key-value cache.

the scan

because Δ, B and C change at every step, the recurrence can't collapse into one convolution. once Ā_t and b_t = B̄_t · x_t exist, each timestep is an affine transformation of the previous state:

h_t = a_t \cdot h_(t-1) + b_t

two transitions compose into one equivalent transition, and the operation is associative — so a sequence of transitions can be grouped into prefix operations instead of a strictly sequential chain:

(a2, b2) \otimes (a1, b1) = (a2 \cdot a1, a2 \cdot b1 + b2)

figure 3the numpy prototype computes states sequentially and through the prefix formulation, then compares them with np.allclose and reports "states match: true".(scroll the figure sideways)

the toy run was 5 steps: 0.000288 s sequential against 0.000112 s for the parallel-style prefix implementation. that timing is a demonstration that the formulation works, not a gpu benchmark — the sequence is length 5 and it's still python-land.

def selective_scan(x, delta, A, B, C):
    # x, delta: (b, L, d)   A: (d, n)   B, C: (b, L, n)
    b, L, d = x.shape
    h = x.new_zeros(b, d, A.shape[1])
    ys = []
    for t in range(L):
        dA = torch.exp(delta[:, t, :, None] * A)          # (b, d, n)
        dBx = delta[:, t, :, None] * B[:, t, None, :] * x[:, t, :, None]
        h = dA * h + dBx
        ys.append((h * C[:, t, None, :]).sum(-1))
    return torch.stack(ys, 1)

mamba_block.py

the shipped implementation runs that recurrence sequentially and compiles it with torch.compile to cut python overhead rather than writing a custom kernel.

the bug i'm still chasing

the pytorch model is trainable, finite and gradient-clean — a random [2, 128, 256] input returns the same shape with no missing or non-finite gradients — but that does not establish exact parity with the prototype, and i'd rather leave the mismatch visible than hide it behind a clean-looking number.

the next correctness gate is isolating the source: shared weights, identical padding and convolution semantics, and step-by-step intermediate comparisons rather than comparing only the final output.

clamp raw dt — to [-20, 10] before softplus, so a stray logit can't produce an absurd step size.
keep A negative — A = -exp(A_log) keeps the continuous transition negative and the dynamics stable.
cap Ā — at 0.999, so the state leaks rather than being remembered forever.
check every build — finite outputs, no missing or non-finite gradients, and gradient norm clipped to 1.0.

the hdfs experiment

each trace is converted into event ids, padded or truncated to a fixed length, embedded, and labeled binary. the split is stratified with a fixed random state of 42, and labels are joined by BlockId.

dataset property	observed value
traces	11,412
unique events	23
vocabulary including pad	24
average trace length	24.99
median trace length	25
maximum trace length	269
normal examples	10,921
anomaly examples	491
anomaly ratio	~4.3%

split	samples	approx. anomaly ratio
train	7,988	4.306%
validation	1,712	4.322%
test	1,712	stratified split

training and results

the classifier is an embedding layer, four mamba blocks, final-timestep pooling and a one-unit linear head, trained with AdamW and BCEWithLogitsLoss at a learning rate of 1e-5 for three epochs.

epoch	train loss
1	0.2870
2	0.0672
3	0.0478

metric	validation
precision	0.9844
recall	0.8514
f1	0.9130
roc-auc	0.9474
pr-auc	0.9063

validation loss 0.06031, all logits finite.

metric	held-out test
f1	0.9504
roc-auc	0.9825
pr-auc	0.9621

pr-auc is the number that matters here, because only about 4.3% of the examples are anomalies — roc-auc alone looks flattering on data that imbalanced. the recall gap between validation (0.8514) and the thresholded test f1 is the honest tension in this run: precision is easy, finding every anomaly is not.

decisions and trade-offs

decision	why	trade-off
numpy prototype first	make the recurrence mathematically inspectable.	extra implementation work before pytorch.
keep A negative	encourage stable state dynamics.	limits the parameterization versus unconstrained A.
input-dependent Δ/B/C	captures mamba's selective behaviour.	more per-token computation.
d_state = 16	small enough to experiment with on colab.	may limit state capacity.
four blocks	enough depth for a meaningful classifier.	not comparable to large pretrained mamba models.
torch.compile the scan	cut python overhead without writing a kernel.	still not a fused selective-scan cuda kernel.
final-timestep readout	simplest sequence-classification head.	may discard information from earlier positions.
hdfs anomaly classification	a concrete long-ish sequence task with real labels.	sequence cap of 64 while traces reach 269.

what worked, what didn't

worked — the selective recurrence is implemented explicitly instead of hidden behind a library, and the associative formulation reproduced the sequential toy states.
worked — MambaBlock returns the expected [batch, sequence, d_model] shape, the gradient checks are clean, and the four-layer stack trained to 0.9504 test f1.
unresolved — numpy and pytorch forwards aren't numerically equivalent yet, and the scan is a compiled sequential recurrence, not a true gpu parallel scan.
unresolved — there's no same-parameter transformer baseline, so this can't make a mamba-vs-transformer claim.
unresolved — only the final timestep is used, and the sequence is truncated at 64 while traces reach 269 — truncation can throw information away.

the paper said “selective state space model.” i said bet — then found out which parts i hadn't actually verified.

next project → repograph all work ↑