Group Normalization¶

import torch
from torch import nn

torch.manual_seed(23)

N, L, C = 8, 16, 32
x = torch.randn(N, L, C)
x_ncl = x.transpose(1, 2)  # (N, C, L)

# BN: channel별로 batch와 sequence 축을 함께 사용함.
# affine=False로 학습 가능한 scale과 bias를 생성하지 않음.
bn = nn.BatchNorm1d(C, affine=False)

# 일반 LN: 각 sample의 (L, C) 전체를 사용하도록 지정함.
# elementwise_affine=False로 학습 가능한 scale과 bias를 생성하지 않음.
ln_all = nn.LayerNorm((L, C), elementwise_affine=False)

# Transformer의 LN: 각 token의 C개 feature만 사용함.
ln_token = nn.LayerNorm(C, elementwise_affine=False)

# GN (1 group): 각 sample의 (C, L) 전체를 사용함.
gn1 = nn.GroupNorm(num_groups=1, num_channels=C, affine=False)

# GN (4 groups): channel을 8개씩 4개 group으로 나눔.
# 각 group의 8개 channel과 전체 sequence 위치를 함께 사용함.
G = 4
gn4 = nn.GroupNorm(num_groups=G, num_channels=C, affine=False)

# 각 layer에 필요한 입력 형태를 지정함.
cases = [
    ("BN", bn, x_ncl),
    ("LN (L, C)", ln_all, x),
    ("LN (token)", ln_token, x),
    ("GN (1 group)", gn1, x_ncl),
    ("GN (4 groups)", gn4, x_ncl),
]

# Training mode에서 출력값을 계산함.
train_outputs = {}
for name, layer, data in cases:
    layer.train()
    train_outputs[name] = layer(data)

# 출력 비교를 위해 (N, L, C) 형태로 맞춤.
y_bn = train_outputs["BN"].transpose(1, 2)
y_ln_all = train_outputs["LN (L, C)"]
y_ln_token = train_outputs["LN (token)"]
y_gn1 = train_outputs["GN (1 group)"].transpose(1, 2)
y_gn4_ncl = train_outputs["GN (4 groups)"]
y_gn4 = y_gn4_ncl.transpose(1, 2)


def report(name, y, dims):
    mean = y.mean(dim=dims)
    var = y.var(dim=dims, unbiased=False)

    # 각 normalization 단위에서 계산한 mean과 variance를 확인함.
    # mean의 절댓값과 variance의 1에 대한 차이의 최댓값을 출력함.
    print(
        f"{name:16s} "
        f"통계 shape={tuple(mean.shape)}, "
        f"max |mean|={mean.abs().max().item():.4f}, "
        f"max |var - 1|={(var - 1).abs().max().item():.4f}"
    )


# (N, L, C) 형태에서 각 방식의 normalization 축을 확인함.
report("BN", y_bn, (0, 1))
report("LN (L, C)", y_ln_all, (1, 2))
report("LN (token)", y_ln_token, (2,))
report("GN (1 group)", y_gn1, (1, 2))

# GN (4 groups)는 group별 통계를 확인하기 위해 reshape함.
# (N, C, L)에서 (N, G, C // G, L)로 변환함.
y_gn4_grouped = y_gn4_ncl.reshape(N, G, C // G, L)
report("GN (4 groups)", y_gn4_grouped, (2, 3))

print(
    "LN (L, C)와 GN (1 group)의 결과 일치:",
    torch.allclose(y_ln_all, y_gn1, atol=1e-6),
)

print(
    "모든 layer의 scale과 bias가 None:",
    all(
        layer.weight is None and layer.bias is None
        for _, layer, _ in cases
    ),
)

# 같은 입력에 대해 training mode와 eval mode의 출력을 비교함.
for name, layer, data in cases:
    layer.eval()
    y_eval = layer(data)

    print(
        f"{name}: train/eval 결과 일치:",
        torch.allclose(train_outputs[name], y_eval, atol=1e-6),
    )

# 예상 출력: 소수점 넷째 자리까지 표시한 예시이며, 실제 실행 확인 전임.
#
# BN               통계 shape=(32,), max |mean|=0.0000, max |var - 1|=0.0000
# LN (L, C)        통계 shape=(8,), max |mean|=0.0000, max |var - 1|=0.0000
# LN (token)       통계 shape=(8, 16), max |mean|=0.0000, max |var - 1|=0.0000
# GN (1 group)     통계 shape=(8,), max |mean|=0.0000, max |var - 1|=0.0000
# GN (4 groups)    통계 shape=(8, 4), max |mean|=0.0000, max |var - 1|=0.0000
# LN (L, C)와 GN (1 group)의 결과 일치: True
# 모든 layer의 scale과 bias가 None: True
# BN: train/eval 결과 일치: False
# LN (L, C): train/eval 결과 일치: True
# LN (token): train/eval 결과 일치: True
# GN (1 group): train/eval 결과 일치: True
# GN (4 groups): train/eval 결과 일치: True
#
# 통계 shape은 normalization에 사용한 축들을 축약한 후 남는 shape임.
# BN의 (32,)는 channel별 32쌍의 통계를 의미함.
# LN (L, C)와 GN (1 group)의 (8,)은 sample별 8쌍의 통계를 의미함.
# LN (token)의 (8, 16)은 sample별·token별 총 128쌍의 통계를 의미함.
# GN (4 groups)의 (8, 4)는 sample별·group별 총 32쌍의 통계를 의미함.
#
# max |mean|이 약 0이면 각 normalization 단위의 mean이 약 0임.
# max |var - 1|이 약 0이면 각 normalization 단위의 variance가 약 1임.
# 0.0000은 반올림된 표시이며, 정확히 0이라는 의미는 아님.
# eps 때문에 variance는 정확히 1이 되지는 않음.
#
# LN (L, C)와 GN (1 group)은 scale과 bias가 비활성화된 상태에서
# 동일한 값들을 사용하여 normalization하므로 결과가 허용 오차 내에서 같음.
#
# BN은 training mode에서 현재 batch의 통계를 사용하고,
# eval mode에서는 running statistics를 사용하므로 이 예제의 출력이 다름.
# LN과 GN은 두 mode 모두 현재 입력의 통계를 사용하므로 출력이 같음.