ScaleOp — Pythia size-conversion checkpoints
Research artifacts for the paper "Wiring Beats Blending: Structure-Aware Compensation for
Transformer Downscaling" (arXiv:2608.02829, v3; v1–v2 carried the
title "What Transfers Between Transformer Sizes — and What Doesn't").
The paper asks whether a pretrained large model can be converted into a smaller sibling instead
of training the small size from scratch. These checkpoints are the endpoints of matched-budget
recovery races: every arm gets the same target shape, data order, schedule, and token budget,
and only the starting weights differ. The initializations compared:
| Init |
Construction |
random |
from scratch |
projection |
dense linear projection of the donor weights |
subclone |
structured selection (magnitude pruning of the donor) |
subclone_rs |
selection + variance-preserving rescale |
hybrid |
selection + least-squares compensation |
hybrid_rs |
selection + compensation + rescale (the paper's method) |
What's in this repo
All files are plain PyTorch state_dicts for GPTNeoXForCausalLM with the target size's
stock Pythia config. They are research checkpoints trained for 30M–1B tokens (0.01–0.3% of
Pythia's pretraining), useful for studying conversion initializations — not production or
instruction-tuned models.
Pair A: Pythia-1.4B → 410M
| File |
Budget |
Final WikiText-103 ppl |
pairA_1.4b-to-410m/hybrid_30M.pt |
30M |
114.1 |
pairA_1.4b-to-410m/subclone_30M.pt |
30M |
355.8 |
pairA_1.4b-to-410m/projection_30M.pt |
30M |
1,054.5 |
pairA_1.4b-to-410m/random_30M.pt |
30M |
1,519.0 |
pairA_1.4b-to-410m/hybrid_100M.pt |
100M |
62.3 |
pairA_1.4b-to-410m/subclone_100M.pt |
100M |
103.1 |
pairA_1.4b-to-410m/hybrid_rs_1B.pt |
1B |
40.0 |
pairA_1.4b-to-410m/subclone_rs_1B.pt |
1B |
40.0 |
pairA_1.4b-to-410m/random_1B.pt |
1B |
57.4 |
Reference points: the real Pythia-410M scores ≈15.6 on the same eval; transfer initialization
beats from-scratch by up to 18× at 30M tokens, and both rescale arms converge to parity at 1B.
Pair B: Pythia-410M → 160M (depth-dominated, 3 data-draw seeds)
| Files |
Budget |
Final ppl (mean over seeds) |
pairB_410m-to-160m/hybrid_30M_s{0,1,2}.pt |
30M |
113.2 ± 4.5 |
pairB_410m-to-160m/subclone_30M_s{0,1,2}.pt |
30M |
118.5 ± 3.3 |
pairB_410m-to-160m/random_30M_s{0,1,2}.pt |
30M |
1,505.2 ± 57.9 |
pairB_410m-to-160m/{hybrid,subclone}_100M_s0.pt |
100M |
77.1 / 74.4 |
Per-seed numbers and the full ablation ladder are in the paper's appendix.
Usage
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoTokenizer, GPTNeoXForCausalLM
path = hf_hub_download(
"YenugulaAIML/scaleop-pythia-conversions",
"pairA_1.4b-to-410m/hybrid_rs_1B.pt",
)
config = AutoConfig.from_pretrained("EleutherAI/pythia-410m") # target size's stock config
model = GPTNeoXForCausalLM(config)
model.load_state_dict(torch.load(path, map_location="cpu"))
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-410m")
out = model.generate(**tokenizer("The capital of France is", return_tensors="pt"), max_new_tokens=10)
print(tokenizer.decode(out[0]))
For pair B, use the EleutherAI/pythia-160m config and tokenizer instead.
Training data
Continued pre-training on a deterministic stream of
pile-uncopyrighted, with a frozen
10,000×128 evaluation corpus that is contamination-guarded against the training stream. The full
recipe, configs, and seeds are in the code repository.
Limitations
These are controlled-experiment artifacts, not usable assistants: budgets stop far short of full
pretraining, all models are base LMs with no instruction tuning or safety training, and they
inherit the biases of the Pile and the Pythia family. Single-seed cells are noted in the paper.
Links
Citation
@article{yenugula2026wiring,
title = {Wiring Beats Blending: Structure-Aware Compensation for Transformer Downscaling},
author = {Yenugula, Ravi Satya Durga Prasad},
journal = {arXiv preprint arXiv:2608.02829},
year = {2026}
}
v3 checkpoints (added 2026-09-25, paper v3 / AISTATS submission)
All files are state_dicts for GPTNeoXForCausalLM with the Pythia-410M config, 1.4B→410M pair, in pairA_1.4b-to-410m/. Seeds are data-draw seeds (skip_docs offset by seed); seed 0 is the run in the original tables.
| file pattern |
what it is |
paper table |
hybrid_rs_1B_s{1,2}.pt, subclone_rs_1B_s{1,2}.pt |
1B-token convergence re-runs (seeds 1–2; seed 0 = *_rs_1B.pt above); WikiText ppl 40.3±0.3 vs 40.3±0.5 over 3 seeds |
tab:conv1b |
subclone_rs_act_30M_s{0,1,2}.pt, hybrid_rs_act_30M_s{0,1,2}.pt |
information-matched control: V/O and MLP selection scored by activation energy √Σ_dd·‖col‖ instead of weight norm; hybrid_rs_act (80.2±2.3) is the best 30M endpoint |
tab:actsel, tab:seeded |
{subclone_rs,hybrid_rs,random}_30M_b16_s{0..2}.pt |
30M tokens, batch 16, no distillation (batch-matched control for the distillation arms); random is seed 0 only |
tab:distill |
{subclone_rs,hybrid_rs,random}_30M_distill_s{0..2}.pt |
same, with a frozen Pythia-1.4B teacher (KD, T=2, α=0.5); subclone_rs + distill is the prune-and-distill recipe |
tab:distill |
Sizes: 1.62 GB each (fp32). Code: experiments/m5_train.py (--init subclone_rs_act|hybrid_rs_act, --distill), lib/subclone.py::select_indices(moments=...) in the linked repository.
