가중치 초기화

Weight initialization; 가중치 초기화
학습을 시작하기 전에 신경망 가중치의 처음 값을 정하는 방법. 층을 지나도 신호와 기울기의 크기가 유지되도록 분포의 분산을 고르는 것이 핵심이다.

가중치 초기화(weight initialization)는 신경망의 가중치를 학습 시작 전에 어떤 값으로 채울지 정하는 일이다. 단순해 보이지만 깊은 신경망에서는 초기값 분포에 따라 학습이 아예 안 되기도 한다. 층을 하나 지날 때마다 활성값의 분산이 조금씩 줄거나 늘면, 수십 층을 지나면서 지수적으로 0에 가까워지거나 폭발한다. 역전파되는 기울기도 마찬가지라 기울기 소실이나 기울기 폭주로 이어진다. 좋은 초기화는 순전파 활성값과 역전파 기울기의 크기가 깊이와 상관없이 비슷하게 유지되도록 분산을 정한다.

대표적인 방법은 tanh·시그모이드용 세이비어 초기화(Xavier initialization, 글로럿 초기화), ReLU용 He 초기화(카이밍 초기화), 직교 행렬을 쓰는 직교 초기화, 그리고 모델 너비가 바뀌어도 하이퍼파라미터가 그대로 옮겨지도록 초기화와 학습률을 함께 정하는 μP(최대 갱신 매개변수화)다. 배치 정규화·레이어 정규화와 잔차 연결이 널리 쓰이면서 초기화에 덜 민감해졌지만, 수천억 매개변수 모델에서는 초기화가 학습 안정성과 하이퍼파라미터 선택에 여전히 영향을 준다.

왜 0이나 아무 값으로 시작하면 안 되나

모두 0(또는 같은 값): 한 층의 뉴런들이 같은 입력을 받아 같은 출력과 같은 기울기를 내므로 영원히 같은 값으로 갱신된다. 대칭이 깨지지 않아 뉴런 여러 개가 하나와 다를 바 없다. 그래서 가중치는 무작위로 뽑는다. 편향은 0으로 두어도 된다.
너무 작은 값(예: 표준편차 0.01): 층을 지날 때마다 활성값이 줄어 깊은 층에서는 0에 가깝고, 기울기도 사라진다. He 등은 표준편차 0.01로 초기화하면 깊은 모델이 학습을 멈춘다고 지적했다.^[1]
너무 큰 값: tanh·시그모이드는 포화 구간에 들어가 기울기가 0이 되고, ReLU 신경망은 활성값이 폭발해 오버플로가 난다.

분산 보존 원리

선형층 $y = W x$ 에서 입력 개수를 $n_{in}$ (fan-in), 출력 개수를 $n_{out}$ (fan-out)이라 하자. 가중치와 입력이 서로 독립이고 평균이 0이면 $Var (y_{j}) = n_{in} Var (w) Var (x)$ 이다. 순전파에서 분산을 유지하려면 $n_{in} Var (w) = 1$ , 역전파에서 기울기 분산을 유지하려면 같은 논리로 $n_{out} Var (w) = 1$ 이어야 한다. 비선형 함수는 이 값에 보정 계수(gain) ${gain}^{2}$ 를 곱한다.

세이비어(글로럿) 초기화

2010년 글로럿과 벤지오는 시그모이드·tanh 신경망에서 활성값과 기울기의 분산이 층을 지나며 변하는 양상을 분석하고, 두 조건을 절충한 분산을 제안했다.^[2] $Var (w) = \frac{2}{n_{in} + n_{out}}, w \sim U [- \sqrt{\frac{6}{n_{in} + n_{out}}}, \sqrt{\frac{6}{n_{in} + n_{out}}}]$ 균등분포 $U [- a, a]$ 의 분산이 $a^{2} / 3$ 이므로 경계가 $\sqrt{6 / (n_{in} + n_{out})}$ 가 된다. 원점 근처에서 기울기 1인 tanh를 가정한 방법이라 ReLU에는 맞지 않는다. 입력 쪽 조건만 쓴 $Var (w) = 1 / n_{in}$ 은 1998년 르쿤 등의 역전파 실무 지침으로 알려져 르쿤 초기화(LeCun initialization)라 부르며, SELU 신경망에 쓴다. 입력과 출력 개수가 같으면 세이비어 초기화와 같다.

