babylm-paragru-20m

ParaGRU BabyLM-10M (Strict-Small)

19.7M causal LM with ParaGRU (mix='diag', Cho gates, eq. 3.1a / 3.3) under BabyLM 2026 Strict-Small (~10M words). Same 6×384 SwiGLU residual stack as the other ParaRNN BabyLM arms; fused Newton (K{=}3).

Report PPL 95.80 after 3 epochs (~16.3 min on one RTX 2080 Ti, ~38.5k tok/s, peak 3.77 GiB). Zero-shot: BLiMP 63.31 (GPT-2 Strict-Small 65.23) — highest BLiMP among the four matched arms so far.

Library: bugkira/pararnn-torch · Sisters: babylm-paralstm-20m, babylm-paranlru-19m, babylm-paraslstm-20m · Paper: ParaRNN, arXiv:2510.21450

---

Model details

Gates: update (z), reset (r), candidate (n) (Cho et al. 2014). Diagonal recurrent vectors (a_z, a_r, a_n).

---

Training

Data & schedule

• Corpus: BabyLM 2026 Strict-Small (~10M words).
• Packing: contiguous T=512 (19 275 train rows).
• Epochs: 3 · Steps: 1809 (603 / epoch) · Seed: 0 · per-epoch row shuffle.
• Batch: 16 × grad-accum 2 → 16 352 CE tokens / step.
• Optim: AdamW lr=6e-4, cosine over 1809 steps, warmup 50, β=(0.9, 0.95), wd=0.01, clip 1.0.

Hardware & speed

Learning curves

Report val PPL (256 packed sequences):

---

Evaluation (BabyLM 2026 Strict zero-shot)

Pipeline: babylm-org/babylm-eval · causal · temperature 1.0.

Reading: eye-tracking 0.65, self-paced reading 0.03.

Skipped in this card: GlobalPIQA, full EWoK (gated), SuperGLUE finetune, AoA.

---

Intended use

Matched BabyLM cell-zoo arm for ParaGRU fused Newton. Small English LM under the Strict-Small budget. Out of scope: chat, long context, multilingual tracks.

---

How to load

Needs pararnn-torch and transformers with trust_remote_code=True.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "bugkira-ai/babylm-paragru-20m"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
model.eval()

ids = tok("The cat sat on the mat.", return_tensors="pt").input_ids
with torch.no_grad():
    logits = model(input_ids=ids).logits
    gen = model.generate(ids, max_new_tokens=16, do_sample=False)
print(logits.shape)  # [1, T, 16000]
print(tok.decode(gen[0], skip_special_tokens=True))

Local export:

uv run --extra lm --with transformers python scripts/export_babylm_hf.py \
  --config configs/train/babylm_gru.yaml \
  --ckpt checkpoints/babylm/gru_shuf_ep3.pt \
  --out checkpoints/babylm/hf_gru_shuf_ep3 \
  --cell_type gru

---

Reproduction

uv run --extra lm --extra train python scripts/train_babylm.py \
  --config configs/train/babylm_gru.yaml \
  --cell_type gru --epochs 3 --gpu 2080

bash scripts/run_babylm_zeroshot.sh checkpoints/babylm/hf_gru_shuf_ep3 1

Config: configs/train/babylm_gru.yaml.
Metrics: results/babylm_gru_shuf_ep3.json, results/babylm_zeroshot_gru_shuf_ep3.json.

---

Citation

@inproceedings{danieli2026pararnn,
  title     = {{ParaRNN}: Unlocking Parallel Computation in Nonlinear RNNs
               through Symbolic Algebra},
  author    = {Federico Danieli and Miguel Sarabia and Aviv Navon and
               Amos Storkey and Aaron van den Oord},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2026},
  note      = {Oral. arXiv:2510.21450},
  url       = {https://arxiv.org/abs/2510.21450}
}
@misc{choshen2026babylm,
  title         = {BabyLM Turns 4 and Goes Multilingual},
  year          = {2026},
  eprint        = {2602.20092},
  archivePrefix = {arXiv}
}

---

License

MIT for these weights and this card. BabyLM eval data and baselines keep their upstream licenses.