Abstract
Sequence modeling layers in modern language models typically face a trade-off between storage capacity and computational efficiency. While Softmax attention offers unbounded storage at prohibitive quadratic costs, linear variants provide efficiency but suffer from limited, fixed-size storage. We propose Fast-weight Product Key Memory (FwPKM), a novel architecture that resolves this tension by transforming the sparse Product Key Memory (PKM) from a static module into a dynamic, "fast-weight" episodic memory. Unlike PKM, FwPKM updates its parameters dynamically at both training and inference time via local chunk-level gradient descent, allowing the model to rapidly memorize and retrieve new key-value pairs from input sequences. Experiments reveal that FwPKM functions as an effective episodic memory that complements the semantic memory of standard modules, yielding significant perplexity reductions on long-context datasets. Notably, in Needle in a Haystack evaluations, FwPKM generalizes to 128K-token contexts despite being trained on only 4K-token sequences.
Abstract (translated)
现代语言模型中的序列建模层通常在存储容量和计算效率之间存在权衡。虽然Softmax注意力提供了无界的存储能力，但其成本是高昂的二次复杂度；而线性变体虽提供高效性却受到有限且固定大小的存储限制。我们提出了一种新的架构——快速权重产品键记忆（FwPKM），它通过将稀疏的产品键记忆（PKM）从静态模块转化为动态的“快权重”情景记忆，解决了这一矛盾。不同于PKM，FwPKM在训练和推理过程中都能通过局部片段级别的梯度下降动态更新其参数，使模型能够快速存储并检索输入序列中的新键值对。 实验表明，FwPKM作为一种有效的情景记忆模块，在补充标准模块的语义记忆方面表现出色，并且在长上下文数据集上实现了显著的困惑度降低。值得注意的是，在Needle in a Haystack（大范围搜索）评估中，尽管仅基于4K令牌长度序列进行训练，FwPKM仍然能够泛化到128K令牌长度的情境中。
URL
https://arxiv.org/abs/2601.00671
