[논문 리뷰] MCU AI 학습 시, 효율적인 메모리 사용 (On-Device Training Under 256KB Memory)
본 글은 다음의 논문을 기반으로 작성되었습니다.
https://arxiv.org/pdf/2206.15472.pdf
제가 생각한 중요 개념과 수식이 주된 내용이며, 자세한 내용은 위의 사이트를 참고해주세요.
또한 논문에서 사용된 그림과 표를 블로그에서 사용하기에는 저작권 등의 문제가 있기 때문에, 논문을 옆에 띄워 놓고, 글 중간에 언급되는 그림(Figure)와 표(Table)을 논문에서 참고해주세요.
현 글은 MCUNetv3에 관한 내용으로, 초기 버전 또는 MCUNet의 기반을 알고싶으시면, 아래 글을 참고해주세요.
Quantization-aware scaling (QAS)
양자화가 gradient update를 왜곡시킨다는 가정 하에, |W| / |G|를 비교하여 양자화하기 전의 값(using fp32)과 양자화 후의 값(int8)을 비교해 보았다.
(논문의 page 3의 Figure 2. 참고)
Figure 2.에서 알 수 있는 것처럼 푸른색 선(no quantization)과 녹색 선(quantization)의 양상이 다르다는 것을 알 수 있으며, 이는 곧 quantization이 gradient update를 왜곡한다는 뜻이다. 논문에서 제시된 양자화 기법인 QAS를 사용하면, 그림의 노란 선에서 알 수 있듯 no quantized 환경과 비슷한 양상을 얻을 수 있다.
논문에서 소개한 QAS(일반 QAS와는 조금 다릅니다.)는 hyper-paremeter가 없는 기법(aka hyper-parameter-free)이다. QAS의 원리는 다음과 같다.
위의 수식에 따르면, W는 기존의 Quantization으로 양자화하고, Gradient에 해당되는 G만 마지막 식을 통해 바꾸는 것 같다.
Memory-Efficient Sparse Update
QAS로 효육적인 Quantization을 했을 지라도, 여전히 training 시, memory consumption이 많다. 논문에서는 모든 층의 weight를 backpropagation으로 업데이트하는 것이 아닌, 나름의 weight 또는 layer 분석을 기반으로 중요한 weight만 또는 layer만 update한다. 부분적으로만 backpropagation을 진행하기 때문에 forward 연산에서 어떤 activation은 사용되지 않으며, 이로써 memory를 아낄 수 있다. 또한 backpropagation을 하지 않는 node들의 weight는 frizing(값을 변경하지 않음)한다.
Automated selection with contribution analysis
Sparse Update할 때, 각 weight 또는 bias가 accuract에 얼마나 기여하는지를 기반으로 update할 layer를 정한다.
where i is a collection of layer indices whose weights are updated, and $\mathbf{r}$ is the corresponding update ratios (1/8, 1/4,1/2,1), r은 전체 W에서 업데이트할 부분의 비율이다.
Memory-Efficient Sparse Update를 그림으로 보면 다음과 같다.
(논문의 page4에 Figure 3. 참고)
(b)에서 backward graph를 위한 연산자들을 만들고, Memory-Efficient Sparse Update를 (c) 단계에서 진행한다. 이 과정에서Automated selection with contribution analysis로 얻어진 짙은 푸른 색 부분의 영역만이 update이 된다.
Tiny Training Engine (TTE)
- Compile-time differentiation and code generation.
Python 등에서 작성된 코드에서 필수적인 부분만 code gneration(python으로 구현됨)으로 C파일로 생성한 후, compile을 통해 필요없는 코드와 overhead를 제거하는 기법이다.
- Backward graph pruning or sparse update.
앞서 설명한 Sparse update와 Dead-code elimination (DCE) 과정을 수반한다.
- Operator reordering.
연산의 순위를 바꾸면서 tensor의 활용도를 높히는 방식이다. 즉, tensors의 life cycle을 조정하여 메모리를 효율적으로 관리하는 방식이다.
Evolutionary Search vs. Random Search (for Automated selection with contribution analysis)
논문에서는 좋은 조건의 acc를 찾기 위한 방법으로 Evolutionart Search와 Random Search를 수행했다. 결과는 다음과 같다.
(논문의 page17에 Figure 13 참고)
그림에서 알 수 있는 것처럼 Evolutionary Serach 방법이 더 나은 \Delta acc를 찾는 것을 볼 수 있다. (구체적인 방법은 기재되어 있지 않았다.)
Single-batch training
MCU는 메모리의 제약으로 인해 2 이상의 batch를 사용하는 건 무리가 있다. 논문에서는 1 batch를 사용했을 때의 문제점을 기술하며, Gradient accumulation 기법으로 이를 해결한다.
(논문의 page16에 Table 4 참고)
위의 표는 하드웨어만을 통제 변인으로 삼아, batch size가 128 때와 1일 때를 비교한 테이블이다.
여기서 논문은 VWW를 사용했을 때를 주목했다. VWW에서 batch size 128 때와 1일 때 정확도 면에서 차이가 발생했다. 이는 논문의 말을 빌리면 (링크) “VWW only has two categories (binary classification), so the information from each label is small, leading to unstable gradients.”라고 표현했다. 즉, class의 수가 줄면, batch size와 정확도 사이에 gradient unstable로 인한 상관관계가 생긴다는 의미이다.
이는 Gradient accumulation 기법으로 극복이 가능한다. computation performance와 memory를 희생하고, stable gradient를 얻는 기법이다.
Gradient accumulation은 batch 하나마다 gradient를 update하는 것이 아닌, n개의 gradient를 누적했다가 이들의 평균을 이용하여 update라는 방식이다. 이렇게 하면 n-batch size를 갖는 학습 효과를 얻을 수 있다고 한다. <참고> (OOM은 Out of Memory입니다)