Test-Time Scaling Makes Overtraining Compute-Optimal

Code and synthetic datasets for the paper:

Test-Time Scaling Makes Overtraining Compute-Optimal (COLM 2026) Nicholas Roberts, Sungjun Cho, Zhiqi Gao, Tzu-Heng Huang, Albert Wu, Gabriel Orlanski, Avi Trost, Kelly Buchanan, Aws Albarghouthi, Frederic Sala arXiv:2604.01411

We present Train-to-Test (T²) scaling laws that jointly optimize model size, training tokens, and the number of inference samples under fixed end-to-end budgets. When test-time compute via repeated sampling is accounted for, the compute-optimal model is substantially smaller and more overtrained than standard Chinchilla scaling prescribes.

Pretrained checkpoints: all checkpoints (Chinchilla-scaling suite extensions and our overtrained models) are on the Hugging Face Hub in this collection.

Setup

pip install -r requirements.txt
pip install -e ./lm-evaluation-harness

For pretraining and checkpoint conversion, additionally install open_lm pinned to commit b1ebb62:

pip install git+https://github.com/mlfoundations/open_lm.git@b1ebb62c11ae56740a76a14cd05adfc5f2abb92c

Synthetic datasets

The four synthetic evaluation tasks from the paper (~1,000 fill-in-the-blank / short-completion questions each, generated with GPT-5 and Claude Opus 4.6) ship inside the vendored evaluation harness:

Task	Dataset
Simple knowledge recall (simple_knowledge_nll)	lm-evaluation-harness/lm_eval/tasks/custom_dataset_knowledge_nll/
Multi-step arithmetic reasoning (simple_reasoning_nll)	lm-evaluation-harness/lm_eval/tasks/custom_dataset_reasoning_nll/
Commonsense causal reasoning (commonsense_causal_nll)	lm-evaluation-harness/lm_eval/tasks/custom_dataset_commonsense_causal_nll/
Spatial reasoning (spatial_reasoning_nll)	lm-evaluation-harness/lm_eval/tasks/custom_dataset_spatial_reasoning_nll/

Each dataset is a JSON list of {"id", "question", "answer"}. The four real tasks used in the paper are lambada_openai_nll, arc_easy_nll, sciq_nll, and openbookqa_nll. All eight tasks report per-example NLL over the reference answer, and pass@k / negative-log-pass@k up to k=1024 (pass_at_k_max in the task configs).

scripts/generate_dataset.py contains the seed-and-expand generation loop (requires scripts/api_keys.json; see scripts/api_keys.json.example).

Evaluation

Run from the repository root (the synthetic task configs use repo-relative dataset paths):

python lm_harness_eval.py --model hf \
    --model_args pretrained=nick11roberts/SL-discrep-chinchilla-rw-params5M_maxstep760-flop_1_25e16_step_767 \
    --tasks simple_knowledge_nll \
    --batch_size auto:16 \
    --output_path eval_results \
    --log_samples

To evaluate the full checkpoint suite (85 Chinchilla-scaling + 22 overtrained checkpoints) on all eight tasks:

OUTPUT_BASE=./eval_results bash scripts/eval_chinchilla_overtrained_all.sh

Checkpoint manifests live in scripts/core_checkpoints.sh, scripts/core_checkpoints_overtrained.sh, and scripts/checkpoint_grid.csv.

Fitting T² scaling laws

After evaluation, macro-average across tasks and fit both T² approaches (Approach 1: additive NLL model; Approach 2: Beta-regression pass@k accuracy model):

OUTPUT_BASE=./eval_results bash run_both_fits.sh

This runs scripts/macro_average_tasks.py followed by plot_passk_isoflop_additive_fit.py and plot_passk_isoflop_chinchilla_fit.py for each task (and the macro average), with and without the overtrained checkpoints. Figures and fitted parameters are written under plots/.

Pretraining overtrained checkpoints

We pretrain with open_lm on tokenized RefinedWeb (webdataset shards of 2048-token sequences), using the fixed hyperparameters of Porian et al. (2024) — see Appendix C of the paper.

# Single run: model config + max steps
DATA_DIR=/path/to/tokenized/2048-v1 bash scripts/train_isoflop.sh params_7M 8941

# Full iso-FLOP band sweeps (see also train_overtrained.sh, train_sanity_check.sh)
DATA_DIR=/path/to/tokenized/2048-v1 bash scripts/train_2.56e19_sweep.sh

Model architecture configs for all 16 sizes (5M–901M parameters) are in scripts/model_configs/. Convert an open_lm checkpoint to Hugging Face format with:

python scripts/convert_to_hf.py \
    --checkpoint ./training_outputs/params_7M_maxstep8941/checkpoints/epoch_1.pt \
    --model-config scripts/model_configs/params_7M.json \
    --out-dir ./hf_checkpoints/params_7M

Post-training (SFT)

Fine-tune checkpoints on ARC-Easy, SciQ, or OpenBookQA (standard fine-tuning in training/sft.py, completion-only SFT in training/sft_completion.py):

python training/sft.py --dataset arc_easy \
    --model nick11roberts/SL-discrep-chinchilla-rw-params220M_maxstep36840-flop_2_56e19_step_36844

# Full pipelines
bash training/allckpts_run.sh arc_easy 6   # fine-tune the whole checkpoint suite
bash training/isoflop_run.sh               # iso-FLOP SFT across model sizes

Evaluate fine-tuned models with training/allckpts_eval_single.py or scripts/eval_sft_completion_only.sh, and plot with training/isoflop_plot.py / training/allckpts_plot.py.

Citation

@inproceedings{roberts2026testtime,
    title={Test-Time Scaling Makes Overtraining Compute-Optimal},
    author={Roberts, Nicholas and Cho, Sungjun and Gao, Zhiqi and Huang, Tzu-Heng and Wu, Albert and Orlanski, Gabriel and Trost, Avi and Buchanan, Kelly and Albarghouthi, Aws and Sala, Frederic},
    booktitle={Third Conference on Language Modeling},
    year={2026},
    url={https://openreview.net/forum?id=0v6BUeQe67}
}

License

MIT (see LICENSE). The vendored lm-evaluation-harness fork retains its original MIT license.