[논문 리뷰]TinyTL: Reduce Activations, Not Trainable Parameters for Efficient On-Device Learning

들어가며

본 논문은 TinyTL: Reduce Activations, Not Trainable Parameters for Efficient On-Device Learning을 리뷰한 글입니다.

논문의 주요 배경

On-Device에서 학습의 어려움

메모리 부족

엣지 장비의 경우 메모리가 매우 제한되어 있다.(라즈베리파이1 A의 경우 256MB) 이는 Fig. 1의 왼쪽 그림에서 보이듯 추론만 하기엔 충분하지만 학습을 하기에는 MobileNetV2 같은 가벼운 모델을 쓰더라도 충분하지 않다. 이 부족한 메모리조차 다른 응용프로그램이나 운영체제와 공유하기 때문에 이런 메모리는 더욱 부족해진다.

에너지 부족

엣지 장비는 에너지도 제한되어 있다. DRAM에 접근할 때 on-chip SRAM에 비해 두 배 더 많은 에너지를 소비한다. 활성화 함수(activations)는 메모리를 많이 사용하기 때문에(The large memory footprint) SRAM에는 적합하지 않아 DRAM을 사용하여야만 한다.

 

따라서 메모리 사용량을 줄여야 한다.

※ SRAM과 DRAM의 특징의 차이점을 잘 정리해 놓은 표가 있어 첨부한다.

 

기존 경량화 기법들의 한계

기존의 기법들은 학습 가능한 파라미터나 FLOPs를 줄였다. 이러한 방법들은 학습 메모리를 직접적으로 절약하진 못한다. Fig. 1의 오른쪽 그림에서 볼 수 있듯 학습 메모리에 병목을 유발하는 것은 파라미터가 아니라 활성 함수(activation)이다. 그림을 보면 MobileNetV2-1.4의 파라미터 사이즈는 ResNet-50에 비해 4.3배 더 적지만 활성 함수의 경우 거의 같다.

 

따라서 역전파에 필요한 중간 활성 함수의 크기를 줄여야 한다.

 

다시 정리하자면 엣지 장비에서 학습을 위해선 메모리 측면이나 에너지 측면에서 메모리 사용량을 줄여야 하는데, 이를 위해 중간 활성 함수의 크기를 줄여야 한다는 것이다.

 

역전파에서의 메모리 사용량

다음 식에서 \(M\)을 일반적인 신경망이라고 할 때, 이는 레이어의 시퀀스로 이루어져있다.

\(\mathbf{w}_{i}\)는 \(i\)번째 레이어의 파라미터를 가리키고, \(\mathbf{a}_{i}\)와 \(\mathbf{a}_{i+1}\)를 각각 \(i\)번째 레이어의 입력과 출력 활성 함수라고 하고, \(L\)을 손실이라고 하자. 역전파에서 \(\frac{\partial L}{\partial \boldsymbol{a}_{i+1}}\)이 주어질 때 \(\frac{\partial L}{\partial \boldsymbol{a}_{i}}\)와 \(\frac{\partial L} {\partial \boldsymbol{w}_{i}}\)를 계산하여야 한다.

 

\(i\) 번째 레이어를 선형 레이어라고 할 때 순전파는 \(\boldsymbol{a}_{i+1}=\boldsymbol{a}_{i} \boldsymbol{W} + b\)이고, 배치가 1이라고 할 때, 역전파는 다음과 같이 표현된다.

 

위 식에서 볼 수 있듯 \(\boldsymbol{a}_{i}\)같은 중간 활성 함수는 바이어스를 계산할 때는 필요하지 않고 \(\frac{\partial L}{\partial \boldsymbol{W}}\)같은 가중치를 계산할 때만 필요하다. 따라서 학습시 바이어스만 업데이트 한다면 학습 메모리를 크게 아낄 수 있을 것이다. 합성곱 레이어나 정규화 레이어도 일종의 선형 레이어기 때문에 이들에도 저런 성질을 적용할 수 있다.

 

아래 테이블에서 볼 수 있듯 sigmoid, h-swish 같은 비선형 활성 함수들은 \(\frac{\partial L}{\partial \boldsymbol{a}_{i}}\)를 계산하기 위해 \(\boldsymbol{a}_{i}\)를 저장하여야 한다. 따라서 메모리 측면에서 효율적이지 못하다. 반면 ReLU 기반 활성 함수는 값이 0보다 작은지를 표현하는 이진 마스크만 저장하면 되기 때문에 다른 활성 함수에 비해 32배 더 작다.

 

 

따라서 앞으로 소개될 방법론에서 본 논문은 ReLU 기반 활성 함수를 사용한다.

 

제안된 모델

Lite Residual Learning

