Schaars deltageheugen: het opschalen van de toestand van lineaire RNN's door schaarste

Samenvatting

Lineaire aandachtsmodellen maken een vaste toestandsgrootte en een vaste hoeveelheid rekenwerk per token mogelijk. Door hun beperkte toestandsgrootte presteren lineaire aandachtsmodellen echter slechter bij het oproepen van lange contexten dan op softmax-aandacht gebaseerde transformatorarchitecturen. Het vergroten van de toestandsgrootte van lineaire aandacht verbetert de herinneringsprestaties, maar ten koste van hogere FLOPs. In dit werk introduceren we Sparse Delta Memory (SDM), een architectuur die de verborgen toestand van gegate lineaire RNN's opschaalt naar een orde van grootte hogere capaciteit met behulp van een schaars adresseringsschema. SDM breidt de Gated DeltaNet-architectuur uit door het dichte sleutel-waarde buitenproduct te vervangen door schaarse lees- en schrijfbewerkingen naar een groot expliciet geheugen. We tonen aan dat, onder een isoFLOP-beperking en met een identiek aantal parameters, een hogere geheugentoestandcapaciteit de prestaties op in-context leren en taken voor het ophalen van lange contexten aanzienlijk verbetert. Bovendien, door de begintoestand van het SDM-geheugen te leren en het dus als parametrisch geheugen te gebruiken, laten we zien dat het model verder verbetert op een breed scala aan algemene kennis- en redeneertaken.

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.