fhe-ssm-pbd430a — a 405M recurrent language model built to run under fully homomorphic encryption
This model is served in https://github.com/xelananv/fhe-ssm: it is a 404,715,520 parameter recurrent language model whose every (!) nonlinearity is a
low-degree polynomial with trained coeffs, so that the CKKS encrypted evaluation computes precisely the model's own
float64 arithmetic without any post-hoc approximation of softmax or GELU. It succesfully generated text for 64 prompts at once under encryption,
autoregressively, for 148 consecutive steps (=> 9,472 tokens) at 128 bit classical parameters, matching its own plaintext
top-1 choice 99.7 % of the time. Those are properties of the circuit; this card is about the model, which is small and
was trained on a few hundred dollars of compute. It writes fluent English; it is not a reasoning model.
Model details
|
|
| parameters |
404,715,520 (embedding and output head are a single tied 50,277 × 1,024 matrix = 51.5M; the 24 blocks hold 353.2M) |
| architecture |
24 blocks, each a time-mix (token shift and a per-channel recurrence S ← a·S + (1−a)·x with a trained, fixed decay a ∈ [0.9, 0.9995]; learned cubic gate) followed by a channel-mix (learned cubic activation × learned quadratic gate), residual; RMSNorm by a Newton inverse-square-root iteration; no attention and no KV cache: the state that's carried between tokens is 2 vectors of 1,024 per block |
| width / hidden |
d = 1,024; channel mix hidden 4,096 |
| tokenizer |
GPT-NeoX (50,277 tokens; EleutherAI/gpt-neox-20b) |
| training context |
1,024 tokens, recurrent state carried across shard boundaries |
| precision of the reference |
float64 (spec_decode/plain_recurrent.py in the github repo is the reference every encrypted run is compared against) |
| closest published relatives |
RWKV / Griffin-style recurrences (fixed per-channel decay); not Mamba (no input-selective decay) |
As for code: the trainer ml-eval/train_fhe_native_ssm.py, the loader spec_decode/_native_loader_copy.py (load_native), the exact recurrent
reference spec_decode/plain_recurrent.py, all in the github repo. Config, seeds and the training log are in this repository and in
the GitHub tree (ml-eval/artifacts/).
Files
| file |
bytes |
sha256 |
what |
pbd430a_weights.npz |
1,824,931,038 |
3ec3f1d8b9b71e7eee38055a7621213991c6417a89576f1c02590963293e784b |
the weights (float32 arrays, keys emb.weight, head.weight, blocks.{0..23}.tm.*, blocks.{0..23}.cm.*, g_out); what every plaintext script loads |
bundle_pbd430a.bin + bundle_pbd430a.index.txt |
3,448,050,888 + 18,332 |
67069799192af9e2447abca0e707c8464576d35f0f06aaae83029bd03b45b5fc |
the flat float64 dump the GPU server reads (--bundle-dir); the index lists tensor offsets |
pbd430a_config.json, pbd430a_seeds.json, pbd430a_train_log.jsonl |
training configuration, seeds, loss/perplexity log |
Training
Two runs on one rented machine (2× NVIDIA RTX PRO 6000 Blackwell 96 GB, 16-core Threadripper, 503 GB RAM; $1.53 per hour for the
whole box), each on one of the two cards while a sibling model trained on the other. Records: the run stamps and console logs in
the GitHub tree, results/dense-demo-s31/pod/state/runs/s31_train_pbd430/ and s31_anneal_pbd430a_a2/; the trainer is
ml-eval/train_fhe_native_ssm.py at the commits named in the stamps.
|
pretraining (pbd430) |
anneal (pbd430a, this checkpoint) |
| started |
2026-08-28 19:10 UTC |
2026-08-31 16:45 UTC, resumed from pbd430 (the optimizer state was forked; pbd430 itself is untouched) |
| data |
FineWeb-Edu only, staged as a token stream (ml-eval/stage_fineweb_tokens.py), 12,144,001,024 tokens = 30 tokens per parameter |
a 419 M-token mix, round-robin in 1 M-token chunks: WikiText-103 train (119.3 M) + Project Gutenberg fiction (150 M) + a FineWeb-Edu replay (150 M) (results/dense-demo-s31/pod_tools/anneal_staged_provenance.json); 409.6 M tokens consumed |
| steps |
1,482,422 steps of 8,192 tokens (batch 8 × 1,024), gradient accumulation 4 → 32,768 tokens per update |
+50,000 steps (to 1,532,422) |
| optimizer |
AdamW-style, lr 6e-4 with 2,000 warm-up steps and a cosine decay to 6e-5, weight decay 0.1, gradient clip 1.0, bf16 autocast, torch.compile, seed 1234; the recurrent state carried across shard boundaries (--window-carry) |
lr at the floor, 6e-5, otherwise unchanged |
| wall time |
61.44 h on one card |
2.08 h on one card |
| final training loss |
3.2004 (ppl 24.5 on the stream) |
2.969 |
| rental cost |
61.44 h × $1.53 = $94 for the box (half of it attributable to this model: the other card trained the sibling) |
2.08 h × $1.53 = $3.2 |
30 tokens/parameter is pretty low and was chosen only because of the limited budget for this project. Before the anneal, WikiText-103 ppl measured 60.74 and 28.55 after.
Evaluation (quality — this is separate from the encrypted-circuit fidelity)
Measured with ml-eval/anchor_eval.py (harness validated on Pythia-410M first; records in the GitHub tree, results/dense-demo-s31/pod/state/runs/):
| checkpoint |
WikiText-103 test, token ppl (285,830 tokens, ctx 1024) |
pretraining-corpus holdout, token ppl (2,000,000 tokens) |
LAMBADA accuracy (5,153 items) |
| pbd430a (this repository; annealed on WT103 train + Gutenberg + FineWeb-Edu replay) |
28.55 |
27.91 |
9.0 % |
| pbd430 (before the anneal) |
60.74 |
— |
10.1 % |
| Pythia-410M, final (for scale, same harness) |
18.17 |
— |
51.6 % |
Read the LAMBADA number as the honest one: the anneal lowers WikiText perplexity because the anneal corpus contains WikiText's own train
split; the model does not know many facts and cannot reason. Under greedy decoding it loops within ~20 words on every prompt tried; the
encrypted long run used a deterministic anti-repetition rule (repetition penalty 1.3 over the last 64 tokens + no-repeat-4-gram; spec_decode/decode_policy.py).
Fidelity is not quality. The encrypted circuit reproduces this model's plaintext top-1 choice in 9,445 of 9,472 comparisons
(99.7 %), with the 27 disagreements at near ties of the plaintext logits, which says that the circuit computes what the model computes,
not that the model is good.
How to run it (plaintext)
git clone https://github.com/xelananv/fhe-ssm && cd fhe-ssm
# put pbd430a_weights.npz + pbd430a_config.json + pbd430a_seeds.json into ml-eval/artifacts/
# one-shot completion, deterministic (the rule the encrypted long run used):
python3 ml-eval/serve_ollama.py --prompt "The most important thing to know about rivers is that" --temperature 0 --repeat-penalty 1.3 --num-predict 60
# or chat with it through any Ollama client (the model is not a GGUF; this serves the same forward pass behind the Ollama API):
python3 ml-eval/serve_ollama.py --port 11435 --chat-style plain
OLLAMA_HOST=127.0.0.1:11435 ollama run fhe-ssm-430m
The chat server samples at temperature 0.8 by default, so every run differs; --temperature 0 --repeat-penalty 1.3 reproduces the
encrypted run's text token for token up to the first near-tie (the server has no 4-gram ban and runs float32, so the two can part
later). Needs Python 3.9+, torch, numpy, transformers (for the tokenizer); ≈ 2 GB of RAM; 50 tokens/s on an Apple laptop.
How to run it under encryption
The GPU server, as well as the client too and the build scripts and the records of the recorded sessions are in the github repo.
(docs/RUNBOOK.md). Two 96 GB cards at ring 2^17 & 128-bit classical parameters, ≈ 176 s per step of 64 lanes.
Intended use and limitations
A research artifact for encrypted inference: the point is that this exact model runs under FHE with measured fidelity. As a
language model it is quite small and undertrained by today's standards, it is also not instruction-tuned, not safety-tuned, and states false facts
fluently ("Leonardo da Vinci painted the Mona Lisa in 1452"). Do not use it for anything where the text has to be true.
Trained on public web text (FineWeb-Edu) and WikiText-103 (as well as Project Gutenberg fiction); it can and does reproduce their biases.
Citation
@software{fhe_ssm_2026,
author = {Ananyev, Alexander},
title = {fhe-ssm: a recurrent language model served under CKKS fully homomorphic encryption with free-running generation},
year = {2026},
url = {https://github.com/xelananv/fhe-ssm}
}
MIT. Alexander Ananyev — [email protected]. X, GitHub, Website
