Context-gated memory slots fix the biggest blind spot in token embeddings

Every existing memory-embedding method for transformers shares one structural flaw: retrieval is a deterministic function of the token's surface form. The same row in the memory table gets activated whether "python" means a scripting language or a snake. MoME fixes this by replacing each token's single memory row with a small set of slots and using the model's hidden state at that position to gate which slots actually get read.

The mechanism is deliberately lightweight. A learned router projects the hidden state into slot logits, selects the top-k slots (typically 2 out of 5–12), and aggregates them with a sigmoid-normalized weighting. The result is injected into the attention value stream in parallel with the standard value projection—not sequentially—which keeps latency overhead small. A hidden-state-to-hidden-state injection variant nearly doubles that overhead compared to the value-stream path, which is why the value-stream design is the default.

Across three backbone families—nanochat-style, Llama/MobileLLM, and Qwen3—MoME beats Value Embedding, Bigram, and STEM on validation bits-per-byte and the CORE aggregate in most matched comparisons, while requiring only marginally more wall time than the no-memory base. The memory-size scaling result is particularly pointed: at every tested memory budget on the nanochat d12 backbone, MoME achieves lower validation bpb than a parameter-matched Bigram run, and its training curves show noticeably tighter run-to-run variance.

The routing analysis is the most interesting part. On the Word-in-Context dataset, filtered to 670 pairs where the target token is a single subword with sufficient preceding context, different-sense pairs show greater routing divergence than same-sense pairs at 117 of 144 layer-head sites. The effect is not uniform: some heads with large routing differences have nearly closed injection gates, meaning they route differently but contribute little to the output. The gate-weighted analysis identifies layer 7, head 6 as having both a strong routing effect and an open injection gate.

The authors are honest about the limits. All experiments are sub-billion-parameter scale, the 100B-token run uses a single seed, and the WiC analysis is descriptive—it shows that routing correlates with word sense but does not establish that this correlation causes better predictions. The kNN token-grouping option, which clusters semantically similar tokens into shared rows to reduce redundancy, is only swept in one ablation and is off by default elsewhere.

For practitioners building small-to-mid-scale language models, MoME is a concrete, low-overhead way to add context-sensitive parametric memory without restructuring the backbone.