ryankim17920/mechanistic-sps

Which State Should Prediction Read? A Mechanistic Analysis of State–Prediction Separation

Checkpoints accompanying the paper Which State Should Prediction Read? A Mechanistic Analysis of State–Prediction Separation (Ryan Kim, 2026): 86 trained language models (Transformer, SPS, Two-tower and Sequential families, 50M–960M parameters), the self-contained model code, and a one-call loader.

Paper: arXiv link coming soon · Code: github.com/RyanKim17920/mechanistic-sps · Intermediate checkpoints / embedding ladders: ryankim17920/mechanistic-sps-extra

The models

A language model keeps a running memory of the text (its state) and also produces the next token (its prediction). The families here differ in how those two jobs are laid out:

• Transformer: one stream does both (Llama-style blocks: RMSNorm, RoPE, SwiGLU).
• SPS (State–Prediction Separation, Monea et al. 2026): every token gets two slots that share all weights, a state slot and a <predict> slot; prediction slots see state keys at any distance and prediction keys within a window of 64.
• Two-tower: a state tower and a prediction tower with separate weights; prediction block i reads state level i (aligned depths).
• Sequential: the same two towers, but every prediction block reads the final state (read_map=final). Multi-level read (read_map=mix): prediction block i reads a learned mixture of all state levels, initialised one-hot on the final level.

Every model uses the GPT-2 BPE tokenizer (vocabulary padded to 50,304; <|endoftext|>=50256 separates documents, 50257 is the <predict> id), context 4,096 (32,768 for long_context_32k/), and was trained from scratch on FineWeb-Edu (sample-100BT) with AdamW, 8×H100.

Quickstart

pip install torch==2.9.0 numpy huggingface_hub tiktoken   # torch 2.9.0 brings triton 3.5.0; Python >= 3.10

import sys, torch, tiktoken
from huggingface_hub import snapshot_download

code = snapshot_download("ryankim17920/mechanistic-sps", allow_patterns=["load.py", "modeling/**"])
sys.path.insert(0, code)
from load import load_model

# a run name is found in any folder; "multilevel_read/s_two_tower_seq12_mix_20b" works too
model = load_model("ryankim17920/mechanistic-sps", "s_two_tower_seq12_mix_20b")     # nn.Module, eval mode, on CUDA

enc = tiktoken.get_encoding("gpt2")
ids = torch.tensor([enc.encode("The mitochondria is the powerhouse of the cell. It produces ATP")], device="cuda")
with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
    logits, loss, stats = model(ids[:, :-1], ids[:, 1:].clone())   # targets are masked in place
print(f"mean next-token NLL: {loss.item():.3f} nats")

load_model(path_or_repo, subpath) takes a Hub repo id, a local copy of this repo (snapshot_download(..., local_dir=...)), a run folder, or a .pt file. It rebuilds the model class and config from the checkpoint's own model_args, loads the weights, and raises if any parameter is missing or unexpected. python load.py ryankim17920/mechanistic-sps <run> --val-bin val.bin scores a tokenized (uint16) validation file exactly as the Val. NLL numbers below were measured (every 4,096-token window; --n-seq N scores only the first N windows, which read about 0.02 nats lower than the full set).

Runtime notes:

• GPU. Trained and evaluated on H100 (bf16 autocast). Everything needs CUDA except building the model; SPS and the standard Transformer run custom Triton kernels (tested with triton 3.5.0).
• Standard Transformer: use the Triton kernel. All Transformer numbers come from the Triton document-masked attention kernel, i.e. use_triton_full_attention=True. load.py sets it for every Transformer checkpoint; if you build the model yourself, set it too (the config default is False, a flex-attention path used only by small test models and on CPU).
• SPS uses its Triton kernel by default; overrides={"attn_backend": "flex"} selects the flex backend instead (benchmark path; no run trained on it).
• Two-tower / Sequential use torch.nn.attention.flex_attention (compiled on first call; older checkpoints record attn_backend="auto", which load.py maps to flex, the backend they trained with). overrides={"attn_backend": "sdpa"} avoids compilation at a memory cost.
• Triton cache. Point TRITON_CACHE_DIR (and TORCHINDUCTOR_CACHE_DIR) at node-local disk. A cache shared across machines on a network filesystem silently returned wrong kernels once (val NLL 9.71 instead of 3.11 on the same checkpoint).
• The forward pass is model(idx_BxT, targets_BxT) -> (logits, loss, stats) (int64 GPT-2 ids, T <= block_size); there is no KV-cache generation API.

