[Pytorch] 가중치 초기화

이번 포스팅에서는 pytorch에서 Tensor의 가중치를 초기화하는 방법을 정리했다. 신경망 모델의 성능은 가중치 초기 값에 민감할 수 있기 때문에, 가중치 초기화 과정은 모델 훈련에서 중요한 역할을 수행한다. 기본적으로 torch.nn.init module을 통해 다양한 가중치 초기화 방법을 사용할 수 있는데, 아래에 간단한 예시와 함께 살펴보자.

 

# 가중치 초기화 정의(Xavier)
import torch
import torch.nn as nn

def init_weights(m):
    if type(m) == nn.Linear:
        torch.nn.init.xavier_uniform_(m.weight)
        m.bias.data.fill_(0.01)
        
model = MyModel()
model.apply(init_weights)

# 함수 사용(주어진 텐서를 평균이 0, 표준편차가 1인 정규 분포로 초기화)

linear = nn.Linear(2, 2)
nn.init.normal_(linear.weight, mean=0, std=1)

 

신경망에서 원하는 Layer에 원하는 방식으로 가중치를 초기화하거나, 주어진 텐서를 직접 초기화 할 수도 있다. 하지만, 가중치를 아무 값이나 사용하면 학습 과정에서 다양한 문제가 발생할 수 있다. 어떤 문제가 생길 수 있는지 case 별로 알아보자.

 

Gradient Vanishing 문제 발생(Sigmoid)

 

• 가중치 초기값이 0이거나 모두 동일한 경우 : 가중치 초기값을 0이가 모두 동일한 값으로 초기화하면, 신경망의 모든 뉴런이 동일한 출력을 생성한다. 이렇게 되면 역전파 과정에서 모든 가중치에 대한 Gradient가 동일하게 계산되므로, 가중치의 업데이트가 동일하게 이뤄진다. 이는 실질적으로 신경망이 깊이에 상관없이 하나의 뉴런을 가진 것과 같은 효과를 만들어낸다. 즉, 뉴런의 개수를 전혀 활용하지 못하는 문제가 발생한다.
  
  
• 가중치 초기값이 너무 큰 경우 : Activation function에 따라 다르지만, sigmoid를 사용하는 경우 가중치를 너무 큰 값으로 초기화 하면 0과 1로 수렴하기 때문에 Gradient Vanishing(기울기 소실) 문제가 발생한다. ReLU를 사용하는 경우, 절대값이 크면 음수일 때는 Dead ReLU 문제가 발생하고, 양수일 때는 Gradient Exploding(기울기 폭주) 문제가 발생한다. 이는 신경망이 깊어지는 경우, 더욱 심각한 문제를 유발한다.
  
  
• Dead ReLU : 순전파 과정에서 어떤 뉴런에서든 음수값이 들어가는 경우 역전파 시에 0이라는 실제 값이 가중치에 곱해지면서 해당 노드가 죽어버리는 현상
  
  

따라서, 가중치는 일반적으로 작은 난수로 초기화하는 것이 좋다. 이렇게 하면, 개별 뉴런이 다양한 출력을 생성하여 원활하게 학습이 진행된다. 그러나, 이 방법 역시 DNN(심층 신경망)에서는 Gradient Vanishing 문제가 발생한다. 이러한 문제를 해결하기 위해, Xavier 초기화나 He 초기화와 같은 방법이 탄생했다. 해당 방법은 가중치의 초기값을 신경망의 Layer 크기에 따라 조절하여, DNN에서도 Gradient Vanishing 혹은 Gradient Exploding 문제를 완화하는데 도움을 준다. 

---

Xavier initialization

 

작은 난수(왼쪽) vs xavier 초기값(오른쪽)

 

Xaiver 초기화의 핵심은 신경망의 각 층에서 출력의 분산이 입력의 분산과 같아야하며, 역전파 과정에서 Gradient의 분산도 입력의 분산과 같아야 한다는 것 이다. 이를 통해 각 층의 출력이 고르게 분포하도록 하여 Gradient Vanishing 문제를 완화하고, 신경망의 학습을 안정화시키는 것을 목표로 한다. 위의 그림에서 tanh 함수를 사용한 결과를 비교해보면, 작은 난수로 초기값을 설정했을 때 보다 xavier 초기값을 사용한 가중치가 더 넓게 분포되어 있음을 확인할 수 있다.

Xaiver 정규분포( $n_{in}$ : 이전 layer의 노드 수, $ n_{out}$ : 현재 layer의 노드 수)

 

Xavier 초기화는 가중치의 초기 값을 설정하는 방법으로서, 평균이 0이고 분산이 $\sqrt{\frac{2}{n_{in}+n_{out}}}$인 정규분포로 가중치를 초기화한다. 이런 방식을 취하면, 각 층을 통과하는 출력 값들의 분산이 입력 값들의 분산과 비슷한 수준으로 유지되는 효과를 얻을 수 있다.

 

Xaiver 균등분포

균등 분포를 사용해 가중치를 초기화할 수도 있다. 이 경우 가중치는 입력 노드와 출력 노드의 개수에 따른 범위에서 추출된다. 이 방법은 활성화 함수가 선형(linear)이거나 대칭적인 S 형태(Sigmoid, tanh)인 경우에 특히 유용하다. 하지만, ReLU 함수를 사용하는 경우에는 다른 방법을 활용해야 한다.

---

He initialization

 

Xavier 초기값(왼쪽) vs He 초기값(오른쪽)

 

Xavier 초기화는 무작위 초기화 방법에 비해 더욱 우수한 성능을 보이지만, ReLU 활성화 함수를 사용할 때에는 그래디언트 소실 문제를 완벽하게 해결하지 못할 수 있다. 왼쪽의 그림에서 레이어가 깊어질수록, Xavier 초기화를 사용한 모델의 출력 값은 0으로 수렴하는 경향을 보인다.

 

 

정규 분포, 균등 분포

이 문제를 해결하기 위해 He 초기화 방법이 도입되었다. He 초기화는 ReLU 활성화 함수와 그 변형체들(예: Leaky ReLU, Parametric ReLU 등)에 적합하도록 설계되었다. Xavier 초기화와 마찬가지로, He 초기화는 각 레이어의 출력 분산이 입력 분산과 같아지도록 하지만, 가중치 출력의 분산은 $\sqrt{\frac{2}{n_{in}}}$($n_{in}$은 입력 뉴런의 개수)으로 설정한다는 차이가 있다. 이는 ReLU 함수가 음수 입력에 대해 0을 출력하기 때문에, 출력의 분산이 입력 분산의 절반에 가까워지는 특성을 보정하기 위함이다. 또한, 균등 분포를 사용하는 경우에는 Xavier 초기화와 유사하지만, 입력 노드의 개수만을 활용한다는 차이점이 있다.

 

 

가중치의 분포가 넓다면, 모델은 입력 데이터의 복잡한 표현을 더욱 잘 학습할 수 있다. 이는 넓은 분포가 가중치 간의 다양성을 증가시키므로, 다양한 특성을 더 잘 포착할 수 있기 때문이다. 그러나 가중치가 너무 크면 과적합(Overfitting)의 위험이 있으므로, 활성화 함수와 모델 구조에 따라 적절한 가중치 초기화 방법을 선택하는 것이 중요하다.

 

 

 

사진 출처

https://yngie-c.github.io/deep%20learning/2020/03/17/parameter_init/

https://www.kaggle.com/code/kalelpark/eng-kor-weight-initialization