기술 요약: 베이지안 어텐션 메커니즘 (BAM)

문제 제기

트랜스포머 기반 언어 모델은 토큰 순서를 처리하기 위해 위치 인코딩 (PE) 에 의존하지만, 기존 방법들은 컨텍스트 길이 외삽 (훈련 중 본 것보다 긴 시퀀스에 일반화하는 능력) 과 관련하여 두 가지 치명적인 한계에 직면해 있습니다:

1. 이론적 불투명성: 많은 PE 기법들 (예: Sinusoidal, RoPE, ALiBi, NoPE) 은 경험적으로 동기부여되었으며 이론적 기반이 제한적이어서, 외삽 중의 그 행동 양상이 잘 이해되지 않습니다.
2. 부족한 평가: 주된 지표로 퍼플렉시티 (perplexity) 에 의존하는 것은 불충분합니다. 모델은 슬라이딩 윈도우 평가에서 지적된 결함처럼, 장기 정보를 진정으로 통합하지 않고도 국소적 어텐션 패턴 (최근 토큰에만 집중) 을 통해 낮은 퍼플렉시티를 달성할 수 있습니다.

방법론: 베이지안 어텐션 메커니즘 (BAM)

저자들은 토큰 내용과 위치의 결합 확률 모델 하에 계산된 값들에 대한 기대값으로 자기 어텐션을 재해석하는 이론적 프레임워크인 베이지안 어텐션 메커니즘 (BAM) 을 제안합니다.

이론적 프레임워크

• 결합 확률 공식화: BAM 에서 어텐션 가중치 $p_{ij}$ (토큰 $j$ 가 토큰 $i$ 에 어텐션할 확률) 는 결합 확률 분포로 정의됩니다:
  $$p_{ij} = p(f_{cont}(q_i, k_j) | g_{pos}(i, j)) \times p(g_{pos}(i, j))$$
  여기서 $f_{cont}$ 는 내용 유사성을, $g_{pos}$ 는 위치 관계를 나타냅니다.
• 위치 인코딩으로서의 사전 분포: 이 프레임워크는 위치 인코딩 전략이 토큰 위치에 대한 특정 사전 분포 $p(g_{pos}(i, j))$ 에 해당한다고 가정합니다.
  • NoPE 는 (인과적 마스크에 의해 제한된) 균일 사전 (Uniform prior) 으로 유도됩니다.
  • ALiBi 는 균일 사전과 라플라스 사전 (거리 기반 지수 감쇠) 의 곱으로 유도됩니다.
• 정규화 스칼라 ($Z$): 유도 과정은 내용과 위치 확률의 곱을 정규화하는 스칼라 $Z$ 를 도입하며, 이는 내용과 위치 간의 상호작용 자유 에너지 또는 정보 결합을 나타냅니다.

제안된 해결책: GGD-BAM

이 프레임워크를 활용하여 저자들은 위치 사전으로 일반화된 가우시안 분포 (GGD) 를 사용하는 GGD-BAM 을 소개합니다. 어텐션 점수에 추가되는 편향 행렬 $B$ 는 다음과 같이 정의됩니다:
$$b_{ij} = -\left| \frac{j - i - \mu}{\alpha} \right|^\beta$$
여기서 $\mu$, $\alpha$, $\beta$ 는 학습 가능한 매개변수 (또는 특정 동작을 위해 고정됨) 입니다.

• $\beta \in (0, 1)$: 라플라스 사전 (ALiBi) 보다 느리게 감소하는 두꺼운 꼬리 (heavy-tailed) 분포를 생성하여, 국소성을 유지하면서도 ALiBi 보다 더 멀리 어텐션을 확장할 수 있게 합니다.
• $\beta < 0$: 표준 GGD 정의의 새로운 완화입니다. 이는 쿼리 토큰 근처에서는 확률 질량이 억제되고 컨텍스트 윈도우의 먼 끝 부분에 집중되는 "역전된" 사전을 생성합니다. 이론적으로 이는 모델이 국소 컨텍스트를 무시하고 임의의 장기 의존성에 집중할 수 있게 합니다.
• 확장 가능한 소프트맥스 (SSMax): 저자들은 BAM 이 컨텍스트가 커짐에 따라 "어텐션 소실 (균일 분포)"에 취약할 수 있다고 지적합니다. 그들은 분포의 날카로움을 보존하기 위해 동적으로 로그잇을 재조정하는 SSMax 와 BAM 을 결합합니다.