Repository layout

<run>/                    paper finals (s size, 20B tokens), unchanged paths
scaling/<run>/            size ladder xxs / xs / m / l: Transformer, SPS, Two-tower, Sequential (+ seeds)
multilevel_read/<run>/    Sequential with the multi-level read (mix, attnmix) at xxs, xs, s, m, l
iso_flop/<run>/           Sequential trained to the training FLOPs of SPS (s, m)
long_context_32k/<run>/   32K-context continuations of six s models
controls/<run>/           architecture controls and depth splits
modeling/                 model code (Transformer, SPS + Triton kernels, Two-tower/Sequential)
load.py                   load_model(), val_nll()
SHA256SUMS                sha256 of every checkpoint

Each run folder holds ckpt_tokens_<N>_final.pt (the final weights after N training tokens) and config.json (model arguments, full training config, params, metrics, file size and sha256). Each .pt is a torch.save dict: model (fp32 state dict), model_args, the training config, and training metadata; the optimizer state is removed. Run names are the experiment names of the code repo (conf/experiment/<run>.yaml); xxs/xs/s/m/l = d_model 384/512/768/1024/1280.

Columns. Params: all parameters. Non-emb.: params minus the vocabulary-sized matrices (token table, <predict> table where separate, untied LM head). Val. NLL: full-sweep validation cross-entropy in nats/token on the held-out FineWeb-Edu split (13,085,928 tokens, 3,198 non-overlapping 4,096-token windows; the iso-FLOP runs use the same split). Downstream acc.: mean zero-shot accuracy (%) over LAMBADA (acc), HellaSwag, PIQA, ARC-Easy, ARC-Challenge, OpenBookQA, SciQ (acc_norm) and WinoGrande (acc), lm-evaluation-harness; — = not evaluated. Seed 1 is the config default (1337); seed n is the run with suffix _seed<n>.

Scaling ladder (scaling/)

