Mémoire à Clés de Produit à Poids Rapides

papers.abstract

Les couches de modélisation séquentielle dans les modèles de langage modernes sont généralement confrontées à un compromis entre capacité de stockage et efficacité computationnelle. Alors que l'attention Softmax offre une capacité de stockage illimitée au prix de coûts quadratiques prohibitifs, les variantes linéaires fournissent de l'efficacité mais souffrent d'une capacité de stockage limitée et de taille fixe. Nous proposons la Mémoire à Clés Produit à Poids Rapides (FwPKM), une nouvelle architecture qui résout cette tension en transformant la Mémoire à Clés Produit (PKM) sparse d'un module statique en une mémoire épisodique dynamique à "poids rapides". Contrairement à la PKM, la FwPKM met à jour ses paramètres dynamiquement, à la fois pendant l'entraînement et l'inférence, via une descente de gradient locale au niveau des segments, permettant au modèle de mémoriser et de récupérer rapidement de nouvelles paires clé-valeur à partir des séquences d'entrée. Les expériences révèlent que la FwPKM fonctionne comme une mémoire épisodique efficace qui complète la mémoire sémantique des modules standards, entraînant des réductions significatives de perplexité sur des ensembles de données à contexte long. Notamment, dans les évaluations "Aiguille dans une botte de foin", la FwPKM généralise à des contextes de 128K tokens bien qu'elle n'ait été entraînée que sur des séquences de 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.