기술 요약: 온라인 벡터 양자화 어텐션

1. 문제 제기

대규모 언어 모델 (LLM) 의 표준 시퀀스 믹싱 레이어는 계산 효율성과 장문맥 성능 사이에서 근본적인 트레이드오프에 직면해 있습니다.

• 자기 어텐션 (Self-Attention): 장문맥 작업 (인컨텍스트 학습, 추론, 검색) 에 매우 효과적이지만, 이차적인 계산 비용 ($O(T^2)$) 과 선형적으로 증가하는 메모리 비용 ($O(T)$) 을 수반하여 매우 긴 시퀀스에는 비용이 많이 듭니다.
• 선형 어텐션 및 상태 공간 모델 (SSM): 이들은 선형 계산 비용 ($O(T)$) 과 상수 메모리 ($O(1)$) 를 제공합니다. 그러나 제한된 메모리 용량으로 인해 정밀한 장거리 정보 처리가 필요한 작업에서는 상당한 성능 저하를 겪습니다.
• 벡터 양자화 (VQ) 어텐션: 키를 고정된 중심점 (centroid) 사전으로 양자화하여 선형 계산과 상수 메모리를 달성하는 이전 접근 방식 (Lingle, 2023) 입니다. 그러나 예비 테스트 결과, 원래의 VQ-어텐션은 장문맥 작업에서 어려움을 겪는 것으로 나타났습니다. 주요 병목 현상은 **정적이며 사전 훈련된 키 사전 ($D_k$)**으로 확인되었는데, 이는 순전파 (forward pass) 동안 키를 효과적으로 모델링하지 못해 성능을 저하시키는 양자화 오류를 초래했습니다. 사전 크기를 늘리는 것은 체감 효과를 보였습니다.

2. 방법론: 온라인 벡터 양자화 (OVQ) 어텐션

저자들은 키 사전 ($D_k$) 과 값 사전 ($D_v$) 을 모두 순전파 동안 온라인으로 학습하는 새로운 시퀀스 믹싱 레이어인 온라인 벡터 양자화 (OVQ) 어텐션을 제안합니다.

이론적 기반: 가우시안 혼합 회귀 (GMR)

이 논문은 OVQ-어텐션과 가우시안 혼합 회귀 (GMR) 사이의 이론적 연결을 확립합니다.

• GKR 로서의 자기 어텐션: 특정 가정 하에서 표준 자기 어텐션은 가우시안 커널 회귀 (GKR) 와 동일합니다.
• GMR 로서의 VQ-어텐션: 저자들은 VQ-어텐션 (및 확장하여 OVQ) 이 GMR 과 더 밀접하게 관련되어 있음을 보여줍니다. GMR 에서 가우시안 혼합 모델 (GMM) 은 입력 - 출력 쌍 (키와 값) 데이터셋에 적합됩니다.
• 연결: 단위 노름 쿼리/키 및 등방성 공분산 가정 하에서 GMR 의 예측 공식은 선형 시간 VQ-어텐션 방정식으로 단순화됩니다:
  $$O_{ovq} = \text{softmax}(\beta Q D_k^\top + \log(c)) D_v$$
  여기서 $c$는 각 중심점에 할당된 키의 수를 나타냅니다.
• 함의: 사전 훈련 대 온라인 방식으로 $D_k$와 $D_v$를 별도로 학습하는 표준 VQ-어텐션과 달리, GMR 은 동일한 데이터셋에 적합하도록 이를 동시에 학습합니다. 이는 동일한 조건부 분포 $P(V|K)$를 정확하게 모델링하기 위해 두 사전 모두를 온라인으로 학습하는 OVQ 접근 방식을 동기화합니다.

온라인 학습 공식화

OVQ-어텐션은 청크 병렬 공식을 사용하는 테스트 시간 훈련 프레임워크 내에서 작동합니다:

1. 예측: 각 쿼리 청크에 대해 모델은 사전 ($D_k, D_v$) 과 카운트 벡터 ($c$) 의 현재 상태를 사용하여 어텐션을 계산합니다. 원래 VQ-어텐션과 달리, 현재 청크의 키와 값은 어텐션 계산에 사용되기 전에 양자화되지 않으며, 사전에 직접 연결됩니다.
2. 상태 업데이트: 예측을 생성한 후, 모델은 현재 청크의 키와 값을 사용하여 사전들을 업데이트합니다.
   • 사전 성장: 중심점의 수 ($N_t$) 는 디리클레 과정에 영감을 받은 포화 성장 함수를 사용하여 최대 $N$으로 점근적으로 성장합니다. 이를 통해 모델은 메모리 용량을 동적으로 적응시킬 수 있습니다.
   • 중심점 초기화: 새로운 중심점은 기존 중심점과의 가장 작은 내적을 가진 데이터 포인트를 선택하여 다양성을 보장하는 "최대 확산" 방식 (k-means++ 와 유사) 으로 초기화됩니다.
   • 업데이트 규칙: 기존 중심점은 배치 K-평균 업데이트의 온라인 버전을 통해 업데이트됩니다. 저자들은 가정 하에서 GMM 에 대한 첫 번째 배치 EM 업데이트가 음수 로그 가능도 (NLL) 를 최소화하는 뉴턴 업데이트와 동등함을 증명하여, 손실 최소화를 위한 2 차 근사를 제공합니다.