Folder	Model	Seed	Size	Shape	Params	Non-emb.	Tokens	Val. NLL	Downstream acc.	File
scaling/xxs_full_attention_20b_fw100_untied/	Transformer (untied)	1	xxs	6L, d=384, 6h	50.1M	11.5M	20.00B	3.3086	39.14	0.20 GB
scaling/xxs_full_attention_20b_fw100_untied_seed2/	Transformer (untied)	2	xxs	6L, d=384, 6h	50.1M	11.5M	20.00B	3.3032	39.68	0.20 GB
scaling/xxs_sps_w64_20b_fw100_untied/	SPS (untied, w=64)	1	xxs	6L, d=384, 6h	50.1M	11.5M	20.00B	3.2512	40.18	0.20 GB
scaling/xxs_sps_w64_20b_fw100_untied_seed2/	SPS (untied, w=64)	2	xxs	6L, d=384, 6h	50.1M	11.5M	20.00B	3.2492	40.51	0.20 GB
scaling/xxs_two_tower_w0_equal3_20b/	Two-tower 3+3	1	xxs	3+3, d=384, 6h	69.5M	11.5M	20.00B	3.3051	39.34	0.28 GB
scaling/xxs_two_tower_w0_equal3_20b_seed2/	Two-tower 3+3	2	xxs	3+3, d=384, 6h	69.5M	11.5M	20.00B	3.2972	39.72	0.28 GB
scaling/xxs_two_tower_seq3_20b/	Sequential 3+3	1	xxs	3+3, d=384, 6h	69.5M	11.5M	20.00B	3.2301	40.44	0.28 GB
scaling/xxs_two_tower_seq3_20b_seed2/	Sequential 3+3	2	xxs	3+3, d=384, 6h	69.5M	11.5M	20.00B	3.2339	39.37	0.28 GB
scaling/xs_full_attention_20b_fw100_untied/	Transformer (untied)	1	xs	8L, d=512, 8h	78.8M	27.3M	20.00B	3.1138	41.77	0.32 GB
scaling/xs_full_attention_20b_fw100_untied_seed2/	Transformer (untied)	2	xs	8L, d=512, 8h	78.8M	27.3M	20.00B	3.1119	41.29	0.32 GB
scaling/xs_sps_w64_20b_fw100_untied/	SPS (untied, w=64)	1	xs	8L, d=512, 8h	78.8M	27.3M	20.00B	3.0440	42.64	0.32 GB
scaling/xs_sps_w64_20b_fw100_untied_seed2/	SPS (untied, w=64)	2	xs	8L, d=512, 8h	78.8M	27.3M	20.00B	3.0465	41.94	0.32 GB
scaling/xs_two_tower_w0_equal4_20b/	Two-tower 4+4	1	xs	4+4, d=512, 8h	104.5M	27.3M	20.00B	3.0771	41.63	0.42 GB
scaling/xs_two_tower_w0_equal4_20b_seed2/	Two-tower 4+4	2	xs	4+4, d=512, 8h	104.5M	27.3M	20.00B	3.0989	41.21	0.42 GB
scaling/xs_two_tower_w0_equal4_20b_seed3/	Two-tower 4+4	3	xs	4+4, d=512, 8h	104.5M	27.3M	20.00B	3.1084	—	0.42 GB
scaling/xs_two_tower_seq4_20b/	Sequential 4+4	1	xs	4+4, d=512, 8h	104.5M	27.3M	20.00B	3.0360	43.15	0.42 GB
scaling/xs_two_tower_seq4_20b_seed2/	Sequential 4+4	2	xs	4+4, d=512, 8h	104.5M	27.3M	20.00B	3.0332	42.24	0.42 GB
scaling/xs_two_tower_seq4_20b_seed3/	Sequential 4+4	3	xs	4+4, d=512, 8h	104.5M	27.3M	20.00B	3.0326	—	0.42 GB
scaling/m_full_attention_20b_fw100_untied/	Transformer (untied)	1	m	24L, d=1024, 16h	430.2M	327.2M	20.00B	2.6526	49.24	1.72 GB
scaling/m_sps_w64_20b_fw100_untied/	SPS (untied, w=64)	1	m	24L, d=1024, 16h	430.2M	327.2M	20.00B	2.5843	51.15	1.72 GB
scaling/m_sps_w64_20b_fw100_untied_seed2/	SPS (untied, w=64)	2	m	24L, d=1024, 16h	430.2M	327.2M	20.00B	2.5812	50.97	1.72 GB
scaling/m_two_tower_w0_equal12_20b/	Two-tower 12+12	1	m	12+12, d=1024, 16h	481.8M	327.2M	20.00B	2.6394	49.06	1.93 GB
scaling/m_two_tower_seq12_20b/	Sequential 12+12	1	m	12+12, d=1024, 16h	481.8M	327.2M	20.00B	2.6155	50.06	1.93 GB
scaling/m_two_tower_seq12_20b_seed2/	Sequential 12+12	2	m	12+12, d=1024, 16h	481.8M	327.2M	20.00B	2.6057	50.13	1.93 GB
scaling/m_two_tower_seq12_slim_20b/	Sequential 12+12 slim (constant input, no prediction table)	1	m	12+12, d=1024, 16h	430.2M	327.2M	20.00B	2.6085	—	1.72 GB
scaling/l_two_tower_seq18_20b/	Sequential 18+18	1	l	18+18, d=1280, 20h	960.1M	766.9M	20.00B	2.5030	53.24	3.84 GB

