statehead-10m

StateHead-10M

BabyLM 2026 Strict-Small submission. A 9.97M-parameter, attention-free causal language model whose entire sequence memory is 4.2 KiB of fixed-size recurrent state — constant at any context length, versus a ~36 MiB KV cache for a comparable transformer at 1,024 tokens.

Architecture

StateHead replaces attention with a bank of recurrent heads. Each head keeps a small state vector s updated by learned elementwise gates:

s_t = sigmoid(a_t) * s_{t-1} + sigmoid(b_t) * tanh(c_t)      # retain + write
y_t = sigmoid(o_t) * s_t                                     # output gate

A gated diagonal linear recurrence in the QRNN / RG-LRU family — scan-parallel in training, O(1) state in inference, no position embeddings (order enters only through the recurrence).

Training

10 epochs over the BabyLM 2026 Strict-Small corpus (100M word exposures, the permitted maximum; 161.2M tokens), batch 32,768 tokens/step, cosine schedule, hybrid Muon+AdamW with head-wise Muon on matrix parameters, bf16 with fp32 scan accumulation, single seed (1337). Total: 1.1 hours on one Apple M5 MacBook Air (~9.7×10¹⁵ training FLOPs). Trained in MLX; exported to this HuggingFace/PyTorch implementation with scan-level parity tests.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("haybales/statehead-10m")
model = AutoModelForCausalLM.from_pretrained("haybales/statehead-10m", trust_remote_code=True)

ids = tok("The child put the apple in the", return_tensors="pt")
print(tok.decode(model.generate(**ids, max_new_tokens=10)[0]))

Optional: set STATEHEAD_COMPILE=1 for a bit-exact torch.compiled scan (inference only, ~1.5–1.8× faster forward; falls back to eager automatically).