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

Final checkpoints of every model reported in the paper Which State Should Prediction Read? A Mechanistic Analysis of State–Prediction Separation.

Paper: arXiv link coming soon · Code: github.com/RyanKim17920/mechanistic-sps

Related repos: ryankim17920/mechanistic-sps-extra holds the intermediate checkpoints and embedding ladders read by two analyses (training trajectory, embedding divergence) and the final of s_two_tower_afsps_20b.

Summary

A language model has to keep a running memory of the text (its state) and also produce the next token (its prediction). State–Prediction Separation (SPS) gives these two jobs separate token slots that still share all weights. This paper pulls the two jobs fully apart into two streams (the Two-tower and Sequential models) so we can trace what each stream writes and what the prediction side reads. Previous-token information concentrates in the state stream, induction matching emerges on the prediction side, and letting every prediction layer read the final state with its own weights (Sequential) lowers loss at matched parameters and FLOPs.

All models: d_model 768, 12 heads, GPT-2 BPE (vocab 50304 padded), context 4096, trained on 20B tokens of FineWeb-Edu (sample-100BT), AdamW, 8×H100.

Models

Layers is state+prediction depth for the Two-tower/Sequential families; SPS runs both roles through the same 12 shared layers. Val. loss is the full-sweep validation cross-entropy (nats/token, FineWeb-Edu validation split) exactly as recorded in the paper's ledger (scripts/analysis/results/ledger.jsonl, field val_nll_full_sweep); Params from the same ledger. Seed 1 is the config's default seed (1337); seed n is the run with suffix _seed<n>.

Folder	Paper arm	Seed	Layers (state+pred)	Params	Training tokens	Val. loss	Size	sha256
`s_full_attention_20b_fw100/`	Transformer, tied head	1	12	130.7M	20,000,145,408	2.851023	0.52 GB	`5bea5f214a73…`
`s_full_attention_20b_fw100_untied/`	Transformer	1	12	169.3M	20,000,145,408	2.862126	0.68 GB	`e9b82b5bbbcf…`
`s_full_attention_20b_fw100_untied_seed2/`	Transformer	2	12	169.3M	20,000,145,408	2.867694	0.68 GB	`b2f16f891aea…`
`s_sps_w64_20b_fw100_seed2/`	SPS, tied head	2	12 (shared slots)	130.7M	20,000,145,408	2.806231	0.52 GB	`8a41fb9983d6…`
`s_sps_w64_20b_fw100_seed3/`	SPS, tied head	3	12 (shared slots)	130.7M	20,000,145,408	2.808532	0.52 GB	`66fb8413c1b3…`
`s_sps_w64_20b_fw100_untied/`	SPS	1	12 (shared slots)	169.3M	20,000,145,408	2.797983	0.68 GB	`c3aa06cb4210…`
`s_sps_w64_20b_fw100_untied_seed2/`	SPS	2	12 (shared slots)	169.3M	20,000,145,408	2.795271	0.68 GB	`59734ac3271e…`
`s_two_tower_afsps_faithful_20b/`	AF-SPS	1	12+12	130.7M	20,000,145,408	2.826122	0.52 GB	`f47e6e84b349…`
`s_two_tower_w0_equal_20b/`	Two-tower 12+12	1	12+12	300.0M	20,000,145,408	2.732093	1.20 GB	`46bd8b565129…`
`s_two_tower_w0_equal_20b_seed2/`	Two-tower 12+12	2	12+12	300.0M	20,000,145,408	2.733251	1.20 GB	`b6c9bce1425f…`
`s_two_tower_w0_equal_tiedattn_20b/`	Two-tower 12+12 (shared attn. only)	1	12+12	271.7M	20,000,145,408	2.767096	1.09 GB	`8d93fdee6561…`
`s_two_tower_w0_s1152p3456_20b/`	Two-tower 12+12 (MLP realloc.)	1	12+12	300.0M	20,000,145,408	2.736785	1.20 GB	`c60157e8beee…`
`s_two_tower_w0_equal6_20b/`	Two-tower 6+6	1	6+6	207.9M	20,000,145,408	2.838113	0.83 GB	`6f783623ebcb…`
`s_two_tower_w0_equal6_20b_seed2/`	Two-tower 6+6	2	6+6	207.9M	20,000,145,408	2.838036	0.83 GB	`9f4bcde0d62e…`
`s_two_tower_w0_shared_20b/`	Shared Two-tower	1	12+12	130.7M	20,000,145,408	2.814383	0.52 GB	`4959e93733a9…`
`s_two_tower_asym_20b/`	Two-tower 12+6 (top-aligned)	1	12+6	251.3M	20,000,145,408	2.756741	1.01 GB	`23c640071b02…`
`s_two_tower_seq12_20b/`	Sequential 12+12	1	12+12	300.0M	20,000,145,408	2.701871	1.20 GB	`c5ae4cc6601e…`
`s_two_tower_seq12_20b_seed2/`	Sequential 12+12	2	12+12	300.0M	20,000,145,408	2.701394	1.20 GB	`69a3fd70646a…`
`s_two_tower_seq12_tied_20b/`	Shared Sequential	1	12+12	131.8M	20,000,145,408	2.826953	0.53 GB	`1df22a630fdd…`
`s_two_tower_seq12_tied_20b_seed2/`	Shared Sequential	2	12+12	131.8M	20,000,145,408	2.825576	0.53 GB	`9a1b77436436…`
`s_two_tower_seq3p9_20b/`	Sequential 3+9	1	3+9	207.9M	20,000,145,408	2.812965	0.83 GB	`fc977521c967…`
`s_two_tower_seq6_20b/`	Sequential 6+6	1	6+6	207.9M	20,000,145,408	2.803027	0.83 GB	`d3954c022a5a…`
`s_two_tower_seq6_20b_seed2/`	Sequential 6+6	2	6+6	207.9M	20,000,145,408	2.799752	0.83 GB	`8a67d184ec5c…`
`s_two_tower_seq6_cpause_20b/`	Sequential 6+6 (⟨predict⟩ input)	1	6+6	169.3M	20,000,145,408	2.804941	0.68 GB	`b4c8e820c176…`
`s_two_tower_seq6_slim_20b/`	Sequential 6+6 (tied head, ⟨predict⟩ input)	1	6+6	130.7M	20,000,145,408	2.815133	0.52 GB	`bbd068ebec38…`
`s_two_tower_seq6_slim_20b_seed2/`	Sequential 6+6 (tied head, ⟨predict⟩ input)	2	6+6	130.7M	20,000,145,408	2.819392	0.52 GB	`0290edf2052b…`
`s_two_tower_seq9p3_20b/`	Sequential 9+3	1	9+3	207.9M	20,000,145,408	2.814565	0.83 GB	`0379bdea6702…`

