Fast-weight Product Key Memory
计算与语言
2026-02-24 v2 人工智能
摘要
现代语言模型中的序列建模层通常面临存储容量与计算效率之间的权衡。虽然softmax 注意力提供无限制的存储，但计算成本高达二次；线性变体更高效，但受限于固定大小的存储。我们引入Fast-weight Product Key Memory（FwPKM），一种稀疏快速权重记忆层，以解决这一矛盾。FwPKM 在训练和推理期间使用块级梯度下降对局部记忆重写目标进行稀疏激活参数的更新。这对激活槽位执行 Test-Time Training（TTT）风格的梯度更新，使其能够在保持每个 token 计算量低且固定的同时，快速记忆和检索大量新的 key-value 关联。实验表明，FwPKM 作为一种有效的情景记忆，能够补充标准模块的语义记忆，在长上下文数据集上显著降低困惑率。值得注意的是，在 Needle-in-a-Haystack 评估中，FwPKM 能够推广到128K token 的上下文，尽管仅在4K token 序列上进行训练。
引用
@article{arxiv.2601.00671,
title = {Fast-weight Product Key Memory},
author = {Tianyu Zhao and Llion Jones},
journal= {arXiv preprint arXiv:2601.00671},
year = {2026}
}