
PSSA: A non-transformer language model written from scratch in Rust
PSSA is a small language model that is not a transformer. It reads text one
以下正文同步自 Hacker News,版权归原站所有,已转换为易读排版。
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.
At matched parameters and on the same corpus, it learns faster than a transformer and generates text about twelve times quicker on the same CPU.
A 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.
Two 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:
PSSA finished at 3.98 training cross-entropy, the transformer at 4.43. That is a gap of 0.45 nats, perplexity 53.7 against 83.7. The transformer spent its entire 12.7M-token budget to reach a loss PSSA had already passed around 2M tokens in.
The two curves never cross, and they never touch:
Training 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:
Every 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.
The held-out gap, 0.43 nats, is essentially the training gap. PSSA is not memorizing harder, it is generalizing better.
Generating 200 tokens on the same CPU, same prompt, same sampler:
A 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.
A recurrent state-space core. Learned continuous state matrices carry information forward in a fixed-size state, instead of attention over the full context window.
An episodic memory bank. 512 slots with hyperbolic (Poincare-style) retrieval and bounded top-4 search, written to and read from during the run.
Plastic weights. Fast updates reinforce what works, novelty drives growth, and a refractory gate rate-limits overwrites so repeated contradictory input does less damage.
Closed-form consolidation. A ridge-regression step folds the fast plastic updates back into the base transition matrix, the way sleep consolidates a day's learning.
No framework. Hand-written linear algebra in Rust, with a CUDA path for training and a scalar CPU reference that every gradient is checked against (max gradient difference 2.98e-8).
What this is not
Being straight about the scale, because the numbers above are easy to over-read:
git clone https://github.com/Sparticle62ops/pssa.git cd pssa cargo build --release ./target/release/oxide_ai_pssa
Running it with no arguments gives you a home screen listing every command plus any checkpoint and corpus it finds in the working directory.
Compute
The 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.
If you have compute to grant, or you work somewhere that does, that is the single highest-leverage thing anyone can offer this project.
Sponsorship 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.
Issues and pull requests are welcome. The parts most in need of hands: kernel performance, a modern recurrent baseline to compare against, and evaluation beyond next-token loss. Validate any branch with cargo test --release before opening a PR.
Solana: 4XPZ9uAa2BMoth6msoHRxTWL4mUrMfq3LGrxbAGja96h
Everything below is for running, training, and working on the project.
Rust toolchain with Edition 2024 support, including Cargo.
Network access only when using an HTTP/HTTPS dataset or a Hugging Face dataset.
Enough memory and disk for larger corpora and serialized models.
Optional: a CUDA device for the GPU training path. The CPU path is the reference and always available.
Direct runtime dependencies are ureq for dataset downloads and tokenizers for byte-level BPE.
How the two runs were matched
Both 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.
The baseline's first session was cut at link 43 by the notebook session limit and its loss CSV did not survive, so links 1 to 43 are read back from that session's own run log instead. The chain resumed from ck43 in a second session and finished all 64 links, and both curves above now cover the full run.
PSSA trained on a Kaggle T4 at roughly 900 tokens/second. The baseline is CPU-only, because train-transformer has no GPU path, and held 212 tokens/second. Those two numbers say nothing about the architectures. On the same CPU-only Kaggle hardware the batched PSSA path measures 375 tokens/second against the baseline's 212, and the loss comparison above is unaffected either way, since it is matched on tokens and updates rather than on time.
The losses are end-of-link training cross-entropy on the stream being fit, not held-out evaluation. For a held-out comparison on an unseen slice, use the compare command described in docs/COMPARISON.md. Generation quality at this scale is poor for both models: PSSA emits "a barget of the Prian Academy", the baseline "a material circulation of the United States".
Two experiments are not yet measured: retention of earlier skills after a corpus switch, and whether ablating the 512 memory slots changes loss.
bash kaggle/kaggle_continue.sh # the PSSA chain bash kaggle/kaggle_transformer_baseline.sh # the parameter-matched baseline
Both read TOTAL, WINDOW and FRESH from the environment and write --loss-csv, so the curve survives a cut session.
oxide_ai_pssa <COMMAND> [OPTIONS]
Commands:
Positional arguments and long/short options can be mixed:
cargo run --release -- train data/downloaded.txt -e 2 -o data/experiment.pssa cargo run --release -- train --data data/downloaded.txt --epochs 2 --out data/experiment.pssa
Chat commands
Inside the REPL:
--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:
cargo run --release -- train data/downloaded.txt -e 1
--skip-tokens 0 --max-tokens 200000 -o chain/ck01.pssa
cargo run --release -- train data/downloaded.txt -e 1
--skip-tokens 200000 --max-tokens 200000 --resume chain/ck01.pssa -o chain/ck02.pssa
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.
DatasetManager accepts one or more comma-separated sources:
cargo run --release -- train science # built-in reference corpus cargo run --release -- train data/downloaded.txt # local text file cargo run --release -- train data/ # every readable file in a directory cargo run --release -- train https://example.org/corpus.txt cargo run --release -- train hf:owner/dataset # Hugging Face repository cargo run --release -- train science,data/downloaded.txt # multiple sources
Local files and directories are read directly; HTTP(S) URLs and explicit hf:owner/dataset sources are downloaded. Structured responses are reduced using common fields such as text, content, article, story, instruction, output, sentence, and summary; structured responses without a supported text field are rejected.
Byte-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.
Download a Hugging Face dataset into a local text file:
cargo run --release -- download wikimedia/wikipedia --out data/downloaded.txt
Network downloads are not validated or curated by Oxide AI. Review licensing, privacy, and content before training on an external corpus.
Clean extracted wikitext-103-raw text before a fresh training run:
./target/release/oxide_ai_pssa clean-wikitext wiki.train.raw --out data/wikitext-clean.txt ./target/release/oxide_ai_pssa train data/wikitext-clean.txt -o data/model.pssa
Also available: oxide_ai_pssa help clean-wikitext
The same command can be used in Kaggle after extracting text from Parquet; it accepts a local UTF-8 text file, not Parquet itself. -o and --out are aliases. The output path is required and must not already exist (including the input path or a link to it). This protects the original corpus; choose a new output name for another run. Read, UTF-8, and write failures exit nonzero through the normal CLI error path, with partial output removed when possible.
oxide_ai_pssa::dataset::clean_wikitext(reader, writer) is the reusable library API (BufRead / Write, returning std::io::Result<()>). The CLI uses buffered file I/O, and the cleaner retains only its input/output line buffers: memory is proportional to the longest line, not the corpus size. Library callers using a buffered writer must flush it themselves; the CLI explicitly checks the flush. No new dependencies are required.
Cleaning is opt-in: existing loaders, tokenizers, training commands, and kaggle/kaggle_continue.sh are unchanged. Do not switch an in-flight resume chain to a cleaned corpus: cleaning changes token IDs/counts and the meaning of --skip-tokens offsets. Prepare and consistently reuse one cleaned corpus for a new chain instead.
Defaults 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.
New 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.
Generation 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 tok
正文由 探How 从来源站点 RSS 同步,内容未经改写;遇到排版缺失或需要图片、视频时请以原文为准。