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

会议: NeurIPS 2026
arXiv: 2609.15126
代码: https://github.com/muchenli/MoME
领域: 其他
关键词: 条件记忆、稀疏查找、多义词消歧、混合记忆嵌入、值流注入

一句话总结¶

本文提出MoME（Mixture of Memory Embeddings），通过将离散词表中每个token的单行记忆槽扩展为多个记忆槽并引入依赖隐层表征的轻量门控路由，打破了现有条件记忆依赖确定性词表/n-gram索引而导致的上下文盲区，在微小计算与延迟开销下显著提升语言模型建模与下游零样本推理表现。

研究背景与动机¶

大规模语言模型的高效扩展一直是拓展智能边界的核心驱动力。密集的参数扩展虽然能直接带来性能跃迁，但不可避免地导致每个前向token的计算量呈线性膨胀。为了解耦模型容量与推理算力，以Mixture-of-Experts（MoE）为代表的条件计算范式已广泛应用于主流前沿模型中。然而，语言建模本质上交织着两类异质子任务：一类是依赖动态计算图的组合推理任务，另一类则是关于命名实体、固定短语和事实三元组的局部静态模式检索。传统Transformer架构欠缺原生的查找基元，只能消耗宝贵的前几层网络容量去隐式模拟静态词表查找，从而造成计算资源的大量浪费。

受此启发，以Engram、STEM、Per-Layer Embedding以及值流记忆（Value Embedding）为代表的“条件记忆”（Conditional Memory）机制逐渐成为继MoE之后解耦容量与算力的第二条稀疏扩展主线。这类方法通过学习外部辅助嵌入表并在前向过程中稀疏查表注入主干，使得模型能够以极低算力增加有效参数量。然而，现有条件记忆方法存在一个致命的结构性限制：查表索引机制通常是输入表面形式（surface form）的确定性函数，例如直接以当前token ID或固定长度的n-gram哈希作为查表行号。这种设计虽然检索极其轻量，却直接导致记忆检索对上下文完全盲目（context-blind）。在自然语言中，同一个表面词在不同语境下具有截然不同的多义语义（例如“python”既可指代编程语言又可指代爬行动物，“bank”既指金融机构又指河岸）；强制让多义词的所有语境模式挤压在唯一固定的记忆行向量中，造成了显著的语义碰撞与表征容量错配。

针对上述多义词表征坍缩与上下文感知缺失的核心痛点，本文提出了MoME（Mixture-of-Memory Embeddings）架构。核心 idea：将记忆表中每个token对应的单个静态记忆向量扩展为由 \(M\) 个记忆槽组成的局部混合池，并引入以当前层隐状态为条件的自适应门控机制动态路由激活 Top-\(K\) 个槽位，在保留 \(O(1)\) 级别极低硬件查找延迟的同时，赋予静态记忆表动态感知深层语义语境的能力。

方法详解¶

整体框架¶

MoME 的核心是在标准 Transformer 架构中交替插入记忆增强层（Memory-Augmented Block），其整体计算流由三部分构成：第一阶段根据输入token标识（或可选的分组聚类映射）确定候选记忆行；第二阶段由 upstream 上下文隐状态通过轻量门控预测槽位对数几率，选出当前语境最相关的 \(K\) 个记忆槽进行加权聚合；第三阶段将聚合得到的动态记忆向量通过残差门控注入到注意力机制的多头值向量流（Value Stream）中。

