稀疏Delta记忆：通过稀疏性扩展线性RNN的状态

摘要

线性注意力模型允许固定状态大小和每token固定计算量。然而，由于状态大小的限制，线性注意力模型在长上下文召回方面落后于基于softmax注意力的Transformer架构。增加线性注意力的状态大小可提升召回性能，但代价是更高的计算量。本文提出稀疏增量记忆（SDM）架构，该架构利用稀疏寻址方案，将门控线性循环神经网络的隐状态扩展至数个数量级更高的容量。SDM通过用对大显存进行稀疏读写操作替代密集的键值外积，扩展了门控DeltaNet架构。我们证明，在等计算量（isoFLOP）约束和相同参数数量的条件下，更高的状态记忆容量能显著提升上下文学习和长上下文检索任务的性能。此外，通过学习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.