레이어 정규화 하에서의 Transformer 안정성

기본 정보

• 논문 ID: 2510.09904
• 제목: Stability of Transformers under Layer Normalization
• 저자: Kelvin Kan (UCLA), Xingjian Li (UT Austin), Benjamin J. Zhang (UNC Chapel Hill), Tuhin Sahai (SRI International), Stanley Osher (UCLA), Krishna Kumar (UT Austin), Markos A. Katsoulakis (UMass Amherst)
• 분류: cs.LG, cs.AI, math.OC
• 발표 시간: 2025년 10월 10일
• 논문 링크: https://arxiv.org/abs/2510.09904

초록

Transformer가 광범위하게 사용되고 있음에도 불구하고, 깊은 Transformer의 훈련은 불안정할 수 있습니다. 레이어 정규화(Layer Normalization)는 표준 구성 요소로서 훈련 안정성을 개선할 수 있지만, 그 위치 선택은 종종 임시적입니다. 본 논문은 서로 다른 레이어 정규화 위치 하에서 Transformer의 전방향 안정성(은닉 상태)과 후방향 안정성(기울기)에 대한 원리적 연구를 수행합니다. 이론적 분석은 훈련 동역학의 핵심 통찰력을 드러냅니다: 훈련이 Transformer를 규칙적인 해로 이끌어가는지 또는 병적 행동으로 이끌어가는지 여부입니다. 전방향 안정성 측면에서, 훈련 후 Transformer의 은닉 상태 증가에 대한 명시적 경계를 도출했습니다. 후방향 안정성 측면에서, 레이어 정규화가 기울기 역전파에 미치는 영향을 분석하여 각 레이어 정규화 위치의 훈련 동역학을 설명합니다. 분석은 또한 Transformer 블록의 잔차 스텝 크기 조정을 지도하며, 적절한 선택은 안정성과 성능을 더욱 개선할 수 있습니다.

연구 배경 및 동기

문제 정의

본 연구가 해결하려는 핵심 문제는: 서로 다른 레이어 정규화 위치가 Transformer 훈련 안정성에 미치는 영향 메커니즘입니다. 구체적으로는:

1. 전방향 안정성 문제: 깊은 네트워크에서 은닉 상태의 증가 제어
2. 후방향 안정성 문제: 역전파 과정에서 기울기의 안정성
3. 아키텍처 설계 지침: 새로운 Transformer 변형에 대한 이론적 지침 제공 방법

중요성 분석

1. 실용적 가치: Transformer는 현대 심층 학습의 기초이며, 그 훈련 안정성은 모델 성능과 훈련 효율에 직접 영향을 미칩니다
2. 이론적 공백: 기존 레이어 정규화 위치 선택은 주로 경험에 기반하며 이론적 근거가 부족합니다
3. 산업 수요: 모델 규모가 계속 증가함에 따라 훈련 안정성 문제가 더욱 심각해집니다

기존 방법의 한계

1. Post-LN: 정교한 최적화 스케줄이 필요하며, 성능이 종종 차선입니다
2. Pre-LN: 초기 훈련 안정성을 개선하지만 과도하게 큰 은닉 상태를 생성하여 수치 불안정성을 초래합니다
3. Peri-LN: 실제로는 좋은 성능을 보이지만 이론적 성질이 충분히 이해되지 않습니다

연구 동기

저자들은 연속 시간 동역학 및 최적 제어 이론의 새로운 관점을 채택하여 Transformer 훈련 문제를 평균장 제어 문제로 모델링함으로써, 초기화 시 행동만이 아닌 훈련 수렴 후 모델 성질을 분석할 수 있습니다.

핵심 기여

1. 이론적 프레임워크 혁신: 최적 제어 이론에 기반한 새로운 프레임워크를 제안하여 서로 다른 레이어 정규화 위치 하에서 Transformer 안정성을 체계적으로 분석합니다
2. 전방향 안정성 분석: 은닉 상태 증가에 대한 명시적 경계를 도출하여 Pre-LN은 무한 증가를 초래하고 Peri-LN은 제어된 증가를 유지함을 증명합니다
3. 후방향 안정성 분석: 레이어 정규화가 기울기 역전파에 미치는 영향 메커니즘을 드러냅니다
4. 잔차 스텝 크기 조정: 안정성과 성능을 개선하는 잔차 스텝 크기 조정 방법을 제안합니다
5. 실험 검증: GPT-2 시리즈 모델에서 이론적 발견을 검증합니다

방법 상세 설명

작업 정의

서로 다른 레이어 정규화 위치 하에서 Transformer의 안정성을 연구합니다:

• 입력: 임베딩 및 위치 인코딩 후의 시퀀스 $X_0 \in \mathbb{R}^{d \times n}$
• 출력: D개 Transformer 블록을 통과한 후의 은닉 상태 $X_D$
• 목표: 전방향 및 후방향 전파의 안정성 분석

