Sphere Retraction Normalizations

기술 요약: 구형 회귀 정규화 (Sphere Retraction Normalizations)

문제 정의
심층 신경망을 안정적으로 학습시키는 것은 여전히 중대한 과제로 남아 있다. 주로 잔차 연결(residual connections)과 정규화(예: Transformer의 Pre-LN 및 Post-LN)를 결합하는 기존 방법론들은 내재적인 한계를 보인다. Pre-LN 구조는 깊은 층이 학습에 기여하는 정도가 점진적으로 감소하는 "깊이의 저주(curse of depth)"를 겪으며, 무한히 커질 수 있는 잔차 스트림(residual streams)을 허용하여 거대한 활성화 값의 축적(attention sinks)을 초래한다. 반면, Post-LN 구조는 초기화 시 출력층 근처에서 큰 그래디언트 문제를 겪으며, 이로 인해 최적화를 늦추는 세심한 학습률 웜업(learning-rate warm-up) 스케줄이 필요하다. 최근의 Geodesic Normalization (GeoNorm)과 같은 연구는 리만 기하학(Riemannian geometry)을 통해 은닉 상태를 초구체(hypersphere) 내에 가둠으로써 학습을 안정화하려고 시도하지만, 업데이트를 오직 지수 사상(exponential map)에만 의존한다는 한계가 있다. 저자들은 지수 사상이 더 넓은 범위의 회귀 사상(retraction maps) 가족 중 단 하나의 구성원일 뿐이라는 점을 지적하며, 초구체 상의 대안적인 회귀(retraction)가 더 우수한 성능이나 안정성을 제공할 수 있는지에 대한 의문을 제기한다.

방법론
본 논문은 **회귀 사상(retraction maps)**에 기반한 초구체 상의 통합된 딥 커넥션 프레임워크인 SpheretNorms를 제안한다.

1. 통합 프레임워크: 저자들은 표준 잔차 연결(RC)과 GeoNorm을 매니폴드 $\mathcal{M}$ 상에 정의된 단일 계산 골격 내로 재구성하였다. 업데이트 규칙은 다음과 같이 추상화된다:
   $$h_l = R_{h_{l-1}}\left[ \alpha_l \cdot \text{proj}_{T_{h_{l-1}}\mathcal{M}}(F_l(h_{l-1})) \right]$$
   여기서 $R$은 회귀 사상이다. GeoNorm이 초구체($S^{d-1}$)와 지수 사상을 선택하는 반면, 이 프레임워크는 다른 회귀 사상들을 허용한다.

2. 구체적 구현:

   • Proj-SpheretNorm: 이 프레임워크를 메트릭 투영 회귀(metric projection retraction)(gnomonic retraction으로도 알려짐)를 사용하여 구현한다. 이는 현재 상태와 업데이트의 벡터 합을 초구체 위에 놓이도록 정규화하는 과정을 포함한다.
   • Cay-SpheretNorm: 현재 상태와 업데이트 벡터로부터 생성된 반대칭 행렬의 케일리 변환(Cayley transform)을 통해 유도되는 케일리 회귀(Cayley retraction)(stereographic retraction)를 사용한다. 이는 폐쇄형 대수적 업데이트를 제공한다.

3. 일반화 (p-SpheretNorm): 저자들은 Proj-SpheretNorm과 Cay-SpheretNorm이 모두 업데이트 크기가 회전 각도로 변환되는 평면 회전 구조를 공유한다는 점을 관찰하였다. 이를 바탕으로, 회전 각도 $\theta_l = p \arctan(\beta_l / p)$ (여기서 $\beta_l$은 정규화된 스텝 크기)로 정의되는 단일 파라미터 가족인 p-SpheretNorm을 도입한다.

   • $p=1$은 Proj-SpheretNorm을 복원한다.
   • $p=2$는 Cay-SpheretNorm을 복원한다.
   • $p \to \infty$는 GeoNorm으로 수렴한다.
   • $p \to 0^+$는 항등 사상(identity mapping)으로 수렴한다.
   • 결정적으로, 유한한 $p$에 대해 회전 각도는 포화(saturate)되어, 지수 사상과 달리 무한한 성장을 방지한다.

4. 구현 세부 사항: 안정성을 보장하기 위해 저자들은 학습 가능한 스텝 크기 $\alpha_l$에 대해 감쇠 인자 $\lambda_l = 1/\sqrt{l}$를 사용하며, 첫 번째 레이어에 수정된 ScaleNorm을 적용하여 은닉 상태의 반지름을 초기화한다.

