mixlab-gpt2-strict-small-replica/README.md

---
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 `<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

```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/).