기술 요약: SimpleGPT: 단순한 정규화 전략을 통한 GPT의 개선

문제 정의

Transformer 기반 대규모 언어 모델(LLM)이 깊이와 너비 측면에서 확장됨에 따라, 최적화 안정성이 성능과 확장성의 병목 현상으로 작 increasingly 중요해지고 있다. 현대적 아키텍처는 잔차 연결(residual connections), 정규화 레이어, 특정 비선형성 등 다양한 구성 요소를 사용하여 학습을 안정화하지만, 이러한 아키텍처적 선택과 고전적 최적화 안정성 조건 사이의 이론적 관계는 여전히 충분히 탐구되지 않았다. 특히, 네트워크 설계, 활성화 스케일, 손실 함수의 헤시안(Hessian) 행렬, 그리고 허용 가능한 최대 학습률 사이의 관계는 명확히 이해되지 않고 있다. 현재의 방법론들은 이들이 2차 최적화 기하학(second-order geometry)에 어떻게 영향을 미치는지 명시적으로 분석하지 않은 채, 경험적 휴리스틱에 의존하는 경우가 많다.

방법론

이론적 프레├──드워크: 2차 기하학

저자들은 2차 기하학의 관점에서 Transformer 최적화를 재검토한다. 그들은 두 번 미분 가능한 목적 함수 $\ell(x)$에 대해, 국소 곡률(local curvature)은 헤시안 행렬 $H_{xx} = \nabla^2 \ell(x)$에 의해 특징지어진다는 것을 확립한다. 경사 하강법(gradient descent)의 안정성은 그래디언트의 리프시츠 상수(Lipschitz constant) $\beta$에 의해 결정되며, 이는 헤시안의 스펙트럼 노름(spectral norm) $\beta = \sup_x \|H_{xx}\|_2$에 해당한다. 최대 허용 학습률 $\eta$는 이 스펙트럼 노름에 반비례한다:
$$\eta \leq \frac{2}{\beta} = \frac{2}{\sup_x \|H_{xx}\|_2}$$
본 논문은 기존의 많은 정규화 전략이 가중치 행렬이 학습 중에 성장함에 따라 헤시안의 스펙트럼 노름을 충분히 억제하지 못한다고 주장한다.

제안된 솔루션: SimpleNorm

저자들은 네트워크 내의 모든 선형 매핑 직후에 정규화 연산자(본 구현에서는 RMSNorm)를 배치하는 정규화 전략인 SimpleNorm을 소개한다.

• 정의: 입력 $x$와 선형 변환 $W$가 주어졌을 때, SimpleNorm은 $\Psi(x) = \text{Norm}(Wx)$로 정의된다.
• 아키텍처 (SimpleGPT): 저자들은 모든 표준 선형 투영(어텐션 메커니즘의 $W_q, W_k, W_v, W_o$ 및 MLP 레이어 $W_1, W_2$ 등)을 SimpleNorm 연산자로 교체한 새로운 아키텍처인 SimpleGPT를 구축한다. 기존의 PreNorm 또는 PostNorm과 달리, SimpleNorm은 "선형 매핑 직후의 정규화"를 하나의 통합된 프리미티브(primitive)로 취급한다.

이론적 메커니즘

본 논문은 SimpleNorm이 최적화를 개선하는 두 가지 이론적 메커니즘을 제공한다.

1. 안정적인 활성화 스케일: 구조적으로 SimpleNorm은 중간 활성화 값의 노름을 $\sqrt{d}$ 차수(여기서 $d$는 차원)로 재조정하여, 깊이가 증가함에 따라 발생할 수 있는 활성화 폭발이나 드리프트를 방지한다.
2. 더 매끄러운 손실 지형 (헤시안 감소): 저자들은 네트워크 활성화에 대한 손실의 헤시안을 분석한다. 그들은 SimpleNorm의 경우 헤시안이 가우스-뉴턴 항($L$)과 곡률 항($C$)으로 분해됨을 증명한다.
   • 고차원에서는 가우스-뉴턴 항이 지배적이다 ($\|C\|_2 \ll \|L\|_2$).
   • 결정적으로, SimpleNorm 헤시안의 스펙트럼 노름은 가중치 행렬의 스펙트럼 노름 $\|W\|_2$에 대해 **스케일 불변(scale-invariant)**이다.
   • 이와 대조적으로, 표준 선형 투영의 경우 헤시안 스펙트럼 노름은 $\|W\|_2$에 따라 제곱으로 스케일링된다.
   • $\|W\|_2$는 일반적으로 학습 중에 성장하므로, SimpleNorm은 훨씬 작은 헤시안 스펙트럼 노름을 유지하여 훨씬 더 큰 안정적인 학습률을 허용한다.

