Pragnosia — a spin-dominant recurrent language model (grown 100M → 1B)

Author: Ashish Kumar

Pragnosia is a transformer/SSM hybrid in which a rotational "spin" recurrence — a diagonal-complex linear-recurrent unit run as a parallel associative scan — is the core token-mixer, with attention only a periodic helper (1 layer in 4). The model was grown while training (function-preserving depth/width growth on probe-confirmed saturation) from 100M to the current 1.02B, on a 176.6B-token corpus, on a single shared H100.

This repo contains every checkpoint of the growth lineage, the full training/eval code, the tokenizer, and the paper. The 1B is still training (checkpoints here will be updated).

Why this architecture

• Causal attribution at scale: ablating the spin carrier in the trained model collapses perplexity 306× (284M) → 1081× (565M) → 1886× (1B) while ablating attention costs only ~3.4× — the carrier is the load-bearing core, and its dominance grows with scale.
• Param-matched wins: at ≈37M (3 seeds, per-arch tuned lr) spin-dominant reaches mean val ppl ≈116 vs ≈361 for attention-only (~3×). A no-phase real-recurrence control also beats attention (3.5×), so the effect is about placement (recurrence as the core mixer), not the complex phase.
• Long context is O(T) (vs attention's O(T²)) and inference is recurrent: O(1)/token, no KV cache. After training with a mixed-window curriculum (up to 8192-token effective context), validation perplexity with the cross-window carry matches or beats the short-context baseline, and generation probes recall the gist of content ~2.5k tokens back. Sharp verbatim ("needle") retrieval is still weak — the carry is a lossy summary, not a lookup.

Checkpoints (the growth lineage)

file	params	arch	val ppl	note
weights/spin_1022M_latest.pt	1022M	29L mlp27 d768	18.0	current — 33 tokens/param (34B tokens). Long-context carry now a decisive, monotonic win: perplexity WITH the carry beats the 256-token baseline at every length, −18.7% at 8192 tokens. Faculties 69%, honest self-persona (accurately names its own weaknesses). This is the best all-round checkpoint.
weights/spin_1022M_multistep.pt	1022M	29L mlp27	17.8	lowest short-context ppl; first multi-step-arithmetic pass (long-context carry pre-re-flip)
weights/spin_706M_converged.pt	706M	24L mlp22	18.8	converged 706M (the 1B's seed)
weights/spin_706M_polish.pt	706M	24L mlp22	20.7	mid-polish
weights/spin_621M_longmix_g1.pt	621M	24L mlp19	20.0	first grow under long-context training
weights/spin_593M_preLongCtx.pt	593M	24L mlp18	21.8	last pure short-context checkpoint
weights/spin_565M.pt	565M	23L mlp18	21.7	the 565M reported in the paper

Coming soon — an instruction-tuned (chat) variant. An SFT pass is in progress on top of the 33 t/p base: a 405M-token blend (15% pretraining replay, 35% CoT reasoning, 25% multi-turn chat, 25% instruction) in the <user>/<assistant> template, to make it conversational and reason step-by-step. Pipeline is in prepare_sft.py. The chat checkpoint will be added here when it completes.

Each checkpoint has a matching .json (architecture config) and, where available, a _state.json (training-state sidecar: cumulative tokens, best val ppl).

Zero-shot comparison (same harness, public models)

model	PIQA	ARC-E	ARC-C	HellaSwag	LAMBADA
Pragnosia-565M	60.7	34.1	24.9	29.3	21.7
GPT-2 (124M)	63.3	34.1	23.1	27.2	32.2
Cerebras-GPT-256M	62.1	32.7	22.1	26.3	29.5
GPT-2-medium (355M)	66.5	37.5	23.3	36.8	42.8
Cerebras-GPT-1.3B (1316M)	67.1	38.9	25.7	35.9	47.0

Honest reading: competitive on ARC/HellaSwag at similar scale, behind on fluency-driven LAMBADA — the models are undertrained for their size (see Limitations). Gaps narrow with training (284M → 565M improved every task except one). The Cerebras-GPT-1.3B row is a token-matched control (a transformer trained at a similar tokens-per-param budget): its multi-hop reasoning is comparably weak, indicating the weakness is scale/training budget, not the spin architecture (details in the paper).

Usage

The architecture is custom (not transformers-compatible) — the loading code is in this repo:

import json, torch, sys
sys.path.insert(0, ".")            # repo root (s6_hybrid.py)
import s6_hybrid as H

cfg = json.load(open("weights/spin_1022M_latest.json"))
H.VOC, H.L = cfg["vocab"], cfg["ctx"]
sd = torch.load("weights/spin_1022M_latest.pt", map_location="cpu", weights_only=True)
m = H.SpinAttentionLM(cfg["vocab"], cfg["d"], cfg["heads"], cfg["layers"],
                      mlp_mult=cfg["mlp_mult"], carrier=cfg["carrier"]).eval()
m.load_state_dict(sd)

from tokenizers import Tokenizer
tok = Tokenizer.from_file("data/bpe.json")
ids = tok.encode("The capital of France is").ids
with torch.no_grad():
    for _ in range(20):
        nxt = int(m(torch.tensor([ids[-cfg["ctx"]:]]))[0, -1].argmax())
        ids.append(nxt)
print(tok.decode(ids))

For long context, thread the carrier state across 256-token windows (logits, S = m(x, state=S, return_state=True)) — see lc_watch.py and train_pragnosia.py::_long_step.

Training

• train_pragnosia.py — GPU-adaptive trainer: grow-as-you-train (function-preserving grow.py), mixed-window long-context curriculum (LONG_MIX=1, up to 8192 effective context), VRAM capping for shared GPUs, all thresholds derived from data (no hand-set constants).
• Corpus: 176.6B tokens (web/books/code/CoT mix), BPE-16384 (data/bpe.json), ctx 256.
• The 1B has seen 24B tokens (23 tokens/param) so far; training continues toward 30+.

Limitations (stated plainly)

• Undertrained for its size — multi-hop reasoning (analogy, multi-step deduction, counterfactuals) is weak; single-step arithmetic, facts, and logic are solid. More tokens is the lever.
• Base model — no instruction tuning or RLHF; it continues text, it does not chat.
• Long-range verbatim recall is weak (lossy recurrent carry); aggregate/gist use of long context works.
• Small-model factuality — it can be confidently wrong; do not use for factual QA without verification.
• The decisive matched carrier="none" transformer baseline at ≥284M is still to run.

Paper

Pragnosia_paper.pdf in this repo — includes the full ablation suite, multi-seed tables, the grow-as-you-train method, long-context curriculum, and all measured limitations.