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

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

Learning curves

[image]

[image]

[image]

Report val PPL (256 packed sequences):

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

[image]

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.

[image]

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.