가중치 초기화와 Kaiming(He) 초기화

2025-09-05 #pytorch #weight-initialization

간단한 Linear Layer를 pytorch로 직접 구현해보면서, 가중치를 어떻게 초기화하느냐에 따라 학습 결과가 달라지는 것을 확인해보았다. 해당 내용에 대한 짧은 글이다.

표준정규분포로 초기화#

torch.randn으로 가중치를 초기화하면 표준정규분포(평균 0, 표준편차 1)에서 값을 뽑는다.

Pasted image 20250905083243.png

이 경우 직접 구현한 선형 모델의 train loss, test loss가 초반에 크게 출렁이는 불안정한 모습을 볼 수 있다. 표준편차 1은 가중치 값의 분산이 꽤 커서, 초기 출력이 들쭉날쭉해지기 때문이다.

정규균등분포로 초기화#

반면에 값의 범위를 좁힌 균등분포로 초기화하면 시작 loss 자체가 훨씬 낮고 안정적이다.

Pasted image 20250905091952.png

Kaiming(He) 초기화 (nn.Linear의 방식)#

nn.Linear는 내부적으로 kaiming_uniform_ 함수로 weight를 초기화한다. 다만 이걸 곧 ReLU용 He 초기화라고 부르기엔 무리가 있다. 기본값은 a=math.sqrt(5)라는 legacy 인자를 쓰기 때문인데, 이 부분은 nn.Linear의 기본 초기화는 He 초기화가 아니다에 따로 정리했다.

스크린샷 2025-09-05 084051.png

여기서 Kaiming 초기화는 He 초기화와 같은 것을 가리킨다. 제안한 사람(Kaiming He)의 이름에서 따온 것이라, Kaiming init이라고도 He init이라고도 부른다.

Kaiming(He) 초기화의 핵심 아이디어는 입력 차원 수(fan-in)를 고려해서 가중치의 분포 범위를 정하는 것이다. 입력이 많이 들어오는 층일수록 가중치를 더 작은 범위에서 뽑아, 여러 입력이 곱해져 더해질 때 출력의 분산이 폭발하지 않도록 맞춰준다.

분산 기준으로 보면, fan-in을 n이라 할 때 가중치의 분산을 대략 2/n 정도로 잡는다. (ReLU 계열을 가정한 값)
균등분포로 구현할 때는 이 분산에 맞는 범위 [-bound, +bound]에서 값을 뽑는다.

이렇게 하면 층을 지나도 출력의 분산이 너무 커지거나 작아지지 않게 유지되고, 학습이 안정적으로 시작된다. ReLU처럼 음수를 잘라내는 활성화 함수와 함께 쓸 때 특히 잘 맞아서, 깊은 신경망에서 표준값처럼 쓰인다. 직접 randn으로 초기화한 것과 nn.Linear를 쓴 것의 loss 곡선이 확실하게 차이 나는 이유가 여기에 있다.

0으로 초기화하면 안 되는 이유#

그 전에 짚고 갈 게 있다. 가중치를 전부 0이나 전부 같은 값으로 두면, 한 층의 모든 뉴런이 같은 입력에 같은 출력을 내고 역전파에서도 같은 gradient를 받는다. 그러면 학습 내내 모든 뉴런이 똑같이 갱신되어 사실상 뉴런 하나짜리 층과 다를 게 없어진다. 이 대칭을 깨려면 가중치를 서로 다른 값으로 흩어줘야 하고, 그래서 무작위 초기화가 필요하다. randn이든 균등분포든 Kaiming이든 공통으로 깔린 전제가 이 대칭성 깨기다.

세 방법의 값 범위 비교#

위의 세 방법이 실제로 어떤 범위의 값을 뽑는지 직접 찍어봤다. fan_in을 바꿔가며 가중치 원소의 min/max, 표준편차, 절댓값이 1을 넘는 비율을 본 것이다.

import math
import torch

torch.manual_seed(0)

def summary(name, w):
    w = w.flatten()
    print(f"{name:<30} "
          f"min={w.min().item():+.3f}  max={w.max().item():+.3f}  "
          f"std={w.std().item():.4f}  "
          f"|w|>1 비율={ (w.abs() > 1).float().mean().item()*100:5.1f}%  "
          f"|w|>3 비율={ (w.abs() > 3).float().mean().item()*100:5.2f}%")

