MLP 모델의 Weight Initialization에 대한 이해

※ 이 글은 사이토 고키 저, 『밑바닥부터 시작하는 딥러닝』의 6.2 단원을 더 깊이 해설한 글입니다.

 

Weight Initialization은 몹시 중요하다.

 

인공지능 모델을 학습시키기 위해서는 패러미터, 특히 가중치(weight)들의 초깃값을 설정해 주어야 한다. 초깃값 설정의 중요성이야 굳이 길게 언급할 필요가 없을 것이다. 애초에 초깃값을 설정해 주지 않고 모델을 학습시키는 것은 불가능할뿐더러, 패러미터의 초깃값을 어떻게 설정하느냐에 따라 모델의 학습은 천차만별로 달라질 수 있다. 비유적으로 표현하자면 공수부대가 깊은 산속에 낙하하여 목적지를 찾아갈 때, 어느 지점에 착륙할 것인가를 결정하는 것과 같다. 다음 그림에서, 아주 약간만 초깃값이 달라져도 전혀 다른 좌표로 최적화되는 사례를 확인할 수 있다.

 

Cousera, Machine Learning - Andrew Ng

 

 그렇다면 어떻게 초깃값을 설정해야 할까?

 

어려운 문제를 단순한 해결책에서부터 시작하는 것은 좋은 자세이므로, 가장 먼저 떠오르는 방법은 그냥 모든 매개변수를 $0$으로 초기화하는 것이다. 때때로 이런 단순한 방법은 기가 막히게 잘 동작하여 우리를 짜릿하게 하지만, 대체로는 잘 동작하지 않으며, 안타깝게도 weight initialization의 문제 역시 여기에 속한다.

 

왜 모든 값을 $0$으로 초기화하는 방법은 잘 동작하지 않는 것일까?  만약 모든 weight 값이 $0$이라면 어떤 값이 입력되든 모두 $0$을 곱하여 zero로 만든 뒤에, bias를 더하여 다음 레이어에 전달할 것이다. 그렇다면 결국 다음 레이어에 있는 노드들에는 입력값과는 무관하게 모두 똑같은 값(bias)이 전달된다는 것인데, 이는 결국 다층 레이어 학습 모델이 갖는 의미를 무색하게 만들 것이다. 스스로를 업데이트하며 손실 함수를 최소화시키는 절묘한 가중치 값을 찾아내는 것이 MLP 모델 학습의 묘미인데, 모든 레이어에서 모두 똑같은 값이 전달되어서야 레이어가 하나인 것과 다름없기 때문이다.

 

이유는 그 뿐만이 아니다. 모델이 역전파를 통해 다시 거꾸로 진행하며 weight 값들을 갱신할 때, 모든 weight 값들이 똑같다면 갱신 역시 똑같은 값으로 이루어질 것이다. 이는 마찬가지로 위에서 언급한 'MLP 모델의 묘미'를 달성하기 어렵다는 것을 의미한다. 결국 패러미터를 $\boldsymbol{0}$으로 초기화한다는 것은 앞에서 보나 뒤에서 보나 문제점이 많다고 할 수 있다.

 

이를 해결하기 위해 고심하다 보면 우리의 목표를 다음과 같이 구체화할 수 있다. 우선 초기화된 패러미터는 $\boldsymbol{0}$이어서는 안 될 것이며, 특히 모든 weight 값들이 똑같지 않도록 패러미터를 초기화해야 한다. 그리고 이렇게 초기화된 가중치로 학습된 모델은 다양한 값들로 이루어진 출력을 반환해야 한다. 다양한 값으로 이루어진 출력은 곧 다양한 입력에 대해 다양한 케이스로 예측이 가능하다는 것을 의미하기 때문이다. 한 마디로 모델의 표현력이 높다는 것이다.

 

이제 랜덤으로 weight 값을 초기화해보자.

 

weight 값이 다양해야 한다면 랜덤하게 weight 값을 초기화할 수 있다. 여기서는 평균이 $0$, 표준편차가 $1$인 표준 정규분포에 따라 임의의 값을 추출해 사용한다고 하자. 그런데 혹시 어떤 이는 다음과 같은 궁금증이 일 수도 있다(적어도, 나는 그랬다). 랜덤으로 값을 넣는 것은 그렇다 치더라도 무작위로 추출하는 방법은 여러 가지가 있을 것인데 왜 하필이면 정규분포를 따라 추출된 값을 넣는다는 것인가? (참고로 균등분포를 사용하는 초기화 방법도 존재하나 이 포스팅에서는 다루지 않을 것이다.)

 

정규분포를 사용하는 이유에 대한 나의 추측은 다음과 같다. 왜냐하면 첫째로, 정규분포가 자연계에서 가장 일반적인 분포이므로 우리가 입력할 값들이 대체로 정규분포를 따를 것이라 예상할 수 있기 때문이다. 그리고 둘째로, 지금부터 우리가 다룰 활성화 함수와 출력 함수를 logistic sigmoid 함수(이하 sigmoid 함수)로 가정할 것이기 때문이다. 지난 글에서 언급하였다시피 sigmoid 함수의 경우 대칭적인 모양을 갖는다. 따라서 표준 정규분포가 잘 들어맞지 않는다고 하더라도 표준편차를 적절히 조절하여 sigmoid 함수 및 입력 데이터의 크기에 정규분포를 절묘히 맞춰 준다면, 뭔가 유의미한 성과를 달성할 것이라 짐작할 수 있다. (대칭적인 정규분포 그래프와 sigmoid 함수의 그래프를 적절하게 포개어 본다고 생각해 보라.)

 

자, 그러면 정규분포 초기화가 과연 잘 동작하는지 실험을 통해 알아보자. 앞서 우선 말했다시피 평균이 $0$, 표준편차가 $1$인 표준 정규분포부터 실험해 볼 것이다. 사용되는 코드는 다음과 같으며, bias는 고려하지 않았다.

 

import numpy as np
import matplotlib.pyplot as plt

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

x = np.random.randn(1000, 100)
node_num = 100
hidden_layer_size = 5
activations = {}

for i in range(hidden_layer_size):
    if i != 0:
        x = activations[i-1]
    w = np.random.randn(node_num, node_num) * 1 # 표준편차 1을 가시적으로 드러내었다.
    a = np.dot(x, w)
    z = sigmoid(a)
    activations[i] = z

for i, a in activations.items():
    plt.subplot(1, len(activations), i+1)
    plt.title(str(i+1) + "-layer")
    plt.hist(a.flatten(), 30, range=(0, 1))

plt.show()

 

위 코드에 대해 설명하자면, $1000 \times 100$의 2차원 입력값 $\boldsymbol{X}$가 주어진다. 그리고 $\boldsymbol {X}$의 각 요소는 $\mathcal {N}(0, 1)$인 표준 정규분포에서 추출되었다. $\boldsymbol {X}$가 거치는 layer의 개수는 총 5개이다. 각 layer에 할당된 weight 값의 행렬 $\boldsymbol {W}$는 $100 \times 100$의 크기를 가지며, 마찬가지로 각 원소 또한 $\mathcal {N}(0, 1)$인 표준 정규분포에서 추출되었다.

 

위 코드는 행렬곱 $\boldsymbol{XW} = \boldsymbol {A}$ 내 각 원소들을 sigmoid 함수에 넣어 새로운 행렬 $\boldsymbol {Z}$를 만든다. 그리고 이와 같은 연산을 다섯 번 반복하며, 각 연산에서 출력되는 행렬 $\boldsymbol {Z}$의 요소 분포를 그림으로 보여준다. 다음과 같이 말이다.

 

$\sigma = 1$인 경우

 

각 분포는 $0$ 또는 $1$에 치우쳐 있는 것을 볼 수 있다(왜 이런 분포가 도출되는지는 후술하겠다).이를 바람직한 분포라고 할 수 있을까? 아쉽게도 바람직한 분포라고는 할 수 없다. 왜 그럴까? 일단 $0$과 $1$에 몰려 있기 때문에 우리가 원하는 다양한 값을 산출한다고 보기에 어렵다고 할 수 있다. 그리고 simoid 함수의 형태에 내재하고 있는 치명적인 문제점을 추가할 수 있다. sigmoid 함수가 $0$과 $1$에 가까운 값을 반환할수록 그 지점에서 gradient는 매우 완만하고 작은 값을 갖는다. 이는 모델의 학습이 매우 지진하게 이루어지는 gradient vanishing problem을 유발한다. 이러한 두 가지 문제점 때문에 $0$과 $1$ 위주로 이루어진 분포를 반환하는 가중치 초기화 방법은 그다지 좋은 방법이라고 할 수 없다.

 

결국 표준편차가 $1$일 경우, 우리가 원하는 결과를 도출할 수 없다. 이제 앞선 계획대로 표준편차를 변화시켜 결과를 살펴보자. $0.01$과 같이 훨씬 작은 값을 표준편차로 설정하면 어떤 분포가 출력될까? 이는 위 코드에서 주석이 달린 라인의 $1$을 $0.01$로 고쳐 주기만 하면 알 수 있다($\because \  \sigma (nX) = n \sigma(X)$). 이렇게 도출된 분포는 다음과 같다.

 

$\sigma=0.01$인 경우

 

