pssa-100m

A 109M-parameter language model with no attention, trained from scratch on one RTX PRO 6000 Blackwell. It is a PyTorch port of the plastic state-space architecture from Sparticle62ops/pssa, scaled from that repo's 1.5M-parameter reference to 109M.

The weights, tokenizer, and a self-contained loader are in this repo. The model writes fluent English in two registers, story and web text. It does not know facts. At 109M parameters it picks up style and local grammar, and the web samples below are invented in a convincing register.

Architecture

18 stacked blocks of:

xn       = rmsnorm(x)
delta    = softplus(W_delta xn)         per-channel step size
h_t      = exp(delta_t·A) h_{t-1} + delta_t·B_t·x_t     diagonal SSM scan, A < 0
y_t      = <h_t, C_t>
q        = poincare(W_qx xn + W_qh y)   query from token AND state
w        = softmax(-d_H(q, k)/tau)      hyperbolic distance to 512 memory slots
m        = w @ v                        bounded episodic read
z        = y/√d_s + gate·W_proj m + rank16-silu-adapter(xn) + MLP
h_i      = h_{i-1} + block_i(h)/√depth

Parameter counts: 101.50M in the core (d_model 768, untied 16k-token embedding and head) and 7.37M across the 18 memory banks. Nothing in the model attends. Compute per token stays flat as context grows because the recurrent state has a fixed size.

Two changes from the Rust reference, both needed at this scale:

• The memory bank trains by gradient instead of surprise-gated inserts. The reference's 3.5-nat novelty threshold stops firing a few hundred steps into a run this size, which would leave the write path idle for the remaining billions of tokens.
• Backprop is truncated at 256 tokens.

The scan runs as a fused Triton kernel with a hand-written backward pass. It matches a plain PyTorch loop to about 1e-7 on every gradient.

Training

4B tokens in two phases:

phase	data	tokens
1	50/50 TinyStories / FineWeb-Edu	3.11B
2	85/15 web/story fine-tune	+0.89B

The learning rate held at 1e-4, then cosine-decayed to 1e-5 over the last 600M tokens. The web share went from 50% to 85% at 3.11B. During that final phase web loss fell from 3.12 to 2.92 nats on the running eval and story loss stayed near 1.19.

Throughput was 127k tok/s at 99% GPU utilization (batch 64, chunk 256, torch.compile on the block). Total GPU time was about 9 hours, spread over three sessions because the machine shut down overnight twice. Checkpoints went to this repo every 40 minutes, and training resumed from them each morning.

Results

Held-out, 1M unseen tokens per corpus, batch 8, sequence length 512:

corpus	cross-entropy	perplexity	next-token acc
TinyStories	1.312	3.71	64.8%
FineWeb-Edu web	2.866	17.6	44.2%

An earlier 2.5B-token checkpoint of the same architecture and tokenizer scored 1.329 on stories and 2.935 on web. The extra 1.5B tokens bought 0.017 nats on story and 0.069 on web.

Web loss was still falling at 4B, so more tokens would help.

Samples, temp 0.7, top-k 40:

Once upon a time there was a little girl named Amy. She was very independent and liked to do things by herself. One day, Amy decided to go to the park and play. As she walked to the park, she saw a pretty flower. She wanted to pick it up and take it home. Amy managed to pick the flower up, but it was too heavy for her. She felt sad and didn't know what to do. Then, Amy had an idea. She asked her friends to help her and they all picked up the flower together.

In the summer of 1969, the U.S. Department of Agriculture and International Cooperation (IAC) announced that the department of agricultural and natural resources was in a position to form a partnership between the Government of India, the United Nations, the Federal Cooperation for Agricultural Development and the World Bank.

Use

# pip install torch safetensors tokenizers
from pssa_model import from_pretrained, generate

model, tok = from_pretrained(".")        # dir containing model.safetensors + tokenizer.json
print(generate(model, tok, "Once upon a time", max_new=100))

pssa_model.py holds the config, model, loader, and sampler. On CPU it runs a plain PyTorch scan. On CUDA it uses the Triton kernel when pssa_kernels.py sits next to it. It does not need transformers.

Weights are bf16 safetensors, 217 MB. The tokenizer is a 16,384-token byte-level BPE trained on TinyStories. It encodes web text without mangling, though rare Unicode costs more tokens than a general-purpose tokenizer would spend.

Files

• model.safetensors: weights, bf16
• pssa_model.py: model + loader + sampler
• pssa_kernels.py: fused Triton scan
• tokenizer.json: 16k byte-level BPE

Credit

The architecture, the Rust reference implementation, and the 1.5M-parameter comparison against a matched transformer come from Sparticle62ops/pssa, released under GPL-3.0. This checkpoint is an independent PyTorch port trained on different data. It is not an official release from that project.