for fan_in in [10, 100, 784]:
    fan_out = fan_in
    n = fan_in * fan_out
    print(f"\nfan_in = {fan_in}  (가중치 원소 {n}개)")

    # 1) torch.randn = 표준정규분포 (평균 0, 표준편차 1)
    w_randn = torch.randn(fan_out, fan_in)
    summary("torch.randn (표준정규 N(0,1))", w_randn)

    # 2) 값의 범위를 좁힌 균등분포 U(-0.5, 0.5)
    w_unif = torch.empty(fan_out, fan_in).uniform_(-0.5, 0.5)
    summary("U(-0.5, 0.5) 좁은 균등분포", w_unif)

    # 3) nn.Linear 기본 = kaiming_uniform_(a=sqrt(5))
    #    bound = sqrt(6 / ((1+a^2) * fan_in)) = 1/sqrt(fan_in)  (a=sqrt(5)일 때)
    w_kaiming = torch.empty(fan_out, fan_in)
    torch.nn.init.kaiming_uniform_(w_kaiming, a=math.sqrt(5))
    bound = 1.0 / math.sqrt(fan_in)
    summary(f"kaiming_uniform a=root5 [+-{bound:.4f}]", w_kaiming)

fan_in = 10  (가중치 원소 100개)
torch.randn (표준정규 N(0,1))      min=-2.613  max=+3.411  std=1.0296  |w|>1 비율= 31.0%  |w|>3 비율= 1.00%
U(-0.5, 0.5) 좁은 균등분포           min=-0.482  max=+0.475  std=0.2713  |w|>1 비율=  0.0%  |w|>3 비율= 0.00%
kaiming_uniform a=root5 [+-0.3162] min=-0.313  max=+0.313  std=0.1721  |w|>1 비율=  0.0%  |w|>3 비율= 0.00%

fan_in = 100  (가중치 원소 10000개)
torch.randn (표준정규 N(0,1))      min=-4.343  max=+4.101  std=1.0077  |w|>1 비율= 32.0%  |w|>3 비율= 0.27%
U(-0.5, 0.5) 좁은 균등분포           min=-0.500  max=+0.500  std=0.2886  |w|>1 비율=  0.0%  |w|>3 비율= 0.00%
kaiming_uniform a=root5 [+-0.1000] min=-0.100  max=+0.100  std=0.0574  |w|>1 비율=  0.0%  |w|>3 비율= 0.00%

fan_in = 784  (가중치 원소 614656개)
torch.randn (표준정규 N(0,1))      min=-4.407  max=+4.668  std=1.0005  |w|>1 비율= 31.7%  |w|>3 비율= 0.28%
U(-0.5, 0.5) 좁은 균등분포           min=-0.500  max=+0.500  std=0.2887  |w|>1 비율=  0.0%  |w|>3 비율= 0.00%
kaiming_uniform a=root5 [+-0.0357] min=-0.036  max=+0.036  std=0.0206  |w|>1 비율=  0.0%  |w|>3 비율= 0.00%

fan_in=100을 기준으로 보면 범위 차이가 분명하다.

randn은 사실상 [-4, +4]까지 퍼지고 표준편차가 약 1이다. 절댓값이 1을 넘는 가중치가 약 30퍼센트나 된다.
U는 범위가 [-0.5, +0.5]로 좁고 1을 넘는 값이 하나도 없다.
kaiming_uniform_(a=√5)는 fan_in=100일 때 [-0.1, +0.1]로 더 좁고, fan_in이 커질수록 범위가 자동으로 더 줄어든다(fan_in=784면 ±0.036).

핵심은 범위의 끝값보다 fan_in에 맞춰 범위가 줄어드느냐다. randn은 fan_in과 무관하게 항상 같은 범위라, 입력이 많은 층일수록 출력이 더 크게 부푼다.

작은 값이 무조건 안정적인 건 아니다#

그러면 0에 가깝게만 두면 항상 안정적일까? 그렇지 않다. 너무 작으면 신호가 층을 지나며 사라진다. 256차원 입력을 ReLU 선형층 12개에 통과시키며 각 층 출력의 표준편차가 어떻게 변하는지 추적해봤다.

import math
import torch

torch.manual_seed(0)

DIM = 256
DEPTH = 12
BATCH = 1024

