[PYTORCH] 레이어 가중치 초기화 방법 5가지와 Xavier vs He 차이 해결책 7가지

가중치 초기화 (Weight Initialization)

 

딥러닝 모델 개발에 있어 많은 개발자들이 아키텍처 설계와 하이퍼파라미터 튜닝에는 수많은 시간을 쏟지만, 정작 학습의 성패를 가르는 첫 단추인 가중치 초기화(Weight Initialization)는 프레임워크의 기본값에 맡겨두곤 합니다. 이는 심각한 실수입니다. 가중치 초기화는 단순한 랜덤 값 채우기가 아닙니다. 그것은 신경망 내부에서 신호(Signal)가 어떻게 전파될지를 결정하며, 학습 속도를 비약적으로 향상시키거나, 반대로 기울기 소실(Vanishing Gradient) 또는 기울기 폭주(Exploding Gradient)의 늪에 빠뜨려 학습 자체를 불가능하게 만들 수도 있는 결정적인 단계입니다. 파이토치(PyTorch)는 강력한 자동 미분 엔진과 유연한 인터페이스를 제공하지만, 최적의 가중치 초기화를 선택하는 것은 여전히 엔지니어의 몫입니다. 본 가이드에서는 시니어 AI 엔지니어의 관점에서 가중치 초기화의 수학적 원리를 이해하고, 파이토치에서 제공하는 5가지 주요 초기화 방법, 그리고 특히 현대 딥러닝에서 가장 중요한 Xavier와 He 초기화의 결정적 차이를 분석합니다. 또한, 실무에서 마주하는 학습 불안정 문제를 즉각적으로 해결할 수 있는 7가지 실전 코드 예제를 제시합니다.

---

1. 가중치 초기화, 왜 해야 하는가? (독창적인 가치 분석)

많은 이들이 가중치 초기화를 단순한 '시작점 설정'으로 오해합니다. 하지만 초기화의 본질은 '신호 전파의 엔트로피 관리'에 있습니다.

1.1. 0으로 초기화했을 때의 재앙 (Symmetry Breaking)

가장 직관적이지만 가장 위험한 방법은 모든 가중치를 0으로 초기화하는 것입니다. 이 경우 모든 뉴런이 동일한 출력값을 가지게 되고, 역전파(Backpropagation) 과정에서 모든 뉴런이 동일한 그래디언트를 업데이트하게 됩니다. 즉, 수만 개의 뉴런이 있어도 실제로는 단 하나의 뉴런만 작동하는 것과 다름없게 됩니다. 이를 대칭성(Symmetry) 문제라고 하며, 초기화는 이 대칭성을 깨는 과정이어야 합니다.

1.2. 신호의 보존: 기울기 소실과 폭주

신경망은 수많은 레이어로 이루어져 있습니다. 입력 신호는 레이어를 거치며 가중치 행렬과 곱해집니다. 만약 가중치가 너무 작게 초기화되면, 신호는 레이어를 거칠 때마다 점점 작아져 마지막 레이어에서는 거의 0에 수렴하게 됩니다. 이는 기울기 소실 문제를 유발하여 학습이 진행되지 않습니다. 반대로 가중치가 너무 크게 초기화되면, 신호는 기하급수적으로 증폭되어 기울기 폭주 문제를 일으키고 모델이 발산해버립니다.

시니어 엔지니어의 핵심 인사이트: 최적의 초기화는 순전파(Forward) 시에는 각 레이어의 출력 분산이 입력 분산과 같게 유지되고, 역전파(Backward) 시에는 각 레이어의 기울기 분산이 레이어를 거쳐도 유지되도록 하는 것입니다. 즉, 신호가 증폭되거나 축소되지 않고 보존되는 지점을 찾는 것이 기술입니다.

---

2. Xavier vs He 초기화: 수학적 원리와 결정적 차이 해결

현대 딥러닝에서 가장 흔히 사용되는 두 가지 표준 초기화 방법인 Xavier(Glorot)와 He 초기화는 활성화 함수(Activation Function)와의 관계에 따라 그 운명이 갈립니다. 이 두 방법의 차이를 명확히 이해하고 적용하는 것이 학습 정체를 해결하는 열쇠입니다.

