선형 어텐션(Linear Attention)을 위한 해마: 순환 상태(Recurrent State)가 망각하는 것에 대한 정확한 메모리
요약
선형 어텐션 모델의 고질적인 문제인 메모리 손실을 해결하기 위해 해마(Hippocampal) 구조를 도입한 HOLA를 제안합니다. HOLA는 압축된 순환 상태와 정확한 KV 캐시를 결합하여, 긴 문맥에서도 정보 망각 없이 높은 검색 성능을 유지합니다.
핵심 포인트
- 선형 어텐션의 정보 압축 과정에서 발생하는 메모리 손실 문제 해결
- 해마 구조를 모방하여 압축 상태와 정확한 KV 캐시를 병행 사용
- Wikitext 및 LAMBADA 벤치마크에서 기존 모델 대비 성능 개선
- 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)를 형성합니다. 즉, 상태는 선형적으로 압축 가능한 구조를 모델링하는 반면, 캐시는 해당 상태를 통해 강제로 처리되지 않아야 할 연관 관계를 저장합니다. 캐시는 학습된 제거 모듈(eviction module) 없이 기록되며, 상태에 실제로 반영된 예측 잔차(prediction residual)인 $\beta \cdot |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 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.AI의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기