bugkira-ai/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

Architecture	6× pre-LN · ParaGRU (mix=diag) · SwiGLU
Width	d_model=384, mlp_mult=4
Parameters	19 656 192
Vocab	16 000 BPE (ByteLevel), Strict-Small
Context	T=512 absolute positions
Dtype (train)	float32
Solver	Newton K=3, Picard P=0, max_recurrent_norm=0.5
Checkpoint tag	gru_shuf_ep3

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

GPU	1× NVIDIA GeForce RTX 2080 Ti (CC 7.5)
Wall clock	2026-09-08 14:29:17 → 14:45:34 (≈16.3 min)
Steady throughput	≈38 518 tok/s (step ≈425 ms)
Peak VRAM	3.77 GiB
Watchdog	0 residual hits · 0 skipped steps

Learning curves

Report val PPL (256 packed sequences):

Step	Epoch end	Report PPL
50	—	1015.07
603	1	131.01
1206	2	99.66
1809	3	95.80

---

Evaluation (BabyLM 2026 Strict zero-shot)

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

Task	ParaGRU	ParaLSTM	ParaNLRU	ParaSLSTM	GPT-2 Strict-Small
BLiMP	63.31	62.86	62.45	62.81	65.23
BLiMP Supplement	56.65	56.64	57.66	55.70	57.25
EWoK	49.45 (fast)	49.36 (fast)	50.18 (fast)	47.36 (fast)	50.63 (full)
Entity Tracking	18.45	18.02	18.15	17.36	19.10
COMPS	50.45	50.87	50.32	50.79	51.81

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.