Multi-level read (multilevel_read/)

Folder	Model	Seed	Size	Shape	Params	Non-emb.	Tokens	Val. NLL	Downstream acc.	File
multilevel_read/xxs_two_tower_seq3_mix_20b/	Sequential 3+3, multi-level read (mix)	1	xxs	3+3, d=384, 6h	69.5M	11.5M	20.00B	3.2237	—	0.28 GB
multilevel_read/xs_two_tower_seq4_mix_20b/	Sequential 4+4, multi-level read (mix)	1	xs	4+4, d=512, 8h	104.5M	27.3M	20.00B	3.0285	—	0.42 GB
multilevel_read/s_two_tower_seq12_attnmix_20b/	Sequential 12+12, input-dependent read over state deltas (attnmix)	1	s	12+12, d=768, 12h	300.0M	184.1M	20.00B	2.6937	—	1.20 GB
multilevel_read/s_two_tower_seq12_mix_20b/	Sequential 12+12, multi-level read (mix)	1	s	12+12, d=768, 12h	300.0M	184.1M	20.00B	2.6900	—	1.20 GB
multilevel_read/s_two_tower_seq6_mix_20b/	Sequential 6+6, multi-level read (mix)	1	s	6+6, d=768, 12h	207.9M	92.0M	20.00B	2.7934	—	0.83 GB
multilevel_read/m_two_tower_seq12_mix_20b/	Sequential 12+12, multi-level read (mix)	1	m	12+12, d=1024, 16h	481.8M	327.2M	20.00B	2.6069	—	1.93 GB
multilevel_read/m_two_tower_seq12_mix_20b_seed2/	Sequential 12+12, multi-level read (mix)	2	m	12+12, d=1024, 16h	481.8M	327.2M	20.00B	2.5876	—	1.93 GB
multilevel_read/l_two_tower_seq18_mix_20b/	Sequential 18+18, multi-level read (mix)	1	l	18+18, d=1280, 20h	960.1M	766.9M	20.00B	2.4891	—	3.84 GB

Iso-FLOP (iso_flop/)

Sequential models trained beyond 20B tokens until their total training FLOPs equal those of the SPS model at the same size (m: 1.343e20, s: 4.533e19), with the same schedule shape (decay over the last 10%), on FineWeb-Edu sample-100BT extended to 60 shards (same validation split).

Folder	Model	Seed	Size	Shape	Params	Non-emb.	Tokens	Val. NLL	Downstream acc.	File
iso_flop/s_two_tower_seq6_isoflop_sps_20b/	Sequential 6+6, iso-FLOP to SPS s untied (36.50B tokens)	1	s	6+6, d=768, 12h	207.9M	92.0M	36.50B	2.7601	—	0.83 GB
iso_flop/m_two_tower_seq12_isoflop_sps_20b/	Sequential 12+12, iso-FLOP to SPS m (38.45B tokens)	1	m	12+12, d=1024, 16h	481.8M	327.2M	38.45B	2.5573	—	1.93 GB

32K long-context continuations (long_context_32k/)

Each model continues the 20B-token 4K checkpoint named in From for 1,048,576,000 tokens at context 32,768 with RoPE base 1e4 → 5e5 (ABF), all parameters trained, identical recipe for every arm (LR 3e-4, batch 16 × 32,768; data: PG19 train books and long FineWeb-Edu documents). NLL is measured on the first 32,768 tokens of each held-out document (79 PG19 test books; 100 FineWeb-Edu documents ≥ 32K tokens from unused shards). model_args stores block_size=32768, rope_theta=5e5.

