빠른 가중치 곱셈 키 메모리

Fast-weight Product Key Memory

현대 언어 모델의 시퀀스 모델링 레이어는 일반적으로 저장 용량과 계산 효율성 간의 균형을 맞추는 어려움에 직면합니다. 소프트맥스 어텐션은 무한한 저장 용량을 제공하지만, 계산 비용이 매우 높고 (이차 함수), 선형 변형은 효율적이지만 저장 용량이 제한적이고 고정되어 있습니다. 본 연구에서는 이러한 긴장을 해소하는 희소 가중치 곱셈 키 메모리 (Fast-weight Product Key Memory, FwPKM) 레이어를 제안합니다. FwPKM은 로컬 메모리 재작성 목표에 기반한 청크 단위 경사 하강법을 사용하여 학습 및 추론 시 희소하게 활성화된 파라미터를 업데이트합니다. 이를 통해 테스트 시간 학습 (Test-Time Training, TTT) 스타일의 경사 업데이트를 수행하여 희소 메모리의 활성화된 슬롯에서 많은 새로운 키-값 연관성을 빠르게 기억하고 검색할 수 있으며, 동시에 토큰당 계산량을 낮게 유지합니다. 실험 결과, FwPKM은 표준 모듈의 의미 기억을 보완하는 효과적인 에피소드 기억 역할을 하며, 긴 컨텍스트 데이터 세트에서 상당한 퍼플렉시티 감소를 가져옵니다. 특히, '바늘 찾기' 평가에서 FwPKM은 4K 토큰 시퀀스로 학습되었음에도 불구하고 128K 토큰 컨텍스트에 대한 일반화 성능을 보입니다.

Sequence modeling layers in modern language models typically face a trade-off between storage capacity and computational efficiency. While softmax attention offers unbounded storage at prohibitive quadratic cost, linear variants are more efficient but suffer from limited, fixed-size storage. We introduce Fast-weight Product Key Memory (FwPKM), a sparse fast-weight memory layer that resolves this tension. FwPKM updates sparsely activated parameters at both training and inference time using chunk-level gradient descent on a local memory-rewrite objective. This performs Test-Time Training (TTT)-style gradient updates on activated slots in a sparse memory, enabling rapid memorization and retrieval of many new key-value associations while keeping per-token compute low and fixed. Experiments show that FwPKM functions as an effective episodic memory that complements the semantic memory of standard modules, yielding significant perplexity reductions on long-context datasets. Notably, in Needle-in-a-Haystack evaluations, FwPKM generalizes to 128K-token contexts despite being trained on only 4K-token sequences.