GeoNorm: 사전·사후 정규화를 하나로 묶는 지오데식 최적화

초록

본 논문은 Transformer의 정규화 위치 문제를 구면 매니폴드 위의 최적화로 재해석하고, 기존의 Pre‑Norm·Post‑Norm을 통합하는 새로운 정규화 기법 GeoNorm을 제안한다. FFN·Attention의 출력은 구면 위의 업데이트 방향으로 보고, 투영 대신 지오데식(지수 지도) 이동을 수행한다. 층별 학습률 감쇠 스케줄을 도입해 안정성을 높였으며, 다양한 모델·데이터셋에서 기존 정규화 방법들을 일관적으로 능가함을 실험적으로 입증한다. 구현 비용은 거의 추가되지 않는다.

상세 분석

GeoNorm은 Transformer 레이어를 구면(ℓ₂‑노름이 일정한 구) 위에서 진행되는 반복 최적화 과정으로 모델링한다. 기존 Post‑Norm은 “x + 모듈(x)” 후 ℓ₂‑노름으로 정규화해 구면에 투영하는 방식인데, 이는 외부적인 투영으로 인해 업데이트 방향이 왜곡되고 수렴 속도가 저하될 위험이 있다. 저자들은 이를 Riemannian 최적화 관점에서 바라보아, 업데이트 벡터 sₖ를 현재 점 xₖ의 접공간 TₓₖΩ에 정사영(vₖ = sₖ – (xₖᵀsₖ/‖xₖ‖²)xₖ) )하고, 구면 위의 지오데식 이동을 구현하는 지수 지도 expₓₖ(αₖvₖ)를 적용한다. 구면에 대한 지수 지도는 닫힌 형태인
expₓ(v)=cos(‖v‖/‖x‖)·x + sin(‖v‖/‖x‖)·(‖x‖/‖v‖)·v
으로, 기본적인 벡터 연산만으로 계산 가능해 계산량이 거의 증가하지 않는다.

이러한 지오데식 업데이트는 두 가지 핵심 장점을 제공한다. 첫째, 구면 위에서의 움직임이 연속적이고 곡률 정보를 보존하므로, Attention·FFN이 생성한 “의도된” 방향을 그대로 전달한다. 둘째, 투영에 비해 수렴 이론이 더 강력하며, 특히 스무스하고 강하게 볼록한 목적함수에 대해 αₖ ≤ 2/L(또는 다항식 감쇠)와 같은 단계 크기 선택이 O(log k/k) 수렴률을 보장한다.

논문은 또한 Pre‑Norm을 GeoNorm의 특수 경우로 해석한다. Pre‑Norm은 정규화를 Residual 앞에 두어, 실제로는 동일한 지오데식 이동을 수행하지만 각 층마다 고정된 각도(비적응적)만을 사용한다는 점에서 차이가 있다. 실험 결과, GeoNorm이 Pre‑Norm과 거의 동일한 구조적 흐름을 가지면서도 학습률 감쇠와 지오데식 이동을 통해 더 안정적인 훈련과 낮은 손실을 달성한다는 것이 확인되었다.

계산 복잡도 측면에서는, 기존 LayerNorm/RMSNorm이 수행하는 ‖·‖ 정규화와 거의 동일한 연산량을 요구한다. 추가적인 파라미터는 없으며, 구현은 기존 Transformer 블록에 “Norm → exp” 단계만 교체하면 된다. 따라서 대규모 LLM에 바로 적용 가능하고, 메모리·시간 오버헤드가 미미하다.

실험에서는 ArXiv, Books3, FinWeb‑Edu 등 다양한 텍스트 코퍼스와 125M, 350M, 1.3B 규모 모델을 대상으로, 학습 길이 512와 1024에서 GeoNorm이 Pre‑Norm, Post‑Norm, DeepNorm, SandwichNorm을 모두 앞선 손실값을 기록했다. 특히 구면 위의 부드러운 이동 덕분에 깊은 네트워크(>48층)에서도 gradient 폭발·소실 현상이 크게 감소했으며, downstream fine‑tuning에서도 일관된 성능 향상을 보였다.

전체적으로 GeoNorm은 정규화 위치 논쟁을 이론적으로 통합하고, 매니폴드 최적화 기법을 실제 LLM에 적용함으로써 기존 방법들의 한계를 극복한다는 점에서 의미가 크다. 다만 구면 외의 다른 매니폴드(예: Stiefel, 저차원 저정밀 행렬)로 확장하거나, 비선형 활성화와 결합했을 때의 동작을 추가로 검증할 필요가 있다.