--- license: mit language: - en library_name: transformers tags: - babylm - babylm-2026 - gpt2 - mixlab - reproduction - causal-language-model datasets: - BabyLM-community/BabyLM-2026-Strict-Small pipeline_tag: text-generation --- # 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`](https://huggingface.co/BabyLM-community/babylm-baseline-10m-gpt2)), trained entirely on a **single Apple M1 Max GPU** with [**mixlab**](https://github.com/mrothroc/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](https://github.com/mrothroc/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`](https://github.com/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](https://github.com/babylm-org/babylm-eval/issues/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`](https://huggingface.co/mrothroc/mixlab-gpt2-strict-small-replica/blob/main/config.json) (the exported HF model config); the full mixlab spec (architecture **and** training recipe) is in [`training_config.mixlab.json`](https://huggingface.co/mrothroc/mixlab-gpt2-strict-small-replica/blob/main/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 `` + chunk + `` (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 + ``/``), warmup 213 + cosine decay - **Tokenizer:** the reference 16,384-vocab tokenizer (``=0, ``=1, ``=2, ``=3, ``=4) ## Use ```python 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](https://github.com/mrothroc/mixlab-babylm-gpt2). 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](https://github.com/mrothroc/mixlab) - Reference: [`BabyLM-community/babylm-baseline-10m-gpt2`](https://huggingface.co/BabyLM-community/babylm-baseline-10m-gpt2) - Companion recipe: [mixlab-babylm-gpt2](https://github.com/mrothroc/mixlab-babylm-gpt2) For the story behind it, read [I Built an ML Architecture Lab in Go](https://michael.roth.rocks/blog/mixlab/).