비교 항목	Xavier (Glorot) 초기화	He 초기화
제안 시기 및 인물	2010년, Xavier Glorot & Yoshua Bengio	2015년, Kaiming He et al.
주요 대상 활성화 함수	S-자형 함수 (Sigmoid, Tanh)	Rectified Linear Unit (ReLU) 및 그 변형 (LeakyReLU)
수학적 근거	선형 활성화 함수를 가정하고 입력과 출력의 분산을 같게 유지.	ReLU가 출력의 절반을 0으로 만든다는 점을 고려하여 분산을 2배로 조정.
가중치 분포 (정규분포 기준)	$$Var(W) = \frac{2}{n_{in} + n_{out}}$$	$$Var(W) = \frac{2}{n_{in}}$$
실무적 특성	층이 깊지 않거나 Tanh를 쓰는 과거 아키텍처에 적합. ReLU와 함께 쓰면 기울기 소실 발생 가능.	현대 깊은 신경망(ReLU 기반)의 표준. 학습을 빠르고 안정적으로 안정화시킴.

2.1. Xavier 초기화의 차이 해결

Xavier 초기화는 활성화 함수가 선형적(Linear)이라고 가정합니다. Sigmoid나 Tanh의 중심 부근은 비교적 선형적이기 때문에 이 함수들과 잘 작동합니다. 하지만 ReLU와 함께 사용하면 문제가 발생합니다. ReLU는 음수 입력을 0으로 차단하므로 출력의 분산을 줄입니다. Xavier는 이 분산 축소를 고려하지 않기 때문에, 층이 깊어질수록 신호가 점점 약해지는 기울기 소실 문제가 여전히 발생할 수 있습니다.</ p>

2.2. He 초기화의 해결책

Kaiming He는 ReLU를 사용할 때 출력 분산이 반으로 줄어든다는 점을 수학적으로 증명했습니다. 이를 해결하기 위해 가중치의 분산을 Xavier보다 2배 더 크게 설정합니다. 즉, $Var(W) = 2/n_{in}$ 으로 설정하여, ReLU에 의해 반토막 날 분산을 미리 2배로 키워 신호 세기를 유지하는 것입니다. 이 간단한 조정이 깊은 ResNet 아키텍처 학습을 가능하게 만든 핵심 비결 중 하나입니다.

---

3. PyTorch에서 가중치를 초기화하는 5가지 방법

파이토치는 torch.nn.init 모듈을 통해 다양한 초기화 함수를 제공합니다. 이 함수들을 사용하여 모델의 레이어 가중치에 직접 접근하여 초기화할 수 있습니다.

3.1. 기본값에 의존 (비추천)

파이토치의 대부분 레이어는 uniform_ 초기화를 기본값으로 사용합니다. 이는 Xavier나 He 초기화가 제안되기 전의 구식 방법으로, 깊은 신경망에서는 성능이 담보되지 않습니다.

3.2. Uniform 초기화 (균등 분포)

지정된 범위 $[a, b]$ 내에서 가중치를 균등하게 선택합니다. torch.nn.init.uniform_ 함수를 사용합니다.

3.3. Normal 초기화 (정규 분포)

평균 $\mu$와 표준편차 $\sigma$를 갖는 정규분포에서 가중치를 선택합니다. 대칭성을 깨는 데 효과적입니다. torch.nn.init.normal_ 함수를 사용합니다.

3.4. Xavier (Glorot) 초기화 해결

앞서 설명한대로 Tanh와 잘 어울립니다. 파이토치는 xavier_uniform_와 xavier_normal_ 두 가지 버전을 제공합니다.

3.5. He 초기화 해결

ReLU 기반의 현대 아키텍처 표준입니다. 파이토치에서는 kaiming_uniform_와 kaiming_normal_ 라는 이름으로 제공됩니다 (Kaiming He의 이름을 땀).

---

4. 실무자를 위한 가중치 초기화 예제 7가지 해결책 (Sample Examples)

실제 프레임워크 설계 및 학습 불안정 문제를 해결하기 위해 즉시 적용 가능한 7가지 파이토치 코드 예제입니다.

Example 1: 특정 레이어만 수동 초기화 해결 방법