输入序列中的token记为 \(x_t\)，进入第 \(\ell\) 个记忆增强块的隐状态记为 \(\mathbf{h}_t \in \mathbb{R}^{d_{\mathrm{model}}}\)。对于每个注意力头 \(i \in [H]\)，模型通过共享的记忆参数张量 \(\mathbf{E}^{\mathrm{mem}} \in \mathbb{R}^{N \times M \times d_{\mathrm{value}}}\) 进行查找，其中 \(N\) 为有效行数（词表大小 \(V\) 或聚类后行数），\(M\) 为每行包含的记忆槽数，\(d_{\mathrm{value}}\) 为每个注意力头的维度。整个计算管线端到端可导，其自上而下的整体架构与数据流向如下所示：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    InToken["输入 token x_t"] --> RowIdx["确定记忆行<br/>词表索引或 kNN 语义分组 f(x_t)"]
    InHidden["隐状态 h_t"] --> SlotGate["上下文感知槽位路由与聚合<br/>W_g 隐状态映射 + Top-K 槽位稀疏选择"]
    RowIdx --> SlotGate
    SlotGate --> AggMem["多头加权聚合记忆向量 m~_t,i"]
    InHidden --> ValGate["值流注入与并行调度<br/>残差门控 gamma(h_t) 调节强度"]
    InHidden --> StdValue["并行分支：标准注意力值投影 v_t,i"]
    AggMem --> FuseValue["门控融合残差 v~_t,i = v_t,i + gamma_t,i * m~_t,i"]
    ValGate --> FuseValue
    StdValue --> FuseValue
    FuseValue --> OutAttn["多头自注意力计算与后续前传"]

关键设计¶

1. 确定记忆行：保持低检索复杂度的第一级索引与分组策略

现有确定性记忆方法直接将词表映射为行，造成低频相似词冗余分配而高频多义词容量不足的问题。MoME 采用两级分解寻址架构，第一级索引函数 \(f: \mathcal{V} \to [N]\) 负责从离散词表中选出对应的记忆行 \(n_t = f(x_t)\)。在标准配置下，采用恒等映射 \(f(x_t) = x_t\)（此时 \(N = V\)），直接继承标准词表查找的高效访存模式；同时，论文进一步设计了一种离线构建的语义分组映射函数 \(f(x) = \mathrm{kNN}(\mathbf{U}^{\mathrm{ref}}_x; \mathbf{U}^{\mathrm{ref}}, c_{\mathrm{grp}})\)，通过对预训练词嵌入矩阵 \(\mathbf{U}^{\mathrm{ref}}\) 执行离线聚类，将语义高度重合的屈折形式（如 cat/cats）聚类到同一行，行数压缩至 \(N = V / c_{\mathrm{grp}}\)。压缩释放的存储空间被全额补偿给每行的槽位数 \(M \leftarrow M \times c_{\mathrm{grp}}\)，在维持总参数量严格不变的前提下将静态物理行转化为更具区分度的动态语义插槽。

2. 上下文感知槽位路由与聚合：利用隐状态解耦多义词槽位

针对同一词在不同上下文中语义坍缩的顽疾，MoME 在第二级索引阶段引入隐状态条件门控。对于第 \(i\) 个注意力头，门控模块将已融合上游全部上下文信息的中间表征 \(\mathbf{h}_t\) 线性投射到 \(M\) 维槽位 logits：

\[ \boldsymbol{\ell}_{t,i} = \mathbf{W}^{(i)}_g \mathbf{h}_t + \mathbf{b}^{(i)}_g \in \mathbb{R}^M. \]

模型在槽位集合 \([M]\) 内执行 \(\mathrm{TopK}\) 稀疏截断，选出激活槽位子集 \(\mathcal{A}_{t,i} = \mathrm{TopK}(\boldsymbol{\ell}_{t,i}, K)\)。当 \(K > 1\) 时，采用 Sigmoid 归一化门控计算归一化权重 \(\alpha^{(i)}_{t,a} = \frac{\sigma(\ell_{t,i,a})}{\sum_{b \in \mathcal{A}_{t,i}} \sigma(\ell_{t,i,b})}\)（若 \(K=1\) 则退化为局部 Softmax）。最终，第 \(i\) 个头从该行选中的记忆槽加权合成聚合记忆向量：

\[ \tilde{\mathbf{m}}_{t,i} = \sum_{a \in \mathcal{A}_{t,i}} \alpha^{(i)}_{t,a} \mathbf{m}_{n_t, a} \in \mathbb{R}^{d_{\mathrm{value}}}. \]

