선형 어텐션(Linear Attention)을 위한 해마: 순환 상태(Recurrent State)가 잊어버리는 것에 대한 정확한 메모리
요약
선형 어텐션 모델의 메모리 손실 문제를 해결하기 위해 해마(Hippocampus) 구조를 도입한 HOLA를 제안합니다. 고정된 순환 상태와 함께 정확한 KV 캐시를 병행 사용하여 긴 문맥에서도 높은 니들 리콜 성능을 유지합니다.
핵심 포인트
- 선형 어텐션의 고정된 순환 상태가 가진 메모리 압축 손실 문제 해결
- 해마 구조를 모방하여 압축 메모리와 정확한 KV 캐시를 결합한 HOLA 제안
- Wikitext perplexity를 16.1% 개선하며 Full-attention Transformer++ 성능 상회
- RULER 테스트에서 32k 토큰까지 강력한 니들 인-어-헤이스택(needle-in-a-haystack) 성능 입증
선형 어텐션(Linear-attention) 및 상태 공간 언어 모델(State-space language models)은 접두사(prefix)를 고정된 크기의 순환 상태(recurrent state)로 압축하여, 손실이 발생하는 정확한 메모리(exact memory)를 대가로 $O(1)$의 메모리를 구현합니다. 즉, 많은 키-값(key--value) 연관 관계가 경쟁할 때 초기 사실들이 덮어씌워지며 니들 리콜(needle recall) 성능이 저하됩니다. 상보적 학습 시스템(Complementary Learning Systems)에서 영감을 받아, 우리는 선형 어텐션에 해마(hippocampal) 보완 장치를 제공합니다. HOLA (Hippocampal Linear Attention)는 기존의 델타 규칙(delta-rule) 상태를 압축 메모리로 유지하면서 제한된 크기의 정확한 KV 캐시(KV cache)를 추가하여, 준매개변수적 테스트 시간 메모리(semiparametric test-time memory)를 형성합니다. 즉, 상태(state)는 선형적으로 압축 가능한 구조를 모델링하는 반면, 캐시(cache)는 해당 상태를 통해 강제로 처리되지 않아야 할 연관 관계들을 저장합니다. 캐시는 학습된 제거 모듈(eviction module) 없이 기록되며, 실제 상태에 반영된 예측 잔차(prediction residual)인 $\beta \times |e|$ 값이 큰 토큰들을 유지합니다. 이후 분리된 RMSNorm-gamma 캐시 읽기(cache read)를 통해 이러한 정확한 KV 쌍을 소프트 평균(soft averaging)이 아닌 날카로운 검색(sharp retrieval)으로 전환합니다. 15B SlimPajama 토큰으로 학습된 340M 파라미터 규모에서, HOLA는 Wikitext perplexity를 27.32에서 22.92(-16.1%)로 낮추어 full-attention Transformer++ (26.88)보다 낮은 성능을 보였으며, LAMBADA perplexity를 30.95에서 30.26으로 개선했습니다. 또한 최고의 선형 인컨텍스트 검색(linear in-context retrieval)을 달성하였으며, RULER needle-in-a-haystack recall 테스트에서 32k 토큰(학습 길이의 16배)까지 GDN이나 매칭된 HOLA+recency 캐시보다 훨씬 더 강력한 견고함을 유지합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 r/LocalLLaMA의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기