모델 정의 후 인스턴스의 레이어 속성에 직접 접근하여 초기화 함수를 적용합니다.

import torch
import torch.nn as nn
import torch.nn.init as init

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(256, 10)

model = MyModel()

# 1. fc1 레이어에 He Normal 초기화 적용 (ReLU 쌍)
init.kaiming_normal_(model.fc1.weight, mode='fan_in', nonlinearity='relu')
# 편향(Bias)은 일반적으로 0으로 초기화
init.constant_(model.fc1.bias, 0)

# 2. fc2 레이어에 Xavier Uniform 초기화 적용 (출력층 예시)
init.xavier_uniform_(model.fc2.weight)
init.constant_(model.fc2.bias, 0)

print("Manual initialization complete.")

Example 2: 모델 내부에서 초기화 정의 해결 (추천 방법 1)

__init__ 메서드 내에서 레이어 생성 직후 초기화를 수행하여 코드를 캡슐화합니다.

import torch.nn as nn
import torch.nn.init as init

class RecommendedModel(nn.Module):
    def __init__(self):
        super(RecommendedModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.relu1 = nn.ReLU()
        self.fc = nn.Linear(16 * 26 * 26, 10) # 예시 크기

        # 생성자 내에서 초기화 수행
        self._initialize_weights()

    def _initialize_weights(self):
        # 모듈 순회하며 초기화
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                # Conv 레이어는 He Normal
                init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                # Linear 레이어도 He Normal (ReLU 전단계)
                init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
                init.constant_(m.bias, 0)

model = RecommendedModel()
print("Internal initialization defined within the model.")

Example 3: apply()를 이용한 모델 외부 전체 초기화 해결 (추천 방법 2)

외부에서 초기화 함수를 정의하고 모델 인스턴스의 apply() 메서드를 사용해 모든 하위 모듈에 재귀적으로 적용합니다. 코드가 깔끔해집니다.

import torch.nn as nn
import torch.nn.init as init

# 초기화 로직을 담은 외부 함수 정의
def weight_init_function(m):
    if isinstance(m, nn.Linear):
        # Linear 레이어에 Xavier Uniform 적용
        init.xavier_uniform_(m.weight)
        if m.bias is not None:
            init.constant_(m.bias, 0.01) # 가끔 작은 상수로 초기화하기도 함
    elif isinstance(m, nn.Conv2d):
        # Conv 레이어에 He Normal 적용
        init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')

model = nn.Sequential(
    nn.Conv2d(1, 10, 5),
    nn.ReLU(),
    nn.Linear(10 * 12 * 12, 50),
    nn.ReLU(),
    nn.Linear(50, 10)
)

# apply 메서드로 모델 전체에 적용
model.apply(weight_init_function)
print("Initialization applied to the entire model instance.")

Example 4: Sigmoid/Tanh 기반 구식 모델의 기울기 소실 해결

과거 아키텍처를 다룰 때 Sigmoid를 쓰면서 He 초기화를 쓰면 학습이 불안정할 수 있습니다. Xavier가 정답입니다.

import torch.nn as nn
import torch.nn.init as init

# Sigmoid 활성화를 쓰는 모델
legacy_model = nn.Sequential(
    nn.Linear(100, 50),
    nn.Sigmoid(),
    nn.Linear(50, 10),
    nn.Sigmoid()
)

# Xavier Normal 초기화로 기울기 소실 억제
for m in legacy_model.modules():
    if isinstance(m, nn.Linear):
        init.xavier_normal_(m.weight, gain=init.calculate_gain('sigmoid'))
        init.constant_(m.bias, 0)

print("Xavier initialization used for Sigmoid activation.")

Example 5: Deep ReLU 신경망의 학습 정체 해결

층이 깊은 ReLU 기반 신경망에서 학습이 진행되지 않는다면 He 초기화가 누락되지 않았는지 점검하십시오. 기본 초기화(Uniform)를 He로 교체하는 것만으로 문제가 해결됩니다.

import torch.nn as nn
import torch.nn.init as init

# 깊은 ReLU 모델
deep_model = nn.Sequential(
    nn.Linear(100, 100), nn.ReLU(),
    nn.Linear(100, 100), nn.ReLU(),
    nn.Linear(100, 100), nn.ReLU(),
    nn.Linear(100, 10), nn.ReLU()
)

# He Uniform 초기화로 분산 보존
def he_init(m):
    if isinstance(m, nn.Linear):
        # nonlinearity='relu' 설정이 중요
        init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='relu')
        init.constant_(m.bias, 0)