由于各注意力头独立执行门控投影，各头既可以共享同一底层记忆槽张量 \(\mathbf{E}^{\mathrm{mem}} \in \mathbb{R}^{N \times M \times d_{\mathrm{value}}}\)（避免参数量随头数 \(H\) 膨胀），又能在不同子空间内分别关注该词不同的多义属性（例如 Head 1 激活技术含义槽位，Head 2 激活常识生物槽位）。

3. 值流注入与并行调度：面向低推理延迟的旁路融合机制

将记忆注入网络中间层若设计不当极易阻塞关键计算路径。若像 Engram 或 STEM 一样将记忆向量注入到块输入隐状态或 SwiGLU 隐藏层中，序列必须串行等待记忆检索与聚合完成才能执行后续算子。MoME 采用双重设计化解该瓶颈：首先，将记忆注入点精准设在注意力机制的值向量流（Value Stream）。由于门控输入仅依赖进入该块的输入 \(\mathbf{h}_t\)，槽位路由计算 \(\mathbf{W}_g \mathbf{h}_t\)、槽位聚合与主干标准值投影 \(\mathbf{v}_{t,i} = \mathbf{W}_v \mathbf{h}_t\) 完全相互独立，可以在 GPU 上完美并发执行。两路计算完成后，仅需执行一个轻量级的元素级门控残差加法：

\[ \tilde{\mathbf{v}}_{t,i} = \mathbf{v}_{t,i} + \gamma_{\theta, t, i} \tilde{\mathbf{m}}_{t,i}, \quad \text{其中} \quad \boldsymbol{\gamma}_\theta(\mathbf{h}_t) = 2\sigma(\mathbf{W}_\gamma \mathbf{h}_t + \mathbf{b}_\gamma) \in \mathbb{R}^H. \]

这里的标量缩放因子 2 保证了门控初始 logits 为 0 时 \(\gamma = 1.0\)，实现平滑的中性值流融合。实测表明，在 Qwen3-4B 模型中，该并行设计将记忆层引入的推理延迟从隐状态注入方案的 0.996 ms 压低至 0.509 ms，降低了近 50% 的额外时延开销。

一个完整示例¶

以多义词 "bank" 在不同语境中的处理流程为例。词表索引 \(f(\text{"bank"})\) 首先从共享表 \(\mathbf{E}^{\mathrm{mem}}\) 中定位到对应行 \(n\)，该行配置有 \(M = 8\) 个记忆槽（维度 \(d_{\mathrm{value}} = 128\)）。

当模型输入句子 "She deposited the check at the bank before noon." 时，前序网络层编码的隐状态 \(\mathbf{h}_t\) 深刻蕴含金融存款语义。在第 7 层第 7 个注意力头处，门控网络计算得到 8 个槽位的 logits，经过 Top-2 筛选激活了槽位 1（权重 0.85）与槽位 4（权重 0.15），两者均对应金融机构表征，加权合成金融语境向量并注入值流；

反之，当输入变为 "The canoe scraped the muddy bank as it landed." 时，水上泛舟的语境使 \(\mathbf{h}_t\) 发生显著语义迁移。相同的词表索引依然检索行 \(n\)，但门控网络输出突变，Top-2 槽位切换为槽位 6（权重 0.92）与槽位 2（权重 0.08），指向地理河岸实体。离散的表面词通过动态槽位分配与连续权重平滑过渡，彻底消除了单向量表征下的义项互相干扰。

损失函数 / 训练策略¶

MoME 保持标准因果自回归交叉熵目标函数，无需设计额外的人工路由正则化损失或辅助负载均衡损失（经消融验证，辅助均衡损失反而会削弱槽位的自然语义特化）。在优化器配置上，Transformer 主干矩阵参数采用 Muon 优化器更新，而嵌入层、输出头以及非矩阵参数（包括值流记忆表 \(\mathbf{E}^{\mathrm{mem}}\) 与门控参数 \(\mathbf{W}_g, \mathbf{W}_\gamma\)）采用 AdamW 优化器进行更新。学习率采用余弦衰减调度，序列长度固定为 2048。

实验关键数据¶

