Transformer LayerNorm

Transformer Layer Normalization 완벽 이해

Transformer를 공부하다 보면 Multi-Head Attention(MHA), Feed Forward Network(FFN)와 함께 반복적으로 등장하는 것이 바로 Layer Normalization(LayerNorm)입니다.

전체 코드는 다음과 같습니다.

import torch
import torch.nn as nn


class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-5):
        super().__init__()

        # γ (gamma): Scale
        self.gamma = nn.Parameter(torch.ones(features))

        # β (beta): Shift
        self.beta = nn.Parameter(torch.zeros(features))

        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        var = x.var(-1, keepdim=True, unbiased=False)

        return self.gamma * (x - mean) / torch.sqrt(var + self.eps) + self.beta


features = d_model = 512
eps = 1e-6

# x = mha_result
# ff = PositionwiseFeedForward(d_model, d_ff, dropout)
# ff_result = ff(x)

x = ff_result

ln = LayerNorm(features, eps)
ln_result = ln(x)

print(ln_result)
print(ln_result.shape)

FFN:https://deepbored.com/transformer-ffn/

MHA:https://deepbored.com/transformer-multi-head-attention-mha/

Input(Embedding+Position Embedding):https://deepbored.com/pytorch-transformer-input/

1. Layer Normalization이란?

Layer Normalization은 입력 벡터의 평균(Mean)과 분산(Variance)을 이용하여 값을 정규화하는 방법입니다.

LayerNorm의 기본 수식은 다음과 같습니다.

LayerNorm(x)=γx−μσ2+ϵ+βLayerNorm(x) = \gamma \frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta

코드와 비교하면 정확하게 다음과 대응합니다.

수식코드의미
xx입력
μmean평균
σ²var분산
εeps수치적 안정성을 위한 작은 값
γgamma학습 가능한 Scale
βbeta학습 가능한 Shift
LayerNorm(x)ln_resultLayerNorm 결과

2. 전체 계산 과정부터 이해하기

LayerNorm의 전체 흐름은 생각보다 간단합니다.

입력 x
   ↓
평균 μ 계산
   ↓
분산 σ² 계산
   ↓
x - μ
   ↓
√(σ² + ε)로 나누기
   ↓
γ 곱하기
   ↓
β 더하기
   ↓
출력 y

코드에서는 다음 세 줄이 이 과정을 모두 수행합니다.

mean = x.mean(-1, keepdim=True)
var = x.var(-1, keepdim=True, unbiased=False)

return self.gamma * (x - mean) / torch.sqrt(var + self.eps) + self.beta

하나씩 살펴보겠습니다.

3. LayerNorm은 어떤 차원을 정규화할까?

Transformer에서 이 부분이 가장 중요합니다.

Transformer 내부 Tensor의 일반적인 shape은 다음과 같습니다.

[batch_size, seq_len, d_model]

예를 들어,

[32, 10, 512]

라면

32  → batch_size
10  → seq_len
512 → d_model

입니다.

LayerNorm은 여기서 마지막 차원인 d_model을 기준으로 정규화합니다.

[ batch_size, seq_len, d_model ]
                         ↑
                    LayerNorm

즉, seq_len 방향으로 정규화하는 것이 아닙니다.

각 Token이 가지고 있는 512개의 feature를 하나의 벡터로 보고 정규화합니다.

4. 각 Token의 feature를 독립적으로 정규화

예를 들어 다음과 같은 문장이 있다고 해보겠습니다.

I    love    Transformer

d_model=512라면 각각의 Token은 512차원 벡터를 가지고 있습니다.

I
→ [x₁, x₂, x₃, ..., x₅₁₂]

love
→ [x₁, x₂, x₃, ..., x₅₁₂]

Transformer
→ [x₁, x₂, x₃, ..., x₅₁₂]

LayerNorm은 다음과 같이 각각 독립적으로 적용됩니다.

Token 1의 512 features
        ↓
    LayerNorm

Token 2의 512 features
        ↓
    LayerNorm

Token 3의 512 features
        ↓
    LayerNorm

따라서 핵심은 다음과 같습니다.

LayerNorm은 Token 간의 값을 정규화하는 것이 아니라, 각 Token 내부의 d_model feature들을 정규화합니다.

5. features = d_model = 512

코드에서는 다음과 같이 설정했습니다.

features = d_model = 512

features는 LayerNorm이 처리해야 하는 마지막 feature dimension의 크기입니다.

Transformer에서는 이것이 일반적으로 d_model과 같습니다.

따라서

ln = LayerNorm(512)

은

