Key-Value Means: Transformer's Expandable Block-Recurrent Compressed Memory

긴 컨텍스트에서 Transformer의 메모리 및 계산 복잡도는 컨텐츠 길이에 비례해 증가한다. KVM은 fixed-size 또는 growable 상태를 활용해 메모리 사용을 부분적으로만 증가시키고, sublinear 메모리 성장과 subquadratic prefill 타임을 달성한다. 또한 block-recurrent attention와 상태 확장을 하나의 Softmax 주의로 처리하여 학습과 추론의 효율성을 높인다.

전통적 Self-Attention은 컨텍스트 길이에 따라 계산 복잡도(N^2)가 커진다. LRNN은 일정 메모리로 빠른 처리지만 긴 컨텍스트 기억이 약하다. KVM은 BSWA를 통해 입력의 창(window) 정보를 처리하고 Overflow를 동적으로 상태로 압축한다. 이때 상태는 expandable하며, 노름 보존을 위한 JIT 정규화와 상태 슬롯의 새 토큰 추가 정책으로 초기 컨텍스트를 넘어서는 정보 보존을 가능하게 한다. 결과적으로 고정 크기 상태의 RNN과 일반 Transformer 사이의 연산-메모리 트레이드오프를 매끄럽게 조정할 수 있다.

KVM 주의는 sink 토큰, BSWA 윈도우, 그리고 상태로 구성된 softmax-attention으로 정의된다. 상태는 chunk 끝에 overflow 토큰 중 일부를 추가(Ae 인덱스)하고, 남은 토큰은 병합된다. 병합은 각 overflow 토큰의 키-값을 현재 상태의 가장 유사한 키에 비례적으로 더해 상태 키와 값에 반영한다. 읽기(readout)는 상태와 BSWA 윈도우를 연결한 a_q, a_k, a_v를 계산하고, 이를 sdpa를 이용해 주의 출력으로 환원한다. 상태의 초기화는 첫 번째 청크의 k¯, v로 구성되며, ρi로 노름을 고정한다. nappend는 e 단계의 오버플로우 토큰에 대해 비중이 높은 토큰을 신규 슬롯에 추가하는 양을 결정한다. RoPE의 부분 적용, Rotational subspace의 제로화, 그리고 토큰의 분리된 가중치 업데이트 등은 안정적 학습과 추론을 돕는다.

메인 벤치마크는 120M / 350M 파라미터에서 Prolong 데이터셋의 8k 컨텍스트를 대상으로 평가되었다. 120M 모델에서 256 고정 상태(KVM 256)는 NIAH-S1에서 4K/8K/16K/32K 토큰에 대해 각각 99.4/97.8/98.4/98.4의 정확도(또는 성능) 수치를 보였고, LB/RULER에서 우수한 순위를 기록했다. KVM sqrt는 비GPTAlpha 비허용 모델 중 최상위 성능을 보였으며, 350M 규모에서도 KVM sqrt가 Long-context에서 강한 성능을 보여주었다. 반면 256 상태를 고정한 경우 NL 길이가 매우 긴 NIAH-S2, S3에서 어려움을 보이기도 했다. ablation 실험은 sink의 제거가 장기 컨텍스트에서 성능 저하를 크게 유발함을 확인했고, v-length normalization, merge gate 제거 등의 변화도 성능에 영향을 주었다. NoPE와 HalfRoPE 변형의 길이 extrapolation 특성 차이가 나타났으며, HalfRoPE가 NoPE보다 긴 문맥에서 더 강한 일반화 성능을 보였다.

전체 아키텍처는 GPTAlpha-2 기반 백본에 KVM 주의를 결합하는 구성이다. KVM 주의는 sink 토큰, BSWA 윈도우, 확장 가능한 상태로 구성되며, 상태 업데이트는 overflow 토큰의 가장 유사한 상태 토큰에 대한 가중 합으로 수행된다. 상태 키 s_k와 상태 값 s_v의 초기화는 첫 청크의 k¯, v를 기반으로 하고, ρi는 슬롯의 readout radius로 남는다. readout은 s_k, s_v를 각각 LN으로 정규화하고, q를 BSWA 윈도우 키(k)와 결합한 A들을 softmax로 처리한다. 이후 y_t를 얻고 잔차에 더한다. KVM의 상태 업데이트는 끝난 chunk에서 overflow block Ωe를 확인하고, nappend에 따라 신규 토큰을 추가하거나 병합한다. nappend는 M⋆(e)−m, |Ωe|의 최솟값으로 결정된다. 병합 시 first S=1의 sink 토큰은 대상에서 제외된다. merge 연산은 각 토큰 j에 대해 πe(j) = argmax i≥S k˘j·LN(s_K+,i)⊤을 사용해 i를 선택하고, 선택된 i에 대해 s_K,new,i와 s_V,new,i를 증가시킨다. 이때 상태 키의 노름은 초기 노름으로 유지되고, readout 시점에만 해당 노름으로 정규화한다. RoPE의 부분 제로화 및 partial RoPE 재병합은 별도의 re-merging 없이 BSWA와 상태 간의 분리된 처리를 가능하게 한다.

KVM-256은 매우 긴 컨텍스트에서 성능이 저하되는 경향이 있다(NIAH-S2, S3에서의 실험 결과). NoPE와 HalfRoPE 간 extrapolation 특성 차이가 존재하며, RoPE 분해/제로화의 표현력 손실 가능성이 있다. 또한 학습/추론 시간은 상태 확장 스케줄, nappend 정책 등에 따라 달라지며, 특정 하이퍼파라미터 의존성이 존재한다.