MoME:把稀疏查表变成上下文感知的记忆混合
核心概要
该工作提出 Mixture of Memory Embeddings(MoME),把每个 token 的单一记忆行扩展为多个记忆槽,并用基于隐藏状态的学习门控在每一位置稀疏选择要读取的槽,从而在保持 token 索引式廉价查表的同时让记忆检索具备上下文感知;在 nanochat、Llama-3/MobileLLM 与 Qwen3 三类骨干的受控预训练中,MoME 在等参数与等训练 FLOP 设定下优于 Value Embedding、Bigram 与 STEM 基线,在十亿参数以下规模显示出更有利的记忆规模扩展趋势,训练与推理仍保持高效,并且对多义词的路由分析显示同一表层 token 在不同语义下被派发到不同记忆槽。
深度剖析
MoME 将 token 索引记忆表从“每 token 一行”改为“每行多个记忆槽”,并用读取隐藏状态的学习门控在每个位置选择激活哪些槽,使记忆寻址不再只是表层形式的确定性函数。 既有记忆嵌入方法(Per-Layer Embedding、Value Embedding、STEM、Engram 等)都通过 token 身份或固定 n-gram 哈希做确定性检索,同一 token 的不同语境义项被压缩进同一个固定向量;MoME 在保留 token 索引查表结构的前提下引入上下文相关的槽选择。 论文给出模块形式化定义(行索引器加槽门控、sigmoid-norm 聚合、门控注入值流),并在三类骨干上做受控预训练对比;消融显示隐藏状态路由优于仅 token 学习路由与固定随机路由。
在三类骨干的受控预训练中,MoME 在等参数与等训练 FLOP 设定下多数匹配比较中优于 Base、Value Embedding、STEM 与 Bigram 基线,且相对 STEM 使用更少记忆参数。 此前记忆嵌入方法的比较多集中在单一骨干或单一规模;该工作把同一机制迁移到 nanochat 风格、Llama/MobileLLM 风格与 Qwen3 风格三类骨干,覆盖 125M、350M 与 0.6B 等规模。 报告了训练/验证 bpb 与 CORE 指标,例如 Llama/MobileLLM 125M 上 MoME-A2/6 验证 bpb 0.8560、CORE 0.1686,对比 VEmbedding 的 0.8620 与 0.1530;Qwen3 0.6B 上 MoME-A2/8 的 CORE 0.2737 高于 VEmbedding 的 0.2530;训练墙钟时间为无记忆基线的约 1.04–1.08 倍。
在十亿参数以下的低算力区间,MoME 随记忆容量增长呈现出比 Bigram 更有利的验证 bpb 轨迹,并且与 Bigram 可组合使用。 既有工作较少系统报告记忆表规模扩展曲线;该工作在 nanochat d12 上固定 FLOP、扫描记忆参数,并额外测试把 Bigram 索引器作为 MoME 的第一阶段索引器。 论文称在每个测试记忆规模上 MoME 的验证 bpb 都低于匹配记忆量的 Bigram,且训练与验证 bpb 曲线的运行间方差更小;组合配置在相同参数预算下优于单独 Bigram。
路由分析显示 MoME 的学习路由与词义相关:同义提示倾向路由到同一记忆槽,而改变语义的提示跳到不同槽,并在 WiC 上给出量化证据。 此前记忆嵌入的检索是确定性的,不存在可分析的“语境路由”;该工作把 MoE 式稀疏路由引入记忆表,并对其语义结构做定性与定量探测。 在 d24 检查点的 144 个记忆层–头位点上评估 670 个严格过滤的 WiC 对,117 个位点的不同义对路由散度高于同义对,110 个位点的同义对槽重叠更高;最大效应出现在 layer 11/head 8,但该头注入门接近关闭,门控加权后最大效应转移到 layer 7/head 6。
启示与展望
该结果面向在受控预训练设定下研究稀疏容量扩展的研究者与工程团队:MoME 适用于保留 token 索引查表、把记忆注入注意力值流的骨干,并已在 nanochat 风格、Llama/MobileLLM 风格与 Qwen3 风格三类架构、125M 至 0.6B 规模以及约 100B token 的 ClimbMix 训练中得到验证;机制可与 Bigram 索引器组合,也可与稀疏 FFN MoE 叠加。其设计目标是让记忆检索随上下文变化,同时把额外延迟控制在较小范围,因此对关注推理延迟与记忆参数预算的部署场景具有参考价值。
论文自身指出主要限制在计算规模:扩展的 100B token 实验总参数约十亿量级,稠密骨干约十亿参数且使用单一种子;CORE 激活分析与 WiC 路由分析提供的是选择性记忆使用与语义敏感路由的描述性证据,仍需因果干预才能确定这些行为是否提升下游准确率或域偏移下的稳健性。WiC 分析中只有 81 个目标 token 组同时包含两种义项标签,合并比较可能受词汇构成差异影响,且未估计留出泛化;自助法区间为逐点区间,最大效应在同一批样本上选出而未做 144 个位点间的多重比较校正。此外,本次载入的正文中部分表格数值以空白形式呈现,若需核对具体逐任务数字,应以原文表格为准。