연속 시간 모델링

Transformer의 연속 시간 표현

표준 Transformer 블록의 스킵 연결 구조를 연속 시간 동역학의 오일러 이산화로 해석합니다:

$\frac{dX(t)}{dt} = \begin{cases} f_{attn}(X(t), t; \theta_{attn}(t)), & t \in [t_i, t_i + \Delta t) \\ f_{ffn}(X(t), t; \theta_{ffn}(t)), & t \in [t_i + \Delta t, t_{i+1}) \end{cases}$

여기서 $\Delta t = \frac{T}{2D}$, $t_i = 2i\Delta t$입니다.

평균장 제어 문제 모델링

훈련 문제를 연속 시간 평균장 제어 문제로 표현합니다:

$\min_\theta \mathbb{E}_{(X_0,y)} G(X(T), y)$$\text{s.t. } \frac{dX(t)}{dt} = f(X(t), t; \theta(t))$

여기서 $f \in \{f_{Pre}, f_{Peri}\}$는 서로 다른 레이어 정규화 위치에 해당합니다.

레이어 정규화의 기하학적 성질

핵심 보조정리 1: 레이어 정규화 출력은 타원면 위에 위치합니다 $\mathcal{E} = \{z \in \mathbb{R}^d : (z - \beta)^T\Gamma^{-2}(z - \beta) = d\}$ 여기서 $\Gamma = \text{diag}(\gamma)$입니다.

전방향 안정성 분석

Pre-LN의 무한성

정리 2: Pre-LN 훈련 문제의 최적해는 크기 면에서 무한합니다.

증명 개요: Hamilton-Jacobi-Bellman (HJB) 편미분방정식을 분석하여 대응하는 해밀턴 함수가 존재하지 않음을 증명하여 훈련 문제가 퇴화됨을 보입니다.

정리 3: 가중치 감소를 사용하더라도 Pre-LN Transformer의 은닉 상태는 지수적으로 증가합니다: $MA(X_D) \leq (1 + C(\lambda))^D \frac{\|X_0\|_F}{\sqrt{nd}} = O(e^D)$

Peri-LN의 제어된 증가

정리 4: Peri-LN Transformer의 은닉 상태는 선형적으로 증가합니다: $MA(X_D) \leq \frac{\|X_0\|_F}{\sqrt{nd}} + 2D(\gamma_{max} + \beta_{max}) = O(D)$

분산은 이차적으로 증가합니다: $\text{Var}(X_D) \leq \frac{(\|X_0\|_F + 2D\sqrt{nd}(\gamma_{max} + \beta_{max}))^2}{nd - 1} = O(D^2)$

후방향 안정성 분석

기울기 계산 공식: $\nabla_{\theta_i} G(X_D) = \nabla_{\theta_i} X_{i+1} \cdot J_{i:D} \cdot \nabla_{X_D} G(X_D)$

여기서 야코비안 행렬: $J_{i:D} = \prod_{j=i+1}^D (I + \nabla_{X_{j-1}} f(X_{j-1}; \theta_{j-1}))$

명제 7: Pre-LN 하에서 민감도 $\nabla_{X_{j-1}} f_{Pre}$는 활성화 값에 비례하여 증가합니다.

명제 8: Peri-LN 하에서 민감도 $\nabla_{X_{j-1}} f_{Peri}$는 활성화 값 크기에 불변입니다.

실험 설정

데이터셋

• OpenWebText 데이터셋: 약 90억 훈련 토큰, 400만 검증 토큰
• GPT-2 시리즈 아키텍처를 사용한 사전훈련

모델 구성

• GPT-2 (1억 2,400만 매개변수)
• GPT-2 Large (7억 7,400만 매개변수)
• GPT-2 XL (15억 매개변수)

평가 지표

• 혼란도 (Perplexity)
• ROUGE 점수 (Rouge1, Rouge2, RougeL)
• BERT 점수 (BertP, BertR, BertF1)
• 훈련 안정성: 발산 실행 횟수 통계

구현 세부사항

• Pre-LN을 위해 조정된 하이퍼파라미터 사용, Peri-LN에 대해 별도로 최적화하지 않음
• 잔차 스텝 크기 조정: $\Delta t \in \{0.1, 1\}$
• 하드웨어: NVIDIA H200 GPU

실험 결과

훈련 안정성 비교

성능 비교 결과

GPT-2 (1억 2,400만) 모델 결과:

• Pre-LN ($\Delta t=1$): 검증 손실 5.43, 혼란도 247.52
• Pre-LN ($\Delta t=0.1$): 검증 손실 3.13, 혼란도 24.43
• Peri-LN ($\Delta t=1$): 검증 손실 3.12, 혼란도 24.17
• Peri-LN ($\Delta t=0.1$): 검증 손실 3.10, 혼란도 23.63

