한 사람의 hidden unit만 보고 균형을 맞추는 법 — Layer Normalization

3단계 여섯 번째 논문. 이 글은 Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton의 Layer Normalization 원문 PDF를 직접 읽고 썼다. 이 글이 확인한 판본은 arXiv 1607.06450v1, 2016-07-21 제출본이다.arXiv 서지 기록 LayerNorm(Layer Normalization, 층 정규화)은 한 training case(학습 사례) 안의 hidden unit(은닉 유닛) 값들을 함께 보고 평균과 분산을 구한 뒤, 그 사례 안에서만 값을 다시 중심화·크기 조절하는 연산이다.

이 글에서 얻을 답과 범위

앞 글의 ResNet은 network를 아주 깊게 쌓았을 때 “더 깊은 plain network가 training error부터 높아지는” degradation problem을 봤다. LayerNorm 논문은 다른 지점에서 시작한다. deep network가 학습 중일 때 다음 layer에 들어가는 activation(활성값)의 scale(크기)와 center(중심)가 계속 움직이고, BatchNorm(Batch Normalization, 배치 정규화)은 이를 mini-batch(미니배치) 단위 통계로 다루려 했다. 하지만 문장 길이가 다른 RNN(Recurrent Neural Network, 순환 신경망), batch size(배치 크기)가 매우 작은 online learning(온라인 학습), train과 test의 다른 통계 처리에서는 불편한 경계가 생긴다.

이 글을 끝까지 읽으면 다음을 답할 수 있어야 한다.

1. BatchNorm은 숫자 표의 어느 축에서 평균·분산을 구하고, LayerNorm은 어느 축에서 구하는가.
2. LayerNorm 식의 (\mu), (\sigma), (\gamma), (\beta), (\epsilon), hidden size (D)가 각각 무엇이며 왜 필요한가.
3. 같은 input인데 batch의 다른 example이 바뀌면 BatchNorm output은 왜 바뀔 수 있고, LayerNorm output은 왜 보통 바뀌지 않는가.
4. RNN의 time step(시간 단계)에 LayerNorm을 적용하면 왜 running mean(이동 평균)을 time step마다 저장하지 않아도 되는가.
5. 논문이 LayerNorm을 CNN(합성곱 신경망)의 만능 대체재라고 말하지 않은 이유는 무엇인가.
6. PyTorch nn.LayerNorm(normalized_shape)에서 normalized_shape가 정하는 축, eps의 역할, train/eval mode의 경계는 무엇인가.
7. Transformer의 LayerNorm, RMSNorm(Root Mean Square Normalization, 평균제곱근 정규화), BatchNorm을 언제 같은 것으로 보면 안 되는가.

범위는 2016년 원 논문의 feed-forward network(순전파 신경망), RNN·LSTM(Long Short-Term Memory, 장단기 메모리), experiment(실험), PyTorch stable LayerNorm의 현재 public API(공개 API)다. 이 글은 Transformer 전체, optimizer(최적화기), mixed precision(혼합 정밀도) kernel(커널)의 모든 구현, 또는 LayerNorm을 포함한 대형 LLM 전체 학습을 재현하지 않는다. 아래 Python 예제는 논문 결과를 재학습한 것이 아니라, normalization axis(정규화 축), affine parameter(아핀 파라미터), batch composition(배치 구성), epsilon, input gradient(입력 기울기)를 작은 수치로 확인한 것이다.

먼저 잡을 한 문장

LayerNorm은 “같은 시점의 한 사례가 가진 feature들”을 한 줄로 보고 그 줄 안에서 평균 0 근처, 분산 1 근처가 되게 만든 뒤, 각 feature가 필요한 scale과 bias를 다시 배우게 한다. 그래서 옆 사람을 같은 batch에 넣거나 빼도 내 줄의 통계는 바뀌지 않는다.

왜 필요한가 — 같은 사람이 아닌 옆 사람 때문에 값이 바뀌는 문제

activation은 다음 층이 읽는 중간 숫자다

neural network(신경망)는 input을 여러 layer로 통과시키며 숫자를 바꾼다. 한 hidden layer에서 linear transformation(선형 변환)을 먼저 하면 각 neuron(뉴런)은 summed input(가중합 입력)을 얻는다.

[ a_i^{(l)}=(w_i^{(l)})^T h^{(l)}+b_i^{(l)}, \qquad h_i^{(l+1)}=f(a_i^{(l)}) \tag{1} ]

학습 중 앞 layer weight가 바뀌면 뒤 layer가 받는 (a)들의 평균과 크기도 함께 바뀐다. 원 LayerNorm 논문은 이 변화가 gradient-based optimization(기울기 기반 최적화)을 어렵게 할 수 있다는 BatchNorm의 문제 설정에서 출발한다. 여기서 “internal covariate shift(내부 공변량 이동)를 완전히 제거한다”는 표현은 조심해야 한다. 논문은 그 표현을 동기로 사용했지만, LayerNorm의 필요성은 한 가지 원인 이론에만 기대지 않는다. 이 글에서는 어느 축의 통계를 읽고 어떤 의존성을 만들거나 없애는가를 확인 가능한 사실로 다룬다.원 논문 1절·2절

BatchNorm은 세로로, LayerNorm은 가로로 읽는다

