Training Infinitely Deep and Wide Transformers

기술 요약: 무한히 깊고 넓은 트랜스포머 학습

문제 제기
트랜스포머가 현대 기계 학습에서 지배적인 아키텍처가 되었음에도 불구하고, 특히 많은 층과 어텐션 헤드를 가진 깊은 아키텍처에 대한 학습 역동에 대한 이론적 이해는 여전히 제한적입니다. 기존 이론적 프레임워크는 종종 가중치가 $1/\sqrt{\text{width}}$로 스케일링된다고 가정하는 신경 탄성 커널 (NTK) 영역에 의존하는데, 이는 네트워크가 초기화 상태에 가깝게 유지되고 특징 학습이 고정되는 '게으른 학습 (lazy training)'으로 이어집니다. 대안적으로, 얕은 네트워크나 신경망 (Ordinary Differential Equations, ODEs 로 모델링된 ResNets) 을 위한 평균장 (mean-field) 영역은 어텐션 메커니즘이 여러 토큰 분포의 진화를 결합한다는 점에서 트랜스포머의 복잡성을 완전히 포착하지 못합니다. 토큰 분포와 매개변수 분포의 동시적 진화를 고려하면서, 무한한 깊이 (층 수) 와 무한한 너비 (어텐션 헤드의 수) 의 결합 극한에서 경계 기반 학습을 분석하기 위한 엄밀한 수학적 프레임워크가 필요합니다.

방법론
본 논문은 ResNets 에 적용되었던 조건부 최적 수송 (conditional optimal transport) 과 Wasserstein 경사 흐름 이론을 확장하여 트랜스포머 설정에 적용함으로써 트랜스포머를 위한 평균장 프레임워크를 개발합니다.

1. 이중 평균장 표현: 이 모델은 트랜스포머를 연속적인 깊이 변수 $s \in [0,1]$로 매개화된 무한히 깊은 구조와 어텐션 매개변수 위의 확률 분포 $\rho$로 매개화된 무한히 넓은 구조로 다룹니다. 핵심적으로, 이 프레임워크는 두 가지 유형의 측도 표현을 사용합니다:

   • 토큰 분포: 층을 통과하는 토큰의 진화는 토큰 공간 위의 확률 측도 $\mu(s)$로 모델링됩니다.
   • 어텐션 매개변수: 각 층의 매개변수는 매개변수 공간 $\Theta$(쿼리, 키, 값 행렬) 위의 측도 $\rho(\cdot|s)$로 표현됩니다.

2. 신경 편미분 방정식 (PDE) 공식화: 특징 벡터에 대한 신경 ODE 를 제어하는 것으로 볼 수 있는 ResNets 와 달리, 이 프레임워크에서 트랜스포머의 학습은 **신경 편미분 방정식 (PDE)**을 제어하는 것에 해당합니다. 토큰 진화는 모든 토큰의 분포를 결합하는 평균장 어텐션 연산자에 의해 결정되는 속도장에 의해 지배되는 연속 방정식에 의해 지배됩니다.
   $$\partial_s \mu(s) + \text{div}(\mu(s) \Phi_{\rho(\cdot|s)}[\mu(s)]) = 0$$
   저자들은 이 PDE 관점을 초기 토큰 분포를 이송 (advect) 하는 흐름 맵 $\Lambda$에 작용하는 Banach 공간 값 ODE 로 재공식화합니다. 이 재매개화는 측도 공간의 미분 가능 구조 부재를 피하고 분석을 위한 향상된 규칙성을 제공합니다.

3. 조건부 Wasserstein 경사 흐름을 통한 학습 역동: 학습 과정은 층 전체의 매개변수 한계 분포가 학습 중에 보존되도록 강제하는 조건부 최적 수송 (COT) 척도를 갖춘 매개변수 분포 공간 위의 경사 흐름으로 모델링됩니다. 저자들은 역방향 ODE 에 의해 지배되는 접변수 (adjoint variables) 를 포함하는 명시적 공식을 얻기 위해 접 민감도 분석을 사용하여 학습 위험의 조건부 Wasserstein 경사를 유도합니다.