def run(name, std_w):
    x = torch.randn(BATCH, DIM)  # 입력 신호 std=1
    layers = [torch.randn(DIM, DIM) * std_w for _ in range(DEPTH)]
    stds = []
    h = x
    for W in layers:
        h = torch.relu(h @ W.T)
        stds.append(h.std().item())
    line = "  ".join(f"{s:7.3f}" if s >= 0.001 else f"{s:7.0e}" for s in stds)
    print(f"{name:<34} {line}")

# 1) 너무 큼: randn std=1
run("randn  std=1.0 (너무 큼)", 1.0)
# 2) 너무 작음: std=0.02
run("std=0.02 (너무 작음)", 0.02)
# 3) 적정: He(Kaiming) std = sqrt(2/fan_in)
he = math.sqrt(2.0 / DIM)
run(f"He  std=sqrt(2/n)={he:.4f} (적정)", he)

입력 std=1.0, 층마다 ReLU. 각 층 출력의 std (12개 층):
                                   L1       L2       L3      ...      L12
randn  std=1.0 (너무 큼)            9.359  103.250  1184.130  ...  2474914873344.000
std=0.02 (너무 작음)               0.188    0.042    0.009  ...  2e-08
He  std=sqrt(2/n)=0.0884 (적정)    0.827    0.810    0.822  ...  0.665

세 경우가 갈린다.

randn(표준편차 1)은 층마다 신호가 약 10배씩 커져 12층 만에 조 단위로 폭발한다.
표준편차 0.02는 반대로 층마다 신호가 줄어 12층 만에 사실상 0으로 사라진다. 신호가 사라지면 gradient도 0이 되어 학습이 안 된다.
He 초기화(표준편차 √(2/n)≈0.088)는 12층을 지나도 표준편차가 0.7~0.9 부근에서 유지된다.

즉 좋은 초기화는 값이 작은 게 아니라, 층을 지나도 신호의 분산이 보존되는 적정 스케일이다. 그 스케일이 fan_in에 따라 정해지고, 이게 Kaiming이 fan_in으로 범위를 정하는 이유다. 처음의 randn 초기화가 폭발까지 가지 않고 출렁이기만 한 건 모델 층이 얕았기 때문이다.

Xavier와 He의 차이#

fan_in으로 분산을 맞추는 초기화에는 크게 두 갈래가 있다.

Xavier 초기화: 활성화 함수가 tanh, sigmoid처럼 0 근처에서 선형에 가깝다고 가정한다. forward와 backward 분산을 둘 다 고려해 분산을 2/(fan_in+fan_out)로 잡는다.
He(Kaiming) 초기화: ReLU를 가정한다. ReLU는 음수 입력을 0으로 잘라 통과 후 분산이 절반으로 줄기 때문에, 이를 보상하려고 분산을 2배로 키워 2/fan_in으로 잡는다.

정규분포로 구현하면 표준편차는 각각 √(2/(fan_in+fan_out))와 √(2/fan_in)이고, 균등분포로 구현하면 범위는 √(6/(fan_in+fan_out))와 √(6/fan_in)이다. 핵심 차이는 ReLU 때문에 붙는 2배 팩터와, fan_out을 함께 보느냐다. ReLU 계열을 쓰면 He, tanh 계열을 쓰면 Xavier를 쓰는 게 보통이다.

정리#

가중치 초기화 방법에 따라 학습 초반의 loss와 안정성이 달라진다.
표준편차 1짜리 randn은 분산이 커서 불안정할 수 있다.
nn.Linear는 kaiming_uniform_로 초기화하지만 a=√5인 legacy 기본값이라, ReLU용 정석 He 초기화와는 다르다(nn.Linear의 기본 초기화는 He 초기화가 아니다). Kaiming init과 He init이라는 이름 자체는 같은 사람(Kaiming He)에서 온 같은 말이다.
가중치를 전부 같은 값으로 두면 대칭성 때문에 학습이 안 되므로, 무작위로 흩어주는 게 초기화의 출발점이다.
좋은 초기화의 기준은 값이 작은 게 아니라 층을 지나도 분산이 보존되는 스케일이다. 너무 크면 폭발하고 너무 작으면 소멸한다.
ReLU 계열엔 He, tanh 계열엔 Xavier를 쓴다. 둘 다 fan_in으로 분산을 맞추되 He는 ReLU 보상으로 분산을 2배 키운다.
최신 LLM은 fan_in 공식 대신 표준편차 0.02 고정 초기화를 쓰고, 잔차 깊이 스케일링과 RMSNorm이 안정성을 함께 받쳐준다.

참고: