⚰️

Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift

Tags
Model Architecture
Created
2021/03/10 15:32
Publication
ICML'15
Rate
5
Summary
딥러닝 모델 학습 시에 각 layer의 input distribution이 계속해서 변동하는 internal covariate shift 현상이 발생하고, 이 영향은 딥러닝 모델이 깊어질수록 나비효과처럼 커져 모델의 학습을 느리게 만들거나 gradient exploding/vanishing과 같이 모델의 학습을 어렵게 만드는 현상으로 나타난다. 이러한 internal covariate shift를 해결하기 위해 각 layer의 input distribution을 mini batch에 기반하여 normalization하는 Batch Normalization 기법을 제안한다. (상세 페이지 참고)

1. Introduction

어떠한 machine learning system을 학습할 때, input distribution이 변동하는 현상을 covariate shift라고 한다. 이러한 covariate shift가 발생했을 때 learning system은 계속해서 이렇게 변동하는 input distribution에 맞추어 학습을 진행해야 하기 때문에 전반적인 학습에 문제가 생길 수 있다.
특히나 딥러닝 모델은 어떤 layer의 output이 다음 layer의 input이 되는 형태이기 때문에 각각의 layer가 sub-network인 거대한 learning system network이며, 이에 따라 layer가 깊어질수록 covariate shift가 나비효과처럼 커지는 현상이 일어날 수 있다.
또 sigmoid function을 activation으로 사용한 경우 절댓값이 매우 작은 몇몇의 노드를 제외하고는 activation value가 0에 수렴해버리는 gradient vanishing 현상이 발생하기도 한다.
이러한 상황에서는 모델의 학습이 매우 길어질 뿐 아니라, 모델의 학습이 learning rate 등의 hyper-parameter 설정 및 weight initialization에 매우 민감해질 수 있다. Batch Normalization은 이렇게 학습 중 internal nodes의 distribution이 변화하는 internal covariate shift 현상을 제거함으로써 더 빠르고 안정적이면서도 효과적인 딥러닝 모델의 학습이 가능함을 보인다.

2. Towards Reducing Internal Covariate Shift

이제 Internal Covariate Shift를 파라미터 학습으로 인한 activation들의 distribution의 변화라고 정의하자. 페이퍼에서는 이러한 internal covariate shift를 제거하고, 각 layer들의 input distribution을 고정함으로써 딥러닝 모델의 학습 속도를 높일 수 있음을 보이고자 한다.
먼저 LeCun 형님의 1998b 페이퍼로 인해 input whitening(input vector를 zero-mean, unit-variance, decorrelation을 가지도록 linear transform 한다)이 모델의 학습 속도를 높여줄 수 있다는 사실이 알려졌다고 한다. 이를 활용하여 각 layer의 input vector들을 whitening하면 internal covariate shift를 없애면 될 것 같다.
하지만 이 과정에서 normalization이 optimization과 독립적으로 이루어진다면 parameter update의 효과가 상쇄될(?) 수 있다. 예컨대 layer input을 input의 training set에 대한 기댓값으로 빼는 normalization(x^=xE[x]\hat{x} = x - E[x])을 수행한다고 하자. 이 때 gradient descent step이 E[x]E[x]bb 에 대한 dependence를 고려하지 않는다면, bias parameter bb 에 update가 일어나도 expectation E[x]E[x] 도 딱 그만큼 변화하게 되어 bias는 무한히 커지면서도(혹은 작아지면서도) loss는 변하지 않는 현상이 발생할 수 있다. 이는 center normalization뿐 아니라 scale normalization에서도 발생할 수 있기 때문에 whitening의 문제는 더욱 심각해진다.
위와 같은 이슈는 normalization이 optimization step과 무관하게 일어나기 때문에 발생한다. 즉 normalization이 gradient 계산에 포함되지 않아서 (다시 말하자면 미분되지 않는 연산이기 때문에) 발생한다. 따라서 이를 해결하기 위해서는 parameter value에 관계 없이 network 자체가 언제나 고정된 distribution을 만들어내도록 해야 한다. 즉 network 자체에서 normalization이 진행되어야 한다.
더군다나 whitening을 위해서는 데이터에 대한 covariance matrix와 그 inverse square root을 구해야 하는데, 이것을 계산하는 비용이 만만치 않다. 따라서 internal covariate shift를 해소하기 위한 differentiable 하면서도 계산 비용이 합리적인 방법이 필요하다.

3. Normalization via Mini-Batch Statistics