주요 기여

1. 이론적 연결: 본 논문은 Transformer의 아키텍처 설계, 활성화 스케일, 헤시안 행렬, 그리고 최대 허용 학습률 사이의 직접적인 연결 고리를 확립했다.
2. SimpleNorm 및 SimpleGPT: 활성화 스케일을 안정화하고 헤시안 스펙트럼 노름을 줄이는 단순한 배치 전략인 SimpleNorm을 도입했다. 저자들은 이 설계가 더 작은 리프시츠 그래디언트 상수를 생성함을 이론적으로 입증했다.
3. 경험적 검증: nanoGPT, LLaMA2, LLaMA3 아키텍처를 사용하여 다양한 모델 규모(1B, 1.4B, 7B, 8B 파라미터)에 걸쳐 광범 \히 실험을 수행했다.
   • SimpleGPT는 표준 관례보다 3배에서 10배 더 큰 학습률을 견딜 수 있다.
   • QKNorm이 적용된 LLaMA2와 같은 강력한 베이스라인과 비교하여 일관된 최적화 안정성과 우수한 성능을 보여준다.
   • 구체적 결과: 7B 규모 모델을 60K 스텝 동안 학습시켰을 때, SimpleGPT는 LLaMA2(QKNorm 포함)의 2.290에 비해 2.208의 학습 손실을 달s성했다 (0.08 감소).
4. 효율성: 이 방법은 계산 오버헤드가 매우 적으며(학습 시간 약 3% 증가), 이는 메모리 제한적인 정규화 연산 때문이지만 torch.compile 퓨전(fusion)을 통해 완화된다.

결과

저자들은 다음 모델들에 대해 검증을 완료했다:

• 1B 파라미터 모델: LLaMA2 기반의 SimpleGPT 1B는 LLaMA2(QKNorm 포함) 대비 학습 손실을 2.478에서 2.466으로 낮추었다.
• 7B 파라미터 모델: SimpleGPT 7B는 20K, 40K, 60K 학습 스텝 전반에 걸쳐 일관된 이득을 보였으며, 학습이 진행됨에 따라 성능 격차가 확대되거나 안정적으로 유지되었다.
• 8B 파라미터 모델: LLaMA3 기반의 SimpleGPT 8B는 3배 더 큰 학습률을 사용했음에도 불구하고 LLaMA3(QKNorm 포함)보다 실질적으로 낮은 학습 손실을 달성했다.
• 1.4B 파라미터 모델: nanoGPT 기반의 SimpleGPT 1.4B는 검증 손실(validation loss)에서 약 0.043의 개선을 보였다.

실험 결과는 SimpleGPT가 표준 아키텍처(PreNorm 또는 PreNorm+QKNorm)가 발산하는 학습률에서도 안정적인 학습을 가능하게 함을 확인시켜 준다.

의의 및 주장

본 논문은 Transformer 설계를 고전적인 2차 최적화 원칙에 근거함으로써, SimpleNorm이 최적화 안정성에 대한 통합된 이론적 설명을 제공한다고 주장한다. 그 의의는 다음과 같다:

• 원칙에 기반한 설계: 경험적 휴리스틱을 넘어 헤시안의 스펙트럼 노름에 의해 유도되는 설계로 나아갔다.
• 확장성: 매우 깊고 거대한 규모의 모델을 훨씬 더 큰 학습률로 학습할 수 있게 하여, 수렴을 가속화하고 최종 성능을 향상시킨다.
• 단순성: 복잡한 아키텍처 수정이나 암시적 유계 비선형성(implicit bounded nonlinearities)에 의존하는 무정규화 방식 대신, 단순한 구조적 변화를 통해 이러한 이득을 달성했다.

저자들은 SimpleNorm이 두 가지 효과를 통해 성능을 개선한다고 결론짓는다: 국소적으로는 곡률 변동성을 줄여 최적화 기하학을 개선하고, 전역적으로는 편재된 정규화 유도 비선형성을 통해 표현 능력을 높인다.