{"slug": "pssa-a-non-transformer-language-model-written-from-scratch-in-rust", "title": "PSSA: A non-transformer language model written from scratch in Rust", "summary": "A from-scratch Rust language model called PSSA, which replaces transformer attention with a recurrent state-space core and a 512-slot episodic memory bank, reached 3.98 training cross-entropy versus 4.43 for a matched transformer over 12.7M tokens of cleaned WikiText-103, a 0.45-nat gap. On a 198,939-token held-out slice, PSSA scored 3.997 cross-entropy and 54.4 perplexity against the transformer's 4.429 and 83.8, with next-token accuracy of 24.1% versus 18.0%, and generated 200 tokens in 226 ms versus 2,735 ms on the same CPU. The author cautions both are 1.5M-parameter research prototypes with poor text quality at this scale, so the result concerns learning efficiency rather than fluency.", "body_md": "PSSA is a small language model that is not a transformer. It reads text one token at a time through a recurrent state-space layer, keeps a bank of episodic memories it can look things up in, and rewrites part of its own weights while it runs. It is written in Rust from scratch, with no PyTorch, no TensorFlow, and no ML framework of any kind underneath it.\n\nAt matched parameters and on the same corpus, it learns faster than a transformer and generates text about twelve times quicker on the same CPU.\n\nA transformer scores every pair of tokens in the context, so its cost per step grows with the square of the sequence length and the whole context is re-read at every step. PSSA carries one fixed-size state along the sequence in a single left-to-right pass, and looks things up in a memory bank instead of re-reading the context, so cost grows linearly with length.\n\nTwo models, same corpus, same tokenizer, same optimizer schedule, same seed, same number of parameters. One is PSSA, one is a standard transformer. Over 12.7M tokens of cleaned WikiText-103:\n\nPSSA finished at **3.98** training cross-entropy, the transformer at **4.43**.\nThat is a gap of **0.45 nats**, perplexity 53.7 against 83.7. The transformer\nspent its entire 12.7M-token budget to reach a loss PSSA had already passed\naround 2M tokens in.\n\nThe two curves never cross, and they never touch:\n\nTraining loss only says a model fit the stream it was fed. So both checkpoints were scored on a 198,939-token slice cut from a part of the corpus neither run ever touched:\n\nEvery checkpoint of both runs, 64 PSSA links and 43 transformer links, scored on a bounded 9,934-token window of that unseen slice. The curves never cross: PSSA is ahead from the first link and finishes 0.51 nats lower. The table below is the final checkpoint of each run on the full slice.\n\n| Held-out slice, 198,939 unseen tokens | PSSA | Transformer | \n|---|---|---|\n| Cross-entropy | **3.997** | 4.429 | \n| Perplexity | **54.4** | 83.8 | \n| Next-token accuracy | **24.1%** | 18.0% | \n\nThe held-out gap, 0.43 nats, is essentially the training gap. PSSA is not memorizing harder, it is generalizing better.\n\nGenerating 200 tokens on the same CPU, same prompt, same sampler:\n\n|  | PSSA | Transformer | \n|---|---|---|\n| 200 tokens | **226 ms** | 2,735 ms | \n| Relative | **12x faster** | baseline | \n\nA recurrent model carries a fixed-size state, so the cost of each new token does not grow with the length of what came before. A transformer re-reads its whole context every step.\n\n- **A recurrent state-space core.** Learned continuous state matrices carry\ninformation forward in a fixed-size state, instead of attention over the full\ncontext window.\n- **An episodic memory bank.** 512 slots with hyperbolic (Poincare-style)\nretrieval and bounded top-4 search, written to and read from during the run.\n- **Plastic weights.** Fast updates reinforce what works, novelty drives growth,\nand a refractory gate rate-limits overwrites so repeated contradictory input\ndoes less damage.\n- **Closed-form consolidation.** A ridge-regression step folds the fast plastic\nupdates back into the base transition matrix, the way sleep consolidates a\nday's learning.\n- **No framework.** Hand-written linear algebra in Rust, with a CUDA path for\ntraining and a scalar CPU reference that every gradient is checked against\n(max gradient difference 2.98e-8).\n\nBeing straight about the scale, because the numbers above are easy to over-read:\n\n- These are **1.5M-parameter models** on 12.7M tokens. That is a research\nprototype, not a competitor to anything you have heard of.\n- Text quality at this scale is poor for both models. PSSA emits \"a barget of the Prian Academy\", the transformer \"a material circulation of the United States\". The comparison is about learning efficiency, not fluency.\n- The speed comparison is CPU-to-CPU, which is fair. The training throughput\nnumbers further down are **not** hardware-matched and should not be read as an\narchitecture result.\n- Two experiments are still unmeasured: retention of earlier skills after a corpus switch, and whether ablating the memory bank changes the loss.\n\n```\ngit clone https://github.com/Sparticle62ops/pssa.git\ncd pssa\ncargo build --release\n./target/release/oxide_ai_pssa\n```\n\nRunning it with no arguments gives you a home screen listing every command plus any checkpoint and corpus it finds in the working directory.\n\nThe whole result above was trained on a free hosted notebook with a single entry-level GPU, in 200,000-token links, because a session gets cut after a few hours. Every interesting question left, whether the gap holds at 10x or 100x these parameters, whether the memory bank matters at scale, how it does against a modern recurrent baseline, needs one thing: a GPU with real VRAM and allocations measured in days instead of hours. Anything meaningfully above the entry-level card this ran on changes what can be asked.\n\nIf you have compute to grant, or you work somewhere that does, that is the single highest-leverage thing anyone can offer this project.\n\nSponsorship funds compute and nothing else. In return you get named here and in the write-up of any result your hardware made possible. Get in touch before sending anything so the details can be agreed.\n\nIssues and pull requests are welcome. The parts most in need of hands: kernel\nperformance, a modern recurrent baseline to compare against, and evaluation\nbeyond next-token loss. Validate any branch with `cargo test --release` before\nopening a PR.\n\nSolana: `4XPZ9uAa2BMoth6msoHRxTWL4mUrMfq3LGrxbAGja96h`\n\nEverything below is for running, training, and working on the project.\n\n- Rust toolchain with Edition 2024 support, including Cargo.\n- Network access only when using an HTTP/HTTPS dataset or a Hugging Face dataset.\n- Enough memory and disk for larger corpora and serialized models.\n- Optional: a CUDA device for the GPU training path. The CPU path is the reference and always available.\n\nDirect runtime dependencies are [`ureq`](https://crates.io/crates/ureq) for\ndataset downloads and [`tokenizers`](https://crates.io/crates/tokenizers) for\nbyte-level BPE.\n\nBoth chains ran 64 links of 200,000 encoded tokens, each link resuming from the previous checkpoint, so the learning-rate schedule and optimizer state continue across the whole run instead of restarting per link.\n\n- Identical corpus: one `clean-wikitext` pass over WikiText-103, reused byte for byte.\n- Identical token IDs: the baseline pins `--tokenizer-from` to the PSSA chain's\nown checkpoint, so neither model sees a different vocabulary.\n- Identical optimization: 30,000-update cosine horizon, no warm-up restart, 512 supervised target tokens per update, seed 42.\n- PSSA: latent 256, recurrent state 16, 512 memory slots, key width 32, vocab 2,048.\n- Baseline: 1,541,120 parameters, 1 layer, width 256, 4 heads, FFN 448, vocab 2,048.\n\nEnd-of-link training cross-entropy:\n\n| Link | Tokens seen | PSSA | Transformer | \n|---|---|---|---|\n| ck01 | 200,000 | 5.733 | 6.461 | \n| ck05 | 1,000,000 | 4.617 | 5.467 | \n| ck10 | 2,000,000 | 4.447 | 5.082 | \n| ck15 | 3,000,000 | 4.292 | 4.858 | \n| ck20 | 4,000,000 | 4.185 | 4.704 | \n| ck25 | 5,000,000 | 4.221 | 4.704 | \n| ck30 | 6,000,000 | 4.070 | 4.561 | \n| ck35 | 7,000,000 | 4.039 | 4.523 | \n| ck37 | 7,400,000 | 3.960 | 4.465 | \n| ck44 | 8,800,000 | 4.004 | 4.480 | \n| ck48 | 9,600,000 | 3.937 | 4.415 | \n| ck52 | 10,400,000 | 3.846 | 4.344 | \n| ck56 | 11,200,000 | 3.887 | 4.375 | \n| ck60 | 12,000,000 | 3.972 | 4.418 | \n| ck64 | 12,800,000 | 3.982 | 4.428 | \n\nThe baseline's first session was cut at link 43 by the notebook session limit\nand its loss CSV did not survive, so links 1 to 43 are read back from that\nsession's own run log instead. The chain resumed from `ck43` in a second session\nand finished all 64 links, and both curves above now cover the full run.\n\nPSSA trained on a Kaggle T4 at roughly 900 tokens/second. The baseline is\nCPU-only, because `train-transformer` has no GPU path, and held 212\ntokens/second. Those two numbers say nothing about the architectures. On the\nsame CPU-only Kaggle hardware the batched PSSA path measures 375 tokens/second\nagainst the baseline's 212, and the loss comparison above is unaffected either\nway, since it is matched on tokens and updates rather than on time.\n\nThe losses are end-of-link training cross-entropy on the stream being fit, not\nheld-out evaluation. For a held-out comparison on an unseen slice, use the\n`compare` command described in [docs/COMPARISON.md](https://github.com/Sparticle62ops/pssa/blob/main/docs/COMPARISON.md).\nGeneration quality at this scale is poor for both models: PSSA emits \"a barget\nof the Prian Academy\", the baseline \"a material circulation of the United\nStates\".\n\nTwo experiments are not yet measured: retention of earlier skills after a corpus switch, and whether ablating the 512 memory slots changes loss.\n\n```\nbash kaggle/kaggle_continue.sh              # the PSSA chain\nbash kaggle/kaggle_transformer_baseline.sh  # the parameter-matched baseline\n```\n\nBoth read `TOTAL`, `WINDOW` and `FRESH` from the environment and write\n`--loss-csv`, so the curve survives a cut session.\n\nGeneral form:\n\n```\noxide_ai_pssa <COMMAND> [OPTIONS]\n```\n\nCommands:\n\n| Command | Purpose | \n|---|---|\n| `train [source]` | Fit a checkpoint on a text corpus and write a `.pssa` file. | \n| `generate <prompt>` | Continue a prompt with a trained checkpoint. | \n| `chat [source]` or`repl [source]` | Interactive prompt loop against a checkpoint. | \n| `evaluate [source]` | Cross entropy, perplexity and accuracy as JSON. | \n| `status` | Checkpoints and corpora in the working directory. Takes no options. | \n| `download <repo>` | Pull a Hugging Face dataset to a local file. | \n| `clean-wikitext INPUT -o OUTPUT` | Stream-clean a raw WikiText file into a new UTF-8 corpus. | \n| `benchmark` | End-to-end smoke test on the built-in corpus. | \n| `gpu-probe` | Check whether a WebGPU compute device is usable. | \n| `help` | Print command and option help. | \n\nOptions:\n\n| Option | Default | Applies to | Description | \n|---|---|---|---|\n| `-d, --data <source>` | `data/downloaded.txt` when present, otherwise`science` | `train` ,`chat` ,`evaluate` | Dataset source, or a comma-separated list. | \n| `-m, --model <path>` | `data/model.pssa` | `chat` ,`generate` ,`evaluate` | Checkpoint to load. | \n| `-o, --out <path>` | Command-specific; required for `clean-wikitext` | `train` ,`download` ,`clean-wikitext` | Output checkpoint or dataset path. Cleaning requires a new file. | \n| `-p, --prompt <text>` | empty | `generate` | Prompt text. Required for generation. | \n| `-e, --epochs <n>` | `4` | `train` | Training epochs. | \n| `-t, --temp, --temperature <float>` | `0.70` | `chat` ,`generate` | Sampling temperature. | \n| `--max-new-tokens <n>` | `64` (maximum 100,000) | `generate` | Generation length cap. | \n| `--latent <n>` | `256` | `train` | Latent dimension. | \n| `--state <n>` | `16` | `train` | Recurrent state dimension. | \n| `--key <n>` | `32` | `train` | Memory-key dimension. | \n| `--memory <n>` | `512` | `train` | Memory bank capacity. | \n| `--chunk <n>` | `64` | `train` | Sequence chunk length. | \n| `--lr <float>` | `1e-3` | `train` | Base learning rate. | \n| `--accumulate <n>` | `8` | `train` | Chunks per optimizer update. | \n| `--warmup-steps <n>` | `0` | `train` | Linear warm-up before cosine decay. | \n| `--seed <n>` | `42` | `train` | Initialization seed. | \n| `--tokenizer <bpe\\|word>` | `bpe` | `train` | Tokenizer family. | \n| `--vocab-size <n>` | `2048` | `train` | BPE vocabulary maximum. | \n| `--max-tokens <n>` | unset | `train` | Global cap across input documents, not per document. | \n| `--skip-tokens <n>` | `0` | `train` | Skip this many tokens before training starts. | \n| `--resume <path>` | unset | `train` | Continue from an existing checkpoint. | \n\nPositional arguments and long/short options can be mixed:\n\n```\ncargo run --release -- train data/downloaded.txt -e 2 -o data/experiment.pssa\ncargo run --release -- train --data data/downloaded.txt --epochs 2 --out data/experiment.pssa\n```\n\nInside the REPL:\n\n- `/exit` or`quit` exits the process.\n- `/info` prints the loaded model path, memory slot count, and adapter count.\n- `/temp <value>` reports a temperature value but does not apply it to later turns. Pass`--temp` when launching`chat` instead.\n\n`--skip-tokens`, `--max-tokens` and `--resume` together let a long corpus be trained as a chain of short runs, so a single run never has to survive a session limit. If a window crosses EOF, selection wraps to the beginning of the corpus. Each link trains its own window and hands its optimizer state to the next:\n\n```\ncargo run --release -- train data/downloaded.txt -e 1 \\\n  --skip-tokens 0      --max-tokens 200000 -o chain/ck01.pssa\ncargo run --release -- train data/downloaded.txt -e 1 \\\n  --skip-tokens 200000 --max-tokens 200000 --resume chain/ck01.pssa -o chain/ck02.pssa\n```\n\n`kaggle/kaggle_continue.sh` drives this pattern end to end: it sets a window size and a link count, walks the corpus offset by offset, and resumes each link from the previous checkpoint. `status` then reports every checkpoint in the chain with its shape and optimizer step count.\n\n`DatasetManager` accepts one or more comma-separated sources:\n\n```\ncargo run --release -- train science                       # built-in reference corpus\ncargo run --release -- train data/downloaded.txt           # local text file\ncargo run --release -- train data/                         # every readable file in a directory\ncargo run --release -- train https://example.org/corpus.txt\ncargo run --release -- train hf:owner/dataset              # Hugging Face repository\ncargo run --release -- train science,data/downloaded.txt   # multiple sources\n```\n\nLocal files and directories are read directly; HTTP(S) URLs and explicit `hf:owner/dataset`\nsources are downloaded. Structured responses are reduced using common fields such as\n`text`, `content`, `article`, `story`, `instruction`, `output`, `sentence`, and `summary`;\nstructured responses without a supported text field are rejected.\n\nByte-level BPE keeps exact UTF-8 case, whitespace, punctuation, and line endings, and has a complete 256-byte fallback alphabet, so valid UTF-8 never collapses to `<unk>`. The previous lowercase word splitter, including its 10,000-word cap and `<unk>` behavior, is available only with `--tokenizer word`.\n\nDownload a Hugging Face dataset into a local text file:\n\n```\ncargo run --release -- download wikimedia/wikipedia --out data/downloaded.txt\n```\n\nNetwork downloads are not validated or curated by Oxide AI. Review licensing, privacy, and content before training on an external corpus.\n\nClean extracted `wikitext-103-raw` text **before a fresh training run**:\n\n```\n./target/release/oxide_ai_pssa clean-wikitext wiki.train.raw --out data/wikitext-clean.txt\n./target/release/oxide_ai_pssa train data/wikitext-clean.txt -o data/model.pssa\n# Also available: oxide_ai_pssa help clean-wikitext\n```\n\nThe same command can be used in Kaggle after extracting text from Parquet; it\naccepts a local UTF-8 text file, not Parquet itself. `-o` and `--out` are aliases.\nThe output path is required and must not already exist (including the input\npath or a link to it). This protects the original corpus; choose a new output\nname for another run. Read, UTF-8, and write failures exit nonzero through the\nnormal CLI error path, with partial output removed when possible.\n\nThe pass:\n\n- Joins `@-@` ,`@.@` , and`@,@` to adjacent text:`guest @-@ starring` →`guest-starring` ,`52 @.@ 9` →`52.9` ,`500 @,@ 000` →`500,000` .\n- Drops balanced heading lines such as `= Title =` and`= = Section = =` .\n- Removes `<unk>` and collapses remaining inline whitespace to single spaces.\n- Removes spaces before `.` ,`,` ,`)` and after`(` ; trims each line.\n- Retains at most one consecutive blank line, including at the start/end. Removing a heading does not introduce a blank line.\n- Writes LF line endings, including a newline on the last retained line.\n\n`oxide_ai_pssa::dataset::clean_wikitext(reader, writer)` is the reusable library\nAPI (`BufRead` / `Write`, returning `std::io::Result<()>`). The CLI uses buffered\nfile I/O, and the cleaner retains only its input/output line buffers: memory is\nproportional to the longest line, not the corpus size. Library callers using a\nbuffered writer must flush it themselves; the CLI explicitly checks the flush.\nNo new dependencies are required.\n\nCleaning is opt-in: existing loaders, tokenizers, training commands, and\n`kaggle/kaggle_continue.sh` are unchanged. **Do not switch an in-flight resume\nchain to a cleaned corpus**: cleaning changes token IDs/counts and the meaning\nof `--skip-tokens` offsets. Prepare and consistently reuse one cleaned corpus\nfor a new chain instead.\n\nThe `train` command performs two phases:\n\n1. **Continuous recurrent ingestion:** token transitions are processed through the PSSA layer. The model updates state, memory, adapters, and routing behavior with a cosine learning-rate schedule.\n2. **Adapter consolidation:** after each epoch, the plastic adapter's fast coefficients are folded into its consolidated coefficients with the configured EMA rate.\n\nDefaults are latent 256, recurrent state 16, memory-key 32, memory capacity 512, chunk length 64, learning rate 1e-3, 8 chunks per update, and seed 42. The resulting binary holds weights, configuration, memory, adapters, and optimizer state. It is not an interchange format for other ML frameworks and should be loaded through `PSSALayer::import_from_pssa_bytes`.\n\nNew saves use **V7**: the full V6 training/resume payload plus a bounded, length-prefixed standard tokenizer JSON. A V7 BPE checkpoint is self-contained and restores its exact ordered vocabulary without access to the training or evaluation corpus. `generate` and `chat` reject `--data` for V7 BPE because retraining a tokenizer on external data would not validate provenance. V7 word checkpoints and V6 checkpoints retain the legacy optional `--data` exact-vocabulary comparison. Checked V5 artifacts remain inference-only and require `--data` because they never contained tokenizer provenance.\n\nGeneration is autoregressive and uses temperature 0.70, a top-24 candidate limit followed by top-p 0.85 filtering, a 1.25 repetition penalty over a recent 64-token window, immediate self-transition suppression, `<unk>` suppression, and a default cap of 64 new tokens, ending early after two generated periods.\n\nV7 BPE inference restores the exact embedded tokenizer and never rebuilds it from a selected dataset. Evaluation supplies its data only as held-out text to the restored tokenizer.\n\n```\ncargo run --release -- benchmark\n```\n\nThe suite exercises synthetic streams for contradictory facts, MQAR-style distractors, burst repetition, model serialization, and short generation prompts. It prints milestone results, is not wired into Cargo's test harness, and is not a quality evaluation on general language tasks.\n\n| Path | Responsibility | \n|---|---|\n| `src/main.rs` | Binary entry point; forwards process arguments to the CLI. | \n| `src/cli.rs` | Argument parsing, home screen, training, chat, generation, evaluation, status, download, and benchmark orchestration. | \n| `src/ui.rs` | Terminal presentation: logo, panels, spinners, progress bars, ANSI-aware width handling. | \n| `src/dataset.rs` | Tokenization, vocabulary construction, built-in corpora, local and remote loading, streaming WikiText cleaning. | \n| `src/pssa.rs` | PSSA layer, forward pass, plastic learning, consolidation, and `.pssa` serialization. | \n| `src/checkpoint.rs` | Checkpoint format versions, resume payloads, and import/export validation. | \n| `src/inference.rs` | Autoregressive sampling and generation constraints. | \n| `src/backend.rs` | GEMM dispatch, CPU reference kernels, and the WebGPU device probe. | \n| `src/memory.rs` | Fixed-capacity hyperbolic memory bank and retrieval/update logic. | \n| `src/adapter.rs` | Low-rank modular adapter projections and updates. | \n| `src/defense.rs` | Refractory rate-limiter primitives for stable updates and overwrite defense. | \n| `src/linalg.rs` | Small allocation-conscious vector, matrix, math, and deterministic RNG utilities. | \n| `src/diagnostics.rs` | CLI banner formatting. | \n| `kaggle/` | Chained-training driver for long corpora on a hosted notebook. | \n| `data/downloaded.txt` | Checked-in corpus used as the default when present. | \n| `data/model.pssa` | Checked-in serialized model artifact. | \n\n```\ncargo fmt --all -- --check\ncargo clippy --release --all-targets\ncargo test --release\n```\n\nIntegration tests live in `tests/`: `allocations.rs`, `bpe_repair.rs`, `checkpoint_repair.rs`, `core_repair.rs`, `linalg.rs`, and `runtime_repair.rs`, with shared artifacts under `tests/fixtures/`. They cover tokenizer round trips, checkpoint import/export across versions, linear-algebra kernels, allocation behavior, and CLI runtime output. Clippy is clean of errors; a number of style warnings in the numeric kernels are left in place deliberately, since rewriting indexed loops there would churn code the gradient tests pin down.\n\n- CPU-oriented prototype with hand-written linear algebra. `gpu-probe` verifies a WebGPU device and a GEMM against the CPU reference, but training and inference still run the layer math on the CPU.\n- The CLI parser is intentionally minimal: no shell-style quoting, and little validation beyond numeric parsing.\n- A missing or unreadable dataset silently falls back to the built-in science corpus in several loading paths.\n- Model and tokenizer vocabularies must remain compatible; a size warning does not repair a mismatch.\n- Model shape cannot change across a resume chain: latent, state, key, memory and vocabulary must match the checkpoint being resumed.\n- Downloaded content can be large and may contain JSON, malformed text, or data unsuitable for training.\n- The REPL temperature command acknowledges a value without changing the active configuration.\n- Benchmark output is milestone-oriented and does not measure perplexity, factuality, latency, or safety.\n- Serialized `.pssa` files are project-specific binary artifacts without version migration tooling.\n\nSee [LICENSE](https://github.com/Sparticle62ops/pssa/blob/main/LICENSE) for the project license.", "url": "https://wpnews.pro/news/pssa-a-non-transformer-language-model-written-from-scratch-in-rust", "canonical_source": "https://github.com/Sparticle62ops/pssa", "published_at": "2026-09-30 03:19:54+00:00", "updated_at": "2026-09-30 03:47:19.498375+00:00", "lang": "en", "topics": ["large-language-models", "machine-learning", "ai-research", "natural-language-processing"], "entities": ["PSSA", "Rust", "WikiText-103", "PyTorch", "TensorFlow", "CUDA"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/pssa-a-non-transformer-language-model-written-from-scratch-in-rust", "markdown": "https://wpnews.pro/news/pssa-a-non-transformer-language-model-written-from-scratch-in-rust.md", "text": "https://wpnews.pro/news/pssa-a-non-transformer-language-model-written-from-scratch-in-rust.txt", "jsonld": "https://wpnews.pro/news/pssa-a-non-transformer-language-model-written-from-scratch-in-rust.jsonld"}}