Key-Value Means

본 논문은 표준 연산과 최소한의 추가 파라미터를 사용하면서도 유연한 상태 성장과 2 차 미만의 프리필 시간을 가능하게 하는 효율적이고 청크 병렬화가 가능한 학습을 통해 트랜스포머와 선형 RNN 의 장점을 통합하는 새로운 블록 재귀 메커니즘인 키-값 평균 (KVM) 을 소개합니다.

핵심

상상해 보세요. 매우 긴 책을 읽으려는데, 읽는 동안 이야기 내용을 머릿속에 간직할 수 있는 '작업 기억 (working memory)'의 양이 제한되어 있다고 가정해 봅시다.

현재 AI 의 문제점
현재의 AI 모델 (Transformer) 은 책에서 읽은 단어 하나하나까지 모두 기억하려는 학생처럼 행동합니다.

• 장점: 전체 이야기가 눈앞에 있기 때문에 매우 정확한 성능을 냅니다.
• 단점: 책이 길어질수록 '작업 기억'이 거대하게 팽창합니다. 100 페이지짜리 책을 읽는 데는 아주 적은 노력이 들지만, 1,000 페이지짜리 책을 읽는 데는 막대한 시간과 에너지가 필요합니다. 이는 걸을 때마다 배낭이 무거워지는 것을 들고 다니려는 것과 같습니다.

순환형 (RNN 스타일) 모델의 문제점
RNN 스타일 모델은 읽은 내용을 고정된 크기의 작은 요약본으로 유지하며, 읽을 때마다 이를 업데이트하는 다른 접근 방식을 취합니다.

• 장점: 매우 빠르고 가볍습니다. 책이 아무리 길어도 배낭은 무거워지지 않습니다.
• 단점: 이야기의 시작 부분을 잊어버릴 수 있습니다. 10 페이지에 등장하는 줄거리 요점에 대해 물어보면, 마지막 몇 페이지만 기억하고 있기 때문에 그 내용을 기억하지 못할 수 있습니다. 이는 다른 가족의 아키텍처로, 다른 장단점을 가지고 있습니다.

새로운 해결책: 키 - 밸류 평균 (Key-Value Means, KVM)
이 논문의 저자들은 **키 - 밸류 평균 (Key-Value Means, KVM)**이라는 새로운 방법을 소개합니다. KVM 은 양쪽 세계의 장점을 모두 결합한 똑똑하고 마법 같은 수첩으로 생각할 수 있습니다.

다음은 간단한 비유를 통해 작동 방식을 설명한 것입니다:

1. "슬라이딩 윈도우" (즉각적인 문맥)

책을 읽고 있는데, 마지막 몇 페이지만 선명하게 보여 주는 확대경이 있다고 상상해 보세요. 이것이 바로 "슬라이딩 윈도우"입니다. KVM 은 표준 AI 가 하듯이 가장 최근의 단어들에 완벽하게 주의를 기울입니다. 이를 통해 즉각적인 문맥을 놓치지 않도록 보장합니다.

2. "압축된 요약" (장기 기억)

그 몇 페이지를 넘겨 읽으면, 오래된 페이지들은 확대경 밖으로 미끄러져 나갑니다. 구형 AI 처럼 버리거나, 현재 AI 처럼 책 전체를 들고 다니려는 대신, KVM 은 다음과 같은 영리한 일을 합니다:

• 방금 미끄러져 나간 페이지들을 살펴봅니다.
• "이 페이지들 중 가장 중요하거나 독특한 것은 무엇인가?"라고 질문합니다.
• 그 중요한 페이지들에 대한 짧고 압축된 요약을 특별한 수첩에 적어 넣습니다.
• 만약 수첩에 이미 있는 내용과 매우 유사한 새로운 페이지가 들어오면, 기존 메모를 업데이트합니다. 만약 완전히 새롭고 놀라운 내용이라면, 수첩에 새로운 줄을 추가합니다.

3. "스마트 병합" (마법 같은 트릭)

이 논문은 "승자 독식 (Winner-Take-All)" 규칙이라는 특정 정보 병합 방식을 설명합니다.

