StateHead
I’ve been working on a small language model called StateHead, which recently became my first published research paper through the BabyLM workshop at EMNLP 2026. Much of my interest in these models has been around how they remember what they’ve read, and this project gave me a way to experiment with that using a fairly small amount of hardware.
When a typical transformer processes text, it keeps representations of earlier tokens in what’s called a key-value cache. It can refer back to that information when predicting what comes next, but the cache grows as the text gets longer. I wanted to experiment with a model that keeps updating the same small set of numbers as it reads. Any information it needs from earlier in the text has to survive in those numbers.
I built StateHead around a set of recurrent heads, each with its own small state vector. As a new token comes in, learned gates control how much of the previous state to retain, what to add, and what to pass onward. The update itself is based on an existing recurrent design used in QRNNs. It doesn’t use attention, and word order comes from the sequence of state updates rather than positional embeddings.
The smaller version has about 10 million parameters and carries 2,160 state values across its six layers. At bf16 precision, those values occupy 4,320 bytes, or about 4.2 KiB, which is where the paper’s ‘four kilobytes’ comes from. The weights and other memory needed to run the model are additional. This small state is the information it carries forward about the text, and it stays the same size as the input gets longer.
I entered it in the Strict-Small track of the BabyLM Challenge, which limits the training corpus to 10 million words and allows up to 100 million word exposures from going through that data multiple times. It was a good setting for a project I could train and experiment with myself. I used MLX on my M5 MacBook Air, and the submitted model’s training loop took about 1.1 hours, excluding setup and evaluation. Being able to finish a run locally made it much easier to keep trying changes.
The strongest result was on BLiMP, a benchmark that checks whether a model gives a grammatical sentence a higher score than a very similar ungrammatical one. In the paper’s evaluation, StateHead scored 68.2, compared with 65.2 for the official GPT-2 baseline, which had about ten times as many parameters. It had picked up quite a bit of grammatical structure with a very small running memory.
The other tests gave me quite a bit more to work on. It struggled with questions about meaning and keeping track of objects as they moved between places, with several results near or below chance. It also did poorly at predicting human reading times. Being able to keep processing text with the same amount of state doesn’t mean it will reliably remember details from much earlier, and the grammar result only describes one part of its ability.
One interesting part of the project was looking more closely at an apparent improvement in those tracking tests. An earlier training recipe for a larger version of the model seemed to produce a large gain. After an evaluator update removed questions with ‘nothing’ among the answers, I went back through the saved predictions. The improvement was concentrated in those removed questions, all of which had ‘nothing’ as the correct answer. Performance on the remaining questions was almost unchanged. I included that comparison in the paper as well.
I’ve made the submitted model and tokenizer available on Hugging Face, with a PyTorch implementation so other people can try it for themselves. The paper goes through the architecture and experiments in more detail, including the results that didn’t turn out as well.
Next I’d like to use StateHead to try different ways of storing and retrieving specific information. I’m particularly interested in whether a more structured memory could help it keep details separate while still being small enough to run locally. Having a model that’s quick to train gives me a useful place to try those ideas and see what actually helps.