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

This repo holds the intermediate checkpoints and embedding ladders used by two analyses in the paper Which State Should Prediction Read? A Mechanistic Analysis of State–Prediction Separation: the training trajectory of gradient conflict between the state and prediction streams, and the embedding divergence between the prediction and state token tables. The final checkpoints of every model reported in the paper are in the main repo, ryankim17920/mechanistic-sps; this repo adds only what those two analyses read beyond the finals.

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

What is here and what it feeds

Run	Model	Files	Analysis (code)	Paper
`s_two_tower_w0_equal_tiedattn_20b`	Two-tower 12+12 (shared attn. only)	`ladder/ckpt_tokens_<N>.pt` at 1, 2, 3, 4, 6, 8, 12, 16B tokens (the 20B point is the final in `ryankim17920/mechanistic-sps`)	`scripts/analysis/a17_trajectory.py partB --arm tiedattn`	App. "Gradient Alignment": state–prediction gradient conflict along training (printed by the Fig. app-gradcos generator; the ladder itself is not drawn)
`s_two_tower_afsps_20b`	AF-SPS before the faithful split (not a paper arm)	`ladder/ckpt_tokens_<N>.pt` at the same 8 points, plus `ckpt_tokens_20000145408_final.pt` (on no other repo)	`scripts/analysis/a17_trajectory.py partB --arm afsps`	App. "Gradient Alignment", as above
`s_two_tower_seq6_20b`, `s_two_tower_seq6_20b_seed2`	Sequential 6+6 (seeds 1, 2)	`emb_ladder.pt`: the state table S (`transformer.wte.weight`), prediction table P (`predict_wte.weight`) and LM head H (`lm_head.weight`) of all 21 checkpoints (every 1B tokens, the pre-decay branch point at 18B, and the final)	`scripts/analysis/a21_emb_divergence.py`	App. "Depth, Input and Pre-Registered Controls", Prediction input: does the prediction table diverge from the state table (the replacement numbers use the final in `ryankim17920/mechanistic-sps`)

The .pt checkpoints use the same format as the main repo's finals: the training checkpoint dict without the optimizer state (model state_dict, model_args, config, iter_num, …). Here the paths in config read ${DATA_ROOT} and provenance omits the host name and job id. Each run folder has a config.json (model arguments, training config, and the size and sha256 of every file); the loss curves of all runs are in the code repo's scripts/analysis/results/ledger.jsonl.

