Slimmable Neural Networks

딥러닝 모델을 실제 어플리케이션에 적용하기 위해 다양한 경량화 기법, 모델 등이 연구되고 있습니다만 deploy를 하기 위한 mobile device 의 종류가 여러 가지고 각각의 computational power와 실행 시간이 다르기에, 목표로 한 device에 맞게 경량화된 모델을 탑재하여야 합니다. 

위의 그림처럼 같은 안드로이드를 사용하는 핸드폰이라 하더라도 image classification을 실행하는 시간은 제각각이기에 하드웨어 성능이 좋은 디바이스에는 더 큰 모델을 deploy 하고, 하드웨어 성능이 안 좋은 디바이스에는 latency 최적화를 위한 모델의 성능 감소가 필연적입니다. (latency-accuracy tradeoff)

MobileNet이나 ShuffleNet 같은 경량화 모델은 latency와 성능의 tradeoff를 조절하는 width multiplier로 다양한 경량화 모델 버젼을 제시하지만 각 모델이 별도로 훈련된 이후에 deploy 되어야 합니다. 또한, 어플리케이션이 실행되는 상황에 따라 같은 디바이스라 하더라도 실행 시간에 유의미한 차이가 날 수 있습니다. 즉, 모델을 탑재하더라도 상황에 따라 더 가볍거나 무거운 모델을 새로 탑재하기 위해서는 훈련, deploy 과정을 똑같이 반복해야 한다는 것이죠. 

이 논문에서는 주어진 리소스 안에서 다양한 경량화 모델을 개별적으로 훈련시키는 것이 아닌 하나의 모델을 훈련시켜 이로부터 더 가벼운 모델을 동시에, 즉각적으로 생성할 수 있는 기법을 제안합니다. 즉, 밑의 그림처럼 하나의 모델을 훈련시켜서 0.75배, 0.5배 가벼운 모델을 새로 훈련시킬 필요 없이 즉각적으로 만들어 낸다는 것이죠.

논문에서는 가벼워지는 것이 가능한 하나의 모델을 slimmable networks, 하나의 모델의 여러 variant 들을 switch 라 부릅니다. 일반적인 스위치처럼 하나의 모델에서 여러 버젼의 모델로 즉각적으로 바뀔 수 있는 것이죠. 이로 인해 하나의 모델만 훈련시키고도 타겟 디바이스 상황에 따라서 channel 수를 맞게 변동시킬 수 있고 다양한 convolution 등의 연산에 쉽게 적용이 가능합니다. 또한, 새로운 모델의 configuration 으로 스위칭 되었을 때 바로 동작이 가능합니다.

하지만 일반적으로 하나의 뉴럴 네트워크는 전체가 같이 동작하며 channel 수를 dynamic 하게 조절할 수는 없습니다. 논문에서는 switchable batch normalization 이란 기법을 사용하여 slimmable networks 를 훈련시킬 수 있는 방법을 제안합니다.

Slimmable neural networks

Naive training

Naive training 은 여러 width configuration 을 가진 공유된 네트워크를 바로 훈련시키는 방법입니다. 그림 1에서 볼 수 있듯이 훈련 데이터에 대한 성능은 안정적이나 테스트 데이터에 대한 에러는 거의 99%에 달할 정도로 성능이 매우 좋지 않습니다.

한 layer의 channel 입장에서 바라볼 때, 전 layer에서 매번 다른 숫자의 입력 channel이 들어오다 보니 feature의 평균과 분산이 상이해집니다. 이는 공유된 batch normalization layer에서 통계치의 불일치를 가져오기 때문에 테스트시의 극악의 성능이 나오는 것이라 예상됩니다.

[그림 1]

Incremental training

Incremental training 은 먼저 작은 모델 (0.35x) A 를 훈련시키고 파라미터를 더해 약간 큰 모델 (0.5x) B 를 훈련시키는 방법입니다. 이때, 추가된 파라미터는 A의 파라미터가 고정되고 fine-tuning 됩니다. 이 방법은 훈련 데이터와 테스트 데이터에서도 괜찮은 성능이 나오나 크기 별로 개별적으로 훈련시킨 모델에 비해서는 성능이 낮습니다. 이는 A 모델에서 B 모델로 확장할 때, 새롭게 연결되는 파라미터들과 기존의 파라미터들이 다같이 훈련되지 않기 때문입니다. 