batch가 3명이고 hidden layer가 4개 unit이라고 하자. 행은 example, 열은 hidden unit이다.

                   hidden unit
              u1      u2      u3      u4
example A    2.0     4.0     6.0     8.0
example B   -2.0     0.0     2.0     4.0
example C    4.0     8.0    12.0    16.0

BatchNorm은 보통 열마다, 즉 같은 hidden unit u1의 A·B·C 값을 모아 평균과 분산을 구한다. LayerNorm은 행마다, 즉 A의 u1부터 u4까지를 모아 평균과 분산을 구한다.

BatchNorm
  u1 column의 A, B, C를 같이 본다
  u2 column의 A, B, C를 같이 본다
  → A output은 B, C가 무엇인지에 따라 training 중 달라질 수 있다

LayerNorm
  A row의 u1, u2, u3, u4를 같이 본다
  B row의 u1, u2, u3, u4를 같이 본다
  → A output은 B, C가 무엇인지에 의존하지 않는다

이 축의 차이가 핵심이다. LayerNorm은 “batch 전체에 normalization을 한다”는 말보다 batch axis를 제외한 한 사례의 normalized feature axis를 정한다고 이해해야 한다. PyTorch도 mean과 standard deviation을 normalized_shape가 정한 마지막 (D)개 dimension(차원)에서 계산한다고 명시한다.PyTorch LayerNorm 공식 문서

RNN은 time step마다 같은 weight를 재사용한다

RNN은 token sequence(토큰 시퀀스)를 한 칸씩 읽는다. t번째 step에서 preactivation(활성화 전 값)은 단순화하면 다음이다.

[ a_t=W_{hh}h_{t-1}+W_{xh}x_t \tag{2} ]

문장은 길이가 서로 다르다. BatchNorm을 RNN에 단순 적용하면 “t=1에서의 batch 통계”, “t=2에서의 batch 통계”를 따로 추정·보관해야 할 수 있다. test sentence가 training에서 본 문장보다 길면 그 뒤 time step의 통계는 어떻게 할지 애매해진다. LayerNorm은 그 순간의 (a_t) vector 내부에서만 (\mu_t), (\sigma_t)를 계산하므로, 별도 example이나 과거 time step의 running statistic을 읽지 않는다. 원 논문은 이 점과 hidden-state dynamics(은닉 상태 동역학)의 안정화를 RNN의 중요한 장점으로 설명한다.원 논문 3.1절

무엇인가 — 한 vector 안에서 평균과 분산을 구하는 연산

원 논문의 식 (15)와 식 (16)

논문 supplementary material(보충 자료)은 LayerNorm을 (D)차원 vector (z)에 적용하는 함수로 다음처럼 쓴다.

[ \operatorname{LN}(z;\alpha,\beta) =\frac{z-\mu}{\sigma}\odot\alpha+\beta \tag{3} ]

[ \mu=\frac{1}{D}\sum_{i=1}^{D}z_i, \qquad \sigma=\sqrt{\frac{1}{D}\sum_{i=1}^{D}(z_i-\mu)^2} \tag{4} ]

논문 본문 3절은 같은 생각을 hidden layer의 summed input (a_i^{(l)})로 쓴다. hidden unit (H)개가 있으면

[ \mu^{(l)}=\frac{1}{H}\sum_{i=1}^{H}a_i^{(l)}, \qquad \sigma^{(l)}=\sqrt{\frac{1}{H}\sum_{i=1}^{H}(a_i^{(l)}-\mu^{(l)})^2} \tag{5} ]

이다. 이름만 (D)에서 (H)로 바뀌었다. (H)는 그 layer의 hidden unit 수다.

왜 gain과 bias를 다시 더하나

정규화만 하면 모든 vector가 비슷한 중심과 scale로 고정된다. 어떤 hidden unit은 큰 값 범위가 유리하거나, 어떤 unit은 0보다 위에서 시작하는 것이 유리할 수 있다. (\alpha_i)와 (\beta_i)는 model이 그 선택을 다시 배우는 통로다.

중요한 순서는 다음이다.

linear sum z
→ 한 사례 안에서 μ와 σ 계산
→ (z - μ) / σ
→ feature별 α를 곱하고 β를 더함
→ ReLU·tanh·sigmoid 같은 nonlinearity

원 논문은 adaptive gain과 bias가 normalization 뒤, nonlinearity 앞에 적용된다고 쓴다.원 논문 초록·3절 LayerNorm은 평균 0·분산 1을 강제로 유지한다는 말은 affine transformation 전의 normalized part만 말한 축약이다. (\alpha)·(\beta)까지 포함한 output의 평균·분산은 일반적으로 0과 1이 아니다.

eps는 원 논문 식에 없는 구현 안전장치다

실제 framework는 보통 분모를 다음처럼 만든다.

[ y_i=\gamma_i\frac{x_i-\mu}{\sqrt{\operatorname{Var}(x)+\epsilon}}+\beta_i \tag{6} ]

(\epsilon)는 아주 작은 positive number(양수)다. vector가 [5, 5, 5]처럼 모든 값이 같으면 variance가 0이고, 식 (3)의 분모는 0이 된다. (\epsilon)를 더하면 0으로 나누지 않는다. PyTorch stable nn.LayerNorm의 default eps는 1e-5다.PyTorch 공식 문서