File	Contents	Training tokens	Size	sha256
`s_two_tower_w0_equal_tiedattn_20b/ladder/ckpt_tokens_1000341504.pt`	ladder, weights only	1,000,341,504	1.09 GB	`da1422ab0f20…`
`s_two_tower_w0_equal_tiedattn_20b/ladder/ckpt_tokens_2000683008.pt`	ladder, weights only	2,000,683,008	1.09 GB	`3c823a20b45a…`
`s_two_tower_w0_equal_tiedattn_20b/ladder/ckpt_tokens_3001024512.pt`	ladder, weights only	3,001,024,512	1.09 GB	`08d6723a7227…`
`s_two_tower_w0_equal_tiedattn_20b/ladder/ckpt_tokens_4001366016.pt`	ladder, weights only	4,001,366,016	1.09 GB	`18534cb00e20…`
`s_two_tower_w0_equal_tiedattn_20b/ladder/ckpt_tokens_6002049024.pt`	ladder, weights only	6,002,049,024	1.09 GB	`7b13c252997c…`
`s_two_tower_w0_equal_tiedattn_20b/ladder/ckpt_tokens_8002732032.pt`	ladder, weights only	8,002,732,032	1.09 GB	`bc5f128ff0d7…`
`s_two_tower_w0_equal_tiedattn_20b/ladder/ckpt_tokens_12004098048.pt`	ladder, weights only	12,004,098,048	1.09 GB	`a092d15828e4…`
`s_two_tower_w0_equal_tiedattn_20b/ladder/ckpt_tokens_16005464064.pt`	ladder, weights only	16,005,464,064	1.09 GB	`38025303cd11…`
`s_two_tower_afsps_20b/ladder/ckpt_tokens_1000341504.pt`	ladder, weights only	1,000,341,504	0.52 GB	`bbadb3ad29a0…`
`s_two_tower_afsps_20b/ladder/ckpt_tokens_2000683008.pt`	ladder, weights only	2,000,683,008	0.52 GB	`a07389c9dc52…`
`s_two_tower_afsps_20b/ladder/ckpt_tokens_3001024512.pt`	ladder, weights only	3,001,024,512	0.52 GB	`a4f243356cbb…`
`s_two_tower_afsps_20b/ladder/ckpt_tokens_4001366016.pt`	ladder, weights only	4,001,366,016	0.52 GB	`1d8387905742…`
`s_two_tower_afsps_20b/ladder/ckpt_tokens_6002049024.pt`	ladder, weights only	6,002,049,024	0.52 GB	`8038ed169786…`
`s_two_tower_afsps_20b/ladder/ckpt_tokens_8002732032.pt`	ladder, weights only	8,002,732,032	0.52 GB	`bdc6f42c4cc5…`
`s_two_tower_afsps_20b/ladder/ckpt_tokens_12004098048.pt`	ladder, weights only	12,004,098,048	0.52 GB	`6a513482b0ab…`
`s_two_tower_afsps_20b/ladder/ckpt_tokens_16005464064.pt`	ladder, weights only	16,005,464,064	0.52 GB	`324fe3807943…`
`s_two_tower_afsps_20b/ckpt_tokens_20000145408_final.pt`	final, weights only	20,000,145,408	0.52 GB	`baf7ecb63bf1…`
`s_two_tower_seq6_20b/emb_ladder.pt`	S/P/H tables of 21 checkpoints		9.74 GB	`83e1a78e011e…`
`s_two_tower_seq6_20b_seed2/emb_ladder.pt`	S/P/H tables of 21 checkpoints		9.74 GB	`e84e06489d51…`

Using them with the code

git clone https://github.com/RyanKim17920/mechanistic-sps && cd mechanistic-sps
export DUALSPS_DATA_ROOT=/path/to/data     # tokenized FineWeb-Edu, see the code README
# a17 part B: the ladder from this repo; the tiedattn 20B final from the main repo
python scripts/dualsps/hf_export.py download --extra s_two_tower_w0_equal_tiedattn_20b s_two_tower_afsps_20b
python scripts/dualsps/hf_export.py download s_two_tower_w0_equal_tiedattn_20b
scripts/run/eval_shim.sh scripts/analysis/a17_trajectory.py partB --arm tiedattn
# a21: the embedding ladders (the --functional pass also needs the finals from the main repo)
python scripts/dualsps/hf_export.py download --extra s_two_tower_seq6_20b s_two_tower_seq6_20b_seed2
scripts/run/eval_shim.sh scripts/analysis/a21_emb_divergence.py --run s_two_tower_seq6_20b

Files land in <DUALSPS_OUT_ROOT>/out/<run>/, where the scripts look for them.

Without the code helpers:

import torch
from huggingface_hub import hf_hub_download

# a ladder checkpoint (same loading as the main repo's finals; see its model card)
path = hf_hub_download("ryankim17920/mechanistic-sps-extra", "s_two_tower_afsps_20b/ladder/ckpt_tokens_8002732032.pt")
ck = torch.load(path, map_location="cpu", weights_only=False)   # the dict holds configs, not just tensors
state_dict = {k.replace("_orig_mod.", ""): v for k, v in ck["model"].items()}

# the embedding ladder: tables keyed by checkpoint name
path = hf_hub_download("ryankim17920/mechanistic-sps-extra", "s_two_tower_seq6_20b/emb_ladder.pt")
emb = torch.load(path, map_location="cpu", mmap=True, weights_only=True)
print(list(emb["checkpoints"]))                      # ckpt_tokens_1000341504, ..., ckpt_tokens_20000145408_final
P = emb["tables"]["ckpt_tokens_20000145408_final"]["P"]   # (50304, 768) float32

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.

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

Inference Providers NEW

This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

ryankim17920
/

mechanistic-sps-extra

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

What is here and what it feeds

Using them with the code

Citation

License

Dataset used to train ryankim17920/mechanistic-sps-extra