Memória Delta Esparsa: Escalonando o Estado de RNNs Lineares através da Esparsidade

Resumo

Modelos de atenção linear permitem um tamanho de estado fixo e uma quantidade fixa de computação por token. No entanto, devido ao seu tamanho de estado limitado, os modelos de atenção linear ficam aquém na recuperação de contexto longo em comparação com arquiteturas de transformador baseadas em atenção softmax. Aumentar o tamanho do estado da atenção linear melhora o desempenho de recuperação, mas ao custo de maior número de FLOPs. Neste trabalho, introduzimos a Memória Delta Esparsa (SDM), uma arquitetura que amplia o estado oculto de RNNs lineares com portas para uma capacidade ordens de grandeza maior usando um esquema de endereçamento esparso. A SDM estende a arquitetura DeltaNet com portas, substituindo o produto externo denso chave-valor por leituras e escritas esparsas a uma grande memória explícita. Mostramos que, sob uma restrição isoFLOP e com um número idêntico de parâmetros, uma maior capacidade de memória de estado melhora significativamente o desempenho em tarefas de aprendizagem em contexto e de recuperação de contexto longo. Além disso, ao aprender o estado inicial da memória SDM e, portanto, usá-la como uma memória paramétrica, mostramos que o modelo melhora ainda mais em uma ampla gama de tarefas de conhecimento comum e raciocínio.

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.