패스트 웨이트 프로덕트 키 메모리

초록

현대 언어 모델의 시퀀스 모델링 계층은 일반적으로 저장 용량과 계산 효율성 사이의 트레이드오프에 직면합니다. Softmax 어텐션은 엄청난 2차 비용으로 무제한 저장 능력을 제공하는 반면, 선형 변형은 효율성을 제공하지만 제한된 고정 크기 저장 용량이라는 한계를 지닙니다. 본 연구에서는 이러한 긴장 관계를 해결하는 새로운 아키텍처인 Fast-weight Product Key Memory(FwPKM)를 제안합니다. FwPKM은 희소 Product Key Memory(PKM)를 정적 모듈에서 동적인 "패스트 웨이트(fast-weight)" 에피소드 메모리로 변환합니다. PKM과 달리, FwPKM은 훈련 및 추론 시점에 로컬 청크 수준 경사 하강법을 통해 매개변수를 동적으로 업데이트하여, 모델이 입력 시퀀스에서 새로운 키-값 쌍을 빠르게 기억하고 검색할 수 있도록 합니다. 실험 결과, FwPKM은 표준 모듈의 의미 메모리를 보완하는 효과적인 에피소드 메모리로 기능하여 장문 컨텍스트 데이터셋에서 현저한 복잡도(perplexity) 감소를 가져옵니다. 특히, Needle in a Haystack 평가에서 FwPKM은 4K-토큰 시퀀스로만 훈련되었음에도 불구하고 128K-토큰 컨텍스트로 일반화되었습니다.

English

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 costs, linear variants provide efficiency but suffer from limited, fixed-size storage. We propose Fast-weight Product Key Memory (FwPKM), a novel architecture that resolves this tension by transforming the sparse Product Key Memory (PKM) from a static module into a dynamic, "fast-weight" episodic memory. Unlike PKM, FwPKM updates its parameters dynamically at both training and inference time via local chunk-level gradient descent, allowing the model to rapidly memorize and retrieve new key-value pairs from input sequences. Experiments reveal 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.