稀疏Delta記憶：透過稀疏性擴展線性RNN的狀態

摘要

線性注意力模型具備固定的狀態大小及每個詞元的固定計算量。然而，由於其有限的狀態容量，這類模型在長上下文回憶能力上不及基於 Softmax 注意力的變壓器架構。增加線性注意力的狀態大小能改善回憶表現，但代價是更高的浮點運算數（FLOPs）。在本研究中，我們提出稀疏增量記憶（Sparse Delta Memory，SDM），這是一種利用稀疏定址機制將門控線性遞迴神經網路的隱藏狀態擴充數個數量級容量的架構。SDM 擴展了 Gated DeltaNet 架構，以對大型明確記憶體的稀疏讀寫取代密集的鍵值外積。我們證明，在等 FLOP 限制下且參數量相同的情況下，更高的狀態記憶容量能顯著提升上下文學習與長上下文檢索任務的表現。此外，透過學習 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.