깊은 트랜스포머 초기화의 두 실패 모드와 해결 방안

초록

본 논문은 트랜스포머 모델에서 초기 가중치 설정이 잘못될 경우 발생하는 “랭크 붕괴”와 “엔트로피 붕괴” 두 가지 실패 모드를 이론적으로 분석한다. 자기‑어텐션을 랜덤 에너지 모델에 매핑하고, 시퀀스 길이가 무한대로 갈 때의 스케일링을 정확히 도출함으로써, 키·쿼리 가중치의 분산 β와 어텐션 잔차 연결 강도 αₛₐ의 조합이 모델의 훈련 가능성을 결정한다는 ‘trainability diagram’을 제시한다. 또한 역전파 시 그래디언트 소실 조건을 밝혀, 실험을 통해 이론의 정확성을 검증한다.

상세 분석

논문은 먼저 기존의 초기화 이론이 완전 연결망에서는 성공했지만, 트랜스포머의 핵심인 자기‑어텐션에서는 두 가지 새로운 병목 현상이 나타난다는 점을 강조한다. 첫 번째인 랭크 붕괴는 어텐션 행렬이 거의 균일해져 모든 토큰이 동일한 표현으로 수렴하면서 토큰 간 차원이 1로 축소되는 현상이다. 이는 깊이가 증가할수록 이중 지수적으로 진행되며, 결과적으로 그래디언트가 급격히 소멸한다. 두 번째인 엔트로피 붕괴는 키·쿼리 가중치가 과도하게 크게 초기화될 경우 소프트맥스가 포화되어 몇 개의 토큰에만 집중되는 불안정한 어텐션 분포가 형성되는 현상이다. 이때 어텐션 행렬의 샤논 엔트로피가 낮아져 학습이 불안정해진다.

핵심 기여는 자기‑어텐션을 통계 물리학의 Random Energy Model(REM)과 정확히 대응시킨 점이다. 초기화된 키·쿼리 가중치가 독립적인 가우시안 변수이므로, 어텐션 스코어 aₜₜ′는 평균 0, 분산 σₐ²인 가우시안이며, 토큰 임베딩 간 내적 qₜₛ에 따라 상관관계를 가진다. 이를 REM의 에너지 E(s)와 동일한 형태로 해석하면, 어텐션 행렬 한 행을 온도 β가 역수인 볼츠만 분포로 보는 것이 가능해진다.

이 매핑을 바탕으로 저자들은 스케일링 법칙 σ_Q² = σ_K² = β·log T / d (β는 상수) 를 제안한다. 여기서 T는 시퀀스 길이이며, 로그 스케일링은 에너지 변동이 O(√log T) 수준이 되도록 보장한다. β가 임계값 β_c 이하이면 어텐션 스코어가 충분히 작아 행렬이 거의 균일해져 랭크 붕괴가 일어나고, β > β_c이면 스코어가 크게 변동해 포화된 어텐션이 나타나 엔트로피 붕괴가 발생한다.

또한 잔차 연결 강도 α_SA가 충분히 크면 토큰 간 코사인 유사도 ρ가 1에 수렴하는 것을 방지할 수 있다. 논문은 α_SA와 β의 조합을 2‑차원 평면에 표시한 ‘trainability diagram’을 제시하고, 각 영역을 (1) 랭크 붕괴, (2) 엔트로피 붕괴, (3) 정상 훈련 가능 영역으로 구분한다. 이 다이어그램은 알고리즘 1을 통해 깊이 L, 레이어 정규화 방식, MLP 스케일 등을 모두 고려해 자동으로 계산할 수 있다.

역전파 분석에서는 키·쿼리 가중치에 대한 그래디언트 노름을 정확히 유도하고, β < β_c 일 때 그래디언트가 지수적으로 소멸함을 증명한다. 따라서 훈련 가능 영역에서는 β가 β_c 위에 있어야 하지만, 동시에 α_SA가 충분히 커야 토큰 다양성이 유지된다.

실험에서는 60‑layer BERT‑style 모델을 TinyStories 데이터에 적용해, 이론이 예측한 α_SA와 β 조합이 실제 훈련 성공 여부와 일치함을 보여준다. 특히, 랭크 붕괴 영역에서는 토큰 코사인 유사도가 층을 거듭할수록 1에 수렴하고, 엔트로피 붕괴 영역에서는 어텐션 행렬이 극단적으로 스파스해져 학습 손실이 발산한다. 반면, 이론이 제시한 ‘안전 영역’에서는 토큰 유사도가 적절히 유지되고, 테스트 손실이 안정적으로 감소한다.

결론적으로, 이 논문은 트랜스포머 초기화 문제를 “스케일링(β) × 잔차 강도(α_SA)”라는 두 차원으로 정량화하고, 통계 물리학적 REM 매핑을 통해 정확한 임계값을 도출함으로써, 실무에서 깊은 트랜스포머를 설계·초기화할 때 명확한 가이드라인을 제공한다.