主实验¶

主干实验在受控预训练环境下评估，覆盖 nanochat、Llama-3/MobileLLM 以及 Qwen3 三大主流架构。下表展示了 Llama/MobileLLM（125M / 350M）以及 Qwen3 0.6B 在相同预训练 token 规模下的代表性下游基准与综合 CORE 分数对比（截取自论文表 3）：

模型架构	方法	记忆参数量	验证集 bpb ↓	CORE ↑	ARC-C	BoolQ	PIQA	训练墙钟相对耗时
Llama/MobileLLM 125M	Base (Dense)	0M	0.8852	0.1382	27.56	55.14	63.82	1.00×
	STEM (1/2层)	755M	0.8633	0.1483	29.35	54.01	65.13	1.00×
	VEmbedding	95M	0.8620	0.1530	28.50	53.67	66.43	1.03×
	MoME-A1/3	95M	0.8623	0.1587	27.30	50.49	66.00	1.07×
	MoME-A2/6	189M	0.8560	0.1686	27.65	57.86	65.18	1.07×
Llama/MobileLLM 350M	Base (Dense)	0M	0.7750	0.1988	34.04	55.66	70.13	1.00×
	STEM	1342M	0.7697	0.2179	35.41	50.52	69.37	1.09×
	VEmbedding	168M	0.7651	0.2214	36.26	47.06	70.62	1.04×
	MoME-A2/5	168M	0.7647	0.2202	36.60	44.65	70.73	1.04×
	MoME-A2/10	336M	0.7610	0.2286	36.77	55.44	71.60	1.08×
Qwen3 0.6B	Base (Dense)	0M	0.7594	0.2491	35.32	59.88	70.78	1.00×
	STEM	1409M	0.7564	0.2556	36.52	56.39	71.98	1.07×
	VEmbedding	470M	0.7427	0.2530	38.05	47.55	71.00	1.06×
	MoME-A2/8	470M	0.7461	0.2737	38.31	63.39	71.76	1.06×

消融实验¶

下表展示了在 nanochat-135M 骨干网络（训练 3.3B tokens，固定 \(3 \times 10^{18}\) FLOPs）上，不同记忆参数规模与 token 聚类分组因子 \(c_{\mathrm{grp}}\) 的对比分析（截取自论文表 2）：

实验配置	记忆表形状定义	记忆参数量	验证集 bpb ↓	综合 CORE ↑	训练吞吐量
Dense Base	–	0M	0.8785 ± 0.0010	0.1447 ± 0.0064	1.15M tok/s
Bigram 哈希	\(N_{\mathrm{hash}} \times d_{\mathrm{model}}\)	151M	0.8636 ± 0.0009	0.1533 ± 0.0116	1.09M tok/s
VEmbedding	\(V \times 6 \times d_{\mathrm{value}}\)	151M	0.8633 ± 0.0006	0.1522 ± 0.0045	1.10M tok/s
MoME (\(c_{\mathrm{grp}}=1\))	\(V \times 6 \times d_{\mathrm{value}}\)	151M	0.8621 ± 0.0003	0.1571 ± 0.0030	1.07M tok/s
MoME (\(c_{\mathrm{grp}}=2\))	\(\frac{V}{2} \times 12 \times d_{\mathrm{value}}\)	151M	0.8611 ± 0.0003	0.1583 ± 0.0020	1.07M tok/s
MoME (\(c_{\mathrm{grp}}=4\))	\(\frac{V}{4} \times 24 \times d_{\mathrm{value}}\)	151M	0.8610 ± 0.0004	0.1554 ± 0.0029	1.07M tok/s
MoME (\(c_{\mathrm{grp}}=1\))	\(V \times 12 \times d_{\mathrm{value}}\)	302M	0.8569 ± 0.0002	0.1635 ± 0.0067	1.03M tok/s
MoME (\(c_{\mathrm{grp}}=2\))	\(\frac{V}{2} \times 24 \times d_{\mathrm{value}}\)	302M	0.8561 ± 0.0003	0.1664 ± 0.0059	1.03M tok/s
MoME (\(c_{\mathrm{grp}}=4\))	\(\frac{V}{4} \times 48 \times d_{\mathrm{value}}\)	302M	0.8565 ± 0.0006	0.1515 ± 0.0110	1.02M tok/s

