スパースデルタメモリ：スパース性による線形RNNの状態拡大

要旨

線形アテンションモデルは、固定の状態サイズとトークンあたりの固定計算量を可能にする。しかし、状態サイズが限られているため、線形アテンションモデルは、ソフトマックスアテンションに基づくトランスフォーマーアーキテクチャと比較して、長文脈の想起において劣っている。線形アテンションの状態サイズを増加させると想起性能は向上するが、より多くのFLOPsを要する。本研究では、スパースアドレッシング方式を用いてゲート付き線形RNNの隠れ状態を桁違いに高い容量に拡張するアーキテクチャであるSparse Delta Memory（SDM）を紹介する。SDMは、Gated 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.