여기에는 판본 경계가 있다. 2016년 원 논문 식 (15)·(16)은 (\epsilon)를 적지 않는다. 따라서 “논문이 eps=1e-5를 정했다”라고 쓰면 틀리다. 1e-5는 현재 PyTorch API의 기본값이며 다른 framework·model config는 다른 값을 쓸 수 있다.

LayerNorm은 무엇을 하지 않는가

선행 개념 — 평균, 분산, 축, population variance

평균과 분산을 직접 계산해 보자

한 사례의 hidden vector가 (x=[2,4,6])이라고 하자.

[ \mu=\frac{2+4+6}{3}=4 ]

[ \operatorname{Var}(x) =\frac{(2-4)^2+(4-4)^2+(6-4)^2}{3} =\frac{8}{3}\approx2.666667 ]

[ \sqrt{\operatorname{Var}(x)+10^{-5}}\approx1.632996 ]

따라서 (\gamma=[1,1,1]), (\beta=[0,0,0])이면

[ \operatorname{LN}(x) \approx[-1.224743,0,1.224743] ]

가 된다. 가운데 4는 평균과 같으므로 0이 된다. 2는 평균보다 작아 음수, 6은 평균보다 커 양수다. 이 계산에서 3으로 나눈 것은 population variance다. PyTorch 현재 문서도 LayerNorm variance가 correction=0과 같은 biased estimator(편향 추정량)를 쓴다고 명시한다.PyTorch 공식 문서

axis는 SQL의 GROUP BY와 비슷한 질문이다

정규화에서 제일 먼저 물어야 할 질문은 “어떤 숫자들을 한 group으로 묶어 평균을 낼 것인가”다. SQL의 GROUP BY를 잘못 잡으면 다른 row가 섞이듯, normalization axis를 잘못 잡으면 다른 token·sample·channel이 섞인다.

위 표는 대표적인 2D·NLP shape 설명이다. image tensor와 framework 설정에 따라 정확한 axes는 달라진다. 이름만 보고 “LayerNorm은 channel만 normalize한다” 또는 “BatchNorm은 무조건 모든 축을 normalize한다”고 외우면 위험하다.

sequence tensor에서는 token 하나씩 hidden dimension을 normalize하는 일이 흔하다

NLP에서 input shape가 (batch, sequence_length, hidden_size)라고 하자. PyTorch에서 nn.LayerNorm(hidden_size)를 적용하면 마지막 hidden_size dimension만 normalize한다. 즉 batch의 각 문장, 각 token 위치마다 자기 hidden vector의 mean·variance를 따로 구한다.

shape = (2 examples, 5 tokens, 768 hidden features)
nn.LayerNorm(768)

각 [example, token, :] 길이 768 vector를 독립적으로 normalize
다른 example의 token, 같은 문장의 다른 token은 통계에 섞지 않음

nn.LayerNorm([C, H, W])처럼 normalized shape를 여러 dimension으로 주면 마지막 여러 축을 함께 묶는다. PyTorch는 이 동작과 affine parameter의 shape가 normalized_shape와 같다고 문서화한다.PyTorch 공식 예제

밑바닥 원리 — BatchNorm과 LayerNorm의 데이터 의존성을 숫자로 비교하기

BatchNorm의 통계는 같은 feature의 여러 example에서 나온다

BatchNorm을 training mode에서 단순화하면, feature (i)에 대해 batch (\mathcal{B})의 여러 example에서

[ \mu_i^{\mathcal{B}}=\frac{1}{|\mathcal{B}|}\sum_{x\in\mathcal{B}}a_i(x) ]

[ \operatorname{BN}(a_i(x)) =\gamma_i\frac{a_i(x)-\mu_i^{\mathcal{B}}} {\sqrt{\operatorname{Var}_{x\in\mathcal{B}}(a_i(x))+\epsilon}}+\beta_i \tag{7} ]

처럼 쓸 수 있다. (a_i(x))는 example (x)의 i번째 feature다. A의 u1 output을 구하는데 B·C의 u1도 분모·분자에 들어간다. batch를 shuffle(섞기)하거나 마지막 작은 batch가 생기는 것 자체가 training activation에 영향을 줄 수 있다는 뜻이다.

반대로 LayerNorm은 example A의 vector만 쓴다.

[ \mu_A=\frac{1}{D}\sum_{i=1}^{D}a_i(A), \qquad \operatorname{LN}(a_i(A)) =\gamma_i\frac{a_i(A)-\mu_A} {\sqrt{\frac1D\sum_j(a_j(A)-\mu_A)^2+\epsilon}}+\beta_i \tag{8} ]

여기서 B·C의 값은 식에 없다. 이것이 batch size 1과 variable-length RNN에서의 실용적인 차이를 만든다.

LayerNorm의 invariance는 무엇이고 무엇이 아닌가

원 논문 5.1절은 BatchNorm, WeightNorm(가중치 정규화), LayerNorm의 invariance(불변성)를 비교한다. LayerNorm은 한 training case를 양의 상수로 rescale(재스케일)하는 변환에 대해 normalized 계산이 변하지 않는 성질을 보인다. 그리고 incoming weight matrix 전체의 scaling·공통 re-centering에 관한 성질도 분석한다.원 논문 표 1·식 6·7

