Mémoire Delta Éparse : passage à l'échelle de l'état des RNN linéaires par la sparsité

Résumé

Les modèles d'attention linéaire permettent une taille d'état fixe et une quantité de calcul fixe par token. Cependant, en raison de leur taille d'état limitée, les modèles d'attention linéaire accusent un retard dans le rappel de contexte long par rapport aux architectures de transformeurs basées sur l'attention softmax. Augmenter la taille d'état de l'attention linéaire améliore les performances de rappel, mais au prix d'un nombre plus élevé de FLOPs. Dans ce travail, nous introduisons la Mémoire Delta Éparse (SDM), une architecture qui porte l'état caché des RNN linéaires à portes à une capacité supérieure de plusieurs ordres de grandeur en utilisant un schéma d'adressage éparse. SDM étend l'architecture Gated DeltaNet en remplaçant le produit extérieur dense clé-valeur par des lectures et écritures éparses dans une mémoire explicite de grande taille. Nous montrons que, sous une contrainte isoFLOP et avec un nombre identique de paramètres, une capacité de mémoire d'état plus élevée améliore significativement les performances sur les tâches d'apprentissage en contexte et de rappel de contexte long. De plus, en apprenant l'état initial de la mémoire SDM et en l'utilisant ainsi comme mémoire paramétrique, nous montrons que le modèle s'améliore encore sur un large éventail de tâches de connaissances générales et de raisonnement.

English

Linear attention models allow a fixed state size and a fixed amount of compute per token. However, due to their limited state size, linear attention models fall behind in long-context recall compared to softmax-attention-based transformer architectures. Increasing the state size of linear attention improves recall performance but at the cost of higher FLOPs. In this work, we introduce Sparse Delta Memory (SDM), an architecture that scales the hidden state of gated linear RNNs to orders of magnitude higher capacity using a sparse addressing scheme. SDM extends the Gated DeltaNet architecture by replacing the dense key-value outer product with sparse reads and writes to a large explicit memory. We show that, under an isoFLOP constraint and with an identical number of parameters, a higher state memory capacity significantly improves performance on in-context learning and long-context retrieval tasks. Moreover, by learning the initial state of the SDM memory and therefore using it as a parametric memory, we show that the model further improves on a wide range of common-knowledge and reasoning tasks.