Snelle-gewicht Product Sleutelgeheugen

Samenvatting

Sequence modeling-lagen in moderne taalmodellen worden doorgaans geconfronteerd met een afweging tussen opslagcapaciteit en computationele efficiëntie. Waar Softmax-attentie onbeperkte opslag biedt tegen prohibitief hoge kwadratische kosten, bieden lineaire varianten efficiëntie maar lijden onder een beperkte, vaste grootte van de opslag. Wij stellen Fast-weight Product Key Memory (FwPKM) voor, een nieuwe architectuur die deze spanning oplost door het spaarse Product Key Memory (PKM) te transformeren van een statische module naar een dynamisch, "fast-weight" episodisch geheugen. In tegenstelling tot PKM werkt FwPKM zijn parameters dynamisch bij, zowel tijdens training als inferentie, via lokale gradient descent op chunkniveau. Hierdoor kan het model snel nieuwe key-value paren uit invoersequenties onthouden en ophalen. Experimenten tonen aan dat FwPKM functioneert als een effectief episodisch geheugen dat het semantisch geheugen van standaardmodules aanvult, wat leidt tot significante perplexiteitsreducties op datasets met lange contexten. Opmerkelijk is dat FwPKM in Needle in a Haystack-evaluaties generaliseert naar contexten van 128K tokens, ondanks dat het enkel getraind is op sequenties van 4K tokens.

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.