End-to-end demonstration that training the real selective-scan Mamba buys
a genuine temporal-modeling win — the payoff of M1–M3.
Task: next-token prediction on a multi-regime sequence (x[t] = μ_regime +
noise). Predicting x[t+1] inside a regime requires integrating recent
history to average out the noise — a memoryless model can't.
Result (same sequence, all three):
persistence (memoryless) MSE = 0.0895
untrained Mamba MSE = 0.2167
trained Mamba (400 Adam) MSE = 0.0002
The trained backbone integrates history to de-noise the regime mean,
beating the memoryless persistence baseline by ~450× and improving
~1000× over its untrained self. (Single-sequence fit: demonstrates the
SSM's temporal-modeling capacity, not held-out generalization.)
This replaces the old non-result ("random Mamba 6.5% vs linear 48%") with
a real "trained SSM exploits temporal structure" demonstration.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Plain randn(0,1) init made A=-exp(A_log) and Δ wildly large, overflowing
the real exp(Δ·A) scan to NaN (the stub never cared — it discarded these
weights). new()/new_seeded() now share a build() with canonical S6 init:
- A_log = ln(1..=d_state) ⇒ A = -(1..=d_state), bounded
- dt_bias so softplus(dt_bias) ≈ 0.01 (small, stable Δ; near-identity
scan at init — intentional for gradient flow)
- D = 1, zero conv bias, projections scaled by 1/√fan_in (capped 0.5)
This fixes the NaN that broke omni-cortex's d231 action-conditioned
predictor training (now green). Seeded determinism preserved.
Tests: active_block helper (Δ overridden to ≈0.69) exercises the
scan-active regime so the liveness check can observe each weight; the
training test asserts a seed-varying backbone weight (conv1d_weight)
moves. All 6 selective-scan tests green incl. the finite-diff grad check.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Demonstrates end-to-end trainability of the real selective scan. A
self-contained Adam loop (the rtx-transformers AdamOptimizer has no
public gradient setter — the spec permits a bespoke loop) fits a teacher
block's output on a fixed input: forward → MSE → analytic backward →
Adam step → rebuild. Over 200 steps the loss drops >50% and the backbone
weight A_log moves, confirming gradients actually train the model (not
just the head). All 6 selective-scan tests green.
The production AdamOptimizer can be wired once it exposes a gradient
setter; the M2 backward already returns grads in its HashMap shape.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Hand-written VJP of the selective-scan forward (Phase-3 spec §5/§7):
MambaBlock::backward(x, d_out) -> per-parameter gradients keyed by
persistence name, summed over the batch. Differentiates the scan
analytically (reverse-time recurrence over the cached h trajectory)
rather than through the immature rtx-tensor autograd tape.
Covers every parameter: in_proj, conv1d_weight, conv1d_bias, A_log
(via A=-exp(A_log) ⇒ dA_log = dA·A), x_proj, dt_proj, dt_bias, D,
out_proj. Adds stable sigmoid_f32 / silu_grad_f32 helpers.
New test analytic_gradients_match_finite_differences: on a small
well-conditioned instance, ≥30 sampled grad elements across all 9
params match central finite differences within (5e-3 + 5e-2·|fd|).
All 5 selective-scan tests green.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
The MambaBlock forward was a stub — SelectiveScan::forward passed input
through, discretize returned zeros, conv1d was a no-op, and B/C were
randn per call, leaving conv1d_weight/dt_proj/A_log as dead weights with
zero temporal mixing. This implements a genuine S6 selective-scan:
- New params: x_proj [d_inner, dt_rank+2*d_state] (data-dependent dt,B,C),
dt_bias [d_inner], D [d_inner] (skip). Added to new()/new_seeded() and
the persistence contract (persistence_tensors/from_persistence_tensors).
- Real forward (CPU f32, looped — backbone is small): in_proj -> causal
depthwise conv1d -> SiLU -> x_proj->(dt,B,C) -> delta=softplus(dt.dt_proj
+dt_bias) -> A=-exp(A_log) -> sequential scan h=dA.h+dBu, y=sum C.h + D.u
-> gate by SiLU(z) -> out_proj. Residual moved OUT (canonical).
Numerically-stable silu_f32/softplus_f32 helpers.
- The scan runs inline (not via the immature rtx-tensor autograd tape);
the analytic backward lands in M2 per docs/phase3_real_ssm_spec.md.
New tests/real_selective_scan.rs (4 cases, all green): liveness (each
formerly-dead weight now moves the output), causality (no future
leakage), seeded determinism, and finite/non-constant output.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