4. 신경 탄성 커널 (NTK) 분석: 수렴을 확립하기 위해, 논문은 어텐션 메커니즘에 대한 NTK 의 단사성 (injectivity) 을 조사합니다. 저자들은 NTK 가 양의 정부호 (positive definite) 가 되기 위한 필요충분조건을 유도하며, 이 속성을 토큰 분포의 로그 - 합 - 지수 (cumulant generating) 함수들이 아핀 함수를 모듈로로 하여 선형 독립인 것과 연결합니다.

주요 기여 및 결과

• 전방향 및 학습 역동의 잘-제시됨 (Well-Posedness): 논문은 전방향 통과 (토큰 진화) 와 학습 진화 (매개변수 분포 흐름) 에 대한 해의 존재성과 유일성을 확립합니다. 토큰 진화가 함수 공간에서 ODE 를 만족하는 흐름 맵으로 특징지어지며, 조건부 Wasserstein 거리 공간에서 경사 흐름 곡선이 존재하고 유일함을 증명합니다.
• 명시적 경사 공식화: 접 민감도 분석을 사용하여, 저자들은 학습 위험의 조건부 Wasserstein 경사에 대한 명시적 공식을 유도합니다. 이는 역방향 ODE 를 만족하는 접변수를 포함하며, 평균장 극한에서 경계 기반 학습을 위한 엄밀한 기초를 제공합니다.
• NTK 단사성 조건: 핵심적인 기술적 기여는 어텐션 메커니즘에 대한 NTK 단사성의 특성화입니다. 논문은 NTK 단사성이 아핀 함수를 모듈로로 한 로그 - 합 - 지수 함수들의 선형 독립과 동치임을 증명합니다. 이 조건은 다음과 같은 다양한 토큰 분포에 의해 만족되는 것으로 나타났습니다:
  • 절대 연속 측도에서 i.i.d. 표본에 대해 거의 확실하게 검증된 구별 조건을 만족하는 이산 분포.
  • 서로 다른 반경을 가진 큐브 위의 균일 분포.
  • 서로 다른 공분산 행렬을 가진 대칭 다변량 라플라스 분포.
  • 서로 다른 분리 파라미터를 가진 두 성분 가우시안 혼합물.
• 전역 최소값으로의 수렴: NTK 가 초기화 시 양수 (특히 NTK 의 $V$-부분) 이고 초기 손실이 충분히 작다는 가정 하에, 논문은 경사 흐름이 선형 수렴 속도로 전역 최소점으로 수렴함을 증명합니다. 이 결과는 지정된 영역에서 최적화 풍경의 인위적 국소 최소값을 제거합니다.

의의 및 주장
이 논문은 무한히 깊고 넓은 트랜스포머 학습을 위한 엄밀한 수학적 기초를 제공한다고 주장합니다. 신경 ODE 인 ResNets 에서 신경 PDE 인 트랜스포머로 평균장 이론을 확장함으로써, 깊은 잔차 네트워크에 대한 이론적 이해와 어텐션 메커니즘의 더 복잡한 역동 사이의 간극을 메웁니다.

저자들은 그들의 프레임워크가 표준 ResNet 분석에는 없는 어텐션 메커니즘을 통한 여러 토큰 분포의 결합을 포착한다고 강조합니다. NTK 단사성 조건 확립은 초기화에 대한 구성적 지침을 제공하여, 초기화가 최적 해에 충분히 가까울 경우 경계 기반 학습이 인위적 국소 최소값에 갇히지 않고 전역 최소점으로 수렴하도록 보장합니다.

논문은 분석이 인코더 스타일 (인과적이지 않은) 어텐션으로 제한된다는 점을 겸손하게 지적합니다. 인과적 마스킹이 있는 디코더 아키텍처로 이 프레임워크를 확장하는 것은 증명에 사용된 대칭성이 깨지기 때문에 상당한 기술적 과제로 식별됩니다. 또한, 현재 공식화는 헤드의 수를 무한대로 증가시키는 동안 헤드의 차원을 고정적으로 유지하며, 임베딩 차원과 헤드의 수의 결합된 스케일링을 포착하는 것은 여전히 해결되지 않은 문제입니다.