• 새로운 정보인 물이 담긴 양동이를, 수첩인 스펀지가 있다고 상상해 보세요.
• 물을 그냥 쏟아붓는 대신, KVM 은 물과 가장 잘 맞는 스펀지의 정확한 위치를 찾아 그곳으로 흡수시킵니다.
• 또한 "Just-in-Time(필요할 때 바로)" 정규화를 사용합니다. 이는 수첩에 정보를 기록할 때는 분모 (개수) 나 합계를 나누지 않은 '원시 상태 (raw sums and counts)'로 누적해 두었다가, 실제로 그 수첩을 읽어볼 때 (쿼리 시점)에만 나누기를 수행하는 방식입니다. 즉, 새로운 항목이 추가될 때마다 매번 다시 나누어 정규화하는 번거로움을 피하고, 필요한 순간에 '지연 (lazy)' 처리된 정규화를 적용하여 정보가 왜곡되지 않도록 합니다.

왜 이것이 중요한가

• 유연한 크기: 속도를 위해 작은 수첩 (고정 크기) 을 유지하도록 KVM 에 지시할 수도 있고, 책이 길어짐에 따라 수첩이 커지도록 (확장 가능 크기) 할 수도 있습니다.
• 속도 대 기억: 중간 지점을 선택할 수 있게 해줍니다. "매우 빠르지만 기억력이 나쁜" 것과 "매우 똑똑하지만 느린" 것 사이에서 선택할 필요가 없습니다. 실시간 사용에 충분히 빠르면서도 전체 이야기를 기억할 만큼 똑똑하도록 조정할 수 있습니다.
• 커스텀 하드웨어 불필요: 실행을 위해 특수하고 비싼 컴퓨터 칩이 필요한 다른 새로운 방법들과 달리, KVM 은 일반적인 소프트웨어 연산을 사용하여 표준 컴퓨터에서 실행할 수 있습니다.

결과

저자들은 이 방법을 언어 모델 (텍스트를 읽고 쓰는 AI) 에서 테스트했습니다.

• 짧은 문맥: KVM 은 최고의 표준 AI 모델과 동등한 성능을 발휘했습니다.
• 긴 문맥: 입력이 수천 개의 토큰으로 길어졌을 때, 확장 가능한 KVM 변형은 '기억력이 나쁜' 고정 메모리 모델보다 훨씬 더 많은 세부 사항을 기억했으며, '전체 어텐션을 사용하는' 모델보다 훨씬 빨랐습니다.
• "건초더미 속의 바늘": AI 가 방대한 텍스트 속에 숨겨진 특정 문장을 찾아야 하는 테스트에서, 수첩이 성장하도록 허용된 KVM 버전은 매우 잘 수행하여 과거의 깊은 곳에서 정보를 실제로 회상할 수 있음을 입증했습니다.

요약하자면, KVM 은 AI 가 피곤해지지 않고, 시작 부분을 잊지 않으며, 무한히 무거워지는 배낭 없이 긴 책을 읽을 수 있는 새로운 방법입니다. 이는 과거에 대한 명확하고 압축된 요약을 유지하면서 현재에 대한 선명한 시야를 유지함으로써 이를 달성합니다.

기술 요약

기술 요약: 키 - 밸류 평균 (KVM)

문제 제기

트랜스포머는 현대 하드웨어에서 효율적인 학습을 제공하지만, 컨텍스트 길이에 비례하여 메모리와 출력 토큰당 추론 시간이 선형적으로 증가하는 ($O(N)$ 메모리, $O(N)$ 추론 시간) 문제를 겪습니다. 반면, 현대의 선형 RNN(LRNN) 은 토큰당 상수 메모리와 시간을 달성 ($O(1)$) 하지만 일반적으로 제한된 장컨텍스트 회상 능력으로 어려움을 겪습니다. 이 격차를 해소하려는 기존 아키텍처들은 종종 고정된 크기의 상태 (회상 능력 제한) 에 의존하거나, 런타임 옵티마이저를 사용한 복잡한 테스트 시간 학습 (속도 저하) 에 의존합니다. 커스텀 커널이나 테스트 시간 학습을 위한 복잡한 하이퍼파라미터 튜닝 없이 메모리 효율성, 속도, 장컨텍스트 회상 능력을 균형 있게 갖춘 아키텍처가 필요합니다.

