提出一种高效稀疏记忆模块，可快速存储和检索大量新信息。
Fast-weight Product Key Memory
- 通过分块梯度更新稀疏激活参数，实现低计算开销的记忆写入。
- 在128K上下文长度中仍能准确定位目标信息，训练仅用4K序列。
- 适合需要快速学习新知识的长文本任务，如问答与记忆增强模型。
现代语言模型中的序列建模层通常面临存储容量与计算效率之间的权衡。虽然Softmax注意力具有无限存储能力但计算代价为二次方，线性变体虽高效却受限于固定大小的存储。本文提出快速权重乘积键记忆（FwPKM），一种稀疏快速权重记忆层，解决了这一矛盾。FwPKM在训练和推理阶段均通过局部内存重写目标进行分块梯度下降，对激活槽位执行类测试时训练（TTT）的梯度更新，实现快速记忆与检索大量新键值对，同时保持每标记计算量低且恒定。实验表明，FwPKM作为有效的情景记忆，可补充标准模块的语义记忆，在长上下文数据集上显著降低困惑度。特别地，在针堆找针测试中，尽管仅在4K序列上训练，仍可泛化至128K上下文长度。
原文摘要 · Abstract (English)
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 cost, linear variants are more efficient but suffer from limited, fixed-size storage. We introduce Fast-weight Product Key Memory (FwPKM), a sparse fast-weight memory layer that resolves this tension. FwPKM updates sparsely activated parameters at both training and inference time using chunk-level gradient descent on a local memory-rewrite objective. This performs Test-Time Training (TTT)-style gradient updates on activated slots in a sparse memory, enabling rapid memorization and retrieval of many new key-value associations while keeping per-token compute low and fixed. Experiments show 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。