각 Token의 512개 feature
          ↓
      LayerNorm

을 의미합니다.

6. 평균 계산 — mean

먼저 평균을 계산합니다.

mean = x.mean(-1, keepdim=True)

여기서 가장 중요한 부분은 -1입니다.

PyTorch에서 -1은 Tensor의 마지막 dimension을 의미합니다.

입력이

[batch_size, seq_len, d_model]

이라면

dim = 0   → batch_size
dim = 1   → seq_len
dim = 2   → d_model

dim = -1  → d_model

입니다.

따라서

x.mean(-1)

은 각 Token의 d_model feature들을 이용하여 평균을 계산합니다.

간단한 숫자로 살펴보기

실제 512차원은 너무 크기 때문에 4차원으로 줄여보겠습니다.

하나의 Token이 다음과 같다고 가정합니다.

x = [1, 2, 3, 4]

평균은

μ=1+2+3+44=2.5\mu=\frac{1+2+3+4}{4}=2.5

입니다.

따라서

mean = 2.5

가 됩니다.

7. keepdim=True는 왜 사용할까?

코드를 보면 다음 옵션이 있습니다.

mean = x.mean(-1, keepdim=True)

keepdim=True는 평균을 계산한 후에도 해당 dimension을 유지합니다.

예를 들어 입력 shape이

[32, 10, 512]

라면

x.mean(-1)

의 결과는

[32, 10]

입니다.

반면,

x.mean(-1, keepdim=True)

은

[32, 10, 1]

이 됩니다.

따라서 이후 다음 연산을 쉽게 수행할 수 있습니다.

x - mean

PyTorch의 Broadcasting을 통해

x      → [32, 10, 512]

mean   → [32, 10,   1]
                    ↓
              Broadcasting

result → [32, 10, 512]

형태로 계산됩니다.

8. 분산 계산 — var

다음으로 분산을 계산합니다.

var = x.var(-1, keepdim=True, unbiased=False)

여기서도 -1을 사용합니다.

따라서 평균과 마찬가지로 각 Token의 d_model feature들을 기준으로 분산을 계산합니다.

분산은 간단히 말하면 값들이 평균에서 얼마나 퍼져 있는지를 나타냅니다.

예를 들어,

[5, 5, 5, 5]

는 모든 값이 동일하기 때문에 분산이 0입니다.

반면,

[1, 5, 10, 20]

은 값들이 넓게 퍼져 있기 때문에 분산이 더 큽니다.

9. unbiased=False의 의미

이번 구현에서는 다음과 같이 작성했습니다.

var = x.var(
    -1,
    keepdim=True,
    unbiased=False
)

unbiased=False를 사용하면 분산 계산에서 N으로 나누는 population variance 방식을 사용합니다.

즉,

σ2=1N∑i=1N(xi−μ)2\sigma^2 = \frac{1}{N} \sum_{i=1}^{N}(x_i-\mu)^2

입니다.

이 설정은 PyTorch의 nn.LayerNorm이 사용하는 분산 계산 방식과 대응합니다.

따라서 직접 LayerNorm을 구현할 때 중요한 부분입니다.

10. x - mean — 중심을 0으로 이동

이제 실제 정규화를 시작합니다.

x - mean

앞의 예제를 다시 사용해 보겠습니다.

x = [1, 2, 3, 4]

mean = 2.5

평균을 빼면

x - mean

= [-1.5, -0.5, 0.5, 1.5]

이 됩니다.

변환된 값의 평균은 0입니다.

즉,

Before

[1, 2, 3, 4]
     중심 ≈ 2.5

        ↓ mean 제거

After

[-1.5, -0.5, 0.5, 1.5]
          중심 = 0

평균을 빼는 과정은 데이터의 중심을 0으로 이동시키는 과정이라고 이해하면 됩니다.

11. torch.sqrt(var + eps)

다음은 분모입니다.

torch.sqrt(var + self.eps)

따라서 실제 Normalization 부분은

(x - mean) / torch.sqrt(var + self.eps)

입니다.

수식으로는

x−μσ2+ϵ\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}

에 해당합니다.

즉,

x
 ↓
평균 제거
 ↓
x - μ
 ↓
분산 기반 Scale 조정
 ↓
(x - μ) / √(σ² + ε)

의 과정을 거칩니다.

12. eps는 왜 필요할까?

코드에서는 다음과 같이 설정했습니다.

eps = 1e-6

이는

0.000001

이라는 아주 작은 숫자입니다.

만약 Token의 feature가 모두 같다면 어떻게 될까요?

x = [5, 5, 5, 5]

평균은 5이고 분산은 0입니다.