He(카이밍) 초기화

ReLU는 입력의 절반(음수)을 0으로 만들어 출력의 2차 모멘트가 입력의 절반이 된다. 2015년 He 등은 이를 보정해 다음 분산을 유도했다.^[1] $Var (w) = \frac{2}{n_{in}}, w \sim 𝒩 (0, \frac{2}{n_{in}})$ 음수 쪽 기울기가 $a$ 인 Leaky ReLU·PReLU에서는 $2 / ((1 + a^{2}) n_{in})$ 이다. 논문 실험에서 30층(합성곱 27층 + 완전 연결 3층) ReLU 모델은 세이비어 초기화로는 학습이 완전히 멈췄고 He 초기화로는 수렴했다. 22층 모델은 둘 다 수렴했지만 He 초기화가 더 빨리 오차가 줄었다.^[1] 같은 해 나온 ResNet도 이 초기화를 썼다.

활성화 함수별 보정 계수

파이토치의 nn.init.calculate_gain이 돌려주는 값이다. 세이비어·He 초기화의 표준편차에 곱한다.^[3]

비선형 함수	gain
선형, 합성곱, 시그모이드	1
tanh	5/3
ReLU	$\sqrt{2}$
Leaky ReLU(기울기 $a$ )	$\sqrt{2 / (1 + a^{2})}$
SELU	3/4 (자기 정규화가 목적이면 'linear' 사용 권장)

직교 초기화

가중치 행렬을 무작위 직교 행렬( $W^{⊤} W = I$ )로 두는 방법이다. 직교 행렬은 모든 특이값이 1이라 벡터 길이를 바꾸지 않으므로, 선형 신경망에서는 몇 층을 쌓아도 신호와 기울기의 크기가 정확히 보존된다. 2013년 색스(Saxe) 등은 깊은 선형 신경망의 학습 동역학을 정확히 풀어, 무작위 직교 초기화가 비지도 사전 학습처럼 깊이와 무관한 학습 시간을 준다고 보였다.^[4] 같은 가중치를 시간 단계마다 반복해 곱하는 순환 신경망의 은닉 상태 행렬에 특히 쓰인다. 비선형 함수를 쓸 때는 보정 계수를 곱한다(ReLU면 $\sqrt{2}$ ). 구현은 가우스 행렬을 QR 분해해 $Q$ 를 쓴다.

잔차 신경망과 트랜스포머의 초기화

잔차 연결 $x_{l + 1} = x_{l} + F (x_{l})$ 에서는 블록마다 분산이 더해지므로 층이 많을수록 출력 분산이 커진다. 이를 막는 방법들이 있다.

잔차 가지 출력 축소: GPT-2는 잔차 경로에 더해지는 층의 가중치를 $1 / \sqrt{N}$ 배( $N$ 은 잔차 층 수)로 줄여 초기화했다.^[5] 허깅 페이스 Transformers의 GPT-2 구현은 기본 표준편차 0.02(initializer_range)를 쓰고, 어텐션·MLP 출력 투영(c_proj)만 $0.02 / \sqrt{2 N_{layer}}$ 로 줄인다.
잔차 블록 마지막 정규화를 0으로: 페이스북 연구진은 ResNet 각 잔차 블록의 마지막 배치 정규화 배율 $γ$ 를 0으로 초기화해, 학습 초기에 신호가 지름길(항등 연결)로만 흐르게 했다. 큰 미니배치 학습에서 특히 도움이 되었다.^[6]
Fixup: 정규화 층 없이 잔차 신경망을 학습하려고, 각 잔차 가지의 마지막 층을 0으로 두고 나머지 층을 잔차 블록 수에 따라 줄여 초기화한다. 저자들은 1만 층짜리 잔차 신경망도 정규화를 쓴 것만큼 안정적으로 학습된다고 보고했다.^[7]

μP: 너비를 바꿔도 하이퍼파라미터가 옮겨지는 초기화

