기술 요약: BHyT – Pre-Layer Normalization의 안정적이고 효율적인 대안

1. 문제 정의

대규모 언어 모델(LLM)은 학습을 안정화하고 효과적인 전이 학습을 가능하게 하기 위해 일반적으로 RMSNorm으로 구현되는 Pre-Layer Normalization(Pre-LN)에 크게 의존한다. 그러나 모델의 깊이가 증가함에 따라, 표준 Pre-LN 설계는 두 가지 결정적인 한계에 직면한다:

1. 깊이의 저주 (The Curse of Depth): 최근의 분석에 따르면, 잔차 연결(residual connections)과 Pre-LN 사이의 상호작용은 레이어를 가로질러 은닉 상태(hidden states)의 크기와 분산을 급격히 증가시킨다. 이러한 상승은 블록 자코비안(block Jacobian)을 항등 행렬(identity matrix)로 밀어내어, 더 깊은 레이어가 비용이 많이 드는 항등 사상(identity mapping)처럼 동작하게 만들고 학습을 불안정하게 만든다.
2. 효율성 대 안정성의 트레이드오프 (Efficiency vs. Stability Trade-off): 기존 솔루션들은 종종 한 가지 측면만을 해결한다. Peri-LN(서브레이어 전후 모두 정규화)이나 LayerNorm Scaling(LNS)과 같은 안정성 지향적 방법은 분산 성장을 억제하지만, 반복적인 통계량 계산으로 인해 상당한 계산 오버헤드를 발생시킨다. 반대로, Dynamic Tanh(DyT)와 같은 효율성 지향적, 정규화 프리(normalization-free) 방식은 통계적 오버헤드는 제거하지만, 깊이 방향의 잔차 성장(depth-wise residual growth)을 명시적으로 제어하지 못해 심층 네트워크에서 포화(saturation) 및 불안정성을 초래할 수 있다.

핵심 과제는 유계 변환(bounded transformations)의 낮은 오버헤드와 잔차 스트림의 스케일 및 분산에 대한 명시적 제어를 결합하여, 처리량(throughput)을 희생하지 않으면서도 깊이 방향의 드리프트(depth-wise drift)를 방지하는 메커니즘을 설계하는 것이다.

2. 방법론: Bounded Hyperbolic Tanh (BHyT)

저자들은 $\tanh(\cdot)$ 쌍곡 탄젠트 비선형성과 명시적이고 데이터 기반인 입력 바운딩(input bounding)을 결합한 Pre-LN의 드롭인 대체제로서 BHyT를 제안한다.

핵심 메커니즘

입력을 재스케일링하기 위해 학습 가능한 글로벌 스칼라에 의존하는 DyT와 달리, BHyT는 데이터 의존적 통계량을 사용하여 $\tanh(\cdot)$ 함수의 입력을 명시적으로 제한한다.

