MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup

Hugging Face Daily Papers Papers
mixture-of-memory context-aware sparse-embeddings llm-scaling ai-research

Summary

Introduces MoME, a context-aware memory mechanism for LLMs that uses a mixture of slots to handle token polysemy, improving over baselines in pretraining experiments.

Scaling large language models efficiently has motivated sparse capacity mechanisms such as Mixture-of-Experts and, more recently, conditional memory: token-indexed embedding tables that augment the backbone with cheap parametric lookups. Existing memory-embedding methods retrieve via a deterministic function of the surface form, which collapses different contextual senses of the same token (e.g., python the language vs. the animal) into a single fixed entry. We introduce Mixture of Memory Embeddings (MoME), a context-aware memory mechanism that replaces each token's single memory row with a mixture of M slots and uses a learned gate over the hidden state to choose which slots to read at each position. In controlled pretraining experiments across nanochat, Llama-3/MobileLLM, and Qwen3 backbones, MoME improves over Value Embedding, Bigram, and STEM baselines in iso-parameter and iso-training-FLOP settings, shows a more promising memory-size scaling trend at sub-billion scale, and remains efficient in training and inference. Qualitative routing analyses on polysemous tokens further suggest that the learned mixture exhibits a degree of semantic interpretability, dispatching the same surface token to distinct memory slots under different senses.
Original Article
View Cached Full Text

Cached at: 09/21/26, 07:19 AM

Paper page - MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup

Source: https://huggingface.co/papers/2609.15126

Abstract

ScalinglargelanguagemodelsefficientlyhasmotivatedsparsecapacitymechanismssuchasMixture-of-Expertsand,morerecently,conditionalmemory:token-indexedembeddingtablesthataugmentthebackbonewithcheapparametriclookups.Existingmemory-embeddingmethodsretrieveviaadeterministicfunctionofthesurfaceform,whichcollapsesdifferentcontextualsensesofthesametoken(e.g.,pythonthelanguagevs.theanimal)intoasinglefixedentry.WeintroduceMixtureofMemoryEmbeddings(MoME),acontext-awarememorymechanismthatreplaceseachtoken’ssinglememoryrowwithamixtureofMslotsandusesalearnedgateoverthehiddenstatetochoosewhichslotstoreadateachposition.Incontrolledpretrainingexperimentsacrossnanochat,Llama-3/MobileLLM,andQwen3backbones,MoMEimprovesoverValueEmbedding,Bigram,andSTEMbaselinesiniso-parameterandiso-training-FLOPsettings,showsamorepromisingmemory-sizescalingtrendatsub-billionscale,andremainsefficientintrainingandinference.Qualitativeroutinganalysesonpolysemoustokensfurthersuggestthatthelearnedmixtureexhibitsadegreeofsemanticinterpretability,dispatchingthesamesurfacetokentodistinctmemoryslotsunderdifferentsenses.

View arXiv pageView PDFGitHub6Add to collection

Get this paper in your agent:

hf papers read 2609\.15126

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper0

No model linking this paper

Cite arxiv.org/abs/2609.15126 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2609.15126 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2609.15126 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

Understanding Is Done Early: A Depth Division of Labor in Large Language Models and Its Use for Unbounded-Context Memory

arXiv cs.CL

This paper introduces CoMem, a method that exploits the depth-wise division of labor in LLMs to cache intermediate residual tensors and recompute only upper layers for retrieval, enabling bounded read compute and memory independent of stored-context length. Evaluated on Qwen3-8B, CoMem achieves strong long-context performance with significant memory savings and prefill speedups.

SimpleMem: Efficient Lifelong Memory for LLM Agents

Papers with Code Trending

Introduces SimpleMem, an efficient memory framework for LLM agents that uses semantic lossless compression to improve accuracy and reduce token consumption, achieving 26.4% F1 improvement and up to 30x reduction in inference-time token usage.

MemTrain: Self-Supervised Context Memory Training

arXiv cs.CL

MemTrain proposes a self-supervised training framework that uses masked reconstruction and intermediate memory recall proxy tasks on Wikipedia corpora to enhance LLM agents' context memory, achieving up to 17.67 point gains on downstream memory-intensive QA benchmarks.