EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention
Metadata
[!meta]- Metadata – PDF
Title:: EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention
Authors:: Xinyu Liu, Houwen Peng, Ningxin Zheng, Yuqing Yang, Han Hu, Yixuan Yuan,
Year:: 2023
ItemType:: preprintKeywords:: Computer Science - Computer Vision and Pattern Recognition
Related::
ABSTRACT
비전 트랜스포머는 높은 모델 능력 덕분에 큰 성공을 거두었습니다. 그러나 이들의 뛰어난 성능은 무거운 계산 비용을 동반하며, 이는 실시간 애플리케이션에는 적합하지 않게 만듭니다. 본 논문에서는 EfficientViT라는 고속 비전 트랜스포머 계열을 제안합니다. 우리는 기존 트랜스포머 모델의 속도가 특히 MHSA의 텐서 재구성 및 요소별 함수에서 메모리 비효율적인 연산에 의해 일반적으로 제한된다는 것을 발견했습니다.
따라서, 우리는 샌드위치 레이아웃을 가진 새로운 빌딩 블록을 설계합니다. 즉, 효율적인 FFN 레이어 사이에 단일 메모리 바운드 MHSA를 사용하여 메모리 효율성을 개선하는 동시에 채널 통신을 강화합니다. 더욱이, 우리는 어텐션 맵이 헤드들 간에 높은 유사성을 공유하며, 이로 인해 계산 중복이 발생한다는 것을 발견했습니다. 이를 해결하기 위해, 우리는 전체 특징의 다른 분할을 어텐션 헤드에 공급하는 연쇄적 그룹 어텐션 모듈을 제시합니다.
이는 계산 비용을 절약할 뿐만 아니라 어텐션 다양성을 향상시킵니다. 포괄적인 실험들은 EfficientViT가 기존의 효율적인 모델들을 능가하며, 속도와 정확도 사이에 좋은 절충점을 이룬다는 것을 입증합니다. 예를 들어, 우리의 EfficientViT-M5는 MobileNetV3-Large보다 정확도에서 1.9% 앞서며, Nvidia V100 GPU와 Intel Xeon CPU에서 각각 40.4%, 45.2% 더 높은 처리량을 달성합니다.
최근의 효율적인 모델 MobileViT-XXS에 비교하면, EfficientViT-M2는 1.8% 더 우수한 정확도를 달성하면서 GPU/CPU에서 5.8배/3.7배 더 빠르게 실행되며, ONNX 형식으로 변환했을 때는 7.4배 더 빠릅니다. 코드와 모델은 링크에서 이용할 수 있습니다. [https://github.com/microsoft/Cream/tree/main/EfficientViT]
Main ideas:
- MHSA의 비효율적인 연산을 해결해보자
- 기존 트랜스포머 구조를 개선하여 빠른 inference를 갖춘 새로운 계층적 모델인 EfficientViT 모델 제안
- 샌드위치 레이아웃(sandwich Layout)
- Cascaded Group Attention - 모든 head에 대해 동일한 feature를 사용하는 이전 self-attention과는 달리 CGA는 각 head에 서로 다른 입력 split을 제공하고 출력 feature를 head 전체에 계단식으로 배열한다.
- Parameter Reallocation - 파라미터 효율성을 높이기 위해 중요한 모듈의 채널 폭을 확장하고 중요하지 않은 모듈은 축소하여 네트워크에서 파라미터를 재할당
Key points:
본 논문에서 언급된 기존 트랜스포머의 문제점을 상세 비교한 내용들.
[self-attention과 residual block 비교] - [https://minyoungxi.tistory.com/95]
[Layer Norm과 Batch Norm의 효율성 차이] - [https://minyoungxi.tistory.com/96]
⭐ Interesting point
To address this issue, in this paper, we explore how to go faster with vision transformers, seeking to find principles for designing efficient transformer architectures. Based on the prevailing vision transformers DeiT [69] and Swin [44], we systematically analyze three main factors that affect model inference speed, including memory access, computation redundancy, and parameter usage. In particular, we find that the speed of transformer models is commonly memory-bound. In other words, memory accessing delay prohibits the full utilization of the computing power in GPU/CPUs [21, 32, 72], leading to a critically negative impact on the runtime speed of transformers [15, 31].
잘 알려진 ViT인 DeiT와 Swin을 기반으로 메모리 액세스, 계산 중복성, 파라미터 사용을 포함한 모델 inference 속도에 영향을 미치는 세 가지 주요 요소를 체계적으로 분석.특히, 저자들은 트랜스포머의 속도가 메모리에 영향이 크다는 것을 알아냄. 즉, CPU/GPU의 컴퓨팅 성능을 완전히 활용하는 것을 막아 트랜스포머의 런타임 속도에 심각한 부정적 영향을 미침.
most memory-inefficient operations are the frequent tensor reshaping and element-wise functions in multi-head selfattention (MHSA). We observe that through an appropriate adjustment of the ratio between MHSA and FFN (feedforward network) layers, the memory access time can be reduced significantly without compromising the performance. Moreover, we find that some attention heads tend to learn similar linear projections, resulting in redundancy in attention maps. The analysis shows that explicitly decomposing the computation of each head by feeding them with diverse features can mitigate this issue while improving computation efficiency. In addition, the parameter allocation in different modules is often overlooked by existing lightweight models, as they mainly follow the configurations in standard transformer models [44,69]. To improve parameter efficiency, we use structured pruning [45] to identify the most important network components, and summarize empirical guidance of parameter reallocation for model acceleration.
메모리를 가장 비효율적으로 사용하는 연산은 muti-head self-attention(MHSA)의 빈번한 텐서 reshaping과 element-wise 함수들임. 저자들은 MHSA와 FFN 레이어 사이의 비율을 적절히 조절해서 성능 저하 없이 메모리 액세스 시간을 크게 줄일 수 있음을 관찰함. 그리고 일부 attention head가 유사한 linear projections를 학습하는 경향이 있어 attention map이 중복되는 것을 확인함. 다양한 feature를 제공하여 각 head의 계산을 명시적으로 분해하면 이 문제를 완화하는 동시에 계산 효율성을 향상시킬 수 있다. 또한 기존 경량 모델은 주로 표준 트랜스포머 모델의 구성을 따르기 때문에 서로 다른 모듈의 파라미터 할당을 간과하는 경우가 많음. 따라서 파라미터 효율성을 향상시키기 위해 구조화된 pruning(가지치기)를 사용하여 가장중요한 네트워크만 식별한다.
Specifically, we design a new block with a sandwich layout to build up the model. The sandwich layout block applies a single memory-bound MHSA layer between FFN layers. It reduces the time cost caused by memorybound operations in MHSA, and applies more FFN layers to allow communication between different channels, which is more memory efficient.
분석의 결과를 바탕으로 본 논문은 EfficientViT라는 새로운 메모리 효율적인 transformer 모델을 제안한다. 구체적으로, 모델을 구축하기 위해 샌드위치 레이아웃으로 새로운 블록을 설계한다. 샌드위치 레이아웃 블록은 FFN 레이어 사이에 하나의 memory-bound MHSA 레이어를 적용한다. MHSA의 memory-bound 연산으로 인해 발생하는 시간 비용을 줄이고 더 많은 FFN 레이어를 적용하여 서로 다른 채널 간의 통신을 허용하므로 메모리 효율성이 향상된다
Then, we propose a new cascaded group attention (CGA) module to improve computation efficiency. The core idea is to enhance the diversity of the features fed into the attention heads. In contrast to prior selfattention using the same feature for all heads, CGA feeds each head with different input splits and cascades the output features across heads. This module not only reduces the computation redundancy in multi-head attention, but also elevates model capacity by increasing network depth. Last but not least, we redistribute parameters through expanding the channel width of critical network components such as value projections, while shrinking the ones with lower importance like hidden dimensions in FFNs. This reallocation finally promotes model parameter efficiency.
그런 다음 계산 효율성을 향상시키기 위해 새로운 cascaded group attention (CGA) 모듈을 제안한다. 핵심 아이디어는 attention head에 공급되는 feature의 다양성을 향상시키는 것이다. 모든 head에 대해 동일한 feature를 사용하는 이전 self-attention과 달리 CGA는 각 head에 서로 다른 입력 split을 제공하고 출력 feature를 head 전체에 계단식으로 배열한다. 이 모듈은 multi-head attention의 계산 중복성을 줄일 뿐만 아니라 네트워크 깊이를 늘려 모델 용량을 높인다. 마지막으로, value projection과 같은 중요한 네트워크 구성요소의 채널 폭을 확장하는 동시에 FFN의 hidden 차원과 같이 중요도가 낮은 구성 요소를 축소하여 파라미터를 재분배한다. 이러한 재할당은 최종적으로 모델 파라미터 효율성을 향상시킨다.
메모리 액세스 오버헤드는 모델 속도에 영향을 미치는 중요한 요소이다. 빈번한 reshaping, element-wise addition, 정규화와 같은 transformer의 많은 연산자는 메모리 비효율적이므로 위 그림에 표시된 것처럼 여러 메모리 장치에 걸쳐 시간이 많이 소요되는 액세스가 필요하다. 계산을 단순화하여 이 문제를 해결하기 위해 제안된 몇 가지 방법이 있지만 표준 softmax self-attention에서는 정확도가 저하되고 가속이 제한되는 대가를 치르는 경우가 많다.
In this work, we turn to save memory access cost by reducing memory-inefficient layers. Recent studies reveal that memory-inefficient operations are mainly located in MHSA rather than FFN layers [31, 33]. However, most existing ViTs [18, 44, 69] use an equivalent number of these two layers, which may not achieve the optimal efficiency. We thereby explore the optimal allocation of MHSA and FFN layers in small models with fast inference.
메모리 비효율적인 레이어를 줄여 메모리 엑세스 비용 절감. 대부분의 ViT 모델들은 동일한 두 개의 레이어를 사용하므로 최적의 효율성을 달성하지 못할 수 있다. 따라서 본 논문은 소형 모델에서 MHSA와 FFN 레이어의 최적 할당을 탐색한다.
위 그림에서 볼 수 있듯이 20%-40% MHSA 레이어를 가진 하위 네트워크는 더 나은 정확도를 얻는 경향이 있다. 이러한 비율은 50% MHSA 레이어를 채택하는 일반적인 ViT보다 훨씬 작다. 또한 저자들은 reshaping, element-wise addition, 복사, 정규화를 포함한 메모리 액세스 효율성을 비교하기 위해 memory-bound 연산에 대한 시간 소비를 측정하였다. Memory-bound 연산은 20% MHSA 레이어가 있는 Swin-T-1.25×에서 전체 런타임의 44.26%로 감소된다. 또한 1.5배 속도 향상을 통해 DeiT와 더 작은 모델로 일반화된다. MHSA 레이어 활용률을 적절하게 줄이면 모델 성능을 향상시키면서 메모리 효율성을 향상시킬 수 있다.
MHSA는 입력 시퀀스를 여러 subspace (head)에 삽입하고 attention map을 별도로 계산하며 이는 성능 향상에 효과적인 것으로 입증되었다. 그러나 attention map은 계산 비용이 많이 들고, 연구에 따르면 그 중 상당수가 그다지 중요하지 않은 것으로 나타났다. 저자들은 계산 비용을 절약하기 위해 소규모 ViT 모델에서 중복 attention을 줄이는 방법을 살펴보았다. 저자들은 1.25배의 inference 속도 향상을 통해 너비가 축소된 Swin-T와 DeiT-T 모델을 학습하고 각 head와 각 블록 내의 나머지 head의 최대 코사인 유사도를 측정하였다
위 그래프에서는 특히 마지막 블록에서 attention head 사이에 높은 유사도가 있음을 관찰할 수 있다. 이 현상은 많은 head가 동일한 전체 feature에 대한 유사한 예측을 학습하고 계산 중복이 발생함을 시사한다. Head가 다양한 패턴을 학습하도록 명시적으로 장려하기 위해 각 head에 전체 feature의 split만 제공하는 직관적인 솔루션을 적용한다. 이는 group convolution 아이디어와 유사하다. 저자들은 수정된 MHSA를 사용하여 축소된 모델의 변형을 학습시키고 attention 유사도를 계산하였다. 위 그래프에서 볼 수 있듯이 모든 head에 대해 동일한 전체 feature를 사용하는 대신 서로 다른 head에서 feature를의 서로 다른 채널별 split을 사용하면 attention 계산 중복을 효과적으로 완화할 수 있다.
EfficientViT의 전체 아키텍처는 위 그림에 나와 있다. 구체적으로 C1 차원의 토큰에 16×16 패치를 삽입하기 위해 overlapping patch embedding을 도입하여 낮은 수준의 시각적 표현 학습에서 모델 용량을 향상시킨다. 아키텍처에는 세 단계가 포함된다. 각 단계에서는 제안된 EfficientViT 빌딩 블록을 쌓고 토큰 수는 각 서브샘플링 레이어에서 4배만큼 줄어든다. 효율적인 서브샘플링을 달성하기 위해 샌드위치 레이아웃을 갖는 EfficientViT 서브샘플 블록을 사용한다. 단, self-attention 레이어는 서브샘플링 중 정보 손실을 줄이기 위해 inverted residual block으로 대체된다. 모델 전반에 걸쳐 Layer Norm (LN) 대신 BatchNorm (BN)을 채택한다는 점은 주목할 가치가 있다. BN은 이전 convolution layer 또는 linear layer로 접힐 수 있으며 이는 LN에 비해 런타임에 이점이 있다. 일반적으로 사용되는 GELU 또는 HardSwish는 훨씬 느리고 때로는 특정 inference 배포 플랫폼에서 잘 지원되지 않기 때문에 ReLU를 activation function으로 사용한다.
본 논문은 위 그림과 같이 ViT를 위한 새로운 효율적인 빌딩 블록을 제안하였다. 이는 메모리 효율적인 샌드위치 레이아웃, cascaded group attention (CGA) 모듈, 파라미터 재할당 전략으로 구성되며 각각 메모리, 계산량, 파라미터 측면에서 모델 효율성을 향상시키는 데 중점을 둔다.
Cascaded Group Attention
Attention head 중복은 MHSA에서 심각한 문제로, 이로 인해 계산 비효율이 발생한다. 효율적인 CNN의 group convolution에서 영감을 받아 ViT를 위한 cascaded group attention (CGA)라는 새로운 attention 모듈을 제안한다. 이는 각 head에 전체 feature를 서로 다르게 분할하여 제공하므로 head 전체에 걸쳐 attention 계산을 명시적으로 분해한다. 그림처럼 계단식 방식으로 각 head의 attention map을 계산한다. 아는 각 head의 출력을 다음 head에 추가하여 feature 표현을 점진적으로 개선한다.
계단식 디자인의 장점은 두 가지가 있다. 우선 각 head에 서로 다른 feature split을 제공하면 attention map의 다양성을 향상시킬 수 있다. Group convolution과 유사하게 cascaded group attention은 QKV 레이어의 입력 채널과 출력 채널이 h배만큼 감소하므로 Flops와 파라미터를 h배만큼 절약할 수 있다. 또한 Attention head를 계단식으로 배열하면 네트워크 깊이가 증가하므로 추가 파라미터를 도입하지 않고도 모델 용량이 더욱 향상된다. 각 head의 attention map계산이 더 작은 QK 채널 차원을 사용하므로 약간의 지연 시간 오버헤드만 발생한다.
샌드위치 레이아웃의 수식.메모리 효율적인 블록을 구축하기 위해 채널 통신에서 memory-bound self-attention 레이어를 줄이고 메모리 효율적인 FFN 레이어를 사용하는 샌드위치 레이아웃 제안.
저자들은 6개의 서로 다른 너비와 깊이 스케일을 사용하여 모델들을 구축하고 각 단계마다 서로 다른 head 수를 설정하였다. MobileNetV3와 LeViT와 비슷하게 나중 단계보다 초기 단계에서 더 적은 수의 블록을 사용한다. 왜냐하면 더 큰 해상도의 초기 단계에서 처리하는 데 더 많은 시간이 소요되기 때문이다. 앞서 분석한 바와 같이, 나중 단계의 중복성을 완화하기 위해 작은 factor (≤2)를 사용하여 단계에 걸쳐 너비를 늘린다. C,L,H는 i번째 너비,깊이, head의 수를 나타낸다.
Parameter Reallocation. To improve parameter efficiency, we reallocate the parameters in the network by expanding the channel width of critical modules while shrinking the unimportant ones. Specifically, based on the Taylor importance analysis in Sec. 2.3, we set small channel dimensions for Q and K projections in each head for all stages. For the V projection, we allow it to have the same dimension as the input embedding. The expansion ratio in FFN is also reduced from 4 to 2 due to its parameter redundancy. With the proposed reallocation strategy, the important modules have larger number of channels to learn representations in a high dimensional space, which prevent the loss of feature information. Meanwhile, the redundant parameters in unimportant modules are removed to speed up inference and enhance the model efficiency.
**파라미터 효율성을 높이기 위해 중요한 모듈의 채널 폭을 확장하고 중요하지 않은 모듈은 축소하여 네트워크에서 파라미터를 재할당한다. 특히 저자들은 Taylor 중요도 분석을 기반으로 모든 단계에 대해 각 head의 Q projection과 K projection에 대한 작은 채널 크기를 설정했다. FFN의 확장 비율도 파라미터 중복으로 인해 4에서 2로 감소된다. 제안된 재할당 전략을 사용하면 중요한 모듈이 고차원 공간에서 표현을 학습할 수 있는 더 많은 수의 채널을 갖게 되어 feature 정보의 손실을 방지할 수 있다. 한편 중요하지 않은 모듈의 중복 파라미터를 제거하여 추론 속도를 높이고 모델 효율성을 향상시킨다.
Experiments
- **데이터셋: ImageNet-1K
구현 디테일
optimizer: AdamW
epoch 수: 300
learning rate: 2.5×10−2, cosine schedule
전체 batch size: 2048
이미지 크기: 224×224
augmentation: Mixup, auto-augmentation, random erasing
8개의 Nvidia V100 GPU
다운스트림 task fine-tuning
optimizer: AdamW
epoch 수: 300
learning rate: 1×10−3, cosine schedule
weight decay: 1×10−8
Conclusion
논문의 주내용은 메모리 연산의 효율성을 높이는 방법입니다. 저자들은 메모리 효율성과 계산 비용을 개선하기 위해 샌드위치 레이아웃을 적용한 새로운 구조의 EfficientViT 모델을 제안합니다. 이 레이아웃은 채널 통신을 강화하고 메모리 연산의 비효율성을 줄입니다. 또한, 다양한 피처 분할을 통해 주의력 헤드를 공급하는 캐스케이드 그룹 주의 모듈을 구현하여 중복을 감소시키고 주의 다양성과 계산 효율성을 향상시킵니다.
ViT의 계산 효율성과 메모리 사용을 개선하면서 정확도를 유지하는 것은 특히 실시간 애플리케이션에 적합한 발전이라고 생각합니다.
하지만 도입된 샌드위치 레이아웃은 FFN의 추가로 인해 모델의 크기가 Efficient CNN에 비해 약간 더 큽니다.