은닉 상태 증가 분석

실험은 이론적 예측을 검증합니다:

• Pre-LN은 더 큰 $\Delta t$에서 빠른 증가를 나타냅니다
• Peri-LN은 더 규칙적인 선형 증가를 유지합니다
• 잔차 스텝 크기 조정은 증가율을 효과적으로 제어합니다

잔차 스텝 크기 조정 효과

1. 성능 향상: Peri-LN + $\Delta t=0.1$은 모든 지표에서 최고의 성능을 보입니다
2. 안정성 개선: Pre-LN은 $\Delta t=0.1$에서 불안정에서 안정으로 변합니다
3. 증가 제어: 은닉 상태의 평균 및 분산 증가율을 효과적으로 감소시킵니다

관련 연구

레이어 정규화 연구

• Post-LN: 가장 초기의 Transformer 설계, 정교한 스케줄이 필요함
• Pre-LN: 훈련 안정성을 개선하지만 큰 활성화 값을 생성함
• Peri-LN: 최근 Gemma2, OLMo2 같은 대규모 모델에서 채택됨

이론적 분석 방법

• 기존 연구는 주로 초기화 시 행동에 초점을 맞추거나 경험적 관찰에 의존합니다
• 본 논문은 훈련 수렴 후 모델 성질을 분석하는 혁신적 접근을 합니다
• 연속 시간 관점은 아키텍처 분석을 위한 새로운 도구를 제공합니다

결론 및 논의

주요 결론

1. Pre-LN 이론적 결함: 훈련 문제는 본질적으로 병적이어서 무한 해를 초래합니다
2. Peri-LN 장점: 잘 정의된 최적화 문제와 제어된 은닉 상태 증가를 제공합니다
3. 잔차 조정 가치: 간단하고 효과적인 안정성 개선 방법입니다

한계

1. 단순화된 가정: 이론적 분석은 연속 시간 근사에 기반합니다
2. 하이퍼파라미터 의존성: 실험은 Pre-LN 조정 하이퍼파라미터를 사용합니다
3. 규모 제한: 실험은 주로 중간 규모 모델에서 수행됩니다

향후 방향

1. 아키텍처 필터링 프레임워크: 새로운 아키텍처 수정을 위한 이론적 필터링 기준 제공
2. 더 큰 규모 검증: 더 큰 모델에서 이론적 발견 검증
3. 다른 정규화 방법: RMSNorm 등의 변형으로 분석 확장

심층 평가

장점

1. 이론적 혁신성 강함: 최적 제어 이론을 사용하여 레이어 정규화 위치 문제를 분석한 첫 시도
2. 수학적 엄밀성: 완전한 이론적 도출 및 증명 제공
3. 실용적 가치 높음: 잔차 스텝 크기 조정 방법이 간단하고 효과적
4. 합리적 실험 설계: 여러 모델 규모에서 이론 검증

부족한 점

1. 이론과 실제의 간격: 연속 시간 가정과 실제 이산 구현 간의 차이
2. 제한된 실험 범위: 주로 GPT-2 시리즈에서 검증, 더 많은 아키텍처 검증 부족
3. 하이퍼파라미터 공정성: Peri-LN에 대한 전문적 하이퍼파라미터 최적화 미실시

영향력 평가

1. 학술적 기여: Transformer 안정성 분석을 위한 새로운 이론적 프레임워크 제공
2. 실용적 가치: 실제 모델 설계 및 훈련 전략 지침
3. 재현성: 코드 및 모델 공개 약속

적용 가능 시나리오

1. 깊은 Transformer 훈련: 특히 대규모 깊은 모델에 적합
2. 아키텍처 설계 지침: 새로운 아키텍처 수정을 위한 이론적 근거 제공
3. 훈련 안정성 개선: 잔차 조정을 통한 훈련 안정성 향상

참고문헌

논문은 다음을 포함한 여러 중요한 연구를 인용합니다:

• Ba et al. (2016): 레이어 정규화 원본 논문
• Xiong et al. (2020): Pre-LN vs Post-LN 비교 연구
• Kim et al. (2025): Peri-LN의 경험적 연구
• He et al. (2016): 잔차 연결의 획기적 연구

종합 평가: 이는 이론과 실제가 잘 결합된 고품질 논문으로, Transformer 안정성 분석을 위한 새로운 수학적 프레임워크를 제공하며 중요한 학술적 가치와 실용적 의의를 지닙니다. 이론적 분석은 엄밀하고 심도 있으며, 실험 검증은 충분하고, 심층 학습 아키텍처 설계에 가치 있는 지침을 제공합니다.