기술 요약: 프리-노름 트랜스포머의 게이트형 정규화 제거 및 스케일 앵커링

문제 제기

정규화 레이어(예: LayerNorm, RMSNorm)는 트랜스포머 아키텍처의 표준 구성 요소로, 활성화 값을 안정화시키고 깊은 네트워크의 훈련을 가능하게 합니다. 그러나 이러한 레이어는 추론 과정에서 상당한 오버헤드를 유발하는 샘플 의존적 비선형 계산을 도입합니다. 이러한 토큰별 통계치가 특히 프리-노름 트랜스포머에서 전체 추론 단계 내내 엄격히 필요한지 여부는 여전히 불분명합니다. 이전 연구들은 초기화 트릭이나 순차적 파인튜닝을 통해 정규화를 제거하려 시도했으나, 내부 정규화 레이어와 최종 정규화 레이어의 고유한 기능적 역할, 특히 프리-로그이트 표현의 스케일을 어떻게 제어하는지에 대한 이해는 부족합니다.

방법론

저자들은 프리-노름 트랜스포머를 위해 설계된 게이트형 정규화 제거 접근법인 TaperNorm을 제안합니다. 이 방법은 모델이 훈련 중 표준 정규화에서 학습된 샘플 독립적 선형 또는 아핀 맵으로 전환할 수 있게 하여, 추론 시 정규화 연산을 제거할 수 있도록 합니다.

핵심 메커니즘: TaperNorm

TaperNorm 은 글로벌 게이트 $g \in [0, 1]$을 사용하여 정규화 브랜치와 고정 스케일링 브랜치 사이를 동적으로 보간합니다. 토큰 벡터 $h$에 대해 다음과 같습니다:
$$\text{TaperNorm}(h; g) = g \cdot \frac{h}{r(h)} D\gamma + (1 - g) \cdot c \cdot h D\tilde{\gamma}$$
여기서:

• $r(h)$는 RMSNorm 스케일링 인자입니다.
• $D\gamma$와 $D\tilde{\gamma}$는 학습 가능한 대각 스케일링 행렬입니다.
• $c$는 레이어별 스칼라입니다.
• $g$는 1(표준 정규화) 에서 시작하여 훈련 기간 동안 코사인 감쇠를 통해 0 으로 감소하는 글로벌 게이트입니다.

훈련 절차:

1. 워머업: 게이트 $g$를 1 에 고정합니다. 모델은 표준 프리-노름 트랜스포머로 훈련됩니다.
2. 보정: 워머업 종료 시, 정규화 브랜치와 스케일링 브랜치 간의 최소제곱 차이를 최소화하도록 스칼라 $c$를 계산하여 분포 이동 없이 원활한 전환을 보장합니다.
3. 테이퍼링: 게이트 $g$를 0 으로 감쇠시킵니다. 스케일링 브랜치 파라미터($\tilde{\gamma}$) 는 학습 가능한 상태로 유지됩니다.
4. 추론 폴딩: $g=0$이 되면 TaperNorm 레이어는 단순한 선형/아핀 맵($c \cdot h D\tilde{\gamma}$) 이 됩니다. 이 연산은 수학적으로 인접한 선형 프로젝션(예: QKV 프로젝션 또는 MLP 가중치) 에 통합되어 추론 그래프에서 명시적 정규화 연산을 제거할 수 있습니다.

스케일 앵커링 분석

이 논문은 내부 및 최종 정규화 레이어 간의 중요한 차이를 규명합니다:

• 앵커로서의 최종 정규화: 최종 정규화 레이어는 암시적 스케일 앵커 역할을 합니다. 수학적으로 이는 0-동차성 (0-homogeneous) 을 가지며, 이는 출력에서 방사형 그라디언트 성분을 제거함을 의미합니다. 이 앵커가 없으면 교차 엔트로피를 최소화하는 과정에서 모델이 손실을 줄이기 위해 로그이트 크기를 무한히 증가시키는 ("로그이트 체이싱") 경향이 있어 불안정성이 발생합니다.
• 명시적 대안: 최종 정규화가 제거된 ("All-Taper" 설정) 경우, 저자들은 고정 목표 스케일 손실($L_{aux}$) 을 도입합니다. 이 보조 손실은 토큰 스케일이 고정 목표 $s_{tgt}$에서 벗어나는 것을 패널티로 부과하여, 로그이트 체이싱 경향을 상쇄하는 명시적 방사형 복원력을 제공합니다.

두 가지 실험 구성이 테스트되었습니다:

