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
1from transformers import AutoModelForCausalLM, AutoTokenizer
2tok = AutoTokenizer.from_pretrained("mrothroc/mixlab-gpt2-strict-small-replica")
3m = 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
