MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup
- URL: http://arxiv.org/abs/2609.15126v1
- Date: Mon, 14 Sep 2026 07:01:52 GMT
- Title: MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup
- Abstract summary: We introduce Mixture of Memory Embeddings (MoME), a context-aware memory mechanism.<n>MoME 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.<n>In pretraining experiments across nanochat, Llama-3/MobileLLM, and Qwen3 backbones, MoME improves over Value Embedding, Bigram, and STEM baselines.
- Score: 35.92693145413963
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: 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.
Related papers
- Dual-Grained Agent Memory and Shapley Context Attribution for Multimodal Agentic Learner [55.80248281761638]
multimodal large language models (MLLMs) deliver impressive perception yet still falter on scientific and mathematical reasoning.<n>We propose DG-Mem, a dual-grained agentic memory framework that augments a frozen MLLM with a non- externally stored memory built once from training-time rollouts and consulted read-only at test time.
arXiv Detail & Related papers (2026-08-24T13:56:01Z) - BitNet Text Embeddings [104.29781473344813]
We present BITEMBED, an extreme low-bit framework for LLM-based text embedding that jointly targets encoding efficiency and vector storage.<n> BITEMBED converts pretrained LLM backbones into BitNet-style embedding encoders with ternary weights, quantized activations, and lightweight normalization refinement.<n> Experiments on MMTEB (eng) with Qwen3-0.6B and Gemma703-2M show that BITEMBED is largely comparable to full precision teacher embedders.
arXiv Detail & Related papers (2026-06-24T10:37:01Z) - Cognitive Memory in Large Language Models [8.059261857307881]
This paper examines memory mechanisms in Large Language Models (LLMs), emphasizing their importance for context-rich responses, reduced hallucinations, and improved efficiency.<n>It categorizes memory into sensory, short-term, and long-term, with sensory memory corresponding to input prompts, short-term memory processing immediate context, and long-term memory implemented via external databases or structures.
arXiv Detail & Related papers (2025-04-03T09:58:19Z) - Wormhole Memory: A Rubik's Cube for Cross-Dialogue Retrieval [4.586907225774023]
This research proposes a wormhole memory module (WMM) to realize memory as a Rubik's cube that can be arbitrarily retrieved between different dialogues.<n> Experimental results show that WMM demonstrated the ability to retrieve memory across dialogues and the stability of quantitative indicators.
arXiv Detail & Related papers (2025-01-24T10:49:45Z) - Towards A Unified View of Sparse Feed-Forward Network in Pretraining
Large Language Model [58.9100867327305]
Large and sparse feed-forward layers (S-FFN) have proven effective in scaling up Transformers model size for textitpretraining large language models.
We analyzed two major design choices of S-FFN: the memory block (a.k.a. expert) size and the memory block selection method.
We found a simpler selection method -- textbftextttAvg-K that selects blocks through their mean aggregated hidden states, achieving lower perplexity in language model pretraining.
arXiv Detail & Related papers (2023-05-23T12:28:37Z) - A Model or 603 Exemplars: Towards Memory-Efficient Class-Incremental
Learning [56.450090618578]
Class-Incremental Learning (CIL) aims to train a model with limited memory size to meet this requirement.
We show that when counting the model size into the total budget and comparing methods with aligned memory size, saving models do not consistently work.
We propose a simple yet effective baseline, denoted as MEMO for Memory-efficient Expandable MOdel.
arXiv Detail & Related papers (2022-05-26T08:24:01Z) - MCML: A Novel Memory-based Contrastive Meta-Learning Method for Few Shot
Slot Tagging [17.205431506768125]
We propose the textbfMemory-based textbfContrastive textbfMeta-textbfLearning (aka, MCML) method to bridge the distribution gap between training episodes and between training and testing respectively.
Experimental results show that the MCML outperforms several state-of-the-art methods on both SNIPS and NER datasets.
arXiv Detail & Related papers (2021-08-26T08:02:21Z) - Semantically Constrained Memory Allocation (SCMA) for Embedding in
Efficient Recommendation Systems [27.419109620575313]
A key challenge for deep learning models is to work with millions of categorical classes or tokens.
We propose a novel formulation of memory shared embedding, where memory is shared in proportion to the overlap in semantic information.
We demonstrate a significant reduction in the memory footprint while maintaining performance.
arXiv Detail & Related papers (2021-02-24T19:55:49Z) - Kanerva++: extending The Kanerva Machine with differentiable, locally
block allocated latent memory [75.65949969000596]
Episodic and semantic memory are critical components of the human memory model.
We develop a new principled Bayesian memory allocation scheme that bridges the gap between episodic and semantic memory.
We demonstrate that this allocation scheme improves performance in memory conditional image generation.
arXiv Detail & Related papers (2021-02-20T18:40:40Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.