Batch Normalization

이 글은 Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift을 읽고 정리한 것이다.

네트워크 학습에 많이 쓰이는 SGD의 경우 각 파라미터의 기울기가 앞쪽 레이어에 의존하기 때문에 조그마한 변화도 네트워크가 깊을 경우 증폭되는 경향이 있다. 또 각 레이어에 들어오는 input은 다른 레이어의 출력이고, 이는 각 층의 파라미터 값에 영향을 받기 때문에 학습이 진행됨에 따라 input의 distribution도 바뀌게 된다. 이에 따라 레이어가 계속해서 변화에 적응해야 하기 때문에 학습 속도가 느려지게 된다.

또한 saturation problem이라는게 있는데, saturation의 정의는 다음과 같다.(https://stats.stackexchange.com/questions/174295/what-does-the-term-saturating-nonlinearities-mean)

• $f$ is non-saturating iff $|\lim_{z \to -\infty}f(z)|=\infty$ or $|\lim_{z \to \infty}f(z)|=\infty$
• $f$ is saturating iff $f$ is not non-saturating.

즉 출력의 범위가 제한된 함수면 saturating, 그렇지 않으면 non-saturating이다. activation function으로 흔히 쓰는 sigmoid의 경우 출력이 [0, 1] 범위로 제한되어 있기 때문에 saturating이다. 이것이 문제가 되는 이유는, input 값이 특정 범위를 벗어나면 함수의 gradient가 0에 가까워지기 때문이다. 예를 들어 sigmoid에 들어오는 input의 절댓값이 커질수록 gradient는 0에 가까워진다. gradient가 작아지면 역전파 과정에서 앞쪽 레이어에 전달되는 값은 더 작아지고(vanishing gradient), 이는 학습을 느리게 한다. 이때까지는 이 문제를 해결하기 위해 activation function을 바꾸거나($ReLU$), 초기값을 잘 설정하거나, learning rate를 낮게 잡았다.

논문에선 네트워크 안쪽 노드의 distribution이 바뀌는 것을 Internal Covariate Shift라고 정의하고, 이를 줄이기 위해 Batch Normalization이라는 방법을 제시했다. 정규화를 mini-batch 단위로 하는 방법인데, 이를 통해 activation function이 적용되는 스칼라 변수의 평균과 분산을 고정시킨다. 이를 적용하면

1. Internal Covariate Shift가 줄어들어 학습이 빨라진다.
2. gradient의 스케일이나 초기값에 대한 dependency도 줄어들어 큰 learning rate를 사용할 수 있다.
3. regularization 효과가 있기 때문에 dropout 등의 기법을 사용하지 않아도 된다.
4. input의 범위가 고정되기 때문에 saturating한 함수를 activation function으로 써도 saturation problem이 일어나지 않는다.

의 장점이 있다.

---

그렇다면 정규화를 어떻게 적용해야 할까?

주의할 점 중 하나는 정규화에 사용되는 파라미터에 관한 연산도 gradient를 계산할 때 모두 고려돼야 한다는 것이다. $x$를 input으로 받아 bias $b$를 더한 값 $y=x+b$를 돌려주는 간단한 레이어를 생각해보자. $y$에서 평균을 뺀 새로운 값 $\hat{y}=y-E[y]$를 돌려주는 방식으로 이 레이어를 정규화 할 것이다. 그럼 $b$에 대한 update는 $b \gets b + \Delta b$로 진행될 것이고, $\Delta b \propto \frac{\partial loss}{\partial b} = \frac{\partial loss}{\partial y} \cdot \frac{\partial y}{\partial b} + \frac{\partial loss}{\partial E[y]} \cdot \frac{\partial E[y]}{\partial b} = 0$ 일 것이다. 그런데 만약 gradient 계산 과정에서 $b$가 $E[y]$에 끼치는 영향을 무시한다면 $\Delta b$는 $0$이 아니게 되고, $b$의 값은 step이 지날 때마다 계속해서 증가할 것이다. 식에서 볼 수 있듯 레이어 구조 상 $b$의 값은 output에 영향을 끼치지 않고(normalization 때문에 상쇄됨) loss 값에도 변화를 주지 않지만 계속해서 $b$의 값이 커지는 것이다. 이는 파라미터의 발산을 가져오게 되고, normalization이 평균을 조정하는데 그치지 않고 input의 스케일도 건드린다면 문제가 생길 여지가 더욱 커지게 된다.

이를 해결하기 위해선 정규화에 쓰이는 파라미터를 계산하는 연산도 gradient descent의 영향을 받는 위치에 두면 된다. tensorflow로 구현한다면 그냥 원래 input으로부터 reduce_mean 등 텐서를 새로 만든 뒤 다시 네트워크에 연결해주면 자연스럽게 적용될 것이다.

또 다른 문제점은 정규화에 쓰일 파라미터, 예를 들어 normalization이 whitening인 경우 mean이나 covariance matrix를 어떻게 계산하느냐이다. 전체 데이터셋에 대해 구하면 가장 좋겠지만 크기가 너무 크기 때문에 실질적으로는 불가능하다. 따라서 전구간 미분 가능하며 모든 데이터셋을 처리하지 않아도 되는 새로운 정규화 방법이 필요하고, 그것이 batch normalization이다.

---

Batch Normalization이 naive한 whitening과 다른 점은 두가지이다.

1. 레이어의 input이 d차원일 때, d개의 스칼라 값 간의 상관 관계를 무시하고 각각 독립적으로 normalize 한다. 즉 $\text{x}=\left(\mathit{x^\mathit{(1)}...x^\mathit{(d)} }\right)$일 때 ($x^i$는 스칼라 값) $Cov[\text{x}], E[\text{x}]$를 구해 정규화 하는게 아니라 $\hat{x}^\mathit{(k)}=\frac{x^\mathit{(k)}-E[x^\mathit{(k)}]} {\sqrt{Var[x^\mathit{(k)}]}}$로 각 차원에 대해 따로 정규화를 진행한다.
2. normalize를 위한 파라미터는 전체 데이터셋이 아닌 mini-batch에서 계산해 사용한다.

1.과 같이 정규화를 진행하면 스칼라 값의 mean과 variance가 각각 0, 1로 고정되는데 이는 레이어가 나타낼 수 있는 범위에 제약을 줄 수 있다. 예를 들어 Sigmoid의 input에 해당 정규화를 진행했다면 정의역이 함수가 거의 선형이 되는 부분만 남기 때문에 표현할 수 있는 정보의 양이 줄어들게 된다. 이를 방지하기 위해 transformation은 identity mapping이 될 수 있는 꼴로 설정한다.

논문에선 각 activation^[1]마다 정규화 후 scale과 bias 변수를 추가해 아핀 변환을 진행한다. 즉 $y^{(k)}=\gamma ^ {(k)} \hat{x}^{(k)}+\beta^{(k)}$ 를 최종적으로 사용한다. 만약 정규화를 하지 않는 것이 더 많은 정보를 나타낼 수 있어 이득이라면 $\gamma^{(k)} = \sqrt{Var[x^{(k)}]}, \beta^{(k)}=E[x^{(k)}]$ 로 학습될 것이다. 뿐만 아니라 원래 네트워크에서 쓰이던 bias는 의미가 없어지는데(정규화 과정에서 mean을 빼므로) 이를 보완해주는 역할도 한다.

2.를 적용하지 않는다면 정규화를 진행할 때마다 해당 activation을 전체 데이터셋에 대해 계산해야 할 텐데, 이는 굉장히 많은 연산이 필요하다. 네트워크의 학습에 GD가 아닌 SGD를 적용하는 것과 같은 맥락으로 이해하면 될 듯하다.

더불어 두번째 simplification은 첫번째 simplification 때문에 가능한데, 만약 joint covariance를 사용해 정규화를 진행한다면 mini-batch 내 데이터 개수 $m$에 비해 activation의 개수 $d$가 더 큰 경우가 많으므로 문제가 생긴다. activation을 모은 벡터가 $\text{x}=\left(\mathit{x^\mathit{(1)}...x^\mathit{(d)} }\right)$일 때 $i$번째 mini-batch에 대한 activation 값을 열벡터로 나타낸걸 $\text{x}_i$ 라고 하고, $\text{x}_i$에서 mean을 뺀 것을 모은 $d \times m$ 행렬을 $\text{F} = [(\text{x}_1-E[\text{x}]) \quad (\text{x}_2-E[\text{x}]) \quad ... \quad (\text{x}_m-E[\text{x}])]$라고 두자.
그럼 Sample Covariance Matrix는 $Cov[\text{x}]=\frac{1}{m-1}FF^T$로 계산된다. 그런데 $rank(F)\le min(d, m)=m$이므로 $rank(Cov[\text{x}]) \le m < d$ 가 성립해 $d \times d$ 행렬인 $Cov[\text{x}]$는 비가역적이어서 정규화에 사용할 수 없다. 그러나 첫번째 simplification에서 joint property를 무시하고 각 스칼라 값에 대해 independent하게 정규화를 진행하므로 mini-batch 단위로 정규화에 쓰이는 파라미터를 계산할 수 있다.

---

이제 실제 Batch Normalization 알고리즘을 살펴보자.
mini-batch에 들어있는 데이터의 개수가 $m$이면, 각각의 activation $x$에 대한 mini-batch의 값은 $B = \{x_{1...m}\}$로 나타낼 수 있다. 이를 앞서 말한 것처럼 정규화한 뒤 아핀 변환한 것을 $y$라고 부르고, 변환 $BN_{\gamma, \beta}:x_{1...m}\rightarrow y_{1...m}$를 Batch Normalization Transform 이라고 부른다. 원래 네트워크에서 $x$가 쓰이는 부분을 모두 $y$로 바꿔주면 BN이 적용된다.

각 연산은 미분 가능하므로 다음과 같이 chain rule을 적용할 수 있다.

$\gamma$와 $\beta$는 네트워크가 학습해야하는 새로운 파라미터이고, 그 외는 non-trainable인 값이다(데이터에 따라 정해지는 값).

학습 과정은 위와 같이 네트워크를 변형해 진행하지만, 학습 도중에 test accuracy를 계산하거나 학습이 끝난 뒤 실제로 네트워크를 사용할 때는 어떤 파라미터를 사용해 정규화를 해야할까?

실제로 네트워크를 사용할 때(inference), 개인적으로는 전체 데이터셋에 대해 mean과 variance를 계산해 사용할 줄 알았는데 알고리즘 의사코드에선 mini-batch에 대해 mean과 variance의 sample mean을 계산해 정규화 파라미터로 사용했다. 전체 데이터셋에 대해 진행하든 mini-batch에 대해 진행하든 expectation은 같을텐데, 이게 실험적으로 더 좋거나 훈련셋이 클 경우도 동작해서 그런 것 같다.

Test accuracy 계산을 할 때는 training이 아직 끝나지 않았을 때이므로 위와 같이 cost가 큰 연산은 하기 힘들다. 따라서 training에서 사용한 파라미터들의 moving average를 계산해 test 과정의 파라미터로 사용한다.

---

다음은 추가적인 내용이다.(논문의 3.2, 3.3 section)

네트워크에서 흔히 나타나는 구조는 input $u$에 대해 아핀변환을 수행하고 non-linear activation function $g$를 적용해 output $z=g(Wu+b)$를 얻는 것이다. 여기서 $W, b$는 각각 학습해야 하는 weight, bias 파라미터이고 $g$는 ReLU나 Sigmoid 등이다. BN은 $x=Wu+b$에도 적용할 수 있을 것이고 $u$에도 적용할 수 있을 것인데, 왜 activation function의 직접적인 input이 되는 $x$에 BN을 적용할까? $u$는 대부분의 경우 non-linear activation function을 거친 다른 레이어의 output 값이므로 distribution의 모양이 Normalazation이 큰 효과를 거두지 못하게 생겼을 것이다. 그에 반해 $x$는 symmetric, non-sparse하게, 즉 좀 더 정규 분포에 가까운 distribution을 가질 확률이 높기 때문에 BN을 적용했을 때 더 좋은 결과를 기대할 수 있다.

CNN의 경우는 parameter sharing이라는 특별한 성질이 있기 때문에 BN도 그에 맞춰서 수정해준다. input이 $H \times W$의 행렬이고 convolution filter가 $P \times Q$의 행렬일 때, convolution을 적용한 output은 $H \times W$의 행렬일 것이고 activation 또한 $H\cdot W$개가 생길 것이다. 그런데 CNN은 원래 데이터의 어떤 위치에 필터를 적용하든 필터의 파라미터는 똑같은 값을 사용하기 때문에, 같은 필터를 적용해서 나온 결과값에 대한 정규화도 같은 방식으로 하는 것이 맞을 것이다. 따라서 데이터의 $H \cdot W$개의 위치에 대해 필터를 적용해 나온 activation을 각각 정규화하는 것이 아니라, "필터를 적용한 결과" 자체를 하나의 스칼라 변수로 보고 한 데이터에 대해 변수의 sample 값이 $H \cdot W$개가 주어진다고 보는 것이다.

다시 정리하면, 각각의 필터에 대해 한 쌍의 $\gamma, \beta$만 설정해 학습하고, mini-batch에 $m$개의 데이터가 있을 때 정규화에 쓰이는 mean과 variance는 $m \cdot H \cdot W$개의 값에 대해 계산한다.

3.3은 BN이 네트워크의 학습 속도에 끼치는 영향에 대해 적혀 있다. 먼저 learning rate가 큰 경우 생기는 문제 중 하나는 weight의 scale이 커져 gradient 또한 증폭돼 발산이 일어난다는 점인데, BN을 적용하면 weight에 스칼라 $a$가 곱해져도 표준 편차로 나눌 때 상쇄되므로 레이어의 input에 대한 gradient의 증폭이 일어나지 않는다는 것이다. 또 weight에 대한 gradient는 오히려 $a$를 곱했을 때 $\frac{1}{a}$배 되기 때문에 더 안정적인 학습이 가능해진다고 한다.

그 뒤의 내용은 저자의 추측이다. layer Jacobian(1차 편미분 행렬)의 singular value들이 1에 가까울수록 학습에 유리하다는 것이 알려져 있는데 몇가지 가정을 하면 BN이 layer Jacobian의 singular value를 1로 만든다고 한다. 두 인접한 레이어에 대해 앞쪽 레이어의 normalized input을 $\hat{x}$, 뒤쪽 레이어의 normalized input을 $\hat{z}$라고 하면 앞쪽 레이어 연산을 transformation $F$로 나타낼 때 $\hat{z} = F(\hat{x})$이 성립한다. 이제 다음 가정을 적용하자.

가정 1. $\hat{x}$와 $\hat{z}$이 Gaussian이다.
가정 2. $\hat{x}$와 $\hat{z}$이 uncorrelated이다(!).
가정 3. $\hat{z} = F(\hat{x}) \approx J\hat{x}$로 transformation을 선형함수로 근사할 수 있다(!!!).

그럼 $\hat{x}$와 $\hat{z}$ 둘다 unit covariance를 가지므로 $I = Cov[\hat{z}] = Cov[J\hat{x}] = JCov[\hat{x}]J^T = JJ^T$가 성립해 $J$의 모든 singular value는 1이 된다. 그런데 가정이 워낙 괴랄하기 때문에 별로 설득력은 없다...
논문에는 다음과 같이 문단을 맺고 있다.

In reality, the transformation is not linear, and the normalized values are not guaranteed to be Gaussian nor independent, but we nevertheless expect Batch Normalization to help make gradient propagation better behaved. The precise effect of Batch Normalization on gradient propagation remains an area of further study.

---

1. 이 논문에선 activation이라는 단어를 activation function의 input으로 쓰이는 하나의 스칼라 변수를 지칭하는데 사용한다. ↩︎