하지만 이 표는 “LayerNorm을 넣으면 model이 모든 입력 변형에 불변”이라는 보증이 아니다. nonlinearity, residual connection, positional encoding(위치 인코딩), downstream layer, learned gain·bias, negative scale, data preprocessing이 추가되면 전체 model의 행동은 별도로 봐야 한다. 이 글의 직접 실행은 더 좁게, batch의 다른 row를 바꿔도 A row LayerNorm result가 같음만 확인한다.

RNN과 LSTM에서는 적용 위치가 중요하다

원 논문 본문은 단순 RNN에 대해 다음처럼 쓴다.

[ h_t=f\left(g\odot\frac{a_t-\mu_t}{\sigma_t}+b\right) \tag{9} ]

여기서 (g), (b)는 hidden state와 같은 차원의 gain·bias이며 time step에 걸쳐 공유된다. (\mu_t), (\sigma_t)는 그 time step의 (a_t)에서 계산한다.원 논문 식 4

paper의 LSTM 실험은 더 구체적이다. forget gate(망각 게이트), input gate(입력 게이트), output gate(출력 게이트), candidate(후보)를 쌓은 vector에 대해 recurrent contribution (W_hh_{t-1})와 input contribution (W_xx_t)을 각각 LayerNorm한 뒤 더한다. 그리고 cell state (c_t)에도 LayerNorm을 적용하는 식을 supplementary material 식 20부터 22에 제시한다.원 논문 supplementary material

잘못된 일반화
  LSTM 전체 출력 하나에 LayerNorm 한 번이면 논문의 LSTM이다

논문이 실제로 제시한 예
  LN(recurrent affine term) + LN(input affine term) + bias
  → gate 계산
  → cell-state 쪽에도 별도 LN 위치가 있다

그러므로 source code나 model config를 읽지 않고 “LSTM에 LayerNorm을 썼다”는 말만 보면 어느 vector, 어느 gate, pre-activation인지 post-activation인지 알 수 없다.

내부 구조와 실제 실행 흐름 — PyTorch API로 축을 확인하기

normalized_shape는 feature axis 계약이다

PyTorch stable API는 다음 signature(함수 모양)를 제공한다.

torch.nn.LayerNorm(
    normalized_shape,
    eps=1e-05,
    elementwise_affine=True,
    bias=True,
)

이 module은 input의 마지막 D dimension에서 mean과 variance를 계산한다. normalized_shape가 integer 하나면 마지막 dimension 하나, tuple 또는 list면 마지막 여러 dimension이 대상이다. output shape는 input shape와 같다. elementwise_affine=True이면 weight와 bias는 normalized_shape의 shape를 가진다.PyTorch LayerNorm 공식 문서

이 표의 첫 행에서 B는 batch size, T는 token 수, H는 hidden size다. LayerNorm(H)는 (B, T)를 group으로 쓰지 않는다. 반대로 (B, C, H, W)에 LayerNorm(C)를 바로 적용할 수 있다고 가정하면 마지막 dimension이 W라서 shape contract가 깨질 수 있다. tensor layout을 먼저 확인해야 한다.

train과 eval에서 LayerNorm 자체의 통계 계산은 같다

BatchNorm은 training 중 현재 mini-batch statistics를 쓰고, inference에서는 running mean·variance를 사용하는 설계가 일반적이다. LayerNorm 논문의 초록은 training과 test에서 정확히 같은 계산을 수행한다고 강조한다. 현재 PyTorch documentation도 LayerNorm이 input data에서 계산한 statistics를 training·evaluation mode 모두에서 쓴다고 명시한다.원 논문 초록 PyTorch 공식 문서

이 말은 model.train()과 model.eval()의 결과가 항상 같다는 뜻이 아니다. LayerNorm 옆에 Dropout(드롭아웃), BatchNorm, sampling(샘플링) layer가 있으면 mode 전환 결과가 달라질 수 있다. 정확한 문장은 LayerNorm module의 mean·variance source는 train/eval 모드에 따라 running statistic으로 교체되지 않는다이다.

framework 구현과 원 논문을 대조할 때의 차이

“LayerNorm은 논문 그대로 framework에 한 줄로 구현된다”는 식의 말은 이 세부 차이를 지운다. 논문의 핵심 축 선택은 유지되지만, production implementation(운영 구현)은 epsilon, dtype(자료형), layout, affine option 같은 API 계약을 함께 가진다.

직접 검증과 재현 — batch에 다른 사람을 넣어 보기

실행 환경과 검증 범위

아래 코드는 macOS의 Python 3.9.6과 표준 라이브러리만으로 직접 실행했다. framework를 호출하지 않고 식 (3)·(6)·(7)을 작은 list로 구현했다. 그래서 PyTorch CUDA kernel 성능이나 원 논문 실험 결과가 아니라 다음 algebra(대수)만 검증한다.

1. [2, 4, 6] 한 vector의 mean과 population variance를 계산한다.
2. feature별 gain·bias가 normalization 뒤에 적용됨을 확인한다.
3. 같은 x에 대해 배치의 나머지 두 row를 바꿔도 LayerNorm output은 같고, training-mode BatchNorm output은 달라짐을 확인한다.
4. time step마다 현재 preactivation vector만 써서 LayerNorm할 수 있음을 보인다.
5. constant vector에서 epsilon이 0 나누기를 막는지 확인한다.
6. LayerNorm input gradient의 analytic formula를 finite difference와 비교한다.

