기술 요약: RMSNorm 의 계산 효율성으로 레이어 정규화 즐기기

문제 제기

레이어 정규화 (Layer Normalization, LN) 는 현대 딥러닝 아키텍처, 특히 트랜스포머 (Transformer) 의 핵심 구성 요소로, 학습을 안정화하고 최적화를 개선합니다. 그러나 추론 (inference) 단계에서 LN 은 샘플별 중심화 (평균을 빼는 작업) 와 스케일링을 필요로 하므로, 정규화 없는 대안들에 비해 무시할 수 없는 계산 오버헤드를 초래합니다.

RMSNorm 은 중심화 작업을 제거하고 스케일링만 유지함으로써 더 효율적인 대안으로 제안되었습니다. RMSNorm 은 기존 연구에서 보고된 바와 같이 실행 시간을 7%~64% 단축하는 등 상당한 속도 향상을 제공하지만, 중심화 작업을 폐기합니다. 이러한 제거는 출력 범위의 불안정성과 더 큰 활성화 변동을 초래할 수 있으며, 기존 문헌 (LeCun et al., 1990; Schraudolph, 1998; Ba et al., 2016) 에서 확립된 중심화와 관련된 이점들을 상실할 수 있습니다. LN 과 RMSNorm 을 동일시하려는 이전 시도들 (예: Jiang et al., 2023) 은 특정 프리-노름 (pre-norm) 트랜스포머 구조로 제한되었으므로, 임의의 심층 신경망 (DNN) 에서 모델 기능이나 학습 역학을 변경하지 않고 LN 을 RMSNorm 으로 대체할 수 있는 일반적인 프레임워크가 존재하는지에 대한 질문은 여전히 열려 있습니다.

방법론

본 논문은 임의의 DNN 에서 LN 레이어가 추론과 학습 모두에서 수학적 동등성을 유지하면서 RMSNorm 으로 대체될 수 있는지 판단하는 일반적인 프레임워크를 제안합니다. 핵심 전략은 LN 의 중심화 작업을 상류 (upstream) 선형 레이어에 "폴딩 (folding)"하는 것입니다.

1. 컬럼 중심화 제약 (CCC) 및 컬럼 기반 가중치 중심화 (CBWC)

저자들은 먼저 LN 이 선형 레이어 직후에 위치하는 경우를 분석합니다. 선형 레이어의 출력이 뉴런 간에 평균이 0 이라면 LN 의 중심화 단계는 불필요해짐을 관찰합니다.

• 컬럼 중심화 제약 (CCC): 선형 레이어의 가중치 행렬에서 각 열의 요소 합이 0 이면 해당 선형 레이어는 CCC 를 만족합니다. 이 제약 하에서 선형 레이어의 출력은 어떤 입력에 대해서도 평균이 0 임이 보장됩니다.
• 컬럼 기반 가중치 중심화 (CBWC): 모델의 용량을 제한하지 않고 학습 중 CCC 를 강제하기 위해, 저자들은 재파라미터화 기법을 도입합니다. 프록시 파라미터 행렬 $W$를 사용하여 투영을 통해 유효 가중치 행렬 $V$로 변환합니다: $V = (I - \frac{1}{m}\mathbf{1}_m\mathbf{1}_m^\top)W$. 이 변환은 유효 가중치가 0 합 열 제약을 만족하도록 보장하면서도 원래 모델의 최적화 역학을 유지합니다.

명제: CBWC 를 적용한 선형 레이어 뒤에 RMSNorm 을 사용하는 모델은, 재파라미터화되지 않은 동일한 선형 레이어 뒤에 표준 LN 을 사용하는 모델과 수학적으로 동등합니다.

2. 임의의 DNN 으로 일반화: 폴딩 가능한 LN

이 프레임워크는 그래프 기반 분석을 통해 임의의 네트워크 구조로 확장됩니다.

• 폴딩 가능한 LN (Foldable LN): 상류 레이어의 제약을 통해 입력이 평균 0 임이 보장될 수 있는 경우, 해당 LN 을 "폴딩 가능"하다고 정의합니다.
• 0 평균 그래프 (Zero-Mean Graph): 저자들은 LN 으로 가는 데이터 흐름을 나타내는 방향 그래프를 구성합니다. LN 에서 역추적하여 "리프 노드 (leaf nodes, 상류 레이어)"를 식별합니다.
  • 모든 리프 노드가 일반 선형 레이어 (공유 가중치를 가진 합성곱 및 RNN 포함) 나 본질적으로 0 평균 출력을 보장하는 레이어에 해당하면, 해당 LN 은 폴딩 가능합니다.
  • 스칼라 연산 (예: 추론 모드에서의 드롭아웃, 온도 스케일링) 은 0 평균 속성을 유지하며 폴딩 가능성을 깨뜨리지 않습니다.
  • 잔차 연결 (addition) 은 각 분기가 개별적으로 폴딩 가능하면 폴딩 가능성을 유지합니다.
  • 비선형 연산 (예: ReLU, Softmax) 이나 연결 (concatenation) 은 일반적으로 0 평균 속성을 깨뜨려 보조 중심화가 삽입되지 않는 한 폴딩을 방해합니다.

3. 탐지 알고리즘

