Induction Heads Interpolate N-Grams

기술 요약: 인덕션 헤드(Induction Heads)의 N-그램 보간

문제 정의
인컨텍스트 학습(ICL)의 주요 회로로 식별된 인덕션 헤드는 전통적으로 "하드 어텐션(hard-attention)" 한계 내에서 특징지어집니다. 이 체제에서 회로는 정확한 패턴 매칭과 최대 우도(maximum-likelihood, ML) 카운팅을 수행합니다. 이는 기계론적으로는 매력적이지만, 고차 의존성(high-order dependencies)에 대해서는 통계적으로 불완전합니다. 즉, 정확한 $k$-gram 매칭은 유한한 샘플 내에서 너무 드물게 발생하여 신뢰할 수 있는 예측을 뒷받침하기 어렵고, ML 추정량은 보이지 않은 연속(unseen continuations)에 대해 0의 확률을 할당합니다. 고전적인 언어 모델은 젤리넥-머서(Jelinek–Mercer) 및 카츠 백오프(Katz back-off)와 같은 스무딩(smoothing) 및 백오프 메커니즘을 통해 이를 해결하며, 이를 통해 확률 질량을 보이지 않은 이벤트로 재분배하고 컨텍스트 길이에 걸쳐 보간합니다. 본 연구가 다루는 핵심 질문은 인덕션 헤드가 엄격한 카운팅을 넘어, 이러한 스무딩을 수행할 수 있는 더 풍부한 클래스의 추정량을 구현하는가 하는 점입니다.

방법론
저자들은 트랜스포머를 훈련시키는 오더-$k$ 마르코프 체인(order-$k$ Markov chains) 환경을 분석하며, 이는 분석적 용이성과 컨텍스트 기반 예측에 내재된 유한 샘플 간의 트레이드오프 사이의 균형을 맞춘 설정입니다.

1. 이론적 프레임워크: 연구는 정밀한 분석을 용이하게 하기 위해 "분리된(disentangled)" 트랜스포머 아키텍처(MLP를 제거하고 잔차 스트림에 연결을 사용)를 활용합니다. 저자들은 상대적 위치 인코딩(RPE)을 사용하는 특정 2층 트랜스포머를 구축하고, 이것이 소프트 컨텍스트 매칭 추정량을 구현함을 증명합니다.
2. 메커니즘 구축:
   • Layer 1 (복사 헤드): 특정 상대적 래그(lag)($t-1, \dots, t-k$)로부터 토큰을 복사하는 $k$개의 어텐션 헤드를 사용하여 현재 컨텍스트의 표현을 구성합니다.
   • Layer 2 (컨텍스트 매칭): 단일 어텐션 헤드를 사용하여 쿼리 컨텍스트를 이력 내의 모든 후보 컨텍스트와 비교합니다. 어텐션 점수는 정확한 매칭을 요구하는 대신, 겹치는 토큰의 수(해밍 오버랩, Hamming overlap)에 따라 가중치를 두는 시프트 정렬된 내적(shift-aligned inner product)을 계산하도록 설계되었습니다.
   • BOS 토큰: 시퀀스 독립적인 어텐션 타겟으로서 가산적 의사 카운트(additive pseudo-counts)를 도입하기 위해 문장의 시작(BOS) 토큰을 포함합니다.
3. 통계적 분석: 저자들은 구현된 폐쇄형(closed-form) 추정량을 도출하고, 이를 두 가지 태스크 사전 확률(prior) 하에서 베이즈 최적 예측량과 비교합니다:
   • 독립 디리클레 사전 확률(Independent Dirichlet Prior): 컨텍스트들이 서로 상관관계가 없는 경우입니다.
   • 계층적 디리클레 사전 확률(Hierarchical Dirichlet Prior): 고차 전이가 하위 차수의 접미사(suffix)를 중심으로 형성되어, 동일한 접미사를 공유하는 컨텍스트 간에 상관관계를 생성하는 경우입니다.

주요 기여
본 논문은 인덕션 헤드가 유한한 어텐션 가중치 규모에서 작동할 때 두 가지 상호 보완적인 스무딩 메커니즘을 구현함을 입증합니다.

