mixlab reproduction of the BabyLM 2026 GPT-2 Strict-Small baseline

A faithful, from-scratch reproduction of the official BabyLM 2026 GPT-2 Strict-Small baseline (BabyLM-community/babylm-baseline-10m-gpt2), trained entirely on a single Apple M1 Max GPU with mixlab, an open Metal/MLX trainer, on the same 10M-word 2026 Strict-Small corpus the organizers use (a data-identical reproduction).

This model reproduces the baseline as a reproducibility and efficiency study. The companion GitHub repo has the full recipe, configs, scripts, and per-task eval reports: mixlab-babylm-gpt2.

Faithfulness: per-component reproduction

For the six zero-shot components, both the reference baseline and this reproduction were scored through one identical eval harness (babylm-org/babylm-eval @ 3bf5142, causal backend); the reference column there is our within-harness re-evaluation of the official baseline, and it reproduces the official published numbers. GLUE is the exception (see †).

Component Reference baseline This reproduction Ξ”
BLiMP 66.35 65.86 βˆ’0.49
BLiMP-supplement 57.07 58.04 +0.97
EWoK 49.23 49.30 +0.07
COMPS 51.72 52.05 +0.33
reading (eye+SPR) 6.50 7.15 +0.65
GLUE (macro) † 63.62 64.15 +0.53
entity tracking ‑ 13.90 25.99 +12.09

Within about one point on every stable component.

† GLUE: the reproduction was fine-tuned and scored here; the reference (63.62) is the official published baseline (GPT-2 model card per-task GLUE), not re-fine-tuned locally. ‑ entity tracking is a length-bias artifact (both models near chance); the gap is not a real capability difference and is excluded from any aggregate claim. AoA (age-of-acquisition) is forfeited (0): the board scores it from a per-checkpoint surprisal trajectory submitted with the predictions, which this reproduction's submission does not include. Separately, the AoA scorer has a confirmed, still-open upstream bug β€” a log2/bits random-chance baseline with a hardcoded ~300k vocabulary instead of the model's actual size (babylm-org/babylm-eval#2) β€” which doesn't zero out small-vocab models (the reference scores 35.68 through it) but makes the metric noise-dominated and extraction-sensitive (one fixed model spans ~12–36). For both reasons we report per-component and exclude AoA.

Architecture (read from the weights)

The exact architecture is in this repo's config.json (the exported HF model config); the full mixlab spec (architecture and training recipe) is in training_config.mixlab.json.

A vanilla GPT-2 Small, β‰ˆ98M params (16,384 vocab Γ— 768 dim, tied to the output head). The reference model card's "124M" is the 50k-vocab number; with the 16k BabyLM vocab the real model is 98,425,344 params.

  • 12 layers, hidden 768, 12 heads, context 1024, tied embeddings
  • Learned absolute position embeddings (GPT-2 WPE)
  • GELU (gelu_new, tanh approximation) feed-forward, intermediate 3072
  • Pre-LayerNorm (eps 1e-5) with affine weight+bias, plus a final LayerNorm
  • Standard causal self-attention with attention + projection biases
  • Dropout 0.1 (embedding, residual, attention), applied during training
  • GPT-2 weight initialization: normal(0, 0.02), residual c_proj scaled by 1/√(2Β·n_layer)

Exported as a native GPT2LMHeadModel; native↔HF logit parity verified (β‰ˆ3e-8).

Training

  • Trainer: mixlab v0.40.0 (-mode arch), Metal/MLX, single Apple M1 Max (32-core GPU, 64 GB unified memory), β‰ˆ11h wall-clock (varies across Apple chip generations), seed 42, fp32
  • Data: the 2026 Strict-Small corpus (β‰ˆ10M words β†’ 18M tokens), full corpus, tokenized whole-block then sliced into 512-token chunks; each chunk framed as <s> + chunk + </s> (the reference convention)
  • Optimizer: AdamW (betas 0.9/0.999, eps 1e-8), lr 5e-5, no weight decay, grad-clip 1.0
  • Schedule: 21,300 steps (~9.7 epochs over the ~18M-token corpus), batch 16 Γ— 514-token sequences (512 content + <s>/</s>), warmup 213 + cosine decay
  • Tokenizer: the reference 16,384-vocab tokenizer (<unk>=0, <s>=1, </s>=2, <pad>=3, <mask>=4)

Use

from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("mrothroc/mixlab-gpt2-strict-small-replica")
m   = AutoModelForCausalLM.from_pretrained("mrothroc/mixlab-gpt2-strict-small-replica")

Reproduce

The full recipe (data prep, train, export, parity, eval), configs, and scripts are in the companion GitHub repo. One command per step; β‰ˆ11h training on an M1 Max.

Implementation Notes

  • This reproduces the 2026 GPT-2 baseline on the 2026 Strict-Small corpus (data-identical).
  • Weight initialization was the remaining mismatch: arch, recipe, corpus, and forward (export) parity all matched the reference as far as verified, yet BLiMP was 5.7 points low until the GPT-2 init scheme (normal 0.02 + residual c_proj scaling) was matched. Validation loss gave no warning (across our runs it did not track BLiMP), so it does not reveal an init-driven generalization gap.
  • Trained in fp32 on Apple silicon; numerical differences from a CUDA run are not characterized.
  • Single training seed; the official seed is unknown, so a residual point or two on BLiMP is expected.
  • Zero-shot faithfulness within one pinned harness; GLUE compares our fine-tune to the official published reference score. Absolute scores can differ across harness versions.

Links

For the story behind it, read I Built an ML Architecture Lab in Go.

Downloads last month
107
Safetensors
Model size
0.1B params
Tensor type
F32
Β·
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Dataset used to train mrothroc/mixlab-gpt2-strict-small-replica