hematteo/readout-recipe-control

Readout recipe control models (31M, five matched arms)

Matched 31M Pythia-style (GPT-NeoX architecture) pretraining runs used as the recipe control in Learning to Read Out: Unembedding Dynamics in Language Model Pretraining (Section 4.5 and Appendix H). All arms train on the first 10B tokens of the Pythia preshuffled Pile (EleutherAI/pile-standard-pythia-preshuffled, GPT-NeoX-20B tokenizer) read in Pythia's released order, with the same parameter seed (0), fp16 precision, global batch 1024 sequences (2,097,152 tokens/step), weight decay 0.1, peak learning rate 1e-3, and a 10B-token budget (4,769 steps). Each arm changes exactly one recipe factor relative to baseline:

Arm	Output readout (W_U) LR multiplier	LR warmup steps	Reported in the paper
baseline/	1.0	1430	yes
wu_lr_0p25/	0.25	1430	yes
wu_lr_4x/	4.0	1430	yes
warmup_short/	1.0	715	yes
warmup_long/	1.0	5720	no (extra arm, same protocol)

Each arm ships config.json (full training config), metrics.csv (train and validation loss, learning rates, W_U row norm statistics, and final LayerNorm gain norm per logged step), a COMPLETE sentinel, and ckpts/step<N>/{model_fp16.pt,metrics.json} at the 16 checkpoint steps the paper analyzes:

0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, ~2480, 4096, 4769

The step near 2480 is where the first 11-hour job slot ended; the trainer checkpoints on slot stop, so that one step differs slightly between arms (2464 to 2497). The others are the dense early steps, the power-of-two token milestones, and the budget end.

The trainer, tokenizer pipeline, slice builders, and the W_U geometry analysis are in the code release under reproduce/recipe_control/ (https://github.com/hematteo/learning-to-read-out; docs/REPRODUCE.md gives the commands and flags per arm). Download into the layout the analysis script expects with hf download hematteo/readout-recipe-control --local-dir "$UM_SSD_ROOT/runs". These are research artifacts for analyzing readout geometry and timing across training, not general purpose language models.

Paper (arXiv:2610.21025) · Project page · Feature atlas · Code · Collection

Revision history: the June 2026 upload held a pilot pull of four arms truncated at step 512 (about 1B tokens) with a long_warmup arm at warmup 2860; this revision replaces it with the complete production runs the paper reports.

Quick start

Each checkpoint is an fp16 GPTNeoXForCausalLM state dict. To read the readout of one arm at one step:

import torch
from huggingface_hub import hf_hub_download

path = hf_hub_download("hematteo/readout-recipe-control", "wu_lr_4x/ckpts/step512/model_fp16.pt")
state = torch.load(path, map_location="cpu", weights_only=True)  # saved on GPU; map to CPU
W_U = state["embed_out.weight"]  # (50304, 256) readout at step 512

To rebuild the full model, construct it with build_model in reproduce/recipe_control/train_control.py of the code release and call load_state_dict(state).

Citation

@inproceedings{he2026learningtoreadout,
  title     = {Learning to Read Out: Unembedding Dynamics in Language Model Pretraining},
  author    = {He, Matteo and Shen, William F. and Iacob, Alex and Jovanovic, Andrej
               and Qiu, Xinchi and Lane, Nicholas D.},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026},
  eprint    = {2610.21025},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
}

MIT. Trained on the first 10B tokens of EleutherAI/pile-standard-pythia-preshuffled; see the Pile's data statement for upstream text provenance.