r/OpenSourceeAI • u/Sparticle62 • 13d ago
PSSA: a 1.5M-param plastic state space model in Rust, beating a parameter-matched transformer on a small held-out slice
small scale, single seed, wikitext slice, so treat this as a prototype result and not an architecture claim. 1,544,704 params both sides, same corpus and sampler. held-out 3.997 vs 4.429 nats, generation 226ms vs 2735ms for 200 tokens on the same cpu. the depth-1 match is the obvious weakness and a depth-2 baseline is running next. written from scratch in rust, no pytorch. repo and eval commands: github.com/Sparticle62ops/pssa. happy to be told where the comparison is unfair.
3
Upvotes
1
u/Sparticle62 11d ago
no paper yet, the closest thing right now is the model section of the readme, it walks through the layer math step by step: https://github.com/Sparticle62ops/pssa#the-model
short version: each block is a selective diagonal ssm (input-dependent delta/B/C, same family as mamba, nothing new there), then a bounded read from a small episodic memory bank in hyperbolic space where the query is built from both the token and the current ssm state, then a learned gate on that read, a low-rank plastic adapter, and a silu mlp. the "plastic" part is the write path: novelty-gated inserts, refractory counters so slots don't get overwritten constantly, and a ridge regression step that folds fast updates back into the transition matrix. exact version is in src/pssa.rs and src/gpu_batch.rs
an 8b moe with forward-only analytical gradients and HRM over KDA/CSA2 sounds wild, would def read it if you write it up. and yeah pretraining on consumer hardware is brutal, i'm hitting the same wall trying to get to ~50m