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 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.
| Subjects: | Computation and Language (cs.CL); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2601.00671 [cs.CL] |
| (or arXiv:2601.00671v2 [cs.CL] for this version) | |
| https://doi.org/10.48550/arXiv.2601.00671 arXiv-issued DOI via DataCite |
Submission history
From: Tianyu Zhao [view email]
[v1]
Fri, 2 Jan 2026 12:37:53 UTC (4,360 KB)
[v2]
Sun, 22 Feb 2026 06:35:55 UTC (8,613 KB)