bugkira-ai/babylm-paralstm-20m

ParaLSTM BabyLM-10M (Strict-Small)

19.7M causal LM with ParaLSTM (CIFG peephole, diagonal \(A,C\)) 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.16 after 3 epochs (~18 min on one RTX 2080 Ti, ~39.1k tok/s, peak 3.84 GiB). Zero-shot: BLiMP 62.86 (GPT-2 Strict-Small 65.23).

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

Model details

Cell: CIFG peephole LSTM (Greff et al. 2017; Danieli et al. 2025 eq. 3.1b / 3.3). State layout (c, h) with diagonal recurrent and peephole vectors.

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

[image]

Learning curves

[image]

Report val PPL (256 packed sequences):

Evaluation (BabyLM 2026 Strict zero-shot)

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

[image]

Reading: eye-tracking 0.80, self-paced reading 0.02.

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

Intended use

Matched BabyLM cell-zoo arm for ParaLSTM 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-paralstm-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_lstm.yaml \
  --ckpt checkpoints/babylm/lstm_shuf_ep3.pt \
  --out checkpoints/babylm/hf_lstm_shuf_ep3 \
  --cell_type lstm

Reproduction

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

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

Config: `configs/train/babylm_lstm.yaml`. Metrics: results/babylm_lstm_shuf_ep3.json, results/babylm_zeroshot_lstm_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.