MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup
Muchen Li, Leonid Sigal, Renjie Liao
cs.CL, cs.AI
2026-09-14
UBC 的 MoME 把每个 token 的单条记忆换成 M 个槽,用当前隐藏态做门控选取,在 nanochat、Llama/MobileLLM 和 Qwen3 三类骨干上以同等参数压过 Value Embedding 与 STEM,100B 训练 CORE-22 到 0.3687。
给 Transformer 加条件记忆,是 MoE 之外另一条稀疏扩容路:用一张按 token 索引的表做廉价查找,专门承接实体名、套话这类静态模式,让骨干少花前几层去「重算一张表」。Gemma 3 的 Per-Layer Embedding、Value Embedding、STEM、Engram 都走这条路,查找函数却几乎都是词面的确定性函数。python 当编程语言和当动物,只能挤进同一行。
容量也按词分配,不按语义分配。cats 和 cat 各占一行,近义重复;bank 一行要同时装银行和河岸。需要一种仍然按 token 寻址、但能按上下文选槽的查找。
MoME 把每个 token 行从单个向量换成 M 个记忆槽。第一级索引仍用 token id(可选地把近义 token 并到同一行),第二级用进入该块的隐藏态做门控,每个 value 头独立 TopK 选出 K 个槽,sigmoid-归一后加权求和,再以残差门控加进注意力 value 流。K=1 时改用 softmax。默认 K=2,配置写成 MoME-A2/M。
注入点选 value 流,是为了让记忆分支和标准 value 投影并行,汇合只做一次门控加法。Qwen3-4B 两层记忆的设定里,这条路径多 0.509 ms,改成隐状态到隐状态注入会多 0.996 ms。多头共享一张表,省参数;门控仍按头分开,不同头可以读不同槽。可选的 kNN 并组把行数缩小 c 倍,省下的容量加到每行槽数上。
因为第一级索引可替换,MoME 还能叠在 Bigram 哈希上一起扩表;同等预算下组合优于单用 Bigram,更大预算时两边一起加最稳。
记忆块和普通 Transformer 块交替放置。
nanochat 风格 135M 骨干、固定 3×10^18 FLOPs(约 3.3B token)、151M 记忆预算下,MoME(cgrp=2) 验证 bpb 0.8611、CORE 0.1583,Value Embedding 是 0.8633/0.1522,Bigram 是 0.8636/0.1533。记忆加倍到 302M 后,同一分组的 CORE 到 0.1664。训练吞吐相对 Bigram 慢不到 2%。
固定 token 预算迁到别的骨干。Llama/MobileLLM 125M、55B token:MoME-A2/6 的 CORE 0.1686,Value Embedding 0.1530,STEM 0.1483,无记忆基线 0.1382。350M、20B token:MoME-A2/10 的 CORE 0.2286,Value Embedding 0.2214。Qwen3 0.6B、20B token:MoME-A2/8 的 CORE 0.2737,STEM 0.2556,Value Embedding 0.2530;BoolQ 上 MoME 63.39,Value Embedding 掉到 47.55。训练墙钟是基线的 1.04–1.08 倍。
100B ClimbMix、nanochat d24(骨干约 0.78B,记忆约 0.61B)上,MoME-A2/12 的 CORE-22 是 0.3687,Value Embedding 0.3484,Dense 0.3441。域内 ClimbMix bpb MoME 0.6504,略差于 Value Embedding 的 0.6463;域外 FineWeb-Edu、enwik9、Shakespeare 都略好。对照 Qwen3-0.6B 用了 36T token 才到 CORE-22 0.3753。
WiC 上 144 个层-头位点里,117 个同义对比异义的路由分布差分为正,110 个离散槽重叠差分为正。bank、bug、drive 的同义提示走同一槽,换义跳槽。
条件记忆这条轴如果只按词面哈希,扩表只会把多义词的冲突写得更大。MoME 证明上下文选槽可以叠在原有 token 索引上,也能和 Bigram 第一级索引一起扩。对已经在 value 流插表的预训练配方,改动面小:多一套槽和一门控,吞吐几乎不动。
增益是真的,也是小步。iso-参数下验证 bpb 只比 Value Embedding 低约 0.002,CORE 更敏感。100B 实验仍在 1.4B 总参数、单一种子。
作者把计算规模写成主要限制:100B 实验骨干约 0.8B,总参数约 1.4B,单一种子。CORE 激活和 WiC 路由是描述性证据,没有做因果干预,不能断定「换槽」本身提高了下游正确率或域外稳健性。并组只在 nanochat 上扫过,cgrp=4 在 302M 时 CORE 掉到 0.1515,分组过猛会伤。STEM 对照用了原论文的 1/2 层插入和大得多的记忆表,参数量不对齐,不能读成全面碾压 STEM。