#!/usr/bin/env python3
"""Small, dependency-free checks for Ba et al. (2016) Layer Normalization."""

from __future__ import annotations

import math


def mean(values):
    return sum(values) / len(values)


def layer_norm(values, gamma=None, beta=None, epsilon=1e-5):
    """Normalize one D-dimensional activation vector using population variance."""
    width = len(values)
    gamma = [1.0] * width if gamma is None else gamma
    beta = [0.0] * width if beta is None else beta
    assert len(gamma) == len(beta) == width
    mu = mean(values)
    variance = mean([(value - mu) ** 2 for value in values])
    reciprocal_std = 1.0 / math.sqrt(variance + epsilon)
    normalized = [(value - mu) * reciprocal_std for value in values]
    output = [g * value + b for value, g, b in zip(normalized, gamma, beta)]
    return output, mu, variance, normalized


def batch_norm_per_feature(batch, epsilon=1e-5):
    """Training-mode BatchNorm over examples, one feature column at a time."""
    width = len(batch[0])
    assert all(len(row) == width for row in batch)
    means = [mean([row[j] for row in batch]) for j in range(width)]
    variances = [mean([(row[j] - means[j]) ** 2 for row in batch]) for j in range(width)]
    return [
        [(value - means[j]) / math.sqrt(variances[j] + epsilon)
         for j, value in enumerate(row)]
        for row in batch
    ]


def dot(left, right):
    return sum(a * b for a, b in zip(left, right))


def layer_norm_input_gradient(values, upstream, epsilon=1e-5):
    """Gradient of dot(upstream, LayerNorm(values)) for gamma=1, beta=0."""
    normalized_output, _, variance, normalized = layer_norm(values, epsilon=epsilon)
    assert normalized_output == normalized
    reciprocal_std = 1.0 / math.sqrt(variance + epsilon)
    centered_upstream = [value - mean(upstream) for value in upstream]
    correction = mean([u * y for u, y in zip(upstream, normalized)])
    return [reciprocal_std * (u - y * correction)
            for u, y in zip(centered_upstream, normalized)]


def finite_difference_gradient(values, upstream, epsilon=1e-5, delta=1e-6):
    gradient = []
    for index in range(len(values)):
        plus, minus = list(values), list(values)
        plus[index] += delta
        minus[index] -= delta
        plus_loss = dot(upstream, layer_norm(plus, epsilon=epsilon)[0])
        minus_loss = dot(upstream, layer_norm(minus, epsilon=epsilon)[0])
        gradient.append((plus_loss - minus_loss) / (2.0 * delta))
    return gradient


def show(label, values):
    print(label + ': [' + ', '.join(f'{value:.6f}' for value in values) + ']')


def main():
    epsilon = 1e-5
    x = [2.0, 4.0, 6.0]

    print('== Eq. (15)/(16): one example across three hidden units ==')
    output, mu, variance, normalized = layer_norm(x, epsilon=epsilon)
    show('input x', x)
    print(f'mean over hidden units: {mu:.6f}')
    print(f'population variance over hidden units: {variance:.6f}')
    show('normalized output', output)
    print(f'normalized mean: {mean(normalized):.6f}')
    print(f'normalized variance with epsilon in denominator: {mean([v * v for v in normalized]):.6f}')

    print('\n== learnable per-hidden-unit gain and bias ==')
    affine, _, _, _ = layer_norm(x, gamma=[1.0, 0.5, 2.0], beta=[0.0, 1.0, -1.0], epsilon=epsilon)
    show('gamma', [1.0, 0.5, 2.0])
    show('beta', [0.0, 1.0, -1.0])
    show('affine LayerNorm output', affine)

    print('\n== LayerNorm ignores unrelated examples, BatchNorm does not ==')
    batch_a = [x, [-2.0, 0.0, 2.0], [4.0, 8.0, 12.0]]
    batch_b = [x, [-2.0, 0.0, 2.0], [100.0, 200.0, 300.0]]
    show('LayerNorm x in batch A', layer_norm(batch_a[0], epsilon=epsilon)[0])
    show('LayerNorm x in batch B', layer_norm(batch_b[0], epsilon=epsilon)[0])
    show('BatchNorm x in batch A', batch_norm_per_feature(batch_a, epsilon=epsilon)[0])
    show('BatchNorm x in batch B', batch_norm_per_feature(batch_b, epsilon=epsilon)[0])
    assert layer_norm(batch_a[0], epsilon=epsilon)[0] == layer_norm(batch_b[0], epsilon=epsilon)[0]
    assert batch_norm_per_feature(batch_a, epsilon=epsilon)[0] != batch_norm_per_feature(batch_b, epsilon=epsilon)[0]

    print('\n== RNN-shaped time steps use current-step feature statistics ==')
    preactivation_t1 = [4.0, 0.0, -4.0]
    preactivation_t2 = [10.0, 8.0, 6.0]
    show('LN(a_t=1)', layer_norm(preactivation_t1, epsilon=epsilon)[0])
    show('LN(a_t=2)', layer_norm(preactivation_t2, epsilon=epsilon)[0])
    print('one gamma/beta set can be reused; no time-step running mean is read here')

    print('\n== epsilon prevents division by zero for a constant vector ==')
    constant_output, constant_mu, constant_variance, _ = layer_norm([5.0, 5.0, 5.0], epsilon=epsilon)
    print(f'constant mean: {constant_mu:.6f}, variance: {constant_variance:.6f}')
    show('output with epsilon', constant_output)
    assert constant_output == [0.0, 0.0, 0.0]

    print('\n== input gradient: analytic result versus finite difference ==')
    upstream = [0.2, -0.4, 0.7]
    analytical = layer_norm_input_gradient(x, upstream, epsilon=epsilon)
    numerical = finite_difference_gradient(x, upstream, epsilon=epsilon)
    show('analytical dL/dx', analytical)
    show('finite-difference dL/dx', numerical)
    assert all(math.isclose(a, b, rel_tol=1e-6, abs_tol=1e-6)
               for a, b in zip(analytical, numerical))
    print('checks passed')


