a-kemeklis/constraining-pretraining-data

Constraining Language Modeling Pretraining Data

Code, data and model for the paper Constraining Language Modeling Pretraining Data (Andrew Kemeklis, 2026). Paper: arXiv link to be added · PDF: paper/constraining_language_modeling_pretraining_data.pdf

The paper restricts a pretraining corpus to the word forms of a fixed set of base words (lemmas). Every other word is replaced with an indexed placeholder (<UNK_n>, <NAME_n>) that preserves coreference, and documents with too many placeholders are dropped. Applied to TinyStories with the 300 most frequent lemmas and at most 8 placeholders per story, this keeps 10.4% of the stories (43.0M tokens). A GPT with 7.0M total parameters (6.3M non-embedding) trained on that data is compared with TinyStories-33M (68.5M total, 28.3M non-embedding parameters).

This repository holds everything in one place:

what	where
the 6.3M-parameter model from the paper	model.safetensors, config.json, tokenizer.json, tokenizer_config.json (repo root)
the filtered dataset (282,603 train / 1,412 validation stories)	data/stories/ (parquet), data/lexicon.json, data/stats.json
the tokenized training blocks used to train the model	data/blocks512/ (512-token blocks, uint16, with meta.json and the tokenizer)
dataset builder	code/build_topk_dataset.py
training, sampling and judge scoring	code/ts300.py, code/paired_stats.py, code/JUDGING.md, code/requirements.txt
judge evaluations behind Tables 3 and 4	eval/
training logs and the hyperparameter sweep	training/, code/sweep_runs.txt
paper source (LaTeX) and figure script	paper/

Results

180 story pairs from 20 in-lexicon openings, both models at the same temperature and seeds. Each packet of 10 pairs was graded by a fresh LLM judge instance (Claude Opus 5.5) with A/B order randomized. Difference is mine minus TinyStories-33M, with paired bootstrap 95% CIs.

condition	grammar errors, mine / ref (diff [CI])	consistency errors, mine / ref (diff [CI])	ending rate, mine / ref	preferred mine : ref
T=0.7, placeholders shown (Table 3)	0.39 / 0.39 (+0.01 [−0.13, +0.14])	1.37 / 2.08 (−0.71 [−0.96, −0.46])	0.97 / 0.84	105 : 72 (3 ties)
T=0.7, placeholders masked (Table 4)	0.74 / 0.36 (+0.39 [+0.23, +0.54])	1.88 / 2.07 (−0.19 [−0.46, +0.08])	0.96 / 0.85	83 : 93 (4 ties)
T=0.8, placeholders masked (Table 4)	0.91 / 0.52 (+0.39 [+0.24, +0.54])	2.04 / 2.26 (−0.22 [−0.48, +0.04])	0.98 / 0.87	81 : 95 (4 ties)

With placeholders shown, the judge can tell which story is mine, so that row's preference count is not blind. Part of the consistency advantage comes from my model's shorter stories (94 vs. 128 words); see the paper for per-word rates.

Using the model

The model is a small custom GPT (pre-RMSNorm, RoPE, GELU MLP, tied embeddings, no biases; 8 layers, d_model 256, 4 heads, vocabulary 2,700, context 512). Load it with load_model from code/ts300.py:

pip install -r code/requirements.txt
hf download a-kemeklis/constraining-pretraining-data --local-dir cpd --exclude "data/*"
python cpd/code/ts300.py sample --ckpt cpd --out samples --n-prompts 1 --samples 1 --temperature 0.7

sample continues the paper's story openings with this model and with roneneldan/TinyStories-33M and writes blind judge packets. Add --ban-placeholders to mask placeholder tokens at decoding (the "masked" condition).

Using the dataset

from datasets import load_dataset
ds = load_dataset("parquet", data_files={
    "train": "hf://a-kemeklis/constraining-pretraining-data/data/stories/train-*.parquet",
    "validation": "hf://a-kemeklis/constraining-pretraining-data/data/stories/validation-*.parquet",
})

Columns: text (the rewritten story), original (the source story), n_words, n_oov, n_unk, n_unk_distinct, n_names, n_dropped, unk_lemmas, status, id. The 300 lemmas, 900 allowed word forms and 50 allowed names are in data/lexicon.json. Placeholders make up 4.6% of training tokens.

Reproducing the paper

pip install -r code/requirements.txt

# 1. rebuild the filtered dataset and tokenizer from the source (optional; the outputs are already in data/)
python code/build_topk_dataset.py --k 300 --max-unk 8 --no-upload

# 2. train the 6.3M model (about 21 minutes on one H100; downloads data/blocks512 from this repo)
python code/ts300.py train --out runs/ts300_6m --n-layer 8 --d-model 256 --n-head 4 --epochs 80 --dropout 0.1 --lr 2e-3

# 3. sample 180 pairs against TinyStories-33M (Table 3: placeholders shown; add --ban-placeholders for Table 4)
python code/ts300.py sample --ckpt runs/ts300_6m/best --out runs/eval_T0.7 --samples 9 --temperature 0.7

# 4. judge each judge_packet_partN.md with a fresh LLM instance (see code/JUDGING.md), concatenate the replies, then
python code/ts300.py score --key runs/eval_T0.7/judge_key.json --judge judge_replies.txt
python code/paired_stats.py runs/eval_T0.7/judge_key.json judge_replies.txt

Steps 3 and 4 can be checked without rerunning anything: each folder in eval/ contains the exact samples.jsonl, judge_packet_part*.md, judge_key.json and judge_replies.txt behind the paper's tables.

eval/ folder	paper
L8_d256_e80_lr2e3_T0.7_n180_withph	Table 3 (placeholders shown, T=0.7)
L8_d256_e80_lr2e3_T0.7_n180	Table 4 (placeholders masked, T=0.7)
L8_d256_e80_lr2e3_T0.8_n180	Table 4 (placeholders masked, T=0.8)

Sampling seeds are fixed, so rerunning step 3 with this checkpoint should reproduce the stories in eval/ (GPU kernels and library versions can cause small differences). The code labels my model ts300-10m in its outputs whatever its size; that label refers to this 6.3M model.

Notes

• code/build_topk_dataset.py is the script that produced the data. It uploads to separate repos when run with a write token; use --no-upload to build locally. Its clitic handling turns can't into ca, so stories containing can't were dropped from the dataset (noted in the paper).
• training/ holds the training log (log.jsonl, full log), run_summary.json, and sweep_summaries.json with the configuration and best validation loss of every run in the sweep (configurations in code/sweep_runs.txt).
• The dataset is derived from karpathy/tinystories-gpt4-clean, a cleaned subset of TinyStories, which are licensed CDLA-Sharing-1.0; the filtered dataset in data/ is shared under the same license.

Citation

@misc{kemeklis2026constraining,
  title  = {Constraining Language Modeling Pretraining Data},
  author = {Kemeklis, Andrew},
  year   = {2026},
  note   = {arXiv identifier to be added}
}