deep_model.apply(he_init)
print("He initialization used for deep ReLU network.")

Example 6: calculate_gain()을 이용한 정밀한 초기화 해결

import torch.nn as nn
import torch.nn.init as init

conv = nn.Conv2d(3, 16, 3)
relu = nn.ReLU()

# 활성화 함수에 따른 이득(Gain) 계산
# Gain은 활성화 함수가 출력 분산을 어떻게 바꾸는지 보정하는 상수
gain = init.calculate_gain('relu') # ReLU의 gain은 sqrt(2)

# Xavier 초기화에 Gain을 곱해 He 초기화와 유사한 효과를 내거나 수동 조정
init.xavier_normal_(conv.weight, gain=gain) 

print(f"Gain calculated for ReLU: {gain}")

Example 7: 사전 학습된 모델(Pre-trained)의 부분 초기화 해결 방법

전이 학습(Transfer Learning) 시 사전 학습된 가중치를 보존하면서, 새로 추가한 마지막 레이어만 초기화해야 합니다.

import torchvision.models as models
import torch.nn as nn
import torch.nn.init as init

# 사전 학습된 ResNet18 로드
resnet = models.resnet18(pretrained=True)

# 1. 기존 가중치 고정 (Freeze)
for param in resnet.parameters():
    param.requires_grad = False

# 2. 마지막 분류층 교체 (새로 생성된 레이어는 requires_grad=True)
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 2) # 클래스 수를 2로 변경

# 3. 새로 추가된 resnet.fc 레이어만 He 초기화
init.kaiming_normal_(resnet.fc.weight, mode='fan_in', nonlinearity='relu')
init.constant_(resnet.fc.bias, 0)

print("Pre-trained model loaded, other layers frozen, and new FC layer initialized.")

---

5. 시니어 엔지니어의 핵심 인사이트: 가중치 초기화와 Batch Normalization

현대 딥러닝 아키텍처는 가중치 초기화에 대한 의존도를 낮추기 위해 배치 정규화(Batch Normalization, BN) 레이어를 적극적으로 도입합니다. BN은 각 레이어의 출력을 강제로 평균 0, 분산 1로 정규화하므로, 가중치 초기화가 다소 잘못되더라도 신호가 폭주하거나 소실되는 것을 강력하게 막아줍니다.

실무적인 권장 사항: BN을 사용하더라도 He 초기화를 기본으로 적용하는 것이 학습 속도를 가장 빠르고 안정적으로 만듭니다. 또한, 사전 학습된 모델을 쓸 때는 이미 최적화된 초기화가 완료된 상태이므로, Example 7처럼 새로 추가하는 레이어에 대해서만 정밀한 초기화를 수행하는 것이 정석입니다.

---

6. 결론 및 요약

가중치 초기화는 파이토치 모델 학습의 성공을 위한 가장 가성비 높은 하이퍼파라미터 튜닝입니다. 기본값에 의존하는 대신, 활성화 함수에 맞춰 최적의 초기화 방법을 선택하십시오.

• Dying ReLU를 피하라: Sigmoid/Tanh 기반의 구식 모델은 Xavier 초기화를 사용하라.
• 깊은 신경망을 학습시켜라: ReLU 기반의 현대 아키텍처는 He (Kaiming) 초기화가 표준이다.
• 캡슐화하라: 모델 내부 _initialize_weights 메서드나 외부 apply() 를 사용하여 초기화 로직을 모델과 통합하라.
• Batch Norm과 함께 사용하라: BN은 초기화 실수에 대한 안전장치이지만, 최적의 초기화와 함께 쓰면 학습 속도를 극대화한다.

---

내용 출처 및 참고 문헌 (Sources)

• PyTorch Official Documentation: torch.nn.init.
• Glorot, X., & Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. Aistats.
• He, K., Zhang, X., Ren, S., & Sun, J. (2015). Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. ICCV.