1. 소프트 컨텍스트 매칭 (보간):

   • 회로는 정확한 매칭과 부분적인 컨텍스트 매칭 모두의 기여를 집계합니다.
   • 매칭은 겹침 크기($|M|$)에 따라 지수적으로 가중치가 부여됩니다.
   • 이는 젤리넥-머서 스무딩과 유사하게 컨텍스트 차원에 대한 데이터 의존적 보간을 유도합니다. 고정된 글로벌 가중치를 사용하는 고전적 JM 스무딩과 달리, 트랜스포머의 가중치는 특정 시퀀스에서 실현된 카운트에 따라 달라집니다.
   • 이 추정량은 혼합 모델(JM 등)과 백오프 전략(Katz 등)을 통합합니다. 정확한 매칭이 드물 때, 모델은 명시적인 규칙 없이도 더 낮은 차수의 부분 매칭으로 자동으로 "백오프"합니다.

2. BOS를 통한 가산적 의사 카운트:

   • BOS 토큰은 시퀀스 독립적인 싱크(sink) 역할을 하여 어텐션을 흡수합니다.
   • 이 메커니즘을 통해 모델은 컨텍스트별 전이 카운트에 상수 의사 카운트를 더할 수 있습니다.
   • 이론적으로, 이는 큰 어텐션 가중치 한계에서 add-$\alpha$ 스무딩(예: 라플라스 스무딩)을 회복하며, 건축적 관습이 어떻게 사전 확률적 정규화를 구현하는지에 대한 회로 수준의 설명을 제공합니다.

3. 경험적 검증:

   • 저자들은 분리된 트랜스포머와 표준(완전 매개변수화된) 트랜스포머를 마르코프 체인 태스크에 훈련시켰습니다.
   • 결과는 훈련된 모델이 예측된 어텐션 패턴을 회복함을 보여줍니다:
     • BOS가 없을 때, 모델은 차원 간 보간을 위해 소프트 컨텍스트 매칭에 의존합니다.
     • BOS가 있을 때, 모델은 의사 카운트를 활용하여 독립 디리클레 사전 확률 하에서 베이즈 최적 성능에 근접합니다.
     • 하위 차수의 컨텍스트가 구조화된 증거를 제공하는 계층적 사전 확률 하에서, 모델은 인터폴레이션 메커니즘을 활용하여 고정된 add-$\alpha$ 베이스라인보다 뛰어난 성능을 보입니다.
   • 기계론적 분석은 훈련된 표준 트랜스포머가 이론이 예측한 특정 가중치 구조(예: 블록 시프트 행렬)를 학습함을 확인합니다.

결과

• KL 다이버전스: add-$\alpha$ 스무딩이 베이즈 최적(독립 디리클레)인 설정에서, BOS를 가진 훈련된 트랜스포머는 베이즈 최적 예측량과 일치합니다. 구조화된 증거가 존재하는 계층적 설정에서는, 인터폴레이션 메커니즘 덕분에 트랜스포머가 고정 스무딩 베이스라인을 능가합니다.
• 어텐션 패턴: 훈련된 모델의 시각화는 2단계 회로를 확인시켜 줍니다: 레이어 1은 래그별 복사를 수행하고, 레이어 2는 겹침에 따라 매칭에 따라 어텐션 질량을 분배하는 소프트 매칭을 수행합니다.
• 비대칭 래그 가중치: 계층적 태스크에서 훈련된 모델은 가장 최근의 토큰(lag 1)에 대한 매칭에 이전 토큰(lag 2)보다 높은 가중치를 할당하는 법을 배웁니다. 이는 최근 토큰이 다음 상태의 더 나은 예측자라는 사전 확률의 구조와 일치합니다.

의의
본 논문은 인덕션 헤드의 기계론적 해석 가능성과 고전적 통계적 스무딩 사이의 간극을 메웁니다. 트랜스포머가 단순히 규제되지 않은 카운팅을 수행하는 것이 아니라, 컨텍스트 내 추정을 정규화하는 법을 배운다는 것을 입증합니다. 데이터 분포에 적응하는 정교한 통계적 절차를 구현하기 위해 컨텍스트 차원을 동적으로 보간하고 건축적 토큰(BOS)을 활용함으로써, 인덕션 헤드가 인컨텍스트 학습 태스크에서 왜 견고한지를 설명합니다.