Adaptive inference and function vectors in deep transformers

문제 정의
트랜스포머는 다양한 도메인에서 복잡한 상관관계를 학습하기 위한 범용적인 기질(substrate)로 부상했으나, 인컨텍스트 러닝(In-Context Learning, ICL)을 위한 내부 메커니즘은 여전히 미스터리로 남아 있다. 경험적 연구를 통해 ICL의 두 가지 주요 모티프—특정 알고리즘(예: 복사, 경사 하강법)을 구현하는 어텐션 회로와 "함수 벡터(function vectors)"(태스크의 내부 표현)의 형성을 식별해 왔으나, 후자에 대한 이론적 이해는 제한적이다. 특히, 딥 트랜스포머가 함수 벡터를 사용하여 어떻게 추론을 수행하는지, 그리고 깊이(depth)와 피드포워드(MLP) 블록이 구체적으로 어떤 계산적 역할을 하는지는 불분명하다. 기존 이론들은 주로 어텐션 기반 회로에 초점을 맞추고 있어, 어텐션과 MLP 간의 다층 상호작용에 의해 가능해지는 메커니즘은 충분히 탐구되지 않았다.

방법론
저자들은 딥 트랜스포머를 분산 베이지안 추론(distributed Bayesian inference)을 구현하는 평균장 상호작용 시스템(mean-field interacting system)으로 모델링하는 이론적 프레임워크를 제안한다. 이 모델은 통신 채널 차원($d$), 국소성(locality), 그리고 깊이($L$)에 대한 구조적 제약을 받는다.

1. 이론적 모델: 프레임워크는 트랜스포머가 잠재 컨텍스트 변수 $\theta$에 의해 결합된 토큰 집합을 처리한다고 가정한다. 추론은 "함수 벡터"($\Phi_\ell$)를 통해 매개되며, 이는 압축된 상태 변수 역할을 한다.

   • 각 층 $\ell$에서, 로컬 비선형성(MLP)은 토큰과 현재 상태 $\Phi_\ell$를 처리하여 임베딩 $\lambda_{\ell+1}$을 생성한다.
   • 어텐션 유사 연산은 이러한 임베딩들을 토큰 전체에 걸쳐 풀링(pooling)하여 새로운 통계량 $\phi_{\ell+1}$을 형성한다.
   • 상태는 직접 합(direct sum)을 통해 업데이트된다: $\Phi_{\ell+1} = \Phi_\ell \oplus \phi_{\ell+1}$.
   • MLP 블록은 정보를 전달할 대상을 선택하는 "라우터(router)"이자 디코더 역할을 한다고 가설을 세우며, 어텐션은 정보를 풀링한다.

2. 최적화 전략: 저자들은 임베딩의 선택을 동적 계획법(dynamic programming) 문제로 프레임화한다. 최적의 전략은 층 $\ell$에서 기대 손실을 최소화하며, 이는 임베딩 선택을 적응적 추론을 위한 실험 설계(experimental design) 전략으로 취급하는 것과 같다.

3. 분석 사례: 이론은 잠재 변수 $\theta$에 대한 두 가지 사전 분포(prior distribution) 하에서 분석된다:

   • 가우시안 사전 분포(Gaussian Prior): 잠재 변수가 표준 정규 분포를 따르는 경우.
   • 트리 사전 분포(Tree Prior): $\theta$가 균형 잡힌 이진 트리를 통과하는 경로를 기반으로 구성되는 계층적 구조를 가진 사전 분포. 이 구조는 점점 더 미세한 스케일에서 잠재 변수를 해결하기 위한 깊이의 이점을 테스트한다.

4. 경험적 검증: 예측은 세 가지 구성의 제약된 선형 어텐션 트랜스포머 아키텍처를 사용하여 테스트된다:

   • 모델 A: 채널 차원 $d=1$이고 모든 층에 MLP 블록이 있는 깊은($L=M$) 모델.
   • 모델 B: 채널 차원이 $M$이고 동일한 MLP를 가진 단일 층 모델.
   • 모델 C: 채널 차원이 $M$이고 MLP가 없는 단일 층 모델(선형 어텐션만 사용).
   • 인과적 섭동(Causal Perturbation): 특정 층에서 "도너(donor)" 컨텍스트의 키(key) 행렬을 "타겟(target)" 컨텍스트에 삽립하여 라우팅 메커니즘을 테스트하는 "키 패칭(key-patching)" 실험이 수행된다.

주요 기여

