희소 델타 메모리: 희소성을 통한 선형 RNN 상태 스케일링

초록

선형 어텐션 모델은 고정된 상태 크기와 토큰당 고정된 계산량을 제공한다. 그러나 제한된 상태 크기로 인해 선형 어텐션 모델은 소프트맥스 어텐션 기반 트랜스포머 아키텍처에 비해 장기 컨텍스트 회상 성능이 떨어진다. 선형 어텐션의 상태 크기를 늘리면 회상 성능이 향상되지만, 더 높은 FLOPs 비용이 발생한다. 본 연구에서는 희소 주소 지정 방식을 사용하여 게이티드 선형 RNN의 은닉 상태 용량을 수십 배 더 높은 수준으로 확장하는 아키텍처인 희소 델타 메모리(SDM)를 소개한다. SDM은 밀집 키-값 외적을 대규모 명시적 메모리에 대한 희소 읽기 및 쓰기로 대체하여 게이티드 델타넷 아키텍처를 확장한다. 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.