mixlab-gptbert-masked-focus-replica

mixlab reproduction of the BabyLM 2025 GPT-BERT masked-focus baseline

A faithful, from-scratch reproduction of the BabyLM 2025 GPT-BERT masked-focus Strict-Small baseline (BabyLM-community/babylm-baseline-10m-gpt-bert-masked-focus), trained entirely on a single Apple-silicon GPU with mixlab, an open Metal/MLX trainer.

This is a faithfulness demonstration for the closed 2025 challenge, not a competition entry. GPT-BERT masked-focus is not a vanilla transformer: it is a masked+causal hybrid with disentangled relative attention, a gated attention output, dense layer aggregation, and a BERT-style MLM head. Reproducing it end to end on commodity hardware is a real test of the trainer.

Faithfulness: per-component reproduction

For the six zero-shot components, both the reference baseline and this reproduction were scored through one identical eval harness. We report the per-component comparison (no aggregate as we did not compete):

† The six zero-shot rows are both re-scored in our harness. For GLUE, the reproduction was fine-tuned and scored here, but the reference value (66.20) is the official published baseline (baselines paper, Table 2) — we did not re-fine-tune the baseline — so the GLUE row alone is not a within-harness re-eval of the reference.

The reproduction stays within ≈2 points on every reported component. The reference column is our within-harness re-evaluation of the official 2025 baseline model; it tracks the official published numbers (baselines paper, Table 2): essentially exact on BLiMP (official 70.4), supplement (63.7), entity (40.0), COMPS (53.5), reading (≈6.4), GLUE (66.2), with EWoK the one outlier at 51.6 vs the official 50.0 (harness variance). AoA is forfeited ≈0 here: the board scores it from a per-checkpoint surprisal trajectory submitted with the predictions, not the final model, which these artifacts don't include. The AoA scorer also 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 makes the metric noise-dominated (one fixed model spans tens of points), so we report per-component, not a macro.

Architecture (read from the weights)

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

≈33M tied parameters (16,384 vocab × 384 dim, tied to the output head). The exported model.safetensors additionally stores a redundant untied lm_head copy.

• 12 layers, hidden 384, 6 heads, seq up to 512, tied embeddings
• DeBERTa-style disentangled relative attention (shared rel-embedding, reuse content QK projections, windowed)
• GeGLU feed-forward (intermediate 1280) with internal LayerNorm
• Affine-free LayerNorm (eps 1e-7), pre- and post-attention (post before the out-projection)
• Attention value-gate + attention biases
• Dense layer aggregation (DWA) — a learned weighted sum over all prior sublayer outputs
• BERT-style MLM head
• Per-example hybrid CLM+MNTP objective (masked-focus = 6.25% causal), mask schedule 0.30→0.15

Exports both MixlabForCausalLM and MixlabForMaskedLM (trust_remote_code=True); native↔HF logit parity verified (≈7e-8).

Training

• Trainer: mixlab (-mode arch), Metal/MLX, single Apple M1 Max GPU, ≈5.5h wall-clock, seed 42
• Data: the BabyLM 2026 detoxified Strict-Small corpus (≈10M words → 15.9M tokens; the current dataset, debiased via the organizers' decontamination pipeline), full corpus, with <s> segment markers inserted in the packed token stream (reference convention)
• Optimizer: LAMB, lr 0.007, weight decay 0.1, z-loss 1e-4, warmup 1.6% + cosine + cooldown
• Schedule: 9,600 steps × 16,384 tokens (≈10 epochs; the reference uses 9,914), sequence curriculum 128→256→512
• Tokenizer: the reference 16,384-vocab tokenizer (<s>=1, </s>=2, <pad>=3, <mask>=4, <unk>=0)

The full reproduction recipe, configs, and scripts are in the companion GitHub repo (see Reproduce). The training config used for this run is linked under Architecture above.

Use

from transformers import AutoModelForMaskedLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("mrothroc/mixlab-gptbert-masked-focus-replica", trust_remote_code=True)
m   = AutoModelForMaskedLM.from_pretrained("mrothroc/mixlab-gptbert-masked-focus-replica", trust_remote_code=True)

Reproduce

The full recipe (data prep with <s> markers, train, export, parity, eval), configs, and scripts are in the companion GitHub repo. One command per step; ≈5.5h training on an M1 Max.

Implementation Notes

• This reproduces the 2025 baseline. The official BabyLM 2026 baselines are GPT-2, a different architecture.
• Corpus: we trained on the 2026 detoxified Strict-Small corpus, not the original 2025 corpus the reference model used (detoxification was introduced in the 2026 round). This reproduces the 2025 architecture and recipe on the current, debiased data rather than a data-identical rerun. The six sources are unchanged and the filter removes only a small fraction of sentences, so the effect on these (non-toxicity) benchmarks should be minimal; we kept the detoxified corpus as the better, current dataset.
• Per-component faithfulness within one identical harness; absolute scores can differ across eval harness versions (BLiMP scoring in particular), so we report per-component, not an aggregate.
• Single training seed; reading and entity tracking are noisy at this scale.
• Reproduced on Apple silicon; numerical differences from a CUDA run are not characterized.

Links

• Trainer: mixlab
• Reference: BabyLM-community/babylm-baseline-10m-gpt-bert-masked-focus
• BabyLM 2025 baselines paper: arXiv:2510.20475

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