Not included: the seed-1 (1337) run of SPS, tied head (s_sps_w64_20b_fw100), whose final checkpoint was lost; the paper's SPS, tied head numbers use seeds 2 and 3, as stated in the paper. The two freeze-and-retrain source runs are s_two_tower_seq6_20b and s_two_tower_w0_equal6_20b.

File layout

<run>/
  ckpt_tokens_20000145408_final.pt   # final weights after 20,000,145,408 tokens
  config.json                        # human-readable: paper arm, model_args, full training
                                     # config, val loss, params, file size + sha256
README.md
SHA256SUMS                           # full sha256 of every .pt

Folder names are the run names used everywhere in the code (conf/experiment/<run>.yaml, the ledger, hf_export.py, common.load), so a folder downloads straight into <DUALSPS_OUT_ROOT>/out/<run>/. The table above maps each folder to its name in the paper.

Each .pt is a torch.save dict: model (fp32 state dict), model_args, the full training config, and training metadata (iter_num, provenance, RNG state, ...). The optimizer state was removed; everything else is as saved by training.

Loading

With the code repo (uv sync, then set DUALSPS_DATA_ROOT in .env), download a run into <DUALSPS_OUT_ROOT>/out/<run>/, where the evaluation, the analyses and the freeze-and-retrain configs look for it. No Hugging Face login is needed:

python scripts/dualsps/hf_export.py download s_two_tower_seq12_20b

Then, in Python (from the repo root; load rebuilds the model from conf/experiment/<run>.yaml, checks it against the checkpoint's model_args, and loads the weights on the GPU):

import sys; sys.path[:0] = ["src", "scripts/dualsps", "scripts/analysis"]
import common
model, cfg, family, ckpt_path = common.load("s_two_tower_seq12_20b")

Or without the helpers, from the checkpoint alone:

import torch
from huggingface_hub import hf_hub_download
from modeling.models.model import config_args
from modeling.models.two_tower.core import TwoTowerConfig, TwoTowerModel

path = hf_hub_download("ryankim17920/mechanistic-sps",
                       "s_two_tower_seq12_20b/ckpt_tokens_20000145408_final.pt")
ck = torch.load(path, map_location="cpu", weights_only=False)   # the dict holds configs, not just tensors
model = TwoTowerModel(TwoTowerConfig(**config_args(TwoTowerConfig, ck["model_args"])))
sd = {k.replace("_orig_mod.", ""): v for k, v in ck["model"].items()}
sd.pop("freqs_cis", None)
model.load_state_dict(sd, strict=False)

(config.json → model_class, or ck["config"]["model"]["_target_"], names the model class of each run: Transformer, SPS or two-tower.)

Citation

@misc{mechanistic_sps_2026,
  title  = {Which State Should Prediction Read? A Mechanistic Analysis of State–Prediction Separation},
  author = {TODO},
  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).

Downloads last month: -; Downloads are not tracked for this model. How to track

ryankim17920
/

mechanistic-sps