본 논문을 읽고 요약한 내용이며 잘못 해석한 내용이 있을 수 있습니다. 설명에 오류가 있다면 댓글로 알려주시면 감사하겠습니다.
[Abstract]
Vision Transformer는 vision task에서 진전을 보이고 있고, 다양한 benchmark에서 유의미한 결과를 얻었다. 그러나 파라미터 수가 많고 attention mechanism과 같이 2차 복잡도를 지니는 구조적인 한계 때문에 경량화된 CNN보다 몇 배는 더 느리다는 단점이 있다. 따라서 실제 어플리케이션에 배포하는게 힘들고, 특히 모바일 디바이스와 같이 제약된 하드웨어에서 이러한 단점이 부각된다.
본 논문의 핵심 질문은 'can transformers run as fast as MobileNet while obtaining high performance?'이다. 이를 위해 본 논문에서는 ViT 기반 모델의 디자인 구조를 다시 살펴보고 새로운 구조를 제안하고자 한다.
[Methodoloy]
On-Device Latency Analysis of Vision Transformers
대부분의 접근 방식은 GPU에서 얻을 수 있는 computation complexity (MACs) 또는 throughput (images/sec)을 통해 transformer의 inference 속도를 최적화하는 것이다. 그러나 이러한 metric은 실제 on-device latency를 반영하지 않는다. Edge device에서 ViT의 속도를 늦추는 원인이 무엇인지 파악하기 위하여 Figure1과 같이 다양한 모델에 대해 latency analysis를 진행하였다.
Observation 1: Path embedding with large kernel and stride is a speed bottleneck on mobile deivces.
Patch embedding은 kernel size와 stride가 큰 non-overlapping convolution layer로 구현되는 경우가 많다. (Figure 2의 모습처럼 겹치지 않게 하기 위해 kernel, stride 사이즈를 16등 크게 설정)
일반적으로 transformer에서 patch embedding layer의 계산 비용이 크지 않다고 생각되어져왔다. 그러나 Figure 1의 DeiT-S 및 PoolFormer-S24(Large kernel & stride)와 LeViT-256을 보면 큰 차이가 있다는 것을 알 수 있다. Large-kernel convolution은 대부분의 컴파일러에서 잘 지원되지 않으며, 가속화될 수 없는 부분이기 때문이다.
*) Reference https://engineering.linecorp.com/ko/blog/neural-network-on-mobile-gpu
TF Lite에서 3x3 필터를 사용하는 합성곱 연산이 Winograd 합성곱 알고리즘(행렬곱 기반의 합성곱보다 수행해야 할 곱셈 연산을 줄여줌)의 가속을 받아 성능이 크게 향상됨.
따라서 저자는 kernel, stride 대신 3x3 사이즈의 kernel을 가지는 CNN + fast downsampling조합으로 Patch embedding을 구현하였다. (Figure 3의 CONV stem)
Observation 2 : Consistent feature dimension is important for the choice of token mixer. MSA is not necessarily a speeed bottleneck.
(LeViT와 PoolFormer 비교) LeViT-256의 대부분은 4D tensor의 CONV으로 구현되며, 패치된 3D tensor에서 attention을 수행해야 하기 때문에 feature를 MSA로 전달할 때 reshape 과정이 여러 번 수행된다. 이렇듯 Reshape를 광범위하게 사용하면, 모바일 장치에서의 LeViT 속도가 제한된다. 반면, Pooling은 4D tensor에 적합하다.
(LeViT와 DeiT 비교) feautre dimension이 일관되고, Reshape가 필요하지 않은 경우 MSA가 오버헤드를 가져오지 않는다.
( 구조적인 차이점을 분석하여 내린 결론 )
따라서 위의 관찰 결과를 바탕으로 저자는 4D feature implementation과 3D MSA를 가지고 비효율적인 빈번한 reshape를 제거하는 dimension-consistnet network를 제안하였다.
Observation 3 : CONV-BN is more latency-favorable than LN (GN)-Linear and the accuracy drawback is generally acceptable.
MLP block은 보통 LN+3D linear projection과 CONV 1x1+BN(Batch Norm) 중 하나가 선택된다. Figure 1을 보면, LN이 network inference time 중 10~20% 정도를 차지하는 것을 볼 수 있다. LN과 같은 dynamic normalization은 inference 단계에서 실행 중인 statistics를 수집하여 latency에 기여한다. 반면, Conv-BN구조는 4D tensor에 대해 연산을 하고 inference 할 때, BN이 Conv 구조에 folding될 수 있으므로 latency를 낮추는 데 더 용이하다. 하지만, Conv-BN은 LN에 비해 성능이 조금 낮게 나온다.
따라서 저자는 token mixer가 pooling일 때는 Conv-BN을 사용하고, MSA일 경우 LN-Linear를 사용하였다. (원래 ViT설계와 일치하게 하기 위해)
Observation 4 : The latency of nonlinearity is hardware and compiler dependent.
GeLU, ReLU, HardSwish와 같은 activation에 대해서도 조사하였다. 선행 연구에서 GeLU가 hardware에 효율적이지 않다고 알려져있으나, iPhone 12에서는 오히려 ReLU보다 GeLU가 더 적합한 것으로 나타났다. (LeViT-256 latency with HardSwish is 44.5 ms while with GeLU 11.9 ms) 따라서 EfficientFormer에서는 GeLU를 사용한다.
Design of EfficientFormer
EfficientFormer는 patch embedding(CONV stem)과 여러 개로 쌓은 meta block(MB)로 구성된다. MB 3D는 transformer기반, MB 4D는 CNN을 기반으로 만들어졌다. 각 단계 사이에는 Embedding으로 표시된 임베딩 차원과 downsample token lenght를 투영하는 임베딩 작업이 있다. 각 Stage는 여러 개의 MB로 구성되어있고 같은 stage내에서는 같은 spatial size를 가진다.
Latency Driven Slimming
Design of Supernet.
dimension-consistent design을 위해 supernet으로부터 NAS를 사용해 최적 구조를 찾는다. Supernet은 위와 같은 MetaPath(MP)를 정의하여 구성된다. MB 3D는 마지막 두 개의 stage에서만 가능하도록 하였다.
이렇게 설계한 이유는 첫째, MSA의 계산은 token length에 quadratic하게 증가하기 때문에 앞쪽에 두게 되면 계산이 너무 증가하게 된다. 두 번째로 global MSA를 뒤쪽에 놓는 것은 앞쪽에서는 low-level feature를 뒤쪽 layer에서는 long-term dependency를 capture한다는 직관과 일치하기 때문이다.
Searching Space. searching space는 각 Stage의 width(C), 각 stage의 block 수 (N), MB 3D에 적용할 마지막 block 이다.
Search algorithm.
supernet의 학습이 완료되면 바로 어떤 path가 best인지 알 수 있는 efficient한 gradient-based search algorithm을 제안한다.
먼저, supernet training 씨에 Gumble Softmax sampling을 함께 사용해 선택된 MP의 importance score를 측정한다. 그 다음 16배수로 나누어진 channel(width)들을 가지는 여럿 MB 4D, MB 3D의 on-deivce latency lookup table을 구축한다. 그리고 single-width를 가지는 Supernet 기준으로 채널 수를 조절하는 gradual slimming을 진행한다.
[Result]
EfficientFormer는 PyTorch 1.11과 Timm library를 사용해 구현되었다. 모델은 NVIDIA A 100, V 100 GPUs에서 학습되었다. mobile speed는 A 14 bionic chip이 장착되고 NPU 사용이 가능한 iPhone 12에서 1000번 inference하고 평균을 내어 결과를 냈다.
[Conclusion & Limitations]
기존의 모델들에 비해 모바일 환경에서 더 낮은 latency를 달성하면서 성능도 acceptable한 결과를 나타냈다. 그러나 여전히 MobileViT와 같은 모델에 비해 큰 파라미터를 나타낸다. 또한, general-purposed하게 디자인되었지만, GeLU와 같이 플랫폼과의 호환성으로 인해 다른 플랫폼에서는 다른 경향을 띈다.