이제 분포는 극단적으로 $0.5$에 몰리고 말았다. 즉 MLP 모델이 다양한 결과를 표현하지 못하고 매우 한정된 결과만 표현하여 단조로운 모델이 돼버린 것이다. 또한, 여러 레이어에서 같은 값만을 표현하기 때문에 레이어를 여러 개로 설계하여 얻는 이점이 상실되었다.

 

표준편차를 $1$로 두었을 때는 분포가 가장자리로 너무 갈려 있었고, 표준편차를 $0.01$로 두었을 때는 분포가 한 가운데에 너무 몰려 버렸다. 따라서 직관적으로 표준편차를 $1$과 $0.01$ 사이의 어떤 값으로 설정하면, 다양한 결과를 표현하는 바람직한 모델을 달성할 수 있을 것이다.

 

 그건 그렇고 표준편차에 따라 왜 저런 분포가 도출되는 것일까?

 

어떤 표준편차가 적절한지 탐구해 보기에 앞서, 표준편차와 MLP 모델의 출력 사이에는 과연 어떤 관계가 존재하는 것인지 생각해 보자. 표준편차가 $1$일 때, 출력값 $\boldsymbol {Z}$에서부터 계산 과정을 한 단계씩 거꾸로 거슬러 올라가며 살펴볼 것이다. 일단 $\boldsymbol {Z}$를 도출하기 위해 sigmoid 함수에 입력되는 행렬은 $\boldsymbol {A}$이다. 행렬 $\boldsymbol {A}$의 요소 $a_{mn}$들이 $0$ 근처에 몰려 있기 보다는, 다분히 크거나 작은 값을 갖는 빈도가 높아야 sigmoid 함수를 거친 $\boldsymbol {Z}$의 분포가 $0$ 또는 $1$로 몰리게 될 것이다. 다시 말하자면 $\boldsymbol{XW} = \boldsymbol {A}$이기 때문에, $\boldsymbol {X}$의 행과 $\boldsymbol {W}$의 열을 곱한 $a_{mn}$이 다분히 크거나 작은 값을 갖는 빈도가 높다는 의미이다.

 

$\boldsymbol {X}$의 행과 $\boldsymbol {W}$의 열을 곱한 값은 표준정규분포에서 임의로 두 원소 $x_i$와 $w_i$를 100번 추출하여 더한 $\sum_{i=1}^{100} x_i w_i$라고 할 수 있다. $x_i$와 $w_i$ 모두 정규분포를 따르므로 $0$ 근방에서 추출될 확률이 높은 값이다. 따라서, 이 둘을 곱한 $x_i w_i$는 훨씬 더 $0$에 가까운 값일 것이다. 그런데 비록 $x_i w_i$이 $0$에 아주 가까운 값이지만, 이러한 값을 백 개 씩이나 더하면 이는 과연 어느 정도로 큰 값이 나올까? 첫 번째 실험 결과는 $\sum_{i=1}^{100} x_i w_i$이 sigmoid 함수에서 $0$ 또는 $1$을 도출할 정도로 큰 값임을 말해준다.

 

※ $\sum_{i=1}^{100} x_i w_i$가 어떤 분포를 따를 것인지 엄밀히 분석하는 것은 내 수학적 지식을 초월하지만 다음의 링크와 관련이 있는 것으로 보인다.

https://math.stackexchange.com/questions/161757/what-is-the-distribution-of-a-random-variable-that-is-the-product-of-the-two-nor

 

이제 표준편차가 $0.01$인, 훨씬 더 $0$에 몰려 있는 정규분포에서 $x_i$와 $w_i$를 추출하여 $\sum_{i=1}^{100} x_i w_i$을 계산했다고 생각해 보자. 비록 $x_i w_i$를 백 개 씩이나 더하더라도, 애초에 추출된 값들이 표준편차가 $1$일 때 보다 훨씬 더 $0$에 가깝기 때문에, $\sum_{i=1}^{100} x_i w_i$은 $0$에 더 가까운 값을 갖는다. 이는 두 번째 실험 결과를 바탕으로 확인할 수 있다.

 

좌측은 $\sigma = 1$일 때의 (표준) 정규분포 곡선이고, 우측은 $\sigma = 0.01$일 때의 정규분포 곡선이다. 확률 밀도 함수를 고려헤보면 $y$ 절편이 100배 상승하였음을 알 수 있다.

 

그러면 어떤 표준편차를 써야 할까?

 

$\sigma = 1$일 때는 출력 분포의 그래프가 너무 양쪽에 치우쳐져 문제였고, $\sigma = 0.01$일 때는 출력 분포의 그래프가 너무 가운데에 몰려 있어서 문제였다. 출력 분포 그래프가 다양한 값을 갖도록 하는 표준편차는 $1$과 $0.01$의 사이에 존재할 것이다. 정확히는 우리가 위 코드블럭에서 node_num으로 설정한 노드의 개수 $100$에 영향을 받을 것이다. 왜냐하면 행렬곱을 통해 가중치를 반영하는 과정에서, $x_i$와 $w_i$를 곱한 값을 노드의 개수만큼 더해줄 것이므로, 노드의 개수와 협응할 수 있는 표준편차를 선택해야 출력 분포 그래프가 좋은 표현력을 가질 수 있기 때문이다.

 

