Memory Attention
Jiale Kang
cs.LG
2026-09-24
把注意力的 V 从稠密投影改成「键+词表查表」,等损失训练省 13.8%~29.6% token,查表放 CPU 预取后 GPU 常驻参数反而比基线少 7.38%。
标准注意力里,Q、K、V 三个向量都由同一份隐状态经三张稠密矩阵投影而来。V 是「被聚合的内容」,但同一个 token 在不同上下文里要写进 V 的东西有相当一部分是重合的,论文摘要里的说法是「内容可跨上下文复用」。每层每步都靠一次矩阵乘法从头算 V,这份计算和这份常驻显存的权重,是注意力里一笔固定开销。
此前已有工作往注意力里加 token 级信息:value embeddings 把词向量混进 V,PLE 在每层注入词表表示,Engram 用 n-gram 索引查记忆。它们都保留了原来的价值投影,新东西是「加在投影之上」。这篇问的是一个更激进的问题:那张 WV 能不能整个删掉,token 身份靠查表,上下文依赖靠本来就有的 K 来补?
Memory Attention 的构造一行就能写完:V = K + Norm(E[s])。
设计的关键是分工:记忆提供「这个 token 是什么」,K 提供「此情此景它意味着什么」。K 来自隐状态,而隐状态已经吸收了前文与前层的计算,所以上下文依赖没有丢。同一个词在不同句子里查到同一行记忆,但加上的 K 不同。WK 和记忆表联合训练,一起学着撑起被删掉的 V。
推理时还有一步折叠:归一化参数固定后可以预先算进表里,在线构造 V 只剩「查一行、加一下」,矩阵乘法没有了。又因为查表地址只由 token id 和层号决定、与隐状态无关,整张表可以搬进 CPU 内存,靠预取流水线在计算间隙把数据送上 GPU,参数不再必须常驻显存。论文还分析了 MA-Recall 扩展:缓存 K 和 token id,历史 V 用 V=K+Ē[s] 现场重建,KV cache 理论上省一半,但这一部分只做了分析,没有实测。
与 MLA 的关系一句话可以说清:两者都在聚合一份共享表示,MLA 聚合潜在向量 C 之后还要做一次学到的值映射,MA 直接聚合 K 加上可加的记忆。
等训练 token 预算、同种子,FineWeb-10BT 预训练,下游取七项 zero-shot 任务的平均分:
| 构型(24 层,10B token) | 标准 | MA | 参数量 |
| MHA | 40.68 | 41.39 | 373M→1,135M |
| GQA | 40.90 | 41.51 | 349M→729M |
| MQA | 41.11 | 41.34 | 336M→526M |
| 门控变体 | 41.40 | 42.88 | 398M→1,160M |
WikiText 困惑度上 MHA 从 31.55 降到 28.64。大一号的模型(24 层、d=2048、20B token)平均分 49.06→50.22,OBQA +3.00、ARC-C +2.65,但 ARC-E 和 PIQA 略降。检索任务差距更大:2K 训练窗口内的单针大海捞针,MA 均值 96.8%,标准注意力 82.1%;外推到 4K 时 41.9% 对 25.9%,单项最大提升 62.4%→98.2%。
等损失 token 效率:小模型 1.42 倍、大模型 1.16 倍,即省 29.6% 与 13.8% 的训练 token。
推理侧(H800,batch 8,prefill 2048 加一步带 2048 历史的 decode):
| 配置 | prefill / decode(ms) | GPU 参数存储 | 总参数 |
| 标准 | 90.97 / 17.64 | 2,602.38 MiB | 1,364M |
| MA(表在 GPU) | 88.20 / 17.79 | 5,410.38 MiB | 2,836M |
| MA-Offload(表在 CPU) | 90.69 / 17.19 | 2,410.38 MiB | 2,836M(其中 3,000 MiB 在 CPU) |
MA-Offload 以比基线低 7.38% 的显存参数驻留,装下了 2.08 倍的总参数,延迟与基线打平。
如果「V 可以靠查表加 K」在大规模上成立,它给出两样直接可用的东西:注意力里少一个稠密乘法;以及一个新的参数落脚点,词表索引的表,天然可以放进 CPU 内存靠预取供数,模型容量与显存占用解耦。对推理成本敏感的线上模型,这是与量化、KV cache 压缩正交的另一条轴。检索任务的提升幅度也值得展开:token 身份直接进 V,可能缓解了多层抽象把稀有 token 信息洗掉的问题。
也要诚实:这更像一份「方向可行」的验证,不是能直接搬进生产的配方。实验停在 2.8B 参数、20B token,单一作者、gmail 邮箱的独立预印本,尚未经同行评审。