Folder	Model	From	Params	PG19 test NLL @32K	FineWeb-Edu long NLL @32K	File
long_context_32k/lc32k_s_full_attention/	Transformer 12L	s_full_attention_20b_fw100_untied	169.3M	2.674	3.175	0.68 GB
long_context_32k/lc32k_s_two_tower_w0_equal6/	Two-tower 6+6	s_two_tower_w0_equal6_20b	207.9M	2.665	3.149	0.83 GB
long_context_32k/lc32k_s_two_tower_seq6/	Sequential 6+6	s_two_tower_seq6_20b	207.9M	2.631	3.101	0.83 GB
long_context_32k/lc32k_s_sps_w64/	SPS 12L (w=64)	s_sps_w64_20b_fw100_untied	169.3M	2.616	3.087	0.68 GB
long_context_32k/lc32k_s_two_tower_w0_equal/	Two-tower 12+12	s_two_tower_w0_equal_20b	300.0M	2.569	3.018	1.20 GB
long_context_32k/lc32k_s_two_tower_seq12/	Sequential 12+12	s_two_tower_seq12_20b	300.0M	2.538	2.970	1.20 GB

Controls (controls/)

Folder	Model	Seed	Size	Shape	Params	Non-emb.	Tokens	Val. NLL	Downstream acc.	File
controls/xxs_two_tower_seq2p4_20b/	Sequential 2+4	1	xxs	2+4, d=384, 6h	69.5M	11.5M	20.00B	3.2505	—	0.28 GB
controls/xxs_two_tower_seq4p2_20b/	Sequential 4+2	1	xxs	4+2, d=384, 6h	69.5M	11.5M	20.00B	3.2390	—	0.28 GB
controls/s_full_attention_24L_20b_untied/	Transformer 24L (untied), non-emb params/FLOPs matched to 12+12	1	s	24L, d=768, 12h	261.3M	184.1M	20.00B	2.7561	—	1.05 GB
controls/s_two_tower_top12p6_widepm_20b/	Two-tower 12+6 (top-aligned), pred MLP 5824 (non-emb params matched to 12+12)	1	s	12+6, d=768, 12h	300.0M	184.1M	20.00B	2.7223	—	1.20 GB
controls/s_two_tower_top12p6_widefm_20b/	Two-tower 12+6 (top-aligned), pred MLP 8576 (FLOPs matched to 12+12)	1	s	12+6, d=768, 12h	338.0M	222.1M	20.00B	2.6991	—	1.35 GB
controls/s_two_tower_seq4p8_20b/	Sequential 4+8	1	s	4+8, d=768, 12h	207.9M	92.0M	20.00B	2.8080	—	0.83 GB
controls/s_two_tower_seq6_initdet_20b/	Sequential 6+6, pred tower initialised from the detached state output	1	s	6+6, d=768, 12h	169.3M	92.0M	20.00B	2.8457	—	0.68 GB
controls/s_two_tower_seq6_init_20b/	Sequential 6+6, pred tower initialised from the state output (with gradient)	1	s	6+6, d=768, 12h	169.3M	92.0M	20.00B	2.8264	—	0.68 GB
controls/s_two_tower_seq6_skv_wide_20b/	Sequential 6+6, shared state K/V read, pred MLP widened to 2752	1	s	6+6, d=768, 12h	208.2M	92.3M	20.00B	2.8165	—	0.83 GB
controls/s_two_tower_seq8p4_20b/	Sequential 8+4	1	s	8+4, d=768, 12h	207.9M	92.0M	20.00B	2.8045	—	0.83 GB
controls/s_two_tower_seq8p4_pti_20b/	Sequential 8+4, per-tower residual-projection init	1	s	8+4, d=768, 12h	207.9M	92.0M	20.00B	—	—	0.83 GB
controls/s_two_tower_seq8p4_initdet_20b/	Sequential 8+4, pred tower initialised from the detached state output	1	s	8+4, d=768, 12h	169.3M	92.0M	20.00B	2.8583	—	0.68 GB
controls/s_two_tower_seq9p3_pti_20b/	Sequential 9+3, per-tower residual-projection init	1	s	9+3, d=768, 12h	207.9M	92.0M	20.00B	—	—	0.83 GB
controls/s_two_tower_shared12_20b/	Shared-weight Sequential (one 12-block stack run as state then pred)	1	s	12+12, d=768, 12h	207.9M	92.0M	20.00B	2.8057	—	0.83 GB
controls/s_two_tower_yoco6_20b/	YOCO-style 6+6 (pred tower continues the state residual; one shared K/V from the final state)	1	s	6+6, d=768, 12h	163.4M	86.1M	20.00B	2.8505	—	0.65 GB
controls/m_two_tower_seq16p8_pti_20b/	Sequential 16+8, per-tower residual-projection init	1	m	16+8, d=1024, 16h	481.7M	327.2M	20.00B	—	—	1.93 GB
controls/m_two_tower_seq18p6_pti_20b/	Sequential 18+6, per-tower residual-projection init	1	m	18+6, d=1024, 16h	481.7M	327.2M	20.00B	—	—	1.93 GB