주요 기여

1. 이론적 통합: 본 논문은 유클리드 잔차 연결과 GeoNorm을 회귀 기반 업데이트로 통합하여, GeoNorm이 지수 사상에 의해 구동되는 리만 잔차 연결의 특수한 사례임을 입증하였다.
2. 새로운 정규화 레이어: 지수 사상에 필요한 삼각 함수의 계산 비용을 피하면서 대수적 연산만을 사용하는 두 가지 새로운 노름 보존(norm-preserving) 정규화 레이어를 도입하였다 (Proj-SpheretNorm 및 Cay-SpheretNorm).
3. 파라미터 가족: 이 방법들을 항등 사상, 대수적 회귀, 그리고 지수 사상 사이를 연결하는 단일 파라미터 가족인 p-SpheretNorm으로 일반화하였다. 저자들은 이 가족의 모든 구성원이 동일한 1차 접선 역학(first-order tangent dynamics)을 갖는 2차 회귀(second-order retractions)임을 증명하였다.
4. 실증적 우수성: nanoGPT(최대 36개 레이어, 6.55B 토큰)에 대한 광범위한 실험 결과, 1-SpheretNorm(p=1에 해당)이 OpenWebText 및 FineWeb-Edu 데이터셋에서 Pre-LN, Pre-DyT, Peri-LN, Keel, GeoNorm보다 낮은 사전 학습 및 검증 손실과 가장 높은 제로샷 정확도를 달성하였다. 특히, 1-SpheretNorm은 Pre-LN과 Keel이 발산하는 설정에서도 수렴하였다.

결과

• 사전 학습: 중간 규모(24 레이어) 및 대규모(36 레이어) 모델에서 1-SpheretNorm이 가장 낮은 검증 퍼플렉시티(perplexity)를 달ino했다. 예를 들어, 24 레이어의 FineWeb-Edu에서 검증 퍼플렉시티를 Peri-LN의 19.09에서 16.81로 낮추었다.
• 안정성: Pre-LN(OpenWebText의 중간 규모에서 발산) 및 Keel(FineWeb-Edu의 소규모 설정에서 발산)과 달리, 모든 SpheretNorm 변형은 테스트된 모든 설정에서 수렴하였다.
• 다운스트림 태스크: FineWeb-Edu로 학습된 대규모 모델에서 SpheretNorm 변형들은 12개 지표 중 11개에서 최고의 점수를 기록하였고, OpenWebText에서는 12개 중 10개에서 최고의 점수를 기록하였다. 이러한 이득은 퍼플렉시티 지표(예: LAMBADA 및 WikiText-2)에서 가장 두드러졌으며, 1-SpheretNorm은 GeoNorm 및 기타 베이스라인을 크게 압도하였다.
• 파라미터 민감도: 결과는 최적의 성능이 유한한 $p$(특히 $p=1$)에서 달성됨을 나타내며, 이는 지수 사상(p가 $\infty$로 갈 때의 극한)이 초구형 잔차 스트림을 위한 선호되는 회귀가 아니라 단지 스펙트럼의 한 끝단임을 시사한다.

의의 및 주장
본 논문은 회귀 사상의 선택이 초구체 상의 딥 커넥션에 있어 결정적인 설계 요소라고 주장한다. 대수적 회귀(특히 메트릭 투영)를 통해 지수 사상을 넘어섬으로써, 노름 보존과 학습 안정성을 달성하면서도 더 우수한 성능을 낼 수 있음을 입증하였다. 이 연구는 "지수 사상"이 매니폴드 상의 안정적인 딥 러닝을 위한 근본적인 필수 요소가 아니라, 더 넓은 스펙트럼 내의 특정 지점일 뿐이라는 점을 시사한다. 저자들은 딥 커넥션 메커니즘을 설계하는 것이 그래프 이론, 최적화, 그리고 동역학계(dynamical systems)가 교차하는 수학적 문제라고 결론지으며, 정보 전파와 야코비안 조건수(Jacobian conditioning) 사이의 최적의 트레이드오프를 식별하는 것이 향후 연구를 위한 열린 방향임을 밝힌다. 또한 특화된 아키텍처(예: 등변 신경망)에 대한 고려와 현재 범위 이상의 더 큰 규모의 모델에 대한 평가 필요성 등 한계점도 인정하였다.