bugkira-ai/babylm-paranlru-19m

ParaNLRU BabyLM-10M (Strict-Small)

18.8M causal LM with ParaNLRU (Griffin / RG-LRU-style diagonal nonlinear slot) under BabyLM 2026 Strict-Small (~10M words). Same 6×384 SwiGLU residual stack as the other ParaRNN BabyLM arms; the cell is lighter (~half the recurrent params of diag sLSTM) and trains ~2.5× faster on the same GPU.

Report PPL 95.9 after 3 epochs (~12.3 min on one RTX 2080 Ti, ~54.7k tok/s, peak 3.76 GiB). Zero-shot: BLiMP 62.45 (GPT-2 Strict-Small 65.23; ParaSLSTM 62.81).

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

Model details

Cell recurrence (library ParaNLRU):

\[ at=\sigma(Wa xt),\quad ht = at\odot h{t-1} + (1-at)\odot\tanh(Wc xt + u\odot h{t-1}). \]

Training

Data & schedule

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

Hardware & speed

[image]

Learning curves

[image]

Report val PPL (256 packed sequences):

Matched ParaSLSTM (diag_fused_shuf_ep3) finished at 102.17 report PPL on the same recipe.

Evaluation (BabyLM 2026 Strict zero-shot)

Official pipeline: `babylm-org/babylm-eval` · backend causal · temperature 1.0.

[image]

Reading (human-likeness): eye-tracking score 0.56, self-paced reading 0.00.

Not included: GlobalPIQA, full EWoK (gated), SuperGLUE finetune, AoA.

Intended use

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

How to load

Requires `pararnn-torch` and transformers with trust_remote_code=True.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "bugkira-ai/babylm-paranlru-19m"
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_nlru.yaml \
  --ckpt checkpoints/babylm/nlru_shuf_ep3.pt \
  --out checkpoints/babylm/hf_nlru_shuf_ep3 \
  --cell_type nlru

Reproduction

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

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

Config: `configs/train/babylm_nlru.yaml`. Train log: outputs/babylm_nlru_ep3_2080.log. Metrics: results/babylm_nlru_shuf_ep3.json, results/babylm_zeroshot_nlru_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}
}

Griffin / RG-LRU lineage: De et al., Griffin (2024). This release is our ParaRNN-compatible diagonal nonlinear slot, trained under BabyLM Strict-Small.

License

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