주요 기여

1. 통합된 이론적 프레임워크: BAM 은 기존 PE 방법들의 확률론적 통합을 제공하며, NoPE 와 ALiBi 를 각각 균일 사전과 라플라스 사전에 명시적으로 매핑합니다.
2. 일반화된 가우시안 사전: GGD-BAM 제안은 유연한 위치 사전 계열을 도입합니다. 모양 매개변수 $\beta$ 를 학습함으로써 모델은 국소 어텐션과 장기 검색 모드 간에 적응적으로 전환할 수 있습니다.
3. 최소한의 오버헤드: 이 방법은 각 어텐션 헤드에 대해 $\theta_\alpha$ 와 $\theta_\beta$ 를 학습하고 $\theta_\mu = 0$ 으로 고정함으로써 120M 매개변수 모델의 경우 구체적으로 384 개의 추가 학습 가능한 매개변수 (1,000 개 미만) 만 도입합니다.
4. 해석 가능성: 이 프레임워크는 어텐션 패턴을 위치에 대한 명시적 확률 분포로 시각화하고 해석할 수 있는 원칙적인 방법을 제공합니다.

경험적 결과

실험은 FineWeb 10B 에서 512 토큰 컨텍스트 길이로 훈련된 120M 에서 1.1B 매개변수까지의 모델들에서 수행되었습니다.

• 패키키 (Passkey) 검색: GGD-BAM (특히 SSMax 와 함께) 은 최대 32,000 토큰 (훈련 길이의 64 배) 의 시퀀스에서 특정 토큰을 검색할 때 완벽한 정확도를 달성했으며, 최대 512,000 토큰 (훈련 길이의 1000 배) 까지 80% 이상의 정확도를 유지했습니다. 반면, RoPE, Sinusoidal, NoPE 와 같은 기준 모델들은 훈련 범위를 넘어선 곳에서 무작위 수준의 정확도로 급격히 저하되었습니다. ALiBi 는 퍼플렉시티 외삽에서는 더 좋았지만 극단적인 길이에서 검색 작업에서 크게 실패했습니다.
• RULER 벤치마크 (NIAH): GGD-BAM 은 "Needle In A Haystack" 하위 세트에서 모든 기준 모델 (RoPE 및 ALiBi 포함) 을 능가했으며, Single 1 작업에서 6,000 토큰을 넘어선 높은 정확도를 유지한 유일한 방법이었습니다.
• 퍼플렉시티: GGD-BAM 은 긴 컨텍스트에서 ALiBi 및 RoPE 와 비교 가능한 퍼플렉시티를 유지하여, 개선된 검색이 언어 모델링 유창성 비용을 수반하지 않음을 입증했습니다.
• 하류 작업: MMLU 및 ARC 벤치마크에서 BAM SSMax 는 RoPE SSMax 를 능가했습니다.
• 어텐션 시각화: 시각화는 $\beta \le 0$ 인 헤드가 국소 컨텍스트를 효과적으로 억제하고 먼 토큰 (패키키) 에 집중하는 반면, $\beta \in (0, 1)$ 인 헤드는 긴 꼬리 어텐션을 유지함을 확인시켜 주었습니다.

중요성과 주장

이 논문은 BAM 이 이중 목적을 수행한다고 주장합니다:

1. 이론적: 위치 인코딩에 대한 엄격한 확률론적 기반을 제공하여, 왜 특정 방법들 (예: ALiBi) 은 퍼플렉시티에는 작동하지만 국소성 때문에 검색에는 실패하는지 명확히 합니다.
2. 실용적: 컨텍스트 길이 외삽을 위한 매우 효율적인 방법을 제공합니다. 어텐션 헤드가 음수 $\beta$ 값을 학습할 수 있게 함으로써, 모델은 훈련 중 그러한 길이에 직접 노출되지 않고도 임의의 긴 컨텍스트에서 정보를 검색하기 위해 국소성을 명시적으로 억제할 수 있습니다.

저자들은 퍼플렉시티가 표준 지표이지만 진정한 긴 컨텍스트 능력을 평가하기에는 불충분하며, 모델이 단순히 국소 패턴에 의존하는 것이 아니라 먼 정보를 주시하고 있는지 확인하기 위해 검색 기반 작업이 필요하다고 강조합니다. 이 연구는 위치 사전을 명시적으로 모델링하는 능력이 차세대 긴 컨텍스트 트랜스포머를 위한 중요한 귀납적 편향임을 시사합니다.