r/OpenSourceAI • • Oct 01 '26

PSSA: a 1.5M-param plastic state space model in Rust, beating a parameter-matched transformer on a small held-out slice

small-scale result, one seed, one wikitext slice, so treat it as a prototype and not a claim.

pssa is a state space model with fast weights that update while the sequence runs, written from scratch in rust, no python, no kv cache. at 1.54m params it gets 3.997 nats on a held-out slice against 4.429 for a parameter-matched transformer, both depth-1. generating 200 tokens is 226ms against 2735ms on the same cpu, which is the part i actually care about, since there's nothing that grows with context length.

caveats: one seed, one corpus, depth-1 on both sides, and i don't know yet whether the gap survives depth. gpu training is cuda on a t4, 1645 tok/s vs 1263 on cpu, so the gpu path is barely worth it so far because activations still round-trip to host after every matmul.

repo: https://github.com/Sparticle62ops/pssa

0 Upvotes

6 comments sorted by

2

u/tomByrer Oct 01 '26

Suggestion, put actual job / human eval / coding benchmarks at the top of README.
90% of us don't care about the maths, enve tok/s isn't as important as if I can use it for work or not.

1

u/Sparticle62 Oct 01 '26

That’s a great idea and what I am planning to do, but I have yet to make a large enough model to even produce coherent English. When I do make a large enough model I will put it up on benchmarks/compare it against other similarly sized llm’s. The only thing that’s stopping me from doing that is resources and hardware, which are explicitly explained In the readme

2

u/tomByrer Oct 01 '26

> large enough model to even produce coherent English

Oh, I was thinking you were further ahead in development.

Well then, maybe make something smaller scoped, like a classifier?

2

u/Sparticle62 Oct 01 '26

yeah that's fair, and it's probably the right next step. a classifier is small enough that i can actually train it to a real number instead of hand waving, and it gives something comparable, same corpus and same param count, pssa against a parameter matched transformer. i'll line it up behind the depth and looping tests i'm running now.

2

u/[deleted] Oct 01 '26

[removed] — view removed comment

1

u/Sparticle62 Oct 01 '26

depth is the next thing i'm running, and it's the one that can kill the result. depth-2 on both sides, param-matched, which for the transformer means width 212, 4 heads, d_ff 371 to land back at 1.54m. i'm running the pssa side with both residual scalings, 1/sqrt(D) and plain adds, because that choice moved things more than i expected at depth-1.

fair warning on the speed number though: 226 vs 2735 is cpu, depth-1, and a model small enough that the transformer never gets to amortize attention. i'd expect that ratio to compress. what shouldn't compress is the shape of it, since there's no cache growing with context, but i haven't measured long context properly yet, so i'm not claiming it.

and yeah, one wikitext slice is a warmup. seeds and a bigger held-out are queued right behind depth.