MoME: Context-Aware Mixture-of-Memory Embeddings Outperform Value Embedding Baselines at Iso-FLOPs

vector-institute · hf · 2026-09-21

Vector Institute introduces MoME (Mixture of Memory Embeddings), a context-aware sparse memory mechanism. Existing memory-embedding methods retrieve via a deterministic function of surface form, collapsing contextual senses of the same token into one fixed entry. MoME replaces each token's single memory row with a mixture of M slots and uses a learned gate over the hidden state to select slots at each position.

In controlled pretraining across nanochat, Llama-3/MobileLLM, and Qwen3 backbones, MoME beats Value Embedding, Bigram, and STEM baselines in iso-parameter and iso-FLOP settings, shows a promising memory-size scaling trend at sub-billion scale, and stays efficient at training and inference. Qualitative routing analyses show the learned mixture dispatches polysemous tokens to distinct slots under different senses.

Original post →

More from Research

Research channel →