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 †).

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

• Trainer: mixlab
• Reference: BabyLM-community/babylm-baseline-10m-gpt2
• Companion recipe: mixlab-babylm-gpt2

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