기술 요약: 레이어 정규화(Layer Normalization) 하에서의 트랜스포머 안정성

문제 정의

현대 딥러닝에서 트랜스포머가 지배적인 위치를 차지하고 있음에도 불구하고, 모델의 깊이가 깊어짐에 따라 트랜스포머를 학습시키는 것은 특히 불안정성에 취약한 문제가 되고 있습니다. 이러한 불안정성은 두 가지 핵심 영역인 순방향 불안정성(수치적 문제를 야기하는 은닉 상태의 통제되지 않는 성장)과 역방향 불안정성(역전파 중 그래디언트 폭주)으로 나타납니다.

이러한 문제를 완화하기 위해 사용되는 주요 아키텍처 구성 요소는 **레이어 정규화(Layer Normalization, LN)**입니다. 그러나 트랜스포머 블록 내에서 LN의 배치 방식은 엄격한 이론적 근거보다는 주로 경험적 휴리스틱을 통해 발전해 왔습니다. 본 논문은 세 가지 주요 구성을 식별합니다:

• Post-LN: 잔차 연결(residual connections) 이후에 LN을 적용합니다. 종종 섬세한 최적화 스케줄링을 요구하며 성능이 최적에 미치지 못하는 경우가 많습니다.
• Pre-LN: 어텐션 및 피드포워드 모듈 이전에 LN을 적용합니다. Post-LN보다 학습 안정성이 개선되어 표준적으로 사용되지만, 과도하게 큰 은닉 상태를 생성하여 깊은 네트워크에서 수치적 불안정성과 그래디언트 폭주를 초래할 수 있다는 것이 경험적으로 관찰되었습니다.
• Peri-LN: 모듈의 입력과 출력 모두에 LN을 적용하는 새로운 구성입니다. 경험적 연구에 따르면 Pre-LN보다 더 우수한 안정성과 더 규칙적인 은닉 상태를 제공하는 것으로 제안되지만, 그 이론적 특성은 여전히 제대로 이해되지 않은 상태입니다.

본 논문이 다루는 핵심 문제는 왜 이러한 서로 다른 배치들이 각기 다른 학습 역학 및 안정성 프로파일을 유도하는지에 대한 원칙적인 이론적 프레임워크가 부족하다는 점이며, 학습이 모델을 규칙적인 솔루션으로 이끄는지 아니면 병리적인 동작으로 이끄는지에 대한 문제입니다.

방법론

저자들은 트랜스포머 학습 역학을 분석하기 위해 **이산 시간 최적 제어 이론(discrete-time optimal control theory)**에 기반한 새로운 이론적 프레임워크를 제안합니다.

1. 최적 제어로의 정식화: 트랜스포머 학습 문제를 이산 시간 평균장 제어(mean-field control) 문제로 간주합니다. 레이어 인덱스 $i$는 이산 시간을 나타내고, 은닉 상태 $X_i$는 시스템 상태를 나타냅니다. 잔차 연결은 역학의 속도(velocity) 역할을 합니다.
2. 해밀턴-야코비-벨만(Hamilton-Jacobi-Bellman, HJB) 분석: 저자들은 질량 운송 방정식(mass transport equations)과 결합된 HJB 편미분 방정식(PDE)을 사용하여 학습 문제의 최적 조건을 분석합니다.
   • 이들은 은닉 상태를 위한 최적 속도장을 정의하는 피드백 제어 메커니즘 역할을 하는 **해밀토니안(Hamiltonian)**을 조사합니다.
   • Pre-LN의 경우, 해밀토니안이 유계(unbounded)임을 입증하여 은닉 상태가 임의로 커질 수 있는 부적절한(ill-posed) 최적 조건이 발생함을 보여줍니다.
   • Peri-LN의 경우, 출력 정규화가 속도장을 타원체(ellipsoid)로 제한하여(Lemma 1 기반), 해밀토니안이 유계가 되도록 보장하고 문제가 적절하게 정의(well-posed)되도록 합니다.
3. 그래디언트 분석: 레이어를 가로지르는 국소 민감도(Jacobian 행렬)의 곱을 통해 그래디언트 전파를 조사함으로써 역방향 안정성을 분석합니다.
4. 잔차 스케일링(Residual Scaling): 이 프레임워크는 신호 전파를 제어하기 위해 잔차 단계의 크기를 조절하는 시간 증분 파라미터 $\Delta t < 1$을 포함하도록 확장됩니다.

핵심 기여

1. 학습된 모델을 위한 이론적 프레임워크

