[논문 리뷰] The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit

Lorenzo Noci, C. F. Li외 5명|arXiv (Cornell University)|2023. 06. 30.

Model Reduction and Neural NetworksPhysics and Astronomy인용 수 3

한 줄 요약

이 논문은 깊이와 너비가 무한에 수렴하는 극한에서 깊고 넓은 트랜스포머의 안정성을 확보하기 위해 소프트맥스 출력을 항등원 중심으로 조정하고 너비에 따라 변하는 온도 매개변수로 로짓을 스케일링하는 수정된 어텐션 메커니즘인 '형태를 가진 트랜스포머'(shaped Transformer)를 소개한다. 공분산 행렬에 대한 확률적 미분 방정식(SDE)을 유도하여 잔차 연결과 아키텍처 형태 조절이 질서가 깨지는 현상과 잘 조절된 표현을 유지하는 데 기여함을 보이며, 유한한 모델에서의 시뮬레이션을 통해 SDE의 정확성이 확인된다.

ABSTRACT

In deep learning theory, the covariance matrix of the representations serves as a proxy to examine the network's trainability. Motivated by the success of Transformers, we study the covariance matrix of a modified Softmax-based attention model with skip connections in the proportional limit of infinite-depth-and-width. We show that at initialization the limiting distribution can be described by a stochastic differential equation (SDE) indexed by the depth-to-width ratio. To achieve a well-defined stochastic limit, the Transformer's attention mechanism is modified by centering the Softmax output at identity, and scaling the Softmax logits by a width-dependent temperature parameter. We examine the stability of the network through the corresponding SDE, showing how the scale of both the drift and diffusion can be elegantly controlled with the aid of residual connections. The existence of a stable SDE implies that the covariance structure is well-behaved, even for very large depth and width, thus preventing the notorious issues of rank degeneracy in deep attention models. Finally, we show, through simulations, that the SDE provides a surprisingly good description of the corresponding finite-size model. We coin the name shaped Transformer for these architectural modifications.

연구 동기 및 목표

• 깊고 넓은 트랜스포머의 불안정성, 특히 초기화 단계에서의 질서가 깨지는 현상 해결
• 깊이와 너비가 발산함에 따라 조절이 잘 되는 이론적으로 탄탄한 안정된 어텐션 메커니즘 개발
• 확률적 미분 방정식(SDE)을 사용하여 공분산 구조의 해석 가능한 극한 기술 도출
• 이론적 SDE 모델을 유한한 크기의 신경망 시뮬레이션과 대조하여 검증
• 항등원 중심 소프트맥스와 온도 스케일링을 통한 아키텍처 형태 조절이 깊은 어텐션 모델에서의 안정적 훈련 가능성을 입증

제안 방법

• 소프트맥스 어텐션 메커니즘을 수정하여 출력을 항등원 중심으로 조정하고 너비에 따라 변하는 온도 매개변수 τ로 로짓을 스케일링
• 비례적 무한 깊이-너비 극한(d/n → γ > 0)에서 공분산 행렬에 대한 극한 확률적 미분 방정식(SDE) 유도
• 잔차(스킵) 연결을 사용하여 SDE의 드리프트 및 확산 항을 우아하게 제어하여 안정성 확보
• 스킵 연결을 포함한 형태를 가진 ReLU 피드포워드 네트워크를 기존 SDE 프레임워크에 확장
• 유한한 네트워크에 대한 충실도와 확률성을 유지하기 위해 비례 극한(d, n → ∞, d/n → γ)을 적용
• SDE 예측과 유한한 크기의 네트워크에서의 경험적 공분산 분포를 비교하여 SDE 모델 검증

실험 결과

연구 질문

• RQ1 무한 깊이-너비 극한에서 수정된 어텐션 메커니즘이 깊고 넓은 트랜스포머에서 질서가 깨지는 현상을 방지할 수 있는가?
• RQ2 비례 극한(d/n → γ > 0)에서 어텐션 레이어의 공분산 구조는 어떻게 행동하는가?
• RQ3 확률적 미분 방정식(SDE)이 이러한 모델의 초기화 단계에서 공분산 역학을 정확하게 기술할 수 있는가?
• RQ4 어떤 아키텍처 수정이 SDE를 안정화시키고 잘 조절된 표현을 보장하는가?
• RQ5 SDE는 유한 크기의 실제 트랜스포머의 행동을 어느 정도 정확하게 예측하는가?

주요 결과

• 형태를 가진 트랜스포머는 깊이와 너비가 극한에 도달할 때조차 잘 조절된 공분산 행렬을 유지함으로써 질서가 깨지는 현상을 방지하며, 안정적인 상관계수 분포가 1 이하로 수렴하는 것으로 확인된다.
• 극한 공분산 역학은 깊이-너비 비율 γ에 따라 인덱싱된 잘 정의된 SDE로 기술되며, 잔차 연결을 통해 드리프트 및 확산 항을 제어할 수 있다.
• 시뮬레이션 결과 SDE는 유한 크기 모델의 공분산 구조를 놀랍게 정확하게 기술함을 확인하였으며, 커널 밀도 추정치가 경험적 데이터와 매우 유사하다.
• 제안된 온도 스케일링과 항등원 중심 소프트맥스는 포화를 감소시키고 어텐션 메커니즘을 선형화하여 훈련 불안정성의 주요 원인을 완화한다.
• GLUE 벤치마크에서의 피봇 테이닝 실험에서, 형태를 가진 트랜스포머는 기준 모델보다 COLA와 RTE에서 성능이 뛰어나며, 특히 깊은 아키텍처(d=24)에서 F1 점수 최대 0.211 향상됨.
• 기울기 학습률이 큰 경우에도 형태를 가진 트랜스포머에서는 열정보 붕괴(열정보가 떨어지는 소프트맥스 분포를 나타내는 현상)가 억제되며, 기준 모델과는 달리 안정성을 유지한다.