if __name__ == '__main__':
    main()

실제 출력은 아래다.

== Eq. (15)/(16): one example across three hidden units ==
input x: [2.000000, 4.000000, 6.000000]
mean over hidden units: 4.000000
population variance over hidden units: 2.666667
normalized output: [-1.224743, 0.000000, 1.224743]
normalized mean: 0.000000
normalized variance with epsilon in denominator: 0.999996

== learnable per-hidden-unit gain and bias ==
gamma: [1.000000, 0.500000, 2.000000]
beta: [0.000000, 1.000000, -1.000000]
affine LayerNorm output: [-1.224743, 1.000000, 1.449485]

== LayerNorm ignores unrelated examples, BatchNorm does not ==
LayerNorm x in batch A: [-1.224743, 0.000000, 1.224743]
LayerNorm x in batch B: [-1.224743, 0.000000, 1.224743]
BatchNorm x in batch A: [0.267261, 0.000000, -0.162221]
BatchNorm x in batch B: [-0.664282, -0.685574, -0.692726]

== RNN-shaped time steps use current-step feature statistics ==
LN(a_t=1): [1.224744, 0.000000, -1.224744]
LN(a_t=2): [1.224743, 0.000000, -1.224743]
one gamma/beta set can be reused; no time-step running mean is read here

== epsilon prevents division by zero for a constant vector ==
constant mean: 5.000000, variance: 0.000000
output with epsilon: [0.000000, 0.000000, 0.000000]

== input gradient: analytic result versus finite difference ==
analytical dL/dx: [0.173505, -0.347010, 0.173506]
finite-difference dL/dx: [0.173505, -0.347010, 0.173506]
checks passed

첫 section은 (\epsilon) 때문에 normalized variance가 정확히 1이 아니라 0.999996임을 보인다. 둘째 section은 (\gamma), (\beta)가 feature별로 output을 다시 바꾸는 모습을 보인다. 셋째 section이 이 글의 핵심 반례다. A row x는 그대로인데 C row를 [4, 8, 12]에서 [100, 200, 300]으로 바꾸자 LayerNorm A output은 같고 BatchNorm A output은 달라졌다.

gradient는 feature끼리도 서로 연결한다

LayerNorm은 각 (x_i)를 자기 자신만으로 처리하지 않는다. (\mu)와 variance가 vector 전체에서 나오므로 하나의 feature를 바꾸면 같은 group의 다른 feature normalization에도 영향을 준다. (\gamma=1), (\beta=0), upstream gradient를 (u)라 두면, 위 코드가 검증한 한 형태의 gradient는 다음이다.

[ \frac{\partial L}{\partial x_i} =r\left[(u_i-\operatorname{mean}(u)) -y_i\operatorname{mean}(u\odot y)\right], \qquad r=\frac1{\sqrt{\operatorname{Var}(x)+\epsilon}} \tag{10} ]

여기서 (y)에는 affine 전 normalized output을 쓴다. 이 식은 원 논문에 인쇄된 backward implementation이 아니라, 이 글의 작은 function에서 forward 식을 미분한 해석이다. [2,4,6], upstream [0.2,-0.4,0.7]에 대해 analytic result [0.173505,-0.347010,0.173506]와 finite difference가 일치했다. 즉 LayerNorm은 feature별 독립 scaling이 아니라 같은 normalized group 내부에 coupling(결합)을 만드는 연산이다.

값을 하나 바꿔 보는 실패 사례

위 코드에서 epsilon=0.0으로 바꾸고 [5,5,5]을 넣으면 variance 0으로 나누게 된다. 이 경우 ZeroDivisionError 또는 framework의 nan·inf가 생길 수 있다. production에서는 단순히 eps를 크게 하면 끝이 아니다.

1. input이 거의 constant가 되는 이유가 정상인지 먼저 확인한다.
2. dtype과 mixed precision에서 variance·reciprocal square root가 안정적인지 확인한다.
3. model config의 eps와 checkpoint를 만든 training config가 같은지 확인한다.
4. output에 nan이 처음 생기는 layer·step을 기록한다.

eps는 numerical guard(수치 안전장치)이지 data corruption(데이터 손상)이나 exploding activation의 근본 원인을 숨기는 설정이 아니다.

성능과 트레이드오프 — 배치 독립성에도 대가가 있다

시간·메모리 비용