표준 매개변수화(standard parametrization, SP)에서는 모델 너비를 키우면 최적 학습률이 달라진다. 그래서 큰 모델의 하이퍼파라미터를 작은 모델에서 찾을 수 없었다. 2022년 마이크로소프트와 OpenAI 연구진은 μP(Maximal Update Parametrization, 최대 갱신 매개변수화)에서는 너비가 변해도 최적 하이퍼파라미터가 거의 그대로라는 것을 보이고, 작은 대리 모델에서 하이퍼파라미터를 찾아 큰 모델에 바로 옮기는 μTransfer를 제안했다.^[8]

μP는 초기화 분산과 층별 학습률을 너비(fan-in, fan-out)에 따라 다르게 정한다. 논문 표 3의 Adam 기준 규칙은 다음과 같다.

구분	입력층 가중치·모든 편향	은닉층 가중치	출력층 가중치
초기화 분산(μP)	$1 / f a n_i n$	$1 / f a n_i n$	$1 / {f a n_i n}^{2}$ (SP는 $1 / f a n_i n$ )
Adam 학습률(μP)	1 (너비와 무관)	$1 / f a n_i n$ (SP는 1)	$1 / f a n_i n$ (SP는 1)

트랜스포머에서는 여기에 더해 어텐션 로짓을 $q^{⊤} k / \sqrt{d}$ 대신 $q^{⊤} k / d$ 로 나눈다. 연구진은 1300만 매개변수 모델에서 옮긴 하이퍼파라미터로 BERT-large(3억 5000만)의 공개 성능을 넘었고, 4000만 매개변수 모델에서 옮겨 67억 매개변수 GPT-3의 공개 성능을 넘었으며 이때 조정 비용은 사전 학습 비용의 7%였다고 보고했다. 구현은 pip install mup로 설치하는 PyTorch 라이브러리로 공개했다.^[8]

직접 확인하기

너비 512인 선형층 50개를 쌓고 초기화만 바꿔 마지막 층 활성값의 표준편차를 재 보는 실험이다(편향 없음, 입력 1024개).

import torch
import torch.nn as nn

def final_std(init, act, depth=50, width=512):
    torch.manual_seed(0)
    x = torch.randn(1024, width)
    for _ in range(depth):
        layer = nn.Linear(width, width, bias=False)
        init(layer.weight)
        x = act(layer(x))
    return x.std().item()

inits = {
    "N(0, 0.01)":     lambda w: nn.init.normal_(w, std=0.01),
    "N(0, 1)":        lambda w: nn.init.normal_(w, std=1.0),
    "Xavier":         lambda w: nn.init.xavier_normal_(w),
    "Xavier(5/3)":    lambda w: nn.init.xavier_normal_(w, gain=nn.init.calculate_gain("tanh")),
    "He":             lambda w: nn.init.kaiming_normal_(w, nonlinearity="relu"),
    "Orthogonal":     lambda w: nn.init.orthogonal_(w),
}
with torch.no_grad():
    for name, init in inits.items():
        print(name, final_std(init, torch.tanh), final_std(init, torch.relu))

PyTorch 2.5에서 돌린 결과는 다음과 같다.

초기화	tanh 50층 뒤 표준편차	ReLU 50층 뒤 표준편차	해석
$𝒩 (0, 0.0 1^{2})$	약 $5 \times 1 0^{- 33}$	약 $2 \times 1 0^{- 40}$	신호 소멸
$𝒩 (0, 1)$	0.98	NaN	tanh는 ±1에 포화(기울기 0), ReLU는 오버플로
세이비어	0.10	약 $3 \times 1 0^{- 8}$	tanh에서 서서히 감소, ReLU엔 부적합
세이비어 × 5/3	0.65	약 4200	tanh엔 적절, ReLU에선 폭발
He	0.56	1.14	ReLU에서 크기 유지
직교(gain 1)	0.10	약 $3 \times 1 0^{- 8}$	ReLU면 gain $\sqrt{2}$ 가 필요

활성화 함수에 맞는 보정 계수를 쓰는 것이 분포 종류(정규·균등·직교)보다 중요하다는 것을 볼 수 있다.

PyTorch 기본값과 사용법

PyTorch의 nn.Linear와 nn.Conv2d는 따로 지정하지 않으면 kaiming_uniform_(weight, a=math.sqrt(5))로 초기화한다. 이는 $U [- 1 / \sqrt{n_{in}}, 1 / \sqrt{n_{in}}]$ 와 같아 분산이 $1 / (3 n_{in})$ 로 He 초기화보다 작다. 편향도 같은 범위의 균등분포다.^[9] 다른 초기화를 쓰려면 직접 덮어쓴다.