1. Internal-Taper: 내부 블록 정규화기는 선형 맵으로 테이퍼링되지만, 최종 정규화는 활성화된 상태로 유지됩니다 (암시적 앵커).
2. All-Taper: 최종 것을 포함한 모든 정규화기가 테이퍼링됩니다. 이는 안정성을 유지하기 위해 고정 목표 스케일 손실이 필요합니다 (명시적 앵커).

주요 기여

1. 게이트형 정규화 제거: RMSNorm/LayerNorm 모듈을 고정된 선형/아핀 리스케일링으로 테이퍼링하여, 가중치 폴딩을 통해 추론 시 토큰별 통계를 제거할 수 있는 방법.
2. 출력 스케일 앵커링 이론: 최종 정규화 레이어의 주요 고유한 역할이 방사형 그라디언트 제거임을 보여주는 분석. 이 논문은 이 앵커가 없으면 교차 엔트로피가 무제한 로그이트 성장을 유발하며, 고정 목표 스케일 손실이 이러한 안정화 역할을 효과적으로 대체할 수 있음을 입증합니다.
3. 실증적 검증: TinyStories 프리-훈련 및 GPT-2 파인튜닝 실험을 통해 내부 정규화를 제거하더라도 검증 손실 증가가 미미함을 보였습니다. 이 접근법은 내부 스케일링을 폴딩한 후 KV 캐시 자동회귀 디코딩에서 최대 1.18 배의 처리량 향상을 달성합니다.

결과

• 프리-훈련 (TinyStories):
  • Internal-Taper (+aux): 모델 크기 (1M30M 파라미터) 전반에 걸쳐 베이스라인 대비 0.71.5% 이내의 검증 손실을 달성했습니다.
  • All-Taper (+aux): 베이스라인 대비 0.8~1.8% 이내의 손실을 달성하여, 명시적 스케일 앵커링을 통해 완전히 정규화 없는 모델이 가능함을 입증했습니다.
  • 동역학: 스케일 앵커가 없는 모델은 로그이트 크기의 급격한 증가 ("로그이트 체이싱") 를 보였습니다. 고정 목표 스케일 손실은 이 행동을 성공적으로 억제하여 그라디언트 크기를 앵커링된 베이스라인과 정렬시켰습니다.
• 추론 처리량:
  • NVIDIA H100 에서 Internal-Taper (퓨즈드) 구성은 표준 RMSNorm 베이스라인 대비 1.15~1.18 배 높은 초당 토큰 처리량을 달성했습니다. Internal-Taper (언퓨즈드) 구성 (명시적 스케일링 연산) 은 1.10~1.14 배의 속도 향상을 달성했습니다.
• 파인튜닝 (GPT-2):
  • 제안된 Internal-TaperLN 및 All-TaperLN 방법은 OpenWebText 및 The Pile 데이터셋에서 "FakeLN" 접근법 (Baroni et al., 2025) 과 표준 Vanilla 모델과 경쟁력 있는 교차 엔트로피 손실을 달성했습니다.
  • 순차적 제거 방법과 달리 TaperNorm 은 단일 글로벌 게이트를 사용하여 선택된 모든 레이어를 동시에 제거합니다.

중요성 및 주장

이 논문은 프리-노름 트랜스포머에서 내부 및 최종 정규화 레이어가 서로 다른 역할을 수행한다고 주장합니다. 내부 정규화기는 성능 저하가 미미한 학습된 샘플 독립적 선형 맵으로 대체될 수 있지만, 최종 정규화 레이어는 출력 스케일 앵커로서 고유한 기능을 수행합니다.

이 연구의 중요성은 다음과 같습니다:

1. 효율성: 추론 그래프에서 계산 비용이 많이 드는 정규화 연산을 제거하여 재훈련 없이 측정 가능한 처리량 향상 (최대 1.18 배) 을 제공하는 실용적인 경로를 제시합니다.
2. 이론적 통찰: 최종 정규화 레이어가 제공하는 안정성은 주로 방사형 그라디언트 제거에 기인함을 명확히 합니다. 이 통찰은 암시적 앵커를 명시적 고정 목표 스케일 손실로 대체하여 "완전 정규화 없는" 모델을 설계할 수 있게 합니다.
3. 강건성: 이 접근법은 훈련 초기에는 표준 프리-노름 동역학을 유지하다가 나중에 선형 맵으로 전환하여, 처음부터 정규화 없이 훈련하려는 방법보다 더 안정적으로 보입니다.

저자들은 범위에 대해 겸손하게 언급하며, 실험이 최대 30M 파라미터 모델과 특정 하드웨어 스택으로 제한됨을 지적합니다. 또한, 수십억 파라미터 영역에서의 활성화 아웃라이어에 대한 강건성과 교차 엔트로피 외의 다른 목적 함수에 대한 적용 가능성은 향후 연구의 열린 과제로 남아 있음을 인정합니다.