• 확률적 입력 바운딩 (Probabilistic Input Bounding): 포화를 방지하기 위해, BHyT는 $\tanh$의 전 단계 인자(pre-$\tanh$ argument)가 높은 확률로 사전 정의된 구간 $[-\lambda, \lambda]$ 내에 존재하도록 제한한다. 체비쇼프 부등식(Chebyshev's inequality)을 사용하여, 이 방법은 입력의 평균($\mu_x$)과 표준 편차($s_x$)를 기반으로 한 스케일링 규칙을 도출한다.
• 이상적인 형태 ($BHyT^*$): 이론적으로, 이 변환은 입력 $x$를 다음과 같이 재스케일링한다:
  $$BHyT^*(x) = \gamma \odot \tanh\left(\frac{\lambda}{\kappa s_x + |\mu_x|} x\right)$$
  여기서 $\gamma$는 학습 가능한 스케일 파라미터이고, $\lambda$는 목표 범위를 정의하며, $\kappa$는 목표 확률 $p$에 따라 스케일링의 보수성을 제어한다. 이는 입력이 높은 확률로 $\tanh$의 비포화 선형 영역(non-saturating linear regime)에 머물도록 보장한다.

실질적 구현: 분산 근사 (Variance Approximation)

Transformer 블록 내의 모든 Pre-LN 지점에서 정확한 평균과 분산을 계산하는 것은 효율성 이점을 무효화할 것이다. 이를 해결하기 위해 BHyT는 블록 수준 분산 근사를 도입한다:

1. 첫 번째 지점 (Attention 전): 입력의 정확한 인스턴스별 분산을 계산한다.
2. 두 번째 지점 (MLP 전): 잔차 출력($x + h_{Attn}$)의 정확한 분산을 다시 계산하는 대신, 이전에 계산된 입력 분산과 모델 의존적인 어텐션 출력 분산 추정치를 사용하여 분산을 근사한다.
   • 어텐션 출력 분산은 균일한 어텐션 가중치와 평균이 0인 가우시안 엔트리를 가정하여 $\tilde{s}^2_{hAttn} \approx \frac{1}{Td} \|W_V W_O\|_F^2 \cdot \frac{\lambda_{Attn}^2}{\kappa^2}$로 근사된다.
   • 이 근사법을 통해 두 번째 정규화 과정에서 통계량을 재사용하거나 사전 계산된 모델 파라미터를 사용할 수 있어, 크리티컬한 순전파 경로에서의 두 번째 리덕션(reduction) 연산을 피할 수 있다.

이론적 안정성

본 논문은 **유한 깊이 분산 전파 경계(finite-depth variance propagation bound)**를 제공한다. 정리 3.5에 따르면, 네트워크의 깊이가 $L$일 때, 하이퍼파라미터가 $\lambda/\kappa < 1/\sqrt{L}$를 만족하면, BHyT 하에서의 출력 분산은 깊이 $L$까지의 모든 레이어에서 LNS보다 엄격하게 작다. 이는 BHyT가 LNS보다 깊이 방향의 분산 성장을 더 효과적으로 억제한다는 이론적 보증을 제공한다.

3. 주요 기여

1. BHyT 제안: 깊이 방향의 활성화 성장을 제어하기 위해 유계 $\tanh(\cdot)$ 변환과 데이터 기반 입력 바운딩을 결합한 Pre-LN의 드롭인 대체제를 제안하였다.
2. 이론적 및 실무적 효율성:
   • 특정 조건 하에서 BHyT가 LNS보다 더 작은 분산 전파 경계를 생성함을 증명하는 유한 깊이 분산 경계 도출.
   • 반복적인 정확한 분산 계산을 줄여 레이턴시를 크게 낮추는 블록 수준 분산 근사 메커니즘 도입.
3. 실험적 검증: Llama-374M, Llama-1B, Llama-3B 모델에 대한 종합적인 평가 수행. 연구 결과 BHyT가 안정성-효율성 트레이드오프를 개선하면서도 강력한 다운스트림 성능을 유지함을 입증하였다.

4. 실험 결과

실험은 토큰 예산이 모델 크기에 비례하는(약 파라미터당 20 토큰) C4 코퍼스로 사전 학습된 Llama 스타일 모델(374M, 1B, 3B)을 대상으로 진행되었다.

• 안정성: BHyT는 활성화 크기와 분산의 깊이 방향 성장을 성공적으로 억제하였으며, RMSNorm 및 DyT보다 우수한 성능을 보였다. 활성화 통계량은 Peri-LN과 같은 안정성 지향적 베이스라인과 유사하게 레이어 전반에 걸쳐 제어된 상태를 유지했다.
• 성능:
  • 사전 학습 (Pretraining): BHyT는 모든 모델 규모에서 가장 낮은 사전 학습 손실(loss)과 퍼플렉시티(perplexity)를 달ien했다. 또한 3-shot 벤치마크(예: ARC, MMLU, Winogrande)에서 두 번째로 좋은 방법보다 최대 0.6 포인트 높은 평균 다운스트림 정확도를 기록했다.
  • 사후 미세 조정 (Post-SFT): Lima1K에 대한 지도 미세 조정 후, BHyT 모델은 RMSNorm, Peri-LN, LNS, DyT에 비해 C4 및 WikiText-2에서 더 낮은 퍼플렉시티(망각이 적고 일반화가 잘 됨을 의미)와 더 높은 다운스트림 정확도를 보였다.
  • 규모 (Scale): Llama-3B에 대한 20B 토큰 비교에서, BHyT는 강력한 안정성 지향 베이스라인인 Peri-LN보다 평균 다운스트림 정확도(사전 학습)에서 3.4 퍼센트 포인트, (사후 SFT)에서 2.5 퍼센트 포인트 더 높은 성능을 보였다.
• 효율성:
  • 학습 처리량 (Training Throughput): BHyT는 RMSNorm 대비 평균 1.6% 빠른 학습을 달성했다.
  • 생성 처리량 (Generation Throughput): BHyT는 Llama-1B에서 1,199.9 tokens/s를 기록하여 RMSNorm보다 1.77%, Peri-LN보다 18.0% 높은 성능을 보였다. DyT가 (통계량이 없어) 가장 빨랐으나, BHyT는 속도와 안정성 사이에서 더 우수한 트레이드오프를 제공했다.
  • 절제 연구 (Ablation): 분산 근사 메커니즘은 수렴이나 다운스트림 성능을 저하시키지 않으면서 정확한 형태의 BHyT* 대비 14.9%의 속도 향상을 제공했다.

5. 의의 및 주장

본 논문은 BHyT가 더 깊고 효율적인 파운데이션 모델을 위한 실질적인 경로를 제공한다고 주장한다. 유계 변환의 낮은 오버헤드와 잔차 스트림 분사에 대한 명시적 제어를 통합함으로써, BHyT는 이중 정규화 방식의 계산적 페널티 없이 "깊이의 저주"를 해결한다.

저자들은 확장 가능한 학습이 안정성을 희생하지 않고 효율성을 개선할 수 있다고 단언한다. BHyT는 단순히 이론적인 대안이 아니라, 다양한 모델 규모(374M에서 3B 파라미터까지)에 걸쳐 안정적인 사전 학습과 강력한 전이 성능을 가능하게 하며, 학습 및 추론 모두에서 레이턴시를 줄이는 RMSNorm의 견고한 드롭인 대체제로 제시된다.