import torch.nn as nn

def init_weights(m):
    if isinstance(m, (nn.Linear, nn.Conv2d)):
        nn.init.kaiming_normal_(m.weight, mode="fan_in", nonlinearity="relu")
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Embedding):
        nn.init.normal_(m.weight, std=0.02)

model.apply(init_weights)

mode="fan_in"은 순전파 분산을, "fan_out"은 역전파 분산을 보존한다.^[3] 이 밖에 xavier_uniform_, xavier_normal_, orthogonal_, trunc_normal_(잘린 정규분포, 비전 트랜스포머 계열에서 사용) 등이 있다.

정리: 무엇을 고를까

상황	권장 초기화
tanh·시그모이드 은닉층	세이비어(tanh는 gain 5/3)
ReLU·GELU 계열 CNN·MLP	He(fan_in), 잔차 블록 끝 정규화 $γ = 0$
RNN 은닉 상태 행렬	직교
트랜스포머	작은 정규분포(예: 표준편차 0.02)와 잔차 출력 투영 축소, 또는 μP
너비를 키워 가며 하이퍼파라미터 탐색	μP
사전 학습 모델이 있을 때	사전 학습 가중치로 시작(전이 학습), 새로 붙인 층만 위 규칙

같이 보기

각주

↑ ^1.0 ^1.1 ^1.2 Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun, 〈Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification〉, ICCV 2015. arXiv:1502.01852
↑ Xavier Glorot, Yoshua Bengio, 〈Understanding the difficulty of training deep feedforward neural networks〉, AISTATS 2010, PMLR 9:249-256. proceedings.mlr.press
↑ ^3.0 ^3.1 PyTorch, 〈torch.nn.init〉, 2.14 문서. docs.pytorch.org
↑ Andrew M. Saxe, James L. McClelland, Surya Ganguli, 〈Exact solutions to the nonlinear dynamics of learning in deep linear neural networks〉, ICLR 2014. arXiv:1312.6120
↑ Alec Radford 외, 〈Language Models are Unsupervised Multitask Learners〉, OpenAI, 2019. cdn.openai.com
↑ Priya Goyal 외, 〈Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour〉, 2017. arXiv:1706.02677
↑ Hongyi Zhang, Yann N. Dauphin, Tengyu Ma, 〈Fixup Initialization: Residual Learning Without Normalization〉, ICLR 2019. arXiv:1901.09321
↑ ^8.0 ^8.1 Greg Yang 외, 〈Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer〉, 2022. arXiv:2203.03466
↑ PyTorch 소스 코드, 〈torch/nn/modules/linear.py〉의 reset_parameters. github.com

참고 문헌

[he15-1] 1.0 ^1.1 ^1.2 Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun, 〈Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification〉, ICCV 2015. arXiv:1502.01852

[glorot10-2] Xavier Glorot, Yoshua Bengio, 〈Understanding the difficulty of training deep feedforward neural networks〉, AISTATS 2010, PMLR 9:249-256. proceedings.mlr.press

[ptinit-3] 3.0 ^3.1 PyTorch, 〈torch.nn.init〉, 2.14 문서. docs.pytorch.org

[saxe13-4] Andrew M. Saxe, James L. McClelland, Surya Ganguli, 〈Exact solutions to the nonlinear dynamics of learning in deep linear neural networks〉, ICLR 2014. arXiv:1312.6120

[gpt2-5] Alec Radford 외, 〈Language Models are Unsupervised Multitask Learners〉, OpenAI, 2019. cdn.openai.com

[goyal17-6] Priya Goyal 외, 〈Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour〉, 2017. arXiv:1706.02677

[fixup-7] Hongyi Zhang, Yann N. Dauphin, Tengyu Ma, 〈Fixup Initialization: Residual Learning Without Normalization〉, ICLR 2019. arXiv:1901.09321

[mup-8] 8.0 ^8.1 Greg Yang 외, 〈Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer〉, 2022. arXiv:2203.03466

[ptlinear-9] PyTorch 소스 코드, 〈torch/nn/modules/linear.py〉의 reset_parameters. github.com

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]

[9]