Memoria Delta Dispersa: Escalamiento del Estado de las RNN Lineales mediante Dispersión

Resumen

Los modelos de atención lineal permiten un tamaño de estado fijo y una cantidad fija de cómputo por token. Sin embargo, debido a su limitado tamaño de estado, los modelos de atención lineal se quedan atrás en la recuperación de contexto largo en comparación con las arquitecturas de transformadores basadas en atención softmax. Aumentar el tamaño del estado de la atención lineal mejora el rendimiento de recuperación, pero a costa de un mayor número de FLOPs. En este trabajo, presentamos la Memoria Delta Dispersa (SDM), una arquitectura que escala el estado oculto de las RNN lineales con compuertas a una capacidad órdenes de magnitud mayor mediante un esquema de direccionamiento disperso. SDM extiende la arquitectura DeltaNet con compuertas al reemplazar el producto externo denso clave-valor con lecturas y escrituras dispersas en una memoria explícita grande. Mostramos que, bajo una restricción isoFLOP y con un número idéntico de parámetros, una mayor capacidad de memoria de estado mejora significativamente el rendimiento en tareas de aprendizaje en contexto y recuperación de contexto largo. Además, al aprender el estado inicial de la memoria SDM y, por lo tanto, usarla como una memoria paramétrica, mostramos que el modelo mejora aún más en una amplia gama de tareas de conocimiento común y razonamiento.

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.