Разреженная дельта-память: масштабирование состояния линейных RNN за счет разреженности

Аннотация

Линейные модели внимания обеспечивают фиксированный размер состояния и фиксированный объем вычислений на токен. Однако из-за ограниченного размера состояния такие модели уступают архитектурам трансформеров на основе softmax-внимания при извлечении информации из длинных контекстов. Увеличение размера состояния линейных моделей внимания улучшает качество извлечения, но ценой роста числа операций с плавающей точкой (FLOPs). В данной работе мы представляем разреженную дельта-память (Sparse Delta Memory, SDM) — архитектуру, которая увеличивает емкость скрытого состояния линейных рекуррентных нейронных сетей с вентильными механизмами на порядки за счет использования разреженной адресации. SDM расширяет архитектуру Gated DeltaNet, заменяя плотное внешнее произведение ключ-значение на разреженные чтение и запись в большую явную память. Мы показываем, что при одинаковых ограничениях по FLOPs и идентичном числе параметров более высокая емкость состояния памяти значительно улучшает производительность на задачах обучения в контексте и извлечения информации из длинных контекстов. Кроме того, обучая начальное состояние памяти SDM и, следовательно, используя ее как параметрическую память, мы демонстрируем, что модель дополнительно улучшает результаты на широком спектре задач на общие знания и рассуждение.

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.