방법론: 키 - 밸류 평균 (KVM)

KVM 은 블록 슬라이딩 윈도우 어텐션 (BSWA) 과 동적으로 확장 가능한 압축 상태를 통합한 새로운 블록 순환 어텐션 메커니즘입니다. 이는 단일 소프트맥스 어텐션 레이어 내에서 작동하여 전통적인 트랜스포머 (확장 가능한 컨텍스트, 청크 단위 병렬성) 와 선형 RNN 의 이점을 통합합니다.

핵심 메커니즘

1. 압축 상태를 갖춘 블록 슬라이딩 윈도우:
   KVM 은 입력을 청크 단위로 처리합니다. 최근 토큰을 위한 고정 크기 BSWA 윈도우와 오래된 토큰을 위한 별도로 주기적으로 업데이트되는 상태를 유지합니다. 토큰 블록이 BSWA 윈도우를 초과할 경우, 폐기되는 대신 상태를 업데이트하기 위해 처리됩니다.

2. 상태 압축 및 병합:
   초과 토큰은 "승자 독식 (winner-take-all)" 코사인 유사도 유사 병합 규칙을 사용하여 상태로 압축됩니다.

   • 유사도 지표: 표준 소프트맥스 대신 KVM 은 온라인 벡터 양자화에서 영감을 받은 최대 희소 업데이트 행렬을 사용하며, 각 초과 키는 단일 가장 상관관계가 높은 상태 키에 할당됩니다.
   • 적시 (Just-In-Time, JIT) 재정규화: 직교하거나 반대 방향인 벡터들을 평균화함에 따라 상태 벡터의 노드가 시간이 지남에 따라 축소되는 것을 방지하기 위해 KVM 은 JIT 정규화를 적용합니다. 상태 키는 어텐션 전에 레이어 정규화 (LayerNorm) 를 사용하여 정규화됩니다. 상태 값은 슬롯 생성 시 결정된 고정된 "읽기 반경 (readout radius, $\rho_i$)"으로 정규화되어 방향 변경을 허용하면서 값의 크기를 보존합니다.
   • 병합 게이트: 데이터에 의존하는 스칼라 게이트가 상태에 흡수되는 유입 초과 키/밸류의 양을 조절합니다.

3. 상태 확장 전략:
   고정 크기 RNN 과 달리 KVM 은 확장 가능한 상태를 지원합니다. 가장 "놀라운 (가장 중복되지 않는)" 초과 토큰은 직접 상태에 추가되고 나머지는 병합됩니다. 이를 통해 초기 컨텍스트 회상을 유지하면서 서브선형 메모리 증가 (예: $O(\sqrt{N})$) 를 가능하게 합니다.

4. 위치 인코딩 처리:
   BSWA 윈도우에서 로터리 위치 임베딩 (RoPE) 과의 호환성을 유지하면서 광범위하게 다른 위치에서 집계되는 토큰을 포함하는 압축 상태에서는 RoPE 를 피하기 위해 KVM 은 부분 RoPE 제로화를 사용합니다. 상태 키의 로터리 부분 공간은 제로화되는 반면, BSWA 윈도우는 완전한 RoPE 를 유지합니다. 이를 통해 모델은 동일한 어텐션 패스 내에서 상태에는 회전되지 않은 쿼리를, 윈도우에는 회전된 쿼리를 사용할 수 있습니다.

5. 싱크 토큰 (Sink Tokens):
   초기 상태 행의 보호된 집합 (싱크) 을 보존하여 중요한 초기 컨텍스트 정보의 저하를 방지하며, 이는 싱크 토큰이 고유한 값 크기를 갖는 문제를 해결합니다.

주요 기여

본 논문은 다음과 같은 구체적인 기여를 제시합니다:

• 새로운 블록 순환 공식화: 승자 독식 병합 규칙을 사용하여 초과 토큰을 동적으로 재정규화된 상태로 압축하는 방법으로, 별도의 압축 레이어의 필요성을 제거합니다.
• 상태 확장 전략: 새로운 초과 토큰을 상태에 추가하여 회상을 희생하지 않으면서 서브선형 메모리 증가를 가능하게 하는 메커니즘입니다.
• 적시 (JIT) 재정규화: 벡터 노드를 유지하고 평균화 중 파괴적 간섭을 방지하기 위해 상태 키와 값을 적시에 정규화하는 방식입니다.
• 부분 RoPE 공유: 상태 키의 RoPE 차원을 제로화하여 압축 및 비압축 영역 간 위치 인코딩을 공유하는 방법으로, 추가 파라미터나 복잡한 재병합 메커니즘의 필요성을 피합니다.
• 통합 아키텍처: 고정 상태 RNN 과 풀 어텐션 트랜스포머 사이를 보간하는 단일 어텐션 레이어로, 메모리 효율성과 회상 간의 연속적인 트레이드오프를 제공합니다.

실험 결과

저자들은 8k 컨텍스트 길이의 Prolong 데이터셋에서 1 억 2 천만 및 3 억 5 천만 파라미터 모델을 학습시켰습니다.

• 장컨텍스트 성능:
  • 고정 상태 KVM(256 토큰): 시퀀스 위치 손실 및 단컨텍스트 벤치마크에서 더 큰 상태의 OVQ/SWA 모델보다 우수한 성능을 보였습니다. 그러나 극단적인 길이에서 새로운 방해 요소 (NIAH-S2/S3) 가 포함된 "Needle In A Haystack"(NIAH) 테스트에서는 상태 용량이 병목 현상이 되어 어려움을 겪었습니다.
  • 확장 가능한 KVM(멱함수/포화 스케줄): "KVM sqrt" 변형 (상태 크기 $\propto \sqrt{N}$) 은 8k 학습 컨텍스트를 넘어선 외삽 영역에서 비하이브리드 GPTAlpha 모델과 경쟁하거나 능가하는 장컨텍스트 벤치마크 (RULER, LongBench, NIAH) 결과를 달성했습니다. 장거리에서 새로운 정보를 검색해야 하는 작업에서 고정 상태 KVM 및 순수 LRNN(RWKV-7) 보다 훨씬 우수한 성능을 보였습니다.
• 단컨텍스트 성능: KVM 변형들은 LAMBADA, ARC, HellaSwag 등 단컨텍스트 벤치마크에서 표준 트랜스포머와 일관된 성능을 보였으며, 이는 BSWA 윈도우가 표준 어텐션 능력을 보존함을 확인시켜 주었습니다.
• 애블레이션 연구: 값 길이 정규화를 제거할 때 가장 큰 성능 저하가 발생했습니다. 싱크 보호 및 병합 게이트를 제거하는 것도 장컨텍스트 검색 능력을 현저히 약화시켰습니다.

중요성 및 주장

본 논문은 KVM 이 고정 상태 RNN 과 풀 어텐션 트랜스포머 사이의 격차를 성공적으로 해소했다고 주장합니다.

• 효율성 대 회상: 유연한 상태 크기 선택을 제공하여 메모리 효율성과 회상 간의 트레이드오프를 사용자가 조정할 수 있게 합니다. 고정 상태의 경우 $O(N)$ 청크 순환 행동을 제공하며, 확장 가능한 상태의 경우 강력한 장컨텍스트 검색과 함께 서브선형 메모리 증가를 달성합니다.
• 구현의 단순성: KVM 은 커스텀 커널 없이 표준 연산을 사용하여 구현 가능하며, 청크 단위 병렬 학습 및 프리필을 지원합니다.
• 하이브리드 잠재력: 이 아키텍처는 LRNN 레이어와 함께 하이브리드 솔루션으로 사용되어 개선된 서브선형 메모리 증가 및 장컨텍스트 추론 능력을 보완할 수 있습니다.
• 런타임 옵티마이저 부재: 테스트 시간 학습 (TTT) 접근법과 달리 KVM 은 SGD 나 Adam 과 같은 런타임 옵티마이저 대신 간단한 상태 업데이트 규칙에 의존하여 관련 하이퍼파라미터 도전을 피합니다.

저자들은 KVM 이 단순하고 효과적인 방식으로 고정 상태 RNN 과 풀 어텐션 사이를 매끄럽게 보간할 수 있음을 입증하여, 장컨텍스트 모델링을 위한 통합 패키지를 제공한다고 결론지었습니다.