The *_pti_20b runs (per-tower residual-projection init) have not been scored with the full sweep.

Paper finals (top-level folders)

The original release, at unchanged paths (the code repo's hf_export.py download <run> fetches these into <DUALSPS_OUT_ROOT>/out/<run>/). s_sps_w64_20b_fw100 (SPS tied, seed 1) is not included: its final checkpoint was lost; the SPS-tied numbers use seeds 2 and 3.

Folder	Model	Seed	Size	Shape	Params	Non-emb.	Tokens	Val. NLL	Downstream acc.	File
s_full_attention_20b_fw100/	Transformer, tied head	1	s	12L, d=768, 12h	130.7M	92.0M	20.00B	2.8510	—	0.52 GB
s_full_attention_20b_fw100_untied/	Transformer	1	s	12L, d=768, 12h	169.3M	92.0M	20.00B	2.8621	45.24	0.68 GB
s_full_attention_20b_fw100_untied_seed2/	Transformer	2	s	12L, d=768, 12h	169.3M	92.0M	20.00B	2.8677	44.53	0.68 GB
s_sps_w64_20b_fw100_seed2/	SPS, tied head	2	s	12L, d=768, 12h	130.7M	92.0M	20.00B	2.8062	45.83	0.52 GB
s_sps_w64_20b_fw100_seed3/	SPS, tied head	3	s	12L, d=768, 12h	130.7M	92.0M	20.00B	2.8085	46.02	0.52 GB
s_sps_w64_20b_fw100_untied/	SPS	1	s	12L, d=768, 12h	169.3M	92.0M	20.00B	2.7980	45.80	0.68 GB
s_sps_w64_20b_fw100_untied_seed2/	SPS	2	s	12L, d=768, 12h	169.3M	92.0M	20.00B	2.7953	45.62	0.68 GB
s_two_tower_afsps_faithful_20b/	AF-SPS	1	s	12+12, d=768, 12h	130.7M	92.1M	20.00B	2.8261	—	0.52 GB
s_two_tower_asym_20b/	Two-tower 12+6 (top-aligned)	1	s	12+6, d=768, 12h	251.3M	135.4M	20.00B	2.7567	—	1.01 GB
s_two_tower_seq12_20b/	Sequential 12+12	1	s	12+12, d=768, 12h	300.0M	184.1M	20.00B	2.7019	47.86	1.20 GB
s_two_tower_seq12_20b_seed2/	Sequential 12+12	2	s	12+12, d=768, 12h	300.0M	184.1M	20.00B	2.7014	47.38	1.20 GB
s_two_tower_seq12_tied_20b/	Shared Sequential	1	s	12+12, d=768, 12h	131.8M	93.2M	20.00B	2.8270	—	0.53 GB
s_two_tower_seq12_tied_20b_seed2/	Shared Sequential	2	s	12+12, d=768, 12h	131.8M	93.2M	20.00B	2.8256	—	0.53 GB
s_two_tower_seq3p9_20b/	Sequential 3+9	1	s	3+9, d=768, 12h	207.9M	92.0M	20.00B	2.8130	—	0.83 GB
s_two_tower_seq6_20b/	Sequential 6+6	1	s	6+6, d=768, 12h	207.9M	92.0M	20.00B	2.8030	45.62	0.83 GB
s_two_tower_seq6_20b_seed2/	Sequential 6+6	2	s	6+6, d=768, 12h	207.9M	92.0M	20.00B	2.7998	45.57	0.83 GB
s_two_tower_seq6_cpause_20b/	Sequential 6+6 (⟨predict⟩ input)	1	s	6+6, d=768, 12h	169.3M	92.0M	20.00B	2.8049	45.62	0.68 GB
s_two_tower_seq6_slim_20b/	Sequential 6+6 (tied head, ⟨predict⟩ input)	1	s	6+6, d=768, 12h	130.7M	92.0M	20.00B	2.8151	45.81	0.52 GB
s_two_tower_seq6_slim_20b_seed2/	Sequential 6+6 (tied head, ⟨predict⟩ input)	2	s	6+6, d=768, 12h	130.7M	92.0M	20.00B	2.8194	45.97	0.52 GB
s_two_tower_seq9p3_20b/	Sequential 9+3	1	s	9+3, d=768, 12h	207.9M	92.0M	20.00B	2.8146	—	0.83 GB
s_two_tower_w0_equal6_20b/	Two-tower 6+6	1	s	6+6, d=768, 12h	207.9M	92.0M	20.00B	2.8381	45.32	0.83 GB
s_two_tower_w0_equal6_20b_seed2/	Two-tower 6+6	2	s	6+6, d=768, 12h	207.9M	92.0M	20.00B	2.8380	45.10	0.83 GB
s_two_tower_w0_equal_20b/	Two-tower 12+12	1	s	12+12, d=768, 12h	300.0M	184.1M	20.00B	2.7321	46.69	1.20 GB
s_two_tower_w0_equal_20b_seed2/	Two-tower 12+12	2	s	12+12, d=768, 12h	300.0M	184.1M	20.00B	2.7333	46.82	1.20 GB
s_two_tower_w0_equal_tiedattn_20b/	Two-tower 12+12 (shared attn. only)	1	s	12+12, d=768, 12h	271.7M	155.8M	20.00B	2.7671	—	1.09 GB
s_two_tower_w0_s1152p3456_20b/	Two-tower 12+12 (MLP realloc.)	1	s	12+12, d=768, 12h	300.0M	184.1M	20.00B	2.7368	—	1.20 GB
s_two_tower_w0_shared_20b/	Shared Two-tower	1	s	12+12, d=768, 12h	130.7M	92.0M	20.00B	2.8144	—	0.52 GB

Loading with the code repo

With the code repo, python scripts/dualsps/hf_export.py download <run> downloads a top-level run into <DUALSPS_OUT_ROOT>/out/<run>/, and common.load("<run>") rebuilds it from conf/experiment/<run>.yaml. For the folders added later, download into the same place:

from huggingface_hub import snapshot_download
snapshot_download("ryankim17920/mechanistic-sps", allow_patterns=["scaling/m_two_tower_seq12_20b/*"], local_dir="hf")
# then move hf/scaling/m_two_tower_seq12_20b/ to <DUALSPS_OUT_ROOT>/out/m_two_tower_seq12_20b/

Citation

@misc{kim2026whichstate,
  title  = {Which State Should Prediction Read? A Mechanistic Analysis of State--Prediction Separation},
  author = {Ryan Kim},
  year   = {2026},
  note   = {arXiv preprint coming soon},
  url    = {https://github.com/RyanKim17920/mechanistic-sps}
}

License

MIT, the same as the code repository. Trained on FineWeb-Edu (ODC-By 1.0).