mean = 5
var  = 0

만약 eps가 없다면 분모가 0이 될 수 있습니다.

(x - mean) / √0

따라서

torch.sqrt(var + self.eps)

처럼 아주 작은 값을 추가합니다.

즉,

eps는 0으로 나누는 문제를 방지하여 수치적 안정성(Numerical Stability)을 확보합니다.

13. γ(Gamma)는 무엇일까?

이제 이번 코드에서 가장 중요한 변경 사항입니다.

self.gamma = nn.Parameter(torch.ones(features))

gamma는 LayerNorm 수식의 γ입니다.

gamma = γ = Scale

역할은 정규화된 feature의 값을 확대하거나 축소하는 것입니다.

초기값은 1입니다.

d_model=512라면,

gamma

[1, 1, 1, ..., 1]
 ↑              ↑
γ₁             γ₅₁₂

처럼 512개의 학습 가능한 값이 존재합니다.

초기에는 1이기 때문에

1 × normalized_x

가 되어 정규화된 값을 그대로 유지합니다.

하지만 학습이 진행되면 γ도 함께 업데이트됩니다.

초기

γ = 1

 ↓ Training

γ = 학습된 값

따라서 모델이 각 feature의 적절한 Scale을 직접 학습할 수 있습니다.

14. β(Beta)는 무엇일까?

다음은 beta입니다.

self.beta = nn.Parameter(torch.zeros(features))

beta는 LayerNorm 수식의 β에 해당합니다.

beta = β = Shift

초기값은 0입니다.

beta

[0, 0, 0, ..., 0]
 ↑              ↑
β₁             β₅₁₂

β의 역할은 정규화된 값을 이동(Shift)시키는 것입니다.

따라서 γ와 β를 함께 보면,

γ (Gamma)
    ↓
Scale
    ↓
확대 / 축소


β (Beta)
    ↓
Shift
    ↓
이동

이라고 이해할 수 있습니다.

15. 왜 γ와 β는 nn.Parameter일까?

두 변수 모두 다음과 같이 정의했습니다.

self.gamma = nn.Parameter(torch.ones(features))
self.beta = nn.Parameter(torch.zeros(features))

nn.Parameter를 사용하면 PyTorch가 이 Tensor들을 모델이 학습해야 하는 파라미터로 등록합니다.

즉 Optimizer가 학습 과정에서

gamma
beta

도 함께 업데이트합니다.

따라서 LayerNorm은 단순히 데이터를 정규화하고 끝나는 것이 아닙니다.

Normalization
      ↓
정규화된 표현
      ↓
gamma로 Scale 학습
      ↓
beta로 Shift 학습
      ↓
모델에 필요한 표현

이라는 과정까지 포함합니다.

16. LayerNorm의 핵심 코드 분석

이제 다음 한 줄을 보면 의미가 훨씬 명확합니다.

return self.gamma * (x - mean) / torch.sqrt(var + self.eps) + self.beta

단계별로 나누면 다음과 같습니다.

① x

 ↓

② x - mean

평균을 빼서 중심을 0으로 이동

 ↓

③ (x - mean) / √(var + eps)

분산을 이용해 값의 Scale을 정규화

 ↓

④ gamma × normalized_x

학습 가능한 γ로 Scale 조정

 ↓

⑤ + beta

학습 가능한 β로 Shift

 ↓

⑥ LayerNorm Output

수식과 코드도 매우 직관적으로 대응합니다.

LayerNorm 수식                 Python

x                              x

μ                              mean

σ²                             var

ε                              self.eps

γ                              self.gamma

β                              self.beta

그래서 a2, b2보다 gamma, beta라는 변수명을 사용하면 LayerNorm 수식을 공부할 때 훨씬 이해하기 쉽습니다.

17. FFN 결과에 LayerNorm 적용하기

이번 코드에서는 다음과 같이 사용하고 있습니다.

x = ff_result

즉 Positionwise Feed Forward Network의 결과를 LayerNorm의 입력으로 사용합니다.

FFN
 ↓
ff_result
 ↓
LayerNorm
 ↓
ln_result

코드에서는

x = ff_result

ln = LayerNorm(features, eps)
ln_result = ln(x)

입니다.

LayerNorm 내부에서는 다시

ff_result
    ↓
mean 계산
    ↓
variance 계산
    ↓
Normalization
    ↓
× gamma
    ↓
+ beta
    ↓
ln_result

순서로 계산됩니다.

18. MHA 결과에도 적용할 수 있다

코드의 주석처럼

# x = mha_result