• 함수 벡터 이론: 본 논문은 함수 벡터를 컨텍스트 정보를 예측을 위해 압축하는 내부 상태 표현으로 정의하는 공식 이론을 제공한다. 이 벡터들이 모델이 여러 층에 걸쳐 점점 더 미세한 스케일에서 잠재 컨텍스트 변수를 추론할 수 있게 함을 입증한다.
• 깊이와 MLP의 역할: 이론은 비가우시안, 계층적 사전 분포의 경우, 깊이와 MLP 블록이 비적응적(non-adaptive), 단일 층 접근 방식보다 우수한 적응적 추론 전략을 구현하는 데 필수적이라고 예측한다. 반대로, 가우시안 사전 분포의 경우, 깊이는 동일한 총 통신 예산을 가진 단일 층에 비해 이점을 제공하지 않는다.
• 메커니즘 매핑: 본 연구는 이론적 구성 요소를 트랜스포머 아키텍처에 매핑하여, MLP를 측정할 통계량을 선택하는 "라우터"로, 어텐션 메커니즘을 이러한 통계량을 풀링하는 역할로 할당한다.
• 실험적 확인: 수치 실험을 통해, MLP를 가진 깊은 트랜스포머(모델 A)가 계층적 트리 사전 분포에 대해 최적화된 비적응적 단일 층 전략(모델 B) 및 선형 어텐션 모델(모델 C)보다 훨씬 뛰어난 성능을 보이며 적응적 추론 전략을 학습할 수 있음을 확인하였다.

결과

• 가우시안 사전 분포: 가우시안 사전 분포의 경우, 평균 제곱 오차(MSE) 손실은 함수 벡터의 구체적인 실현(realization)과는 독립적이며 측정 전략에 달려 있다. 이론과 실험 모두 깊은 모델(모델 A)이 동일한 채널 용량을 가진 단일 층 모델(모델 B)보다 나은 성능을 보이지 않음을 보여준다.
• 트리 사전 분포: 계층적 트리 사전 분포의 경우, 이론은 (이전 결과에 기반하여 한 번에 하나의 좌표를 측정하는) 적응적 전략이 MSE 손실이 깊이에 따라 지수적으로($\alpha^M$) 스케일링되는 반면, 비적응적 전략은 $M$의 거듭제곱으로 스스케일링된다고 예측한다.
  • 실험 결과, 모델 A는 모델 B보다 현저히 낮은 MSE 손실을 달연하여, 모델이 적응적 추론 전략을 학습했음을 확인시켜 주었다.
  • MLP가 없는 모델 C는 비적응적 이론적 베이스라인과 유사한 성능을 보였으며, 이는 적응적 라우팅 메커니즘을 위한 MLP의 필요성을 확인시켜 준다.
• 인과적 섭동: 키 패칭 실험은 "도너" 컨텍스트(리프 $\sigma'$)의 키를 "타겟" 컨텍스트(리프 $\sigma$)에 삽입했을 때, 경로가 조기에 갈라지고 개입이 네트워크 초기에 발생하면 손실이 크게 증가함을 보여준다. 이는 모델의 내부 상태(함수 벡터)가 컨텍스트 경로에 따라 적응적으로 진화함을 확인시켜 준다.

의의
본 논문은 ICL을 잠재 변수에 대한 베이지안 추론으로 보는 관점과, ICL을 알고리즘 실행(예: 경사 하강법)으로 보는 메커니즘적 관점을 결면시키고자 한다. 깊은 트랜스포머가 함수 벡터에 의해 매개되는 분산된 적응적 층별 절차로서 베이지안 추론을 구현할 수 있음을 보여줌으로써, 저자들은 깊이와 MLP 블록이 기존에 설명된 것들(예: 프리컨디셔닝된 경사 하강법)보다 더 풍부한 클래스의 ICL 알고리즘을 가능하게 한다는 것을 입증한다. 이 연구는 함수 벡터를 회로 수준의 프레임워크 내에 배치하며, 이러한 구조가 수십억 개의 파라미터를 가진 모델의 조직화에 유리한 이유가 통신 제약 하에서 효율적이고 적응적인 추론을 가능하게 하기 때문임을 시사한다. 저자들은 자신들의 이론이 선형 풀링을 사용하지만 교환 가능성(exchangeability)을 존중하며, 소프트맥스 어텐션의 특정 토큰 간 검색이 ICL의 주요 동인이 아닌 시나리오에서도 충분할 수 있다고 언급한다.