Schnellgewicht-Produktschlüssel-Speicher

papers.abstract

Sequenzmodellierungsschichten in modernen Sprachmodellen stehen typischerweise vor einem Zielkonflikt zwischen Speicherkapazität und Recheneffizienz. Während Softmax-Attention unbegrenzte Speicherkapazität bei prohibitiv quadratischen Kosten bietet, liefern lineare Varianten Effizienz, leiden aber unter begrenztem, festgrößigem Speicher. Wir schlagen Fast-weight Product Key Memory (FwPKM) vor, eine neuartige Architektur, die diesen Konflikt löst, indem sie den spärlichen Product Key Memory (PKM) von einem statischen Modul in einen dynamischen, "Fast-weight"-episodischen Speicher transformiert. Im Gegensatz zu PKM aktualisiert FwPKM seine Parameter dynamisch sowohl während des Trainings als auch während des Inferenzvorgangs durch lokalen Gradientenabstieg auf Chunk-Ebene, was dem Modell ermöglicht, schnell neue Schlüssel-Wert-Paare aus Eingabesequenzen zu memorieren und abzurufen. Experimente zeigen, dass FwPKM als effektiver episodischer Speicher fungiert, der das semantische Gedächtnis standardmäßiger Module ergänzt und zu signifikanten Perplexitätsreduktionen auf Langkontext-Datensätzen führt. Bemerkenswerterweise generalisiert FwPKM in "Needle in a Haystack"-Evaluierungen auf Kontexte mit 128K Tokens, obwohl es nur auf 4K-Token-Sequenzen trainiert wurde.

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 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.