arXiv：cs.LG· Jiale Kang·· 4 小时前
Memory Attention：用 token 级记忆替代注意力 value 投影
Memory Attention
AI 导读
论文提出 Memory Attention（MA），用层级 token 记忆与上下文 key 组合生成 value，替代注意力中专门的 value 投影。推理时归一化可折叠进记忆表，value 构建简化为查表与加法，并支持 CPU offloading 与预取以减少 GPU 参数存储。在匹配训练 token 预算下，多种注意力配置的实验显示语言建模与平均下游性能均有提升。
正文
Abstract:Language models typically construct attention values from contextual hidden states, even when some of their content may be reusable across contexts. We investigate whether token-indexed memory can replace the dedicated value projection when complemented by contextual information. We propose Memory Attention (MA), which forms values by combining layer-specific token memory with contextual keys. The memory supplies token-specific representations, while the keys preserve context dependence. At inference, normalization can be folded into the memory tables, reducing value construction to lookup and addition. Token-indexed retrieval also enables CPU offloading with prefetching, reducing GPU parameter storage. Under matched training token budgets and with additional memory parameters, experiments across attention configurations show improved language modeling and average downstream performance.
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2609.28399 [cs.LG] |
| (or arXiv:2609.28399v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2609.28399 arXiv-issued DOI via DataCite |
Submission history
From: Jiale Kang [view email]
[v1]
Wed, 23 Sep 2026 17:04:30 UTC (286 KB)
[v2]
Wed, 14 Oct 2026 16:08:41 UTC (892 KB)
来源：arXiv：cs.LG · arxiv.org