모델 초기화에 집중하거나 경험적 증거에 의존하는 기존 연구와 달리, 이 프레임워크는 수렴 시점의 학습된 모델을 분석합니다. 이는 학습 역학이 규칙적인 솔루션을 유도하는지 아니면 병리적인 현상을 유도하는지에 대한 체계적인 평가를 제공합니다.

2. LN 배치에 대한 엄격한 분석

• Pre-LN 불안정성: 본 논문은 Pre-LN의 경우 학습 문제의 최적해가 크기 측면에서 유계가 아님을 증명합니다. 가중치 감쇠(weight decay)를 적용하더라도 은닉 상태는 지수적으로($O(e^D)$) 성장할 수 있으며, 그래디dt은 활성화 값에 비례하여 성장하여 학습 불안정성을 초래합니다.
• Peri-LN 안정성: 저자들은 Peri-LN이 은닉 상태의 평균 절대값에 대해 선형 성장($O(D)$)을, 분산에 대해 이차 성장($O(D^2)$)을 보장한다는 명시적인 경계값을 도출합니다. 결정적으로, Peri-LN의 국소 민감도는 **활성화 크기에 불변(invariant)**하여, 활성화 값이 크더라도 그래디언트 폭주를 방지함을 보여줍니다.

3. 불확실성 정량화

이 분석은 Peri-LN이 입력에서 최종 은닉 상태까지(Wasserment distance를 통해 측정된) 불확실성의 전파를 제어하는 반면, Pre-LN은 입력 불확실성을 무한한 수준까지 증폭할 수 있음을 확립합니다.

4. 잔차 단계 스케일링

안정성 분석에 따라, 저자들은 잔차 단계의 크기를 조절하는 잔차 단계 스케일링($\Delta t < 1$) 메커니즘을 도입합니다. 이들은 이 스케일링이 다음을 수행함을 이론적으로 입증합니다:

• 은닉 상태 성장 및 출력 불확실성에 대한 경계값을 날카롭게 만듭니다.
• 국소 민감도를 스케일링하여 그래디언트 폭주를 완화합니다.
• 추가적인 계산 비용 없이 순방향 및 역방향 안정성을 모두 개선합니다.

실험 결과

저자들은 OpenWebText 데이터셋으로 학습된 GPT-2 변체(124M ~ 1.5B 파라미터)를 사용하여 이론적 발견을 검증합니다.

• 학습 안정성: 5회의 시행 실험에서, Pre-LN 모델은 가중치 감쇠를 활성화했을 때 1회, 비활성화했을 때 3회에서 발산했습니다. 반면, Peri-LN 모델은 가중치 감쇠 설정에 관계없이 5회 모두 안정성을 유지했습니다.
• 성능: Peri-LN은 Pre-LN을 위해 하이퍼파라미터를 튜닝했을 때와 대등한 성능을 달 기록했습니다.
• 스케일링의 영향: Peri-LN과 잔차 스케일링($\Delta t = 0.1$)을 결합했을 때 모든 지표(Val Loss, Perplexity, ROUGE, BERT scores)에서 일관되게 가장 좋은 성능을 보였으며, 레이어 전반에 걸쳐 은닉 상태의 모멘트(평균 절대값 및 분산) 성장을 효과적으로 제어하여 이론적 예측과 일치하는 결과를 보여주었습니다.

의의 및 주장

본 논문은 트랜스포머의 레이어 정규화 배치에 관한 경험적 관찰에 대해 최초로 엄격한 이론적 근거를 제공한다고 주장합니다.

• 불안정성 설명: 왜 Pre-LN이 근저에 깔린 HJB PDE의 부적절성(ill-posedness)을 통해 무한한 은닉 상태와 그래디언트 폭주를 유발하는지 설명합니다.
• Peri-LN의 정당성: Peri-LN이 제어된 성장과 불변하는 민감도를 유지함을 보여줌으로써 안정성에 대한 이론적 기반을 제공합니다.
• 설계 가이드라인: 새로운 아키텍처 수정 사항을 비용이 많이 드는 실제 학습 전에 "건전성 검사(sanity-check)"할 수 있는 원칙적인 방법을 제공합니다.
• 실질적 개선: 잔차 단계 스케일링의 도입은 단순하고 비용이 들지 않는 수정 사항임에도 불구하고, 안정성과 성능을 크게 향상시킨다는 것을 실험을 통해 입증하였습니다.

저자들은 본 연구가 경험적 휴리스틱과 이론적 이해 사이의 간극을 메우며, 연속 시간 관점을 통해 이론적 통찰력과 실질적인 아키텍처 개선을 모두 제공한다고 강조합니다.