본 논문은 계산 그래프를 순방향으로 탐색하여 폴딩 가능한 LN 을 자동으로 탐지하고, CBWC 가 필요한 특정 상류 일반 선형 레이어를 식별하는 알고리즘 (알고리즘 1) 을 제시합니다. 이 과정은 계산적으로 경량이며 단일 순방향 전달 (forward pass) 만 필요합니다.

주요 기여

1. 일반 프레임워크: 프리-노름 트랜스포머로 제한된 기존 연구의 한계를 넘어, 임의의 DNN 에서 LN 을 RMSNorm 으로 대체할 수 있는지 판단하는 이론적 프레임워크를 제시합니다.
2. 수학적 동등성: LN 의 중심화 작업을 상류 가중치에 정확히 폴딩할 수 있게 하는 CBWC 를 도입하여, CBWC+RMSNorm 방식이 학습과 추론 모두에서 기존 LN 기반 모델과 동일한 출력과 기울기를 생성하도록 보장합니다.
3. 폴딩 가능한 LN 탐지: "폴딩 가능한 LN"을 정의하고, 주어진 아키텍처에서 변환 가능한 LN 을 식별하는 그래프 기반 탐지 알고리즘을 개발했습니다.
4. 실용적 확장: 정확한 폴딩 조건이 충족되지 않는 경우 (예: 학습 중 드롭아웃으로 인한 경우) 에 대한 처리에 대한 논의를 포함하며, 정확한 동등성이 부분적으로 깨지더라도 CBWC+RMSNorm 이 순수 LN 과 경쟁력 있음을 시사합니다.

실험 결과

저자들은 GPT-2, BERT, OPT, Phi, ViT, Swin-Tiny 등 여러 아키텍처와 텍스트 분류, 번역, 이미지 분류, 파인튜닝 등 다양한 작업에 걸쳐 방법을 검증했습니다.

• 추론 가속화:

  • 이 방법은 모델 예측을 변경하지 않고 GPT-2, BERT, OPT 등 대표 모델에서 **2%~12%**의 엔드투엔드 추론 속도 향상을 달성합니다.
  • 시퀀스 길이가 길어질수록 가속 효과가 더 두드러집니다 (16,384 토큰의 BERT 의 경우 최대 24.95% 속도 향상).
  • 처리량이 크게 증가합니다 (A100 GPU 에서 초당 +1,119~+3,385 토큰).
  • 폴딩 알고리즘의 오버헤드는 최소화되어 (몇 단계의 추론 후 상쇄됨) 메모리 풋프린트를 증가시키지 않습니다.

• 학습 및 파인튜닝 성능:

  • 정확한 폴딩 조건이 성립하는 경우 (예: 추론) 에는 이 방법이 베이스라인과 동일합니다.
  • 드롭아웃과 같은 조건이 정확한 동등성을 깨는 실제 학습 환경에서는 CBWC+RMSNorm이 표준 LN 과 순수 RMSNorm 사이에서 일관된 성능을 보입니다.
  • 이미지 분류 (Swin-Tiny 를 사용한 ImageNet-100) 에서 이 방법은 세 가지 변형 (LN, RMSNorm, CBWC+RMS) 중 가장 높은 테스트 정확도와 손실 값을 달성했습니다.
  • 텍스트 생성 및 파인튜닝 작업 (Alpaca 에 대한 GPT-3) 에서 이 방법은 가장 낮은 테스트 손실과 퍼플렉시티를 달성하면서 효율성 향상을 제공했습니다.

중요성 및 주장

본 논문은 레이어 정규화의 계산 비용을 줄이면서도 이론적 및 실용적 이점을 최대한 유지하는 원칙적인 방법을 제공한다고 주장합니다.

• 정확한 동등성: "폴딩 가능한" LN 의 경우, 이 방법은 수학적으로 정확한 대체를 제공하여, LN 으로 학습된 모델을 재학습이나 성능 저하 없이 원래 훈련된 모델에 RMSNorm 을 배포할 수 있게 합니다.
• 효율성: 대규모 언어 모델 (LLM) 과 비전 - 언어 모델 (VLM) 의 대규모 배포에 중요한, 널리 사용되는 아키텍처에서 측정 가능한 엔드투엔드 추론 속도 향상 (2%~12%) 을 제공합니다.
• 강건성: 정확한 폴딩 조건이 완전히 충족되지 않더라도 (예: 학습 중 드롭아웃 사용 시), 이 방법은 표준 LN 과 경쟁력 있으며, 충실도와 효율성 사이의 유리한 균형을 제공합니다.
• 일반성: 이전 접근법들과 달리, 이 프레임워크는 특정 트랜스포머 변형뿐만 아니라 임의의 DNN 에 적용 가능하며, 프로세스를 자동화하는 탐지 알고리즘을 포함합니다.

저자들은 그들의 분석이 대표적인 아키텍처를 다루지만 모든 현대 대규모 모델을 포괄하지는 않는다고 인정합니다. 또한 학습 중 드롭아웃과 같은 모듈이 있는 아키텍처에서는 정확한 동등성이 실패할 수 있음을 지적하며, 향후 연구에서 이러한 모듈에 대한 처리를 정교화할 수 있음을 시사합니다.