계산적 특성

• 메모리: OVQ 는 전체 시퀀스가 아닌 사전과 카운트에 의해 상태가 정의되므로 시퀀스 길이에 관계없이 상수 메모리($O(N)$) 를 유지합니다.
• 계산: 이 레이어는 선형 계산 비용 ($O(T)$) 을 수반합니다. 업데이트 연산은 희소합니다. 청크 내의 각 토큰은 사전의 단일 행만 업데이트합니다. 이는 상태 업데이트가 상태 차원에 비례하여 확장되는 선형 어텐션과 달리, 상태 크기 ($N$) 를 메모리 풋프린트를 증가시키지 않고 크게 늘릴 수 있게 합니다.

3. 주요 기여

1. OVQ-어텐션 제안: 이전 VQ-어텐션의 정적 사전 제한을 극복하기 위해 키와 값 사전 모두를 온라인으로 학습하는 새로운 시퀀스 믹싱 레이어.
2. 이론적 프레임워크: OVQ-어텐션과 가우시안 혼합 회귀를 연결하는 새로운 유도 과정으로, 사전 성장 및 희소 업데이트 규칙을 포함한 온라인 학습 알고리즘에 대한 원칙적인 기반 제공.
3. 실증적 검증: 합성 작업 (인컨텍스트 회상, 인컨텍스트 학습) 과 장문맥 언어 모델링 (PG19) 에 대한 포괄적인 테스트.

4. 실험 결과

저자들은 전체 자기 어텐션, 선형 어텐션, SSM(예: Mamba, Gated Delta Net) 을 포함한 강력한 베이스라인에 대해 OVQ-어텐션을 평가했습니다.

• 합성 인컨텍스트 회상 (ICR):
  • OVQ-어텐션은 "Basic ICR" 및 "Positional ICR" 작업에서 원래 VQ-어텐션 및 선형 어텐션 베이스라인을 크게 능가했습니다.
  • 전체 자기 어텐션 베이스라인이 사용하는 상태 크기의 10-25% 만 사용함에도 불구하고 64k 시퀀스 길이까지 전체 자기 어텐션 베이스라인과 일치하거나 약간 벗어난 성능을 보였습니다.
  • 테스트 시 더 큰 사전 크기로 성능이 향상되어 효과적인 길이 외삽을 입증했습니다.
• 합성 인컨텍스트 학습 (ICL):
  • 여러 함수가 포함된 장문맥 선형 회귀 작업에서 OVQ-어텐션은 최대 128 개의 서로 다른 함수가 있는 시나리오에서 모든 함수를 성공적으로 학습하여 전체 자기 어텐션의 성능과 일치했습니다.
  • 원래 VQ-어텐션 및 선형 어텐션 베이스라인은 이러한 어려운 장문맥 시나리오에서 단일 함수조차 학습하지 못했습니다.
• 장문맥 언어 모델링 (PG19):
  • 슬라이딩 윈도우 레이어와 OVQ 를 교차 배치한 하이브리드 아키텍처는 순수 슬라이딩 윈도우 또는 Gated Delta Net 모델보다 성능을 획기적으로 향상시켰습니다.
  • OVQ-어텐션은 16k 문맥 길이에서 PG19 데이터셋에서 표준 자기 어텐션 베이스라인과 일치하거나 약간 상회하는 성능을 보였습니다.
• 단문맥 벤치마크:
  • 표준 단문맥 벤치마크 (PIQA, HellaSwag 등) 에서 OVQ-어텐션은 표준 자기 어텐션과 유사하게 작동하여 압축이 필요하지 않은 작업에서 성능이 저하되지 않음을 확인했습니다.

5. 중요성 및 주장

이 논문은 OVQ-어텐션이 장문맥 처리를 위한 효율성과 성능의 균형을 맞추는 중요한 진전이라고 주장합니다.

• 효율성 대 용량: 선형 어텐션 및 SSM 과 달리, OVQ-어텐션은 총 상태 크기와 무관한 희소 메모리 업데이트를 사용합니다. 이는 모델이 상수 메모리 풋프린트와 선형 계산 비용을 유지하면서 매우 큰 메모리 상태 (높은 용량) 를 활용할 수 있게 합니다.
• 연속 학습 잠재력: 저자들은 OVQ-어텐션의 희소하고 온라인 클러스터링 특성이 선형 어텐션/SSM 에서 사용되는 밀집 그래디언트 업데이트에 비해 "파국적 망각 (catastrophic forgetting)"에 본질적으로 더 강건하다고 제안합니다. 이는 데이터 분포가 시간에 따라 변화하는 연속 학습 시나리오에서 OVQ 를 유망한 후보로 위치시킵니다.
• 확장성: 현재 실험은 상대적으로 작은 규모 (<500M 파라미터) 이지만, 아키텍처의 이론적 특성은 자기 어텐션의 이차적 비용 없이 효율적인 장문맥 LLM 을 향한 실현 가능한 경로임을 시사합니다.

이 논문은 OVQ-어텐션이 메모리 요구 사항의 일부로 장문맥 작업에서 경쟁력 있는 성능을 달성하여 자기 어텐션에 대한 매력적인 대안을 제공하며, 효율적인 시퀀스 모델링을 위한 중요한 연구 방향을 표시한다고 결론지었습니다.