메모리 사용량 분석에 근거하여 메모리 소비량을 줄이는 방법 중 하나는 아래 (b) 그림에서 나타나듯 사전 학습된 특징 추출기의 가중치는 고정하고 바이어스만 업데이트 하는 것이다. 그러나 바이어스만 업데이트 할 경우 적응 용량(?)(adaptation capacity)에 한계가 있다. 따라서 Lite Residual Learning을 도입한다.

 

가중치는 고정하고 바이어스만 학습하는 레이어는 다음과 같이 수식화 할 수 있다.

 

앞서 말했듯 단순히 바이어스만 업데이트 하면 한계가 분명하기에 적은 메모리 사용량을 유지하면서 모델 용량(model capacity)을 개선하기 위해 출력을 정제하기 위한 잔차 특징 맵(residual feature map)을 생성하는 lite residual module을 제안한다.

 

위 식에서 \(\mathbf{a'}_{i}=\textrm{reduce}(\mathbf{a}_{i})\)는 저감 활성 함수(reduced activation).이다. lite residual module 학습에는 전체 활성 함수 \(\mathbf{a}_{i}\) 말고 저감 활성 함수 \(\mathbf{a'}_{i}\)만 있으면 된다.

lite residual module에는 mobile inverted bottleneck blocks (MB-block)을 적용했다. 모델의 주요 원칙을 활성 함수를 작게 유지하는 것으로 하고 이에 따라 활성 함수 크기를 줄이기 위해 두 개의 디자인 차원을 탐구했다. (Width, Resolution)

 

Width

널리 사용되는 inverted bottleneck은 depthwise convolution의 작은 용량을 보장하기 위해 6배*에 달하는 엄청난 수의 채널을 요한다. 이는 매개변수 측면에선 효율적이지만 활성 함수 측면에선 비효율적이다. 심지어 1× 채널을 6× 채널로 변환하는 과정에서 두 개의 1×1 정사영(projection) 레이어가 필요한데, 이는 전체 활성 함수를 12배 증가시킨다.

※ linear bottleneck과 비교하는 것 같지만 불확실하기 때문에 더 찾아봐야 할 것 같다.

 

Depthwise convolution은 또한 256 채널에서 1×1 OPs/Byte가 4% 미만으로 매우 낮은 산술 집중도(arithmetic intensity)를 갖고 있다. 그러므로 메모리 측면에서 매우 비효율적이다.

※ 하드웨어의 연산 성능을 충분히 활용하지 못한다는 뜻으로 받아들였다.

 

이러한 한계점을 해결하기 위해 lite residual module은 Grop Convolution을 사용한다. 이는 depthwise convolution보다 훨씬 높은 산술 집중도를 가지고 있으며 FLOPs와 메모리 사이에 좋은 트레이트 오프를 제공한다. 또한 1×1 정사영(projection) 레이어를 제거하여 전체 채널 수를 6.5배 감소시킨다.

 

Resolution

활성 함수의 크기는 해상도에 따라 4배씩 커진다. 그러므로 lite residual module은 2×2 average pooling으로 입력 특징 맵을 다운샘플링하여 해상도를 줄인다. 출력은 다시 메인 브랜치의 출력 특징 맵의 크기와 맞춰주기 위해 bilinear upsampling을 해준다.

 

이 두 Resolution과 Width를 최적화한 결과 lite residual module의 활성 함수는 inverted bottleneck에 비해 약 26배 더 적다.

 

실험 조건

데이터셋

ImageNet으로 사전학습 한 다음 8개의 객체 분류 문제로 전이 학습 하였다.

모델 아키텍쳐

ProxylessNAS-Mobile을 백본으로 잡고 해당 모델의 MB-block을 모두 likte residual module로 바꾸었다. Group의 수는 2개, 커널 크기는 5로 하였고, 활성 함수는 ReLU를 사용하였다. 배치 정규화는 소규모 학습 배치에서 더 잘 작동하도록 그룹 정규화로 바꾸었으며 그룹당 채널의 수는 8로 해주었다.

결과

다음 테이블에서 B는 Bias를 의미하여 L은 LiteResidual을 의미한다. FT-Last는 마지막 레이어만 파인튜닝 한 것을 의미한다. FT-Full은 전체 네트워크가 파인튜닝 되었음을 의미한다.

 아래 표를 살펴보면 알 수 있듯 TinyTL은 지속적으로 FT-Last와 FT-Norm+Last를 비슷하거나 적은 메모리 사용량으로 상회하였다.

 TinyTL은 해상도를 320까지 올리면(기존은 224) FT-Full과 비교해서 6배 더 효율적인 메모리로 비슷한 수준의 정확도에 도달할 수 있다.

 

다음 그림은 입력 해상도를 변경하면서 성능을 비교한 것이다. FT-Full의 경우 입력 해상도가 작아지면 정확도도 상당히 떨어졌지만 TinyTL은 FT-Full에 비교해서 3.9~6.5배 더 적은 메모리를 사용하면서도 비슷하거나 더 높은 정확도를 얻었다.