Multi-Head Attention의 출력 역시 LayerNorm의 입력으로 사용할 수 있습니다.

즉 LayerNorm 자체의 관점에서는

MHA Result
    ↓
LayerNorm

이든

FFN Result
    ↓
LayerNorm

이든 동일한 원리로 동작합니다.

입력이 [batch_size, seq_len, d_model]이라면 마지막 d_model 차원을 기준으로 정규화합니다.

19. 실제 Transformer에서는 Residual Connection도 중요하다

여기서 한 가지 주의할 점이 있습니다.

현재 예제는

x = ff_result
ln_result = ln(x)

으로 LayerNorm 자체의 동작을 확인하는 코드입니다.

하지만 원래 Transformer의 대표적인 Post-Norm 구조에서는 Residual Connection이 함께 사용됩니다.

즉 개념적으로

LayerNorm(x+Sublayer(x))LayerNorm(x + Sublayer(x))

입니다.

FFN을 예로 들면,

이 됩니다.

따라서 현재 코드는 LayerNorm 내부 연산을 이해하기 위한 예제이고, 전체 Transformer Block을 구현할 때는 Residual Connection까지 함께 고려해야 합니다.

20. LayerNorm 전후의 shape은 같다

마지막으로 다음 코드를 실행합니다.

print(ln_result)
print(ln_result.shape)

입력 x의 shape이

[32, 10, 512]

라면 LayerNorm을 통과한 이후에도

[32, 10, 512]

입니다.

즉,

LayerNorm은 Tensor의 shape을 바꾸는 연산이 아닙니다.

각 feature의 값을 정규화할 뿐입니다.

21. nn.LayerNorm과 비교하면?

PyTorch에서는 직접 구현하지 않고 다음처럼 사용할 수도 있습니다.

ln = nn.LayerNorm(
    normalized_shape=512,
    eps=1e-6
)

ln_result = ln(x)

직접 구현한 코드와 비교하면 다음과 같습니다.

직접 구현nn.LayerNorm의미
featuresnormalized_shape정규화할 차원
gammaweightScale γ
betabiasShift β
epsepsNumerical Stability

즉 우리가 직접 작성한

self.gamma
self.beta

는 PyTorch nn.LayerNorm 내부의

weight
bias

와 대응합니다.

22. 전체 구조 한 번에 보기

이번 코드의 전체 흐름을 한 장으로 정리하면 다음과 같습니다.

23. 핵심 정리

Layer Normalization에서 꼭 기억해야 할 핵심은 다음과 같습니다.

① 마지막 차원 d_model을 기준으로 정규화한다.

[batch_size, seq_len, d_model]
                       ↑
                  이 차원

② 각 Token은 독립적으로 정규화된다.

Token 1 → 자신의 d_model features 정규화
Token 2 → 자신의 d_model features 정규화
Token 3 → 자신의 d_model features 정규화

③ 평균을 계산한다.

mean = x.mean(-1, keepdim=True)

④ 분산을 계산한다.

var = x.var(
    -1,
    keepdim=True,
    unbiased=False
)

⑤ 실제 Normalization을 수행한다.

(x - mean) / torch.sqrt(var + self.eps)

⑥ gamma(γ)는 학습 가능한 Scale이다.

self.gamma = nn.Parameter(
    torch.ones(features)
)

⑦ beta(β)는 학습 가능한 Shift이다.

self.beta = nn.Parameter(
    torch.zeros(features)
)

⑧ LayerNorm은 shape을 변경하지 않는다.

[32, 10, 512]

      ↓ LayerNorm

[32, 10, 512]

마무리

LayerNorm을 한 문장으로 정리하면,

각 Token의 마지막 차원(d_model)에 있는 feature들을 평균과 분산으로 정규화한 뒤, 학습 가능한 gamma(γ)와 beta(β)를 이용하여 모델에 필요한 Scale과 Shift를 다시 학습하는 과정입니다.

이번 구현에서는 변수명을 직접

self.gamma
self.beta

로 사용했기 때문에 LayerNorm 공식과 코드의 관계도 훨씬 명확합니다.

결국 가장 중요한 코드는 다음 세 줄입니다.

mean = x.mean(-1, keepdim=True)
var = x.var(-1, keepdim=True, unbiased=False)

return self.gamma * (x - mean) / torch.sqrt(var + self.eps) + self.beta

이 세 줄을 이해하면 Transformer에서 사용하는 Layer Normalization의 핵심 원리를 거의 모두 이해했다고 볼 수 있습니다.

태그

댓글 남기기

이메일 주소는 공개되지 않습니다. 필수 필드는 *로 표시됩니다