다시 한번 강조하자면, 우리의 근본적인 목표는 표현력이 높은 출력을 달성할 수 있는 weight 초기화 방법을 찾는 것이다. 이를 위해서 노드의 개수 $n$이 반영된 표준편차를 따르는 정규분포를 바탕으로 weight 값을 초기화해야 한다는 의미이다.

 

Xavier Initialization & He Initialization

 

위와 같은 아이디어를 통해 도출된 방법이 바로 Xavier Initialization과 He Initialization이다. 두 초기화 방법 모두 정규분포로 초기화하는 방법과, 균등분포로 초기화하는 방법이 존재하지만, 여기서는 정규분포로 초기화하는 방법만 다루도록 하겠다. Xavier Initialization은 표준편차를 $\sqrt{\frac{1}{n}}$로 설정한다. 직관적으로 $x_i w_i$는 두 정규분포의 곱이고 이를 노드의 개수 $n$개만큼 더하였으므로,  $\sqrt {n}$으로 나누는 계산을 통해 출력되는 행렬의 분포를 정규분포로 다시 상쇄시켜 줄  수 있을 것이라 추측할 수 있다. 실제로 산출되는 분포를 보면 다음과 같다.

 

Xavier Initialization

 

레이어가 깊어질수록 분포가 다소 어그러지기는 하지만 대체로 다양한 값을 표현하는 분포를 출력하는 것을 볼 수 있다. 깊어질 수록 분포가 어그러지는 이유는 $\boldsymbol {W}$이 거듭해서 곱해지는 연산을 $\sigma = \sqrt {\frac {1}{n}}$이완벽하게 상쇄하지는 못하기 때문인 것으로 추측된다.

 

He Initialization은 활성화 함수 및 출력 함수가 sigmoid 함수가 아니라 ReLU 함수인 경우에 사용되는 초기화 방법으로, 표준편차가 $\sqrt {\frac {2}{n}}$인 정규분포를 통해 $\boldsymbol {W}$을 초기화한다. 만일 ReLU 함수를 사용하는데 표준편차가 $\sqrt {\frac {1}{n}}$인 정규분포를 통해 초기화한다면, 레이어가 깊어질수록 출력 분포는 모두 $0$으로 수렴하게 된다. 다음 그림을 보자.

 

ReLU 함수를 활성화 함수로 사용하는 경우, 좌측은 $\sigma = \sqrt {\frac {1}{n}}$일 때의 분포이고, 우측은 $\sigma = \sqrt {\frac {2}{n}}$일 때의 분포이다.

 

정규분포로 초기화하는 경우 음수와 양수는 50 대 50의 비율로 추출된다. ReLU 함수의 경우, 음수 입력을 0으로 반환하므로 전체 100,000개의 요소 중 절반인 50,000개의 값은 0으로 반환됨을 확인할 수 있다. 왼쪽 그림에서 Xavier 초기화의 경우, 레이어가 깊어질 수록 분포가 0으로 빨려 들어가는데 반해, 오른쪽 He 초기화를 적용하면 이러한 현상이 사라짐을 확인할 수 있다.

 

$\sigma = \sqrt {\frac {2}{n}}$과 ReLU 함수는 어떤 관계가 있는 것일까?

 

솔직히 말해 둘이 무슨 관계가 있고, 도대체 어떻게 Xavier 초기화의 문제점이 극복되는 것인지 나로서는 아무런 짐작조차 할 수가 없었다. 책을 보면 다음과 같이 쓰여 있다.

 

"ReLU는 음의 영역이 0이라서 더 넓게 분포시키기 위해 2배의 계수가 필요하다고

(직감적으로) 해석할 수 있겠습니다." - p.207

 

불행히도 나에게는 저자가 독자에게 기대하고 있는 '직감'이란 것이 없는 것 같다. 실은 왜 2배의 계수인지도 잘 모르겠다. $\sqrt {\frac {2}{n}}$는 $\sqrt {\frac {1}{n}}$의 $\sqrt {2}$배가 아닌 것인가? 다만 우측 그림에서 저렇게 절묘토록 평탄한 분포를 따르는 것을 보면 필시 어떤 수학적 원리가 뒷받침되었으리라고 의심할 뿐이다.

 

다소 마무리가 아쉽지만 이것으로 가중치 초기화에 대한 글을 마치려고 한다. 많은 도움이 되었길 바란다.