[논문 리뷰] Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse

Lorenzo Noci, Sotiris Anagnostidis외 4명|arXiv (Cornell University)|2022. 06. 07.

Neural Networks and Applications인용 수 4

한 줄 요약

이 논문은 트랜스포머에서 신호 전파를 조사하며, 초기화 시 랭크 붕괴가 쿼리 및 키의 기울기 소실의 근본 원인임을 규명한다. 쿼리와 키의 기울기 흐름을 안정화하고 토큰 표현의 랭크를 유지하기 위해 깊이에 따라 달라지는 잔차 스케일링을 제안하며, 쿼리/키 파라미터와 값 파라미터 간의 비대칭적 기울기 노름을 드러내어 적응형 최적화기의 효용성을 설명한다.

ABSTRACT

Transformers have achieved remarkable success in several domains, ranging from natural language processing to computer vision. Nevertheless, it has been recently shown that stacking self-attention layers - the distinctive architectural component of Transformers - can result in rank collapse of the tokens' representations at initialization. The question of if and how rank collapse affects training is still largely unanswered, and its investigation is necessary for a more comprehensive understanding of this architecture. In this work, we shed new light on the causes and the effects of this phenomenon. First, we show that rank collapse of the tokens' representations hinders training by causing the gradients of the queries and keys to vanish at initialization. Furthermore, we provide a thorough description of the origin of rank collapse and discuss how to prevent it via an appropriate depth-dependent scaling of the residual branches. Finally, our analysis unveils that specific architectural hyperparameters affect the gradients of queries and values differently, leading to disproportionate gradient norms. This suggests an explanation for the widespread use of adaptive methods for Transformers' optimization.

연구 동기 및 목표

• 초기화 단계에서 자기어텐션 레이어에서 발생하는 랭크 붕괴의 원인과 영향을 이해한다.
• 랭크 붕괴가 쿼리 및 키 파라미터의 기울기 소실을 유도함으로써 훈련을 방해하는 방식을 조사한다.
• 토큰 표현의 랭크를 유지하고 신호 전파를 안정화하기 위해 깊이에 따라 달라지는 잔차 스케일링 전략을 제안한다.
• 아키텍처 하이퍼파라미터가 쿼리, 키, 값의 기울기 노름에 미치는 영향을 이질적으로 분석한다.
• 딥 트랜스포머 훈련에서 적응형 최적화기의 광범위한 사용에 대한 이론적 및 실증적 근거를 제공한다.

제안 방법

• 랜덤 행렬 이론을 사용하여 단일 자기어텐션 레이어에서 초기화 시 쿼리, 키, 값의 기대 기울기를 이론적으로 유도한다.
• 토큰 표현 간余弦 유사도를 안정화함으로써 랭크 붕괴를 방지하는 깊이에 따라 달라지는 잔차 스케일링 기법을 제안한다.
• 쿼리 및 키의 기울기 노름이 입력 노름에 대해 세제곱으로 증가하는 반면, 값의 기울기 노름은 선형으로 증가하여 기울기 크기의 극명한 불균형을 초래함을 보여준다.
• 쿼리와 값 간 기울기 노름을 균형 잡기 위해 온도 조정된 소프트맥스를 도입하여 실증적으로 검증한다.
• 이론적 분석과 토이 태스크, IWSLT’14 번역 작업에서의 실증 검증을 통해 기울기 역학과 안정성을 검증한다.
• 학습 가능한 잔차 스케일링을 사용해 모델을 훈련하고 상관관계, 활성화 노름, 기울기 흐름을 모니터링하여 이론적 예측의 타당성을 검증한다.

실험 결과

연구 질문

• RQ1 초기화 시 랭크 붕괴가 딥 트랜스포머의 쿼리 및 키 파라미터 기울기에 어떤 영향을 미치는가?
• RQ2 어떤 아키텍처 메커니즘이 깊이에 걸쳐 신호 전파를 안정화하고 랭크 붕괴를 방지하는가?
• RQ3 왜 Adam과 같은 적응형 최적화기들이 트랜스포머에서 더 잘 작동하는가? 이는 기울기 노름의 불균형으로 설명될 수 있는가?
• RQ4 입력 상관관계, 시퀀스 길이, 임bedding 차원은 쿼리, 키, 값의 상대적 기울기 노름에 어떤 영향을 미치는가?
• RQ5 깊이에 따라 달라지는 잔차 스케일링은 토큰 표현의 랭크를 안정화하고 딥 트랜스포머의 훈련 역학을 향상시킬 수 있는가?

주요 결과

• 초기화 시 랭크 붕괴는 쿼리 및 키 파라미터의 기울기 소실을 유도하며, 이는 깊은 트랜스포머의 학습 능력을 심각하게 저해한다.
• 깊이에 따라 달라지는 잔차 스케일링은 토큰 표현 간 여유도 유사도를 안정화함으로써 랭크를 유지하고, 무한한 깊이의 극한에서도 열악한 상태를 방지한다.
• 쿼리 및 키의 기울기 노름은 입력 노름에 대해 세제곱으로 증가하는 반면, 값의 기울기 노름은 선형으로 증가하여 기울기 크기의 심각한 불균형을 초래한다.
• 실증 결과는 적절한 스케일링이 없을 경우 깊이를 늘릴수록 토큰 정렬도가 높아지고 훈련 성능이 악화됨을 확인하였으며, 이는 그림 1에 나타나 있다.
• 온도 조정된 소프트맥스를 도입하고 $\tau = 8.5$로 설정함으로써 쿼리와 값 간 기울기 노름을 균형 잡을 수 있었으며, 이는 이론적 예측을 검증한다.
• 훈련된 모델에서 잔차 스케일링 파라미터 $\beta_1, \beta_2$는 훈련 중에 적응하며, 토큰 상관관계는 낮게 유지된다. 이는 제안된 스케일링의 이론적 안정성을 지지한다.