VSN sits between snap_feature_assemble and Mamba2 input:
snap_assemble → window_tensor_d [B, K, 40] (raw)
→ VSN fwd → vsn_out_d [B, K, 40] (gated) + vsn_gates_d [B*K, 40]
→ Mamba2 fwd → ...
Forward path: per-position softmax over FEATURE_DIM features, output[i] =
x[i] * gates[i]. Gates initialised near-uniform (W_vsn ~ N(0, 1/√FEATURE_DIM),
b_vsn = 0) so the model starts from "all features matter equally" and
learns regime-conditional gates.
Backward path: Mamba2 backward already computed d_x_from_in (gradient
w.r.t. its input) on its grads_buffers — previously labeled "unused but
allocated", now consumed by VSN bwd as grad_y. Zero refactor to
mamba2_block.rs.
VSN bwd writes:
grad_W_vsn [40, 40] → opt_vsn_w (AdamW, default wd)
grad_b_vsn [40] → opt_vsn_b (AdamW, wd=0)
vsn_grad_x_d [B*K, 40] → discarded (snap_features are non-trainable
transforms of raw MBP-10 data)
Param grads are explicitly memset_zeros before each VSN bwd call (the
kernel uses += semantics like the GRN bwd, but VSN runs ONCE per step
not K times, so zeroing makes the += a clean overwrite — matches
Adam's `step` expectations).
Eval path mirrors training (VSN fwd applied between snap_assemble and
Mamba2 fwd) so eval sees the gated features layers were trained on.
Trainer now manages 19 AdamW: CfC×4 + GRN heads×10 + LN×2 + VSN×2 +
Mamba2 grouped.
Synthetic overfit smoke: stride=1 initial=0.30 → final=0.00 in 50 steps
(faster than pre-VSN's 0.32, consistent with VSN's near-uniform init
giving a small head start). All 8 perception_overfit tests pass.
Note: the smoke proves VSN's W/b actually move via the Mamba2 input
gradient — if `d_x_from_in` were zero, VSN params wouldn't update and
the chain still converges but VSN remains identity. The fact that
initial loss DIFFERS (0.30 vs 0.32) shows VSN is in the forward path
end-to-end.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
73d68ab786