Switchable batch normalization

논문에서는 switchable batch normalization 이라 하여 switch 마다 독립적인 batch normalization layer를 구성합니다. Naive training 방법과 비교했을 시, switch 마다 달라지는 feature 통계의 불일치를 switch 마다 개별적인 batch normalization layer로 해결할 수 있습니다. Incremental training 방법과 비교했을 시, 여러 switch를 동시에 training 시킬 수 있기 때문에 모든 파라미터들을 다같이 훈련시킬 수 있습니다.

Batch normalization의 파라미터는 전체 모델 크기에서 1% 가량만을 차지하고 deploy에 추가적으로 소요되는 실행 시간도 매우 적습니다. 이에 따라 slimmable networks 에서 switch를 바꿀 때마다 switch에 맞는 batch normalization layer를 on the fly 느낌으로 새롭게 구성해줄 수 있습니다. 따라서 새로운 configuration으로 모델을 switch 하더라도 추가적인 메모리나 실행 시간 없이 잘 동작할 수 있습니다.

Training slimmable neural networks

Slimmable networks의 훈련은 모든 switch에서 발생하는 성능의 평균을 최적화시키는 것입니다. 따라서 switch 마다 back-propagation을 통해 gradient 를 누적하고 한 번에 업데이트 합니다. 이때, switchable width list를 미리 정의하여 각 switch 별 configuration을 정의합니다. 전체 알고리즘은 다음과 같습니다.

 

Experiment

ImageNet classification에 대한 결과는 다음과 같습니다. 각 경량화 모델의 버젼을 개별적으로 훈련시킨 것과 slimmable networks를 한 번에 훈련시켜 switch 별로 성능을 비교하였습니다.

Slimmable neural networks의 개별적인 switch는 각 모델을 별도로 훈련시킨 것과 거의 동일한 성능을 얻었습니다. 특히, 극단적으로 slim 한 (0.25x) 네트워크에서는 개별적으로 훈련시킨 것보다 성능 향상이 더 큰 것을 볼 수 있습니다. 이는 slimmable neural networks의 각 switch가 모델 파라미터를 공유하고 여러 switch가 같이 훈련되다 보니 암묵적인 model distillation이 수행되었다고 추측됩니다.

More switches in slimmable networks

Switch의 개수에 따른 성능은 어떨까요? Switch의 개수가 많을 수록 latency-accuracy tradeoff 에 따른 모델 선택의 여지가 더 많아집니다. 밑의 표에서 볼 수 있듯이 switch가 많아도 성능이 비슷하게 유지된다는 것을 알 수 있습니다.

 

Discussion

하나의 channel이 다양한 switch에서 어떤 역할을 맡고 있는지 분석해 보았습니다. 따라서 S-MobileNet v1의 특정 채널 (channel 3_9)에 대해 가장 큰 activation 을 갖고 있는 이미지를 밑의 그림과 같이 분석합니다.

빨간 색 테두리가 된 것은 misclassified 된 이미지입니다. 신기하게도 switch 마다 channel 3_9의 역할이 Width가 커질 수록 channel 3_9 는 흰 색에서 노란 색을 인식하는 역할로 달라지는 것을 볼 수 있습니다. 이는 slimmable networks에서 특정 채널이 비슷한 역할을 수행하지만 switch 별로 살짝 다른 역할을 맡는다는 것으로 해석할 수 있습니다.

또한, 밑의 그림에서 얕은 layer (BN 1_1 to 1_8)과 깊은 layer (BN 12_1 to 12_8)의 switch에 따른 batch normalization 파라미터의 분포를 보면 얕은 층의 파라미터는 깊은 층에 비해 상대적으로 일정하고 층이 깊어질수록 switch 마다 파라미터의 변동성이 더 커집니다. 이것 또한, 특정 채널이 switch 마다 살짝 다른 역할을 맡게 된다는 사실을 방증합니다.

 

참조

• Slimmable Neural Networks