SimpleGPT: 간단한 정규화로 GPT 학습 효율 극대화

초록

본 논문은 Transformer의 최적화 안정성을 2차 기하학(헤시안) 관점에서 재조명하고, 선형 변환 직후 정규화를 적용하는 SimpleNorm 방식을 제안한다. SimpleNorm은 중간 활성화 스케일을 고정하고 헤시안 스펙트럼을 크게 감소시켜 학습률을 310배 확대할 수 있게 한다. 1B8B 규모의 GPT 모델에 적용한 실험에서 기존 LLaMA2‑QKNorm 대비 훈련 손실이 0.08 낮아지는 등 일관된 성능 향상을 확인했으며, 코드는 공개 예정이다.

상세 분석

SimpleGPT 논문은 크게 두 가지 이론적·실험적 기여로 구성된다. 첫째, 저자들은 손실 함수 ℓ(x)의 β‑smoothness 정의를 통해 학습률 η가 0 < η ≤ 2/β, 즉 η ≤ 2/‖H‖₂ (H는 헤시안) 조건을 만족해야 안정적인 Gradient Descent가 가능함을 재확인한다. 기존 Transformer 설계에서는 Residual, LayerNorm, RMSNorm 등 경험적 기법이 사용됐지만, 이들이 헤시안 스펙트럼에 미치는 정량적 영향을 명시적으로 분석하지 못했다.

두 번째 핵심은 SimpleNorm(Ψ)이다. Ψ(x)=Norm(Wx) 형태로, 선형 변환 직후 바로 정규화를 수행한다. 여기서 Norm은 RMSNorm을 사용했으며, 결과는 y=√d · D · u (u는 정규화된 벡터, D는 학습 가능한 스케일 γ) 로 표현된다. 이 구조는 활성화 ‖y‖₂가 √d · γ_min ≤ ‖y‖₂ ≤ √d · γ_max 범위에 머물게 하여, 깊이가 깊어지거나 가중치가 커져도 활성화 폭발을 방지한다.

헤시안 분석에서는 두 항으로 분해된다. Gauss‑Newton 항 L = Jᵀ H_yy J 은 ‖W‖₂²에 비례하지만, SimpleNorm에서는 정규화 행렬 P와 D가 삽입돼 ‖L‖₂ ≈ τ κ² ‖H_yy‖₂ 로, W의 스펙트럼 노름에 독립적인 상수 τ, κ에 의해 제한된다. 반면, 정규화 자체가 만든 곡률 항 C는 ‖C‖₂ ≤ 3 κ² √d ‖g_y‖₂ 로, 고차원에서 L에 비해 무시할 정도로 작다(정리 4.1). 따라서 전체 헤시안 ‖H‖₂는 선형 레이어만 사용했을 때보다 크게 감소한다.

실험에서는 1B, 1.4B, 7B, 8B 규모의 GPT 모델에 SimpleNorm을 적용하고, 기존 LLaMA2‑QKNorm, DeepNorm, PreNorm 등과 동일한 학습 설정에서 비교했다. 7B 모델을 60 K 스텝 학습했을 때 최종 훈련 손실이 2.208(기존 2.290)으로 0.08 감소했으며, 학습률을 3~10배 높여도 발산 없이 안정적으로 수렴했다. 또한, 학습 초기에 그래디언트 노름이 작아져 초기 단계에서의 과도한 파라미터 업데이트를 억제함을 확인했다.

이 논문은 “정규화는 어디에 두는가”라는 설계 질문을 “정규화가 헤시안 스펙트럼을 어떻게 축소하는가”라는 수학적 질문으로 전환함으로써, 기존 경험주의적 설계에 이론적 근거를 제공한다. SimpleNorm은 구현이 간단하고 연산 비용이 기존 RMSNorm과 동일하거나 약간 증가하는 수준이므로, 대규모 LLM 훈련 파이프라인에 바로 적용 가능하다.