먼저 위에서 제시한 1) 미분 가능성에 대한 문제와 2) 계산 비용에 대한 문제를 해결하기 위해 두 가지 simplification을 제시한다.
먼저 whitening처럼 layer input과 output을 jointly normalize하는 게 아니라, input과 output을 각각 독립적으로 normalize한다.
x^(k)=x(k)E[x(k)]Var[x(k)]\hat{x}^{(k)} = \frac{x^{(k)} - E[x^{(k)}]}{\sqrt{Var[x^{(k)}]}}
이러한 normalization은 decorrelation을 포함하지 않더라도 모델의 학습 속도를 높이는 데 도움을 준다.
하지만 이렇게 단순히 normalization을 하는 것은 layer의 representation power를 해칠 수 있다. 예컨대 sigmoid activation을 사용한다면, 이렇게 normalize된 범위 내에서는 sigmoid function이 거의 linear function에 가까우니, 사실상 linear activation을 사용하는 것과 유사한 성능을 가지게 될 수도 있다. 이러한 문제를 해소하기 위해 normalization transform이 identity transform을 학습할 수 있도록 다음과 같이 scale parameter와 location parameter를 도입한다.
y(k)=γ(k)x^(k)+β(k)y^{(k)} = \gamma^{(k)}\hat{x}^{(k)} + \beta^{(k)}
이 두 parameter들은 다른 model parameter와 동시에 학습되며, 언급했듯 normalization function이 identity가 될 수 있는 여지를 남겨준다. ResNet의 skip connection의 포인트 중 하나도 바로 identity function의 학습이 가능하다는 것인데, 유사한 패턴이 보이는 것 같아 신기하다.
그리고 효율적인 학습을 위해 normalization을 위한 expectation과 variance는 각각의 mini-batch를 통해 estimation한다. 이렇게 매우 효율적인 연산으로 미분 가능한 normalization을 이뤄내는 Batch Normalization layer가 완성된다.
이 때 BN transform xyx \rightarrow y 에서 normalized activation x^\hat{x} 의 역할이 다소 숨겨져있지만, normalized feature x^\hat{x} 가 sub-network x^y\hat{x} \rightarrow y 에 input되는 것으로 해석될 수 있으며, x^\hat{x} 의 distribution이 고정되어 있음으로써 internal covariate shift가 해소되어 학습 속도가 높아질 것이라고 기대할 수 있다.

3.1. Training and Inference with Batch-Normalized Networks

이제 이러한 Batch Normalization layer는 딥러닝 모델의 어떠한 layer에도 삽입될 수 있으며, optimizer에 관계 없이 적용될 수 있다.
다만 mini-batch에 기반한 normalization을 inference에 적용하게 되면 stochasticity가 발생하게 되는데, 확실히 이는 undesirable한 상황이다. 따라서 deterministic한 Batch Normalization layer를 얻기 위해 training 시에 exponential average를 통해 normalization statistics를 update하고, 이렇게 구한 statistics를 inference 시에는 deterministic하게 사용하게 된다.

3.2. Batch-Normalized Convolutional Networks

Batch Normalization은 어떤 activation에도 사용될 수 있는데, 통상적으로 affine transform(i.e. fully-connected layer) 뒤, 그리고 nonlinear activation 이전에 삽입하는 것을 권장한다. 그 이유는 affine transform 이전의 input은 다른 nonlinear function의 출력값일 것이고, 이에 따라 단순히 first moment를 normalization하는 것이 covariate shift를 해소하지 못할 가능성이 큰 반면, affine transform을 거친 input은 보다 대칭적이고 Gaussian스러운 분포를 가질 가능성이 크기 때문이라고 한다. 이 경우 어차피 Batch Normalization에 location parameter가 포함되어 있기 때문에 affine transform의 사실상 bias parameter는 아무런 효과가 없다.
Convolution layer에 적용한다고 해도, 똑같이 feature map의 각각의 위치마다 mini-batch normalization을 적용해주면 된다.

3.3. Batch Normalization enables higher learning rates

이러한 BN layer는 activation을 normalization함으로써 parameter의 작은 변화가 나비효과가 되어 gradient explosion 혹은 vanishing을 발생시키는 것을 방지하여 훨씬 안정적인 학습을 가능하게 한다. 이로 인해 훨씬 더 높은 learning rate로 빠른 학습을 유도하는 것이 가능해진다.

3.4. Batch Normalization regularizes the model

학습 시 BN layer는 어떤 한 data sample만이 아니라 mini-batch 내의 다른 data sample들도 고려하여 normalization을 수행하기 때문에 해당 data sample에 deterministic한 결과를 생성해내지 않고, 이로 인해 model regularization 효과를 가져다준다. 이로 인해 BN layer를 사용할 경우 Dropout layer를 사용하지 않거나 강도를 낮추어도 comparable한 성능을 낼 수 있음을 보인다.

4. Experiments

위에서 언급한 효과들 - 1) 안정적인 input distribution, 2) 빠른 학습, 3) model regularization - 이 모두 사실임을 여러 실험을 통해 경험적으로 증명한다. 명료한 아이디어를 합리적으로 구현하고, 실험을 통해 증명한 아름다운 페이퍼이다!
E.O.D.