형상화된 트랜스포머: 무한한 깊이와 너비의 극한에서의 어텐션 모델

The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit

June 30, 2023

저자: Lorenzo Noci, Chuning Li, Mufan Bill Li, Bobby He, Thomas Hofmann, Chris Maddison, Daniel M. Roy

cs.AI

초록

딥러닝 이론에서 표현의 공분산 행렬은 네트워크의 학습 가능성을 검토하기 위한 대리 지표로 사용됩니다. 트랜스포머의 성공에 영감을 받아, 우리는 무한 깊이와 무한 폭의 비례적 극한에서 스킵 연결이 적용된 수정된 소프트맥스 기반 어텐션 모델의 공분산 행렬을 연구합니다. 초기화 시점에서 극한 분포는 깊이 대 폭 비율로 인덱싱된 확률 미분 방정식(SDE)으로 설명될 수 있음을 보입니다. 잘 정의된 확률적 극한을 달성하기 위해, 트랜스포머의 어텐션 메커니즘은 소프트맥스 출력을 항등 행렬로 중심화하고, 소프트맥스 로짓을 폭에 의존하는 온도 매개변수로 스케일링하여 수정됩니다. 우리는 해당 SDE를 통해 네트워크의 안정성을 검토하며, 드리프트와 확산의 스케일이 잔차 연결의 도움으로 우아하게 제어될 수 있음을 보입니다. 안정적인 SDE의 존재는 매우 큰 깊이와 폭에서도 공분산 구조가 잘 동작함을 의미하며, 이는 딥 어텐션 모델에서 흔히 발생하는 랭크 퇴화 문제를 방지합니다. 마지막으로, 우리는 시뮬레이션을 통해 SDE가 해당 유한 크기 모델을 놀라울 정도로 잘 설명함을 보입니다. 이러한 구조적 수정을 '셰이프드 트랜스포머(Shaped Transformer)'라고 명명합니다.

English

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.