关键发现¶

同参数下更优的表征效率：在严格 iso-parameter 约束下（如 151M 记忆预算），MoME (\(c_{\mathrm{grp}}=1\)) 验证集 bpb 达到 0.8621，优于静态 Value Embedding 的 0.8633 与 Bigram 的 0.8636，且训练方差显著降低，证明动态槽位路由有效缓解了单一表征容量竞争。
分组聚类（\(c_{\mathrm{grp}}\)）释放记忆潜能：通过将形态重合的词聚合（\(c_{\mathrm{grp}}=2\)），为每行腾出 2 倍的记忆槽位，模型在 151M 与 302M 预算下均取得了最佳 CORE 得分（分别为 0.1583 与 0.1664）。但过度压缩（\(c_{\mathrm{grp}}=4\)）会导致第一级索引语义混杂，下阶段槽位区分度边际递减。
复合扩展兼容性（Compound Scaling）：将 Bigram 作为第一阶段行索引、MoME 作为第二阶段槽位选择时，在相同参数量下显著超过纯 Bigram，并在扩大容量后取得更优缩放曲线，证明两轴稀疏化可直接正交叠加。
层间记忆使用规律：对网络各层注入量 \(\gamma_{\ell} \tilde{\mathbf{m}}_\ell\) 统计发现，记忆模块在网络极浅层（负责早期事实快速注入）与极深层（负责最终输出预测微调）激活强度最大，中间层则更多依赖主干自注意力与 FFN 计算。

亮点与洞察¶

打破记忆查找“语境盲区”：创新性地在词表静态查找与动态条件计算之间找到了优雅平衡点，仅用一层极轻量的线性门控即可为多义词提供上下文相关的细粒度记忆检索。
与注意力值投影完美并行的硬件友好架构：将注入点定于注意力值向量流（Value Stream），使得门控打分、槽位检索与标准注意力投影可以高度并行化，彻底消除了序列关键路径上的串行等待，在大模型上额外推理延迟低至 1.5%~2.3%。
具备语义可解释性的自发多义词解耦：在没有任何显式词义标注监督的情况下，模型仅凭语言建模自监督损失自发学会将同形多义词（如 bank/bug/drive）在不同上下文中精准路由至不同的记忆插槽，在 Word-in-Context 评测中超过 81% 的注意力头表现出统计显著的义项敏感性。

局限与展望¶

计算规模仍局限于 Sub-Billion 尺度：受算力限制，主干预训练最大验证到 0.8B 主干 + 0.6B 记忆（约 1.4B 总参数量），更大规模百亿/千亿级稠密与 MoE 模型上的扩展规律仍有待系统性验证。
因果干预证据有待强化：目前的 WiC 与层分析主要为相关性统计描述，缺乏针对特定记忆槽进行消融或激活改写的严格因果干预实验，以直接证实下游推理提升与词义解耦的因果关联。
未来方向：模型行为可控引导：利用学到的语义解耦槽位路由机制，未来有望通过在推理期主动偏置/冻结特定记忆槽的选择路径，实现对语言模型生成风格、专业领域偏好或事实知识的高效无损引导。

评分¶

新颖性: ⭐⭐⭐⭐☆ [将 MoE 稀疏门控思想精巧迁移至词表级参数化条件记忆，解决多义词坍缩难题]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖三大模型家族、iso-FLOP/iso-parameter 严格对齐、扩展定律分析及深度的 WiC 语义路由探针]
写作质量: ⭐⭐⭐⭐⭐ [叙述严谨清晰，架构图与动机脉络层层递进，消融与时延开销分析详实透明]
价值: ⭐⭐⭐⭐☆ [为大模型低成本扩展非易失事实记忆与解决知识瓶颈提供了极具工程落地潜力的实用新范式]