normalization group 길이를 (D)라고 하면 forward는 mean을 위한 한 번의 reduction(축소 연산), variance를 위한 한 번의 reduction, normalize·affine을 위한 elementwise pass가 필요하다. 따라서 한 vector당 시간은 (O(D)), output·mean·reciprocal standard deviation 및 backward에 필요한 intermediate(중간값)를 고려한 메모리는 implementation마다 달라진다.

LayerNorm은 batch statistics를 all-reduce(여러 device 통계 합산)할 필요가 없다는 장점이 있다. data-parallel training(데이터 병렬 학습)에서 작은 device-local batch가 되는 상황에도 batch size 자체가 통계 정의를 바꾸지 않는다. 반대로 hidden size가 매우 크고 sequence가 길면 (B × T)개의 group마다 reduction을 해야 한다. “BatchNorm running state가 없으니 LayerNorm은 항상 더 빠르다”는 결론은 나오지 않는다. hardware, fused kernel, tensor layout, sequence length, dtype이 측정값을 결정한다.

원 논문 실험은 무엇을 보였고 무엇을 보이지 않았나

원 논문은 image-sentence ranking, question answering, skip-thought, DRAW, handwriting generation, permutation-invariant MNIST를 다뤘다. RNN 계열에서 training speed와 generalization 개선을 관찰했고, MNIST feed-forward experiment에서 batch size 128과 4를 비교해 LayerNorm이 batch size에 견고하다는 결과를 보였다.원 논문 6절·그림 2·6

그러나 CNN에는 다른 결론을 썼다. preliminary experiment에서 LayerNorm은 normalization 없는 baseline보다 speedup을 보였지만, BatchNorm이 다른 방법보다 성능이 좋았다. 논문은 convolution feature map의 boundary(경계) 근처 hidden unit이 거의 켜지지 않아, layer 전체 unit이 비슷한 기여를 한다는 가정이 깨질 수 있다고 설명한다.원 논문 6.7절

운영에서 볼 지표

LayerNorm에는 BatchNorm running mean·variance buffer가 일반적으로 없지만, checkpoint가 단순하다는 뜻은 아니다. learnable affine parameter, architecture의 pre-LN·post-LN 위치, eps, tensor layout이 모두 contract(계약)다.

실패와 운영 기준 — 축과 위치를 틀리면 다른 모델이 된다

hidden_size만 알면 되는 것이 아니다

아래 둘은 이름은 모두 LayerNorm이지만 normalize하는 숫자 묶음이 다르다.

# token마다 hidden feature 768개를 따로 normalize
token_norm = nn.LayerNorm(768)

# image 하나의 C, H, W를 한 묶음으로 normalize
image_norm = nn.LayerNorm([channels, height, width])

후자의 height, width가 input resolution에 묶이면 image size가 달라질 때 parameter shape와 input contract가 깨질 수 있다. frontend API schema를 잘못 읽어 request field를 섞는 것처럼, tensor의 last dimensions를 확인하지 않고 LayerNorm(C)를 넣으면 원한 normalization이 아닐 수 있다.

pre-LN과 post-LN은 LayerNorm 논문이 정해 주지 않는다

Transformer에서 residual sublayer (S)를 (x)에 붙이는 대표 순서는 둘이다.

[ \text{post-LN}\quad x_{next}=\operatorname{LN}(x+S(x)) ]

[ \text{pre-LN}\quad x_{next}=x+S(\operatorname{LN}(x)) ]

둘 다 LayerNorm을 사용하지만 gradient path·training stability·final normalization 위치가 다르다. 2017년 Attention Is All You Need은 sublayer 뒤 residual addition에 LayerNorm을 둔 post-LN 형태를 제시했다. pre-LN과 post-LN의 차이는 2016 LayerNorm 논문의 결론이 아니라, Transformer architecture를 읽을 때 별도로 확인해야 할 설계 선택이다.

RMSNorm은 LayerNorm의 철자가 짧은 버전이 아니다

RMSNorm은 보통 평균 (\mu)를 빼지 않고 root mean square(RMS, 제곱 평균의 제곱근)만으로 scale을 조절한다.

[ \operatorname{RMSNorm}(x) =\gamma\odot\frac{x}{\sqrt{\frac1D\sum_i x_i^2+\epsilon}} ]

그래서 [2,4,6]의 LayerNorm은 중심을 4로 옮겨 [-,0,+] 모양이 되지만, RMSNorm은 모든 원소의 sign와 평균 위치를 그대로 둔 채 크기를 조절한다. RMSNorm 원 논문은 re-centering이 불필요할 수 있다고 제안한다. 최신 LLM config에서 RMSNorm을 만났다면 “LayerNorm과 똑같고 빠르다”라고 치환하지 말고 mean subtraction, bias, epsilon, parameter shape를 각각 확인한다.

BatchNorm을 대체할지 결정하는 질문

선택은 normalization 이름이 아니라 model family, tensor shape, batch regime, measured metric, deployment contract로 한다.

대안과 비교 — 같은 정규화라는 이름 아래 다른 축

