How heytamjid/babylm25 works
1. Repository purpose
- This is a BabyLM reproduction repo centered on pretraining small language models and evaluating them on BLiMP-style minimal pairs.
- The code supports both standard masked language modeling (MLM) and ELECTRA-style replaced-token detection (RTD), plus an optional hybrid path that adds an MLM head on top of an ELECTRA discriminator.
2. Main execution path
- The primary entrypoint is train_mask.py, which orchestrates training, validation, checkpointing, and optional logging.
- train_mask_basic.py appears to be a simpler or older training variant, likely for baseline MLM runs without the more advanced masking/ELECTRA options.
3. Model architecture layer
- The default backbone is a DeBERTa-v3-style masked LM setup unless overridden.
- electra_model.py contains the custom architecture used for the ELECTRA-derived MLM head (ElectraHeadForMaskedLM per the README), combining a pretrained discriminator backbone with a newly trained prediction head.
- The repository also appears to rely on Hugging Face model loading conventions, including AutoModel / AutoModelForMaskedLM and trust_remote_code for custom classes.
4. Training subsystem
- Training consumes:
- a training text corpus,
- a validation corpus,
- a tokenizer,
- and a configurable model/backbone size.
- The training loop supports common LM knobs: batch size, gradient accumulation, learning rate, weight decay, dropout, seed, CPU data workers, checkpoint cadence, and evaluation cadence.
- It also supports dynamic masking, with scheduled mask updates and optional mask decay over training.
5. ELECTRA branch
- Passing --custom electra switches the objective from MLM to RTD.
- In this mode the model is split into:
- a generator that proposes replacement tokens,
- and a discriminator that learns to detect which tokens were replaced.
- The README makes clear that the discriminator checkpoint is the downstream artifact saved for later fine-tuning or evaluation, and that it is a plain DebertaV2Model at save time.
6. ELECTRA-to-MLM recovery path
- Passing --custom electra_head builds a second-stage MLM model on top of a previously trained ELECTRA discriminator.
- This stage is additive: it reads the discriminator checkpoint, optionally unfreezes the top N transformer layers, and trains a new MLM head with configurable depth.
- This is the repo’s bridge from an RTD-only checkpoint back to a model that can be scored with standard pseudo-log-likelihood methods.
7. Masking and sequence-length control
- The training pipeline supports several masking behaviors:
- standard BERT-style MLM masking probabilities,
- forced full masking for ELECTRA generator inputs,
- soft masking,
- and a mask decay schedule.
- Sequence length can vary by epoch via the --max_seq_len schedule, which is important for BabyLM-style curriculum or compute budgeting.
8. Data and preprocessing boundary
- The repository separates raw/processed corpus files from the training code:
- data/bb24.train
- data/bb25_small.dev
- preprocessing.py is the boundary for turning raw text into model-ready inputs or preparing corpus formats, while tokenizers/bb24.model is the trained tokenizer artifact consumed by the training and scoring scripts.
- This makes the data pipeline largely external to the core training loop: the scripts assume prebuilt text corpora and a SentencePiece-style tokenizer model already exist.
9. Evaluation subsystem
- score_blimp.py is the evaluation entrypoint for BLiMP-style minimal-pair scoring.
- It supports multiple scoring modes:
- MLM pseudo-log-likelihood for standard masked LMs,
- RTD-realness scoring for ELECTRA discriminators,
- and the MLM-head-on-discriminator variant.
- The script consumes a model checkpoint plus a minimal-pair dataset in JSONL-like form, and compares “good” vs “bad” sentences using configurable field names.
10. Model registration and reloadability
- register_model.py exists to make custom architectures reloadable through Hugging Face’s model registry / remote-code mechanism.
- This is important because the repo includes nonstandard model variants beyond vanilla Transformers, and the architecture must survive checkpoint save/load cycles for downstream use.
11. Artifacts and checkpoints
- Training outputs are saved under user-specified output directories, typically with checkpoint-N/ subdirectories.
- The repo distinguishes between:
- a model checkpoint intended for downstream fine-tuning,
- and a separate checkpoint or head-training stage used only for zero-shot scoring conversion.
- The trained model collection is also published externally on Hugging Face, indicating that checkpoints are a first-class output of the pipeline.
12. External dependencies and tooling
- The project is Python-based and depends on the Hugging Face ecosystem for model loading and checkpoint interoperability.
- It also optionally integrates with Weights & Biases for experiment tracking.
- The README mentions LAMB vs AdamW, FLOPs logging, and checkpoint evaluation at BabyLM-required token-count milestones, which suggests the training system is designed for reproducible benchmark submission rather than general-purpose serving.
13. Overall architecture in one line
- The repo forms a pipeline of text corpus + tokenizer → masked/ELECTRA pretraining → checkpoint export → BLiMP-style scoring, with a custom model-registration layer enabling nonstandard architectures to be saved and reloaded cleanly.
Main entry points: train_mask.py Train mask · score_blimp.py Score BLiMP · electra_model.py ELECTRA head · bb24.model Tokenizer
Components 15
Data & preprocessing
Corpus and tokenizer artifacts plus preprocessing boundary.
- Train corpus · Primary pretraining text input.
data/bb24.train - Dev corpus · Development text used for validation during training.
data/bb25_small.dev - Tokenizer · Shared tokenizer artifact consumed by training and scoring.
tokenizers/bb24.model - Preprocess · Boundary for corpus formatting and input preparation.
preprocessing.py
Pretraining
Masked LM and ELECTRA-style training paths.
- Train mask · Primary training entrypoint with masking, checkpoints, and logging.
train_mask.py - Basic train · Simpler MLM training variant.
train_mask_basic.py - ELECTRA model · Custom ELECTRA-derived MLM head and discriminator-backed architecture.
electra_model.py - MLM path · Standard masked language modeling training and scoring route.
- RTD path · Replaced-token detection branch with generator and discriminator.
- ELECTRA head · Adds a new MLM head on top of a trained ELECTRA discriminator.
electra_model.py - Checkpoints · Saved model outputs for downstream use and scoring.
- W&B · Optional experiment tracking integration.
Evaluation
BLiMP-style scoring and model reload support.
- Register model · Makes custom architectures reloadable through Hugging Face conventions.
register_model.py - Score BLiMP · Evaluates checkpoints on BLiMP-style minimal pairs with multiple scoring modes.
score_blimp.py - HF loading · AutoModel and AutoModelForMaskedLM loading plus remote-code support.
Connections 20
- Preprocess → Train corpus: prepares · Turns raw or external text into training-ready corpus files.
- Preprocess → Dev corpus: prepares · Turns raw or external text into validation corpus files.
- Tokenizer → Train mask: consumed by · Tokenizer is required for training inputs.
- Tokenizer → Score BLiMP: consumed by · Tokenizer is used during scoring and tokenization.
- Train corpus → Train mask: input · Primary training corpus.
- Dev corpus → Train mask: validate · Validation corpus for periodic evaluation.
- Train mask → MLM path: orchestrates · Main trainer drives standard MLM runs.
- Train mask → RTD path: orchestrates · Main trainer also drives ELECTRA-style RTD runs.
- Basic train → MLM path: baseline · Simpler MLM-only training route.
- RTD path → Checkpoints: saves · Exports discriminator checkpoints for downstream use.
- MLM path → Checkpoints: saves · Exports masked LM checkpoints.
- Checkpoints → Score BLiMP: evaluated by · Saved checkpoints are scored on minimal-pair tasks.
- Register model → HF loading: enables · Registers custom classes for reloadable HF checkpoints.
- ELECTRA model → HF loading: loads via · Custom ELECTRA-derived models rely on HF remote-code loading.
- Train mask → W&B: optional · Training can emit experiment logs to W&B.
- ELECTRA model → ELECTRA head: extends · Hybrid stage adds an MLM head on top of discriminator weights.
- ELECTRA head → Checkpoints: exports · Produces a scored model artifact for pseudo-log-likelihood evaluation.
- Score BLiMP → MLM path: scores · Uses MLM pseudo-log-likelihood for masked language models.
- Score BLiMP → RTD path: scores · Uses realness scoring for ELECTRA discriminators.
- Score BLiMP → ELECTRA head: scores · Supports MLM-head-on-discriminator evaluation mode.