Sparse Delta Memory: Skalierung des Zustands linearer RNNs durch Sparsität

Zusammenfassung

Lineare Aufmerksamkeitsmodelle ermöglichen eine feste Zustandsgröße und einen festen Rechenaufwand pro Token. Aufgrund ihrer begrenzten Zustandsgröße liegen lineare Aufmerksamkeitsmodelle jedoch bei der Erinnerung langer Kontexte hinter Transformer-Architekturen zurück, die auf Softmax-Aufmerksamkeit basieren. Eine Vergrößerung der Zustandsgröße linearer Aufmerksamkeit verbessert die Erinnerungsleistung, geht jedoch mit höheren FLOPs einher. In dieser Arbeit stellen wir Sparse Delta Memory (SDM) vor, eine Architektur, die den verborgenen Zustand von gegateten linearen RNNs mithilfe eines spärlichen Adressierungsschemas um Größenordnungen auf eine höhere Kapazität skaliert. SDM erweitert die Gated DeltaNet-Architektur, indem es das dichte Schlüssel-Wert-äußere Produkt durch spärliche Lese- und Schreibvorgänge in einen großen expliziten Speicher ersetzt. Wir zeigen, dass unter einer isoFLOP-Beschränkung und bei identischer Parameteranzahl eine höhere Speicherkapazität des Zustands die Leistung bei Aufgaben des In-Context-Lernens und des Abrufens langer Kontexte signifikant verbessert. Darüber hinaus zeigen wir, dass das Modell durch das Erlernen des Anfangszustands des SDM-Speichers und dessen Nutzung als parametrischen Speicher eine breite Palette von Aufgaben zu Allgemeinwissen und logischem Denken weiter verbessert.

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.