| 방법 | 통계의 주된 범위 | running statistic | 대표 장점 | 대표 주의점 | | --- | --- | --- | --- | | BatchNorm | batch의 같은 feature 또는 channel | 일반적으로 필요 | CNN에서 강한 baseline | small batch·RNN time step·train/eval 차이 | | LayerNorm | 한 example의 normalized feature axis | 보통 없음 | batch 독립, RNN·Transformer에 자연스러움 | CNN spatial statistics에는 맞지 않을 수 있음 | | WeightNorm | weight vector의 L2 norm | activation batch 통계 아님 | weight parameterization을 바꿈 | LayerNorm과 같은 input-axis 연산이 아님 | | GroupNorm | 한 image의 channel group과 spatial axis | 없음 | small batch vision에서 대안 | group 수와 layout 선택 필요 | | RMSNorm | 한 example의 feature axis, mean은 생략 | 없음 | mean subtraction을 생략한 LLM 계열 대안 | LayerNorm과 다른 output·parameter semantics |

원 논문은 BatchNorm·WeightNorm과 invariance를 비교하지만, GroupNorm·RMSNorm은 이후 연구다. 역사적으로 같은 표에 두어 비교할 수는 있어도 “논문이 이들을 실험했다”고 쓰면 안 된다.

흔한 오해와 최초 질문에 대한 답

LayerNorm은 batch를 전혀 보지 않나

보통 NLP의 LayerNorm(hidden_size) 설정에서는 그렇다. 하지만 정확한 말은 normalized_shape가 지정한 마지막 dimensions의 statistics를 쓴다는 것이다. batch dimension을 포함해 normalize하도록 tensor를 재구성하거나 다른 API를 쓰면 행동이 달라질 수 있다. axis가 사실이다.

LayerNorm output은 항상 평균 0이고 분산 1인가

affine 전 normalized vector는 epsilon을 무시하면 그렇다. 하지만 epsilon 때문에 정확히 1보다 조금 작을 수 있고, gain·bias 뒤 final output은 그렇지 않다. 위 직접 실행의 0.999996과 affine output이 이 차이를 보여 준다.

LayerNorm이면 train과 eval 결과가 똑같나

LayerNorm이 자기 통계를 고르는 규칙은 같다. 그러나 dropout, sampling, BatchNorm 등 model의 다른 부분이 mode별로 달라질 수 있다. “LayerNorm 자체의 statistics가 같다”가 정확한 범위다.

LayerNorm은 BatchNorm보다 최신이라 항상 좋은가

아니다. 원 논문도 CNN preliminary experiment에서 BatchNorm이 더 좋았다고 적었다. model·data·batch·axis에 맞는지 benchmark로 확인해야 한다.

Transformer의 LayerNorm과 이 논문은 같은가

기본 연산의 조상은 이 논문이다. 하지만 Transformer에서는 residual addition 전후 위치, hidden size 축, pre-LN·post-LN 선택, mixed precision, RMSNorm 등 추가 설계가 있다. LayerNorm 식 하나를 안다고 Transformer training 전체를 안 것은 아니다.

한 문장 답

LayerNorm은 batch의 다른 example을 기준으로 내 activation을 정규화하던 방식 대신, 한 example 내부의 feature들을 기준으로 정규화해 batch size와 RNN time step 통계 의존성을 줄인 연산이다. 단, axis·gain·bias·eps·배치 위치를 잘못 잡으면 같은 이름이어도 다른 model이 된다.

출처 및 검증 경로

1차 자료

1. Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton. Layer Normalization, arXiv 1607.06450v1, 2016-07-21. 이 글은 본문 1절부터 7절, 식 1부터 7, 표 1, 그림 2·6, supplementary material 식 15부터 28을 직접 확인했다.
2. PyTorch. torch.nn.LayerNorm stable API. 이 글은 normalized_shape, 마지막 D dimensions의 mean·variance, biased variance correction=0, eps=1e-5, elementwise_affine, train·evaluation 양쪽에서 input statistics를 쓴다는 현재 API 계약을 확인했다. 이 문서의 동작은 2016 논문 식의 모든 구현 세부를 보증하는 출처가 아니다.
3. Biao Zhang, Rico Sennrich. Root Mean Square Layer Normalization, arXiv 1910.07467. LayerNorm과 RMSNorm의 mean subtraction 차이를 구분하기 위한 이후 1차 논문이다.
4. Ashish Vaswani et al. Attention Is All You Need, NeurIPS 2017. Transformer 원 논문의 post-LN residual 배치를 확인하는 1차 자료다.

직접 실행

• Python 3.9.6, 표준 라이브러리만 사용했다. 한 vector의 mean·population variance·affine output, 배치 다른 row 변경 전후 LayerNorm·BatchNorm 결과, time-step vector, constant vector의 epsilon, analytic gradient·finite difference를 위 코드로 실행했고 본문 출력은 실제 결과다.
• 이 검증은 PyTorch kernel benchmark, GPU mixed-precision 정확도, RNN·LSTM 전체 학습, 원 논문의 MNIST·NLP benchmark 수치 재현이 아니다. 그 결과는 원 논문의 명시된 model·data·optimizer·batch 조건에서 나온 외부 실험 결과로만 인용했다.

이 글의 해석

• SQL GROUP BY 비유와 “옆 사람” 비유는 normalization axis를 설명하기 위한 해석이다. 원 논문이 database query를 제안한 것은 아니다.
• 식 (10)의 gradient는 이 글의 작은 LayerNorm function을 미분해 finite difference로 확인한 계산이다. 원 논문이 제시한 production backward kernel 또는 모든 affine·shape의 완전한 gradient 식이라고 주장하지 않는다.