State-Routed Low-Rank MLP
Implementation & benchmark of arXiv:2608.25088 — The Von-Neumann State-Space Transformer for neural decoding
Source paper: The Von-Neumann State-Space Transformer for neural decoding arXiv:2608.25088 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace the shared Transformer feed-forward matrix with a shared base matrix plus a token-dependent low-rank update synthesized from a recurrent state-space controller. Unlike mixture-of-experts, the token does not merely interpolate expert outputs: it changes the actual matrices used inside both gated projections and therefore creates a continuous family of token-specific operators. The low-rank restriction keeps parameter growth and compute modest while allowing the controller to implement slowly varying computation across a sequence.
Formulas
Mathematical statement
The baseline gated feed-forward map is the paper's \(\mathrm{SwiGLU}(x)=W_{2}(\mathrm{SiLU}(W_{1}x)\odot W_{3}x)\), where \(x\in\mathbb{R}^{d}\), \(W_1,W_3\in\mathbb{R}^{m\times d}\), \(W_2\in\mathbb{R}^{d\times m}\), and \(\odot\) is elementwise multiplication. Adapt this using a state-space memory \(s_t\in\mathbb{R}^{q}\): \(s_t=A s_{t-1}+B u_t\), where \(u_t\) is the current normalized token, \(A\in\mathbb{R}^{q\times q}\) is a stable transition matrix with spectral radius less than one, and \(B\in\mathbb{R}^{q\times d}\). The instruction code is \(\alpha_t=\operatorname{softmax}(P s_t+b)\in\mathbb{R}^{K}\), with \(P\in\mathbb{R}^{K\times q}\), and each dynamic matrix is \(W_j(t)=W_j^{0}+\sum_{k=1}^{K}\alpha_{t,k}U_{j,k}V_{j,k}^{\mathsf T}\). Here \(W_j^0\) is a shared base matrix, \(U_{j,k}\) and \(V_{j,k}\) have rank \(r\), and \(j\in\{1,2,3\}\). The resulting operator is \(W_2(t)[\operatorname{SiLU}(W_1(t)x_t)\odot W_3(t)x_t]\). The low-rank update has only \(r(m+d)\) parameters per instruction for each matrix, instead of \(md\), while the state-space recurrence supplies temporal persistence in the instruction code.
Implementation notes
1. Integration point: in each pre-norm Transformer block, replace the ordinary SwiGLU matrices W1, W2, and W3. Feed the normalized token sequence u[1:T] into a lightweight causal controller before the MLP; use its state s[t] to generate the dynamic low-rank updates. Start with one controller shared across layers, or one controller per layer if capacity is insufficient. 2. Pseudocode: initialize s=0; for t in 1..T, compute s = A @ s + B @ u[t], compute alpha = softmax(P @ s + b), set W1t=W10+sum_k alpha[k]*U1[k]@V1[k].T and analogously W2t,W3t, compute y[t]=W2t @ (silu(W1t @ u[t]) * (W3t @ u[t])), and return h[t]=h[t]+y[t]. Parameterize A as A=diag(tanh(a)) for a diagonal stable SSM in the MVP; later test a low-rank-plus-diagonal A. 3. Compute from the proposed mathematics: the shared-plus-low-rank decomposition, state recurrence, and soft instruction code. Estimate empirically: effective code entropy, state time constant, rank r, and whether codes remain smooth across adjacent tokens. Log alpha entropy and ||Wj(t)-Wj0||F. 4. First experiment: compare a 6-layer decoder-only Transformer with ordinary SwiGLU, standard top-1 or soft MoE, and this module on WikiText-2 or a small sequential neural-decoding dataset. Match trainable parameters and FLOPs as closely as possible, using d=256, m=1024, q=32, K=8, r=8. Train with three random seeds and evaluate validation loss versus training examples, long-context extrapolation, and throughput. Success means lower validation loss at the same number of examples and no worse than 10% throughput, with the strongest expected gain on long contexts or heterogeneous token regimes. Ablate the state recurrence, low-rank updates, and code conditioning separately.
Verification
Stage 1 · Toy mechanism gate: Failed ✗
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built a self-contained state-routed low-rank SwiGLU MVP with a diagonal stable SSM controller, soft routing, weighted rank-3 updates for all three projections, mathematical checks, and a seeded synthetic persistent-regime regression. The math check passed: spectral radius 0.9969<1, maximum update rank 3, softmax sum error 0, and homogeneous state norm contracted to 0.163. The idea reached lower validation MSE than baseline (0.9616 vs 0.9908, about 2.9% better), but required 24.44s versus 1.42s and 6284 versus 3456 parameters, failing the proposed <=10% throughput-overhead criterion.
- Agent confidence
- 8/10
- Baseline
- validation MSE 0.99077; train MSE 1.01413; runtime 1.42s; parameters 3456
- Idea
- validation MSE 0.96160; train MSE 0.97513; runtime 24.44s; parameters 6284; mean update Frobenius norm 1.9405; alpha entropy 0.4227; math spectral radius 0.99689
Limitations: Only one random seed and one small synthetic regression task were tested; no WikiText or language-model evaluation, Transformer integration, FLOP-matched comparison, multi-layer/shared-controller study, long-context extrapolation, or ablations were run. The routed implementation loops over sequence positions and materializes per-token matrices, so its measured speed is not representative of an optimized fused kernel.
How to run: python3 experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- experiment.py 5.2 KB View Raw
- report.md 1.5 KB View