bugkira-ai/babylm-paraslstm-20m

ParaSLSTM BabyLM-10M (Strict-Small)

20.5M causal LM with channelwise Diag-sLSTM under BabyLM 2026 Strict-Small (~10M words). Documents the library’s fused Newton–Picard path on a public LM track: training dynamics, consumer-GPU throughput, and the official zero-shot suite.

Report PPL 102.17 after 3 epochs (29.5 min on one RTX 2080 Ti, **21.9k tok/s**, peak 4.0 GiB). Zero-shot: BLiMP 62.81 (GPT-2 Strict-Small 65.23).

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

---

Model details

Architecture	6× pre-LN blocks · ParaSLSTM (mix=diag, solver=auto / fused Newton) · SwiGLU MLP
Width	d_model=384, 8 heads (diag cell), mlp_mult=4
Parameters	20 545 536
Vocab	16 000 BPE (ByteLevel), trained on Strict-Small
Context	T=512 absolute positions
Dtype (train)	float32 (fused Newton at this width is fp32-stable on this stack)
Solver	Newton K=3, Picard P=3, picard_adapt=true, max_recurrent_norm=0.5
Checkpoint tag	diag_fused_shuf_ep3

---

Training

Data & schedule

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

Hardware & speed

GPU	1× NVIDIA GeForce RTX 2080 Ti (Turing, CC 7.5)
Wall clock	2026-09-07 14:36:53 → 15:06:25 (≈29.5 min, ≈0.49 GPU-h)
Steady throughput	≈21 925 tok/s (mean after warmup; step ≈746 ms incl. periodic short val)
Peak VRAM	4.00 GiB
Watchdog	0 residual hits · 1 Newton-divergent step skipped

Learning curves

Report val PPL (256 packed sequences):

Step	Epoch end	Report PPL
250	—	240.47
500	—	159.52
603	1	144.06
1206	2	107.21
1809	3	102.17

Shuffle vs contiguous packing (same YAML / step budget; no-shuffle was on RTX 3060):

Contiguous HF document order produced once-per-epoch train-loss waves; per-epoch shuffle removed that structure and improved final report PPL (102 vs 204).

---

Evaluation (BabyLM 2026 Strict zero-shot)

Official pipeline: babylm-org/babylm-eval · backend causal · temperature 1.0.
Baseline column: Baseline-GPT2-Strict-Small from the evaluation README.

Task	ParaSLSTM fused (ours)	GPT-2 Strict-Small
BLiMP	62.81	65.23
BLiMP Supplement	55.70	57.25
EWoK	47.36 (fast subset)	50.63 (full)
Entity Tracking	17.36	19.10
COMPS	50.79	51.81

Reading (human-likeness): eye-tracking score 0.45, self-paced reading 0.01 (babylm-eval reading report). Baseline Strict-Small GPT-2 quotes 5.63 Δ%R² on the same track’s human-likeness table — metric definitions differ by report field; compare within one pipeline revision.

Not included in this release card: GlobalPIQA download, full EWoK (gated Hub dataset), SuperGLUE finetune, AoA (needs intermediate word-budget checkpoints).

---

Intended use

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

---

Solver notes (for systems readers)

• Parallel Newton + scan targets O(log T) depth per iteration on sequence length; this run uses fused scan kernels with Picard warm-start (P=3, adaptive bump on high residual).
• Shuffle mixes domain blocks in the packed cache; train loss stays smoother across an epoch and final PPL improves. Keep picard_adapt on for shuffled LM runs; this checkpoint recorded one skipped divergent step out of 1809.
• First training step runs optional sequential↔Newton agreement smoke (verify_first_step); disable for pure timing.

---

How to load

Requires pararnn-torch (for the cell/solver) and transformers with trust_remote_code=True.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "bugkira-ai/babylm-paraslstm-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 (from the library repo):

uv run --extra lm --with transformers python scripts/export_babylm_hf.py \
  --ckpt checkpoints/babylm/diag_fused_shuf_ep3.pt \
  --out checkpoints/babylm/hf_diag_fused_shuf_ep3

Weights ship as model.safetensors (and pytorch_model.bin). Tokenizer: tokenizer.json without training EOS post-processor (eval-style encoding).

---

Reproduction

# Train (example)
uv run --extra lm --extra train python scripts/train_babylm.py \
  --config configs/train/babylm.yaml \
  --cell_type diag_fused --epochs 3 --gpu 2080

# Zero-shot (needs babylm-eval + evaluation_data/)
bash scripts/run_babylm_zeroshot.sh checkpoints/babylm/hf_diag_fused_shuf_ep3 1

Config: configs/train/babylm.yaml.
Train log: outputs/babylm_diag_fused_ep3_shuf_2080.log.
Metrics: results/babylm_diag_fused_shuf_ep3.json, results/babylm_zeroshot_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.