- 이 논문은 VVision Transformer(ViT) 모델에서 발생하는 계산 중복을 해결하고자 메모리 효율적인 설계를 적용한 SHViT모델을 제안한다. - SHViT는 단일 Head-Attention module을 도입하여 성능을 유지하면서 불필요한 계산을 줄이고, 여러 장치에서 뛰어난 성능을 보여준다.
1. Introduction
- ViT는 다양한 비전 작업에서 우수한 성능을 보이지만, 많은 데이터와 계산 비용이 요구된다. - 이를 개선하기 위해 본 논문은 ViT의 매크로 및 마이크로 설계에서 발생하는 중복 문제를 분석하고, 효율적인 패치 임베딩과 단일 헤드 설계를 제안함. - SHViT는 다양한 장치에서 빠른 속도와 높은 정확성을 달성해 실시간 비전 작업에 적합한 솔루션을 제공.
2. Analysis and Method
2.1 Analysis of Redundancy in Macro Design
- 대부분의 효율적인 모델들은 4X4 patchify stem / 4-stage 구성을 채택한다. - 일반적인 ViT모델은 의미있는 입력 토큰을 생성하기 위해 16X16 patchify stem을 사용하여 후속 MHSA 레이어에 전달 - 이러한 차이에 주목하여, 더 큰 stride를 가진 patchify stem이 MHSA 레이어뿐만 아니라, 엄격한 지연 시간 제한 내에서 효과적인 표현 학습에도 중요할 것이라고 가정
- 16X16 patchify stem과 3단계 설계를 사용하여 MetaFormer 블록을 기반을 두가지 모델을 구성 - 두 모델은 동일한 feature 맵 크기를 갖도록 유사한 채널 수로 설정
- 모델 (b)는 GPU/ CPU에서 (a)보다 각각 3.0/ 2.8배 더 빠르지만, 정확도는 1.5% 낮았다. - 해상도가 256X256으로 증가했을때 (b')모델은 (a)와 비교해 속도가 크게 증가, 성능 또한 크게 손실되지 않음
기존의 4×4 패치화 대신 16×16 패치화와 3단계 구조를 적용한 효율적인 매크로 설계를 제안. 이를 통해 초기 단계의 공간 중복을 줄이고 메모리 접근 비용을 감소시키며, 해상도가 높아져도 성능을 유지하는 결과를 확인.
2.2 Analysis of Redundancy in Micro Design
- MHSA 레이어는 여러 하위 공간에서 독립적으로 Attention 맵을 계산하고 적용, 이를 통해 성능 향상을 꾸준히 입증함 - Attention map은 계산 비용이 많이 들며, 많은 head가 반드시 필요하지 않다는 결과가 나옴 - 이 논문은 ViT 모델의 다중 Head중복을 3가지 실험을 통해 탐구 1. attenttion map visualization (attention map 시각화) 2. head similarity analysis (head 유사성 분석) 3. head ablation study (head 소거 실험 )
head similarity analysis
- 동일한 레이어 내에서 각 헤드와 다른 헤드 간의 평균 코사인 유사도를 측정
head ablation study
- 특정 레이어에서 일부 헤드의 출력을 무력화하고 나머지 레이어에서는 전체 헤드를 유지하여 성능 영향을 평가
- 초기 단계에서는 상위 성능 헤드가 컨볼루션처럼 동작하는 경향이 있으며, 제거 시 성능에 미치는 영향이 적은 헤드는 일반적으로 정보를 더 전역적으로 처리 - 후반 단계에서는 특징 및 예측 수준에서 많은 중복이 발생. (EX) DeiT 모델의 후반 단계는 78.3%의 헤드 유사도를 보이며, Swin-T 모델 또한 높은 유사성을 나타냄 Swin-T의 경우, 하나의 헤드를 제거해도 성능 저하가 거의 없으며, 경우에 따라 성능이 약간 향상되기도 함.
- 이전 접근 방식은 전체 네트워크를 학습시킨 후 불필요한 헤드를 가지치기하지만, 이는 훈련 중 계산 자원과 메모리 사용량을 증가시킨다. - 이를 해결하기 위해 단일 헤드 어텐션 모듈을 설계하여 헤드 중복을 본질적으로 피한다. - 이를 통해 훈련 및 추론 과정에서 효율성을 높였다.
MHSA 레이어의 다중 헤드는 성능 향상에 기여하지만, 많은 헤드가 반드시 필요하지는 않다. 후반 단계에서는 높은 중복이 발생하며, 특정 헤드를 제거해도 성능에 큰 영향을 미치지 않는다. 본 연구는 이러한 중복을 피하기 위해 단일 헤드 어텐션 모듈을 설계하여 훈련과 추론의 효율성을 높임.
2.3 Single-Head Self-Attention
- 위 분석에 기반하여 SHSA를 제시함 - SHSA는 입력 채널의 일부에만 단일 헤드로 Attention을 적용하여 공간적 특징을 집계하며, 나머지 채널은 그대로 둔다. (Default 값으로 Γ = 1/4.67 을 사용)
- 여기서 WQ, WK, WV, WO는 투영 가중치/ dqk는 쿼리 및 키의 차원 (Default = 16) - 최종 투영은 전체 채널에 적용되어 Attention 특징의 효과적인 전파를 보장한다.
- Single head attention의 3가지 접근법을 비교 - (a): Bottleneck/ (b): Full channel/ (c): Partial channel(Ours) => (c)가 가장 CPU와 GPU에서 가장 높은 이미지 처리 속도를 보여주며, TOP-1 정확도가 가장 높다. (75.2%)
- Attention 레이어를 효과적으로 사용하기 위해 Layer Normalization이 필수적 - multi head 방식의 경우 reshape와 같은 데이터 이동이 필요 - MHSA의 실행시간 대부분이 reshape와 normalization 같은 메모리 기반 작업에 소비 => 메모리 기반 작업의 사용을 줄이거나 일부 입력 채널에만 적용하여 SHSA모듈은 GPU/ CPU의 컴퓨팅 파워를 최대한 활용
공간적 특징 집계를 위해 일부 채널에만 단일 헤드를 적용하는 SHSA를 제안. 이를 통해 메모리 효율성과 성능을 유지하면서 불필요한 계산을 줄임.
2.4 Single-Head Vision Transformer
<SHViT의 구조> - 입력 이미지가 주어지면 먼저 4개의 3X3 Stride conv를 적용 - 표준 ViT모델에서 사용하는 Stride-16 16X16 Convolution 토큰화 방식과 비교하여 , 중첩된 패치화 스템은 거 나은 로컬 표현을 추출할 수 있다. - 이후 토큰은 계층적 표현 추출을 위해 3단계의 SHViT 블록을 통과 - 각 SHViT 블록은 3가지 주요 모듈로 구성 1. Depthwise Convolution (DWConv) : local feature 집계, conditional position embedding 2. Single-Head Self-Attention (SHSA) : 글로벌 컨텍스트 모델링 3. Feed-Forward Network (FFN) : 채널 간 상호작용 - 정보 손실 없이 토큰을 줄이는 단계가 포함 - 마지막으로, global average polling과 fully-connected layer를 사용해 예측을 출력
3. Experiments
3.1 Implementation Details
- ImageNet-1K에서 이미지 분류를 수행 (1000개 범주의 1.28M 학습 이미지와 50K 검증 이미지 포함) - 모든 모델은 AdamW optimizer를 사용하여 학습 속도 10^-3과 배치 크기 2048로 300 에포크 동안 처음부터 학습 - 코사인 학습률 스케줄러와 5 에포크 동안의 선형 워밍업을 사용 - SHViT-S1에서 S4까지 가중치 감쇠는 각각 0.025/0.032/0.035/0.03으로 설정 - 공정한 비교를 위해, Mixup, 랜덤 지우기, 오토 어그멘테이션을 포함한 동일한 데이터 증강 방법을 사용 - 384²와 512² 해상도의 경우, 모델을 30 에포크 동안 미세 조정하고, 가중치 감쇠 10^-8과 학습 속도 0.004를 적용 - 다양한 하드웨어 플랫폼에서 처리량 성능 평가
3.2 SHViT on ImageNet-1K Classification
다른 모델들과 SHViT를 비교
- SHViT-S1, S2, S3, S4는 SHViT 모델의 네 가지 변형으로, 각각 모델의 깊이와 너비가 다르게 설정된 버전을 나타냄 1. SHViT-S1: 상대적으로 얕은 깊이와 임베딩 차원, 가장 가벼운 버전으로 처리 속도가 가장 빠르다. 2. SHViT-S2: 조금 더 깊은 구조와 큰 임베딩 차원, 성능과 효율성의 균형을 맞춘 변형 3. SHViT-S3: 더 높은 임베딩 차원과 해상도를 가지며, 높은 Accuracy를 목표로 함 4. SHViT-S4: 가장 깊고 넓은 네트워크로 고해상도 입력과 큰 임베딩 차원을 처리할 수 있어, 최상의 정확도를 제공 => 모델의 복잡성과 목표 성능에 따라 "층의 깊이, 임베딩 차원, 해상도, 부분 채널 비율 및 확장 비율"을 조정
SHViT는 다양한 CNN, ViT, 하이브리드 모델과 비교하여 정확도와 처리 속도에서 우수한 균형을 보여준다. 특히 SHViT-S4는 고해상도 및 모바일 플랫폼에서 탁월한 성능을 유지하며, 효율적인 비전 백본으로 활용
3.3 SHViT on downstream tasks
- Object detection(객체 탐지): SHViT-S4는 MobileNetV3와 비교하여 모바일 장치에서 2.3배 더 빠르고 +8.9 AP의 성능 향상을 보여준다. MobileFormer와 비교해도 A100 GPU에서 3.2배, 모바일 장치에서 8.2배 더 빠른 속도를 유지하면서도 더 높은 성능을 달성
- Instance Segmentation(인스턴스 분할): SHViT-S4는 PoolFormer보다 AP 성능에서 더 높은 수치를 보이며, EfficientViT 및 EfficientNet과 비교하여 GPU, CPU, 모바일 장치에서 각각 4.3배, 8.1배, 2.1배 더 빠르다.
SHViT의 3단계 큰 스트라이드 패치 스템이 고해상도에서 유의미한 토큰을 생성하면서도 계산 효율성을 유지하여, 다양한 장치와 작업에서 경쟁력 있는 성능을 보임을 나타낸다.
3.4 Ablation Study
SHSA 레이어가 다중 헤드 자기어텐션(MHSA) 레이어처럼 글로벌 컨텍스트를 효과적으로 캡처할 수 있는지 평가하기 위해, SHSA 레이어를 MHSA 레이어로 대체하거나 제거하는 소거 실험을 수행
- SHSA 레이어는 MHSA 레이어에 비해 더 나은 속도-정확도 균형을 나타낸다. - SHSA 레이어를 제거하면 모델 속도는 빨라지지만 정확도가 크게 떨어진다. - 반면, 모델 (2)는 SOTA 모델과 비교할 때 매우 경쟁력 있는 성능을 달성할 수 있으며, 이는 제안된 매크로 설계가 엄격한 지연 시간 제약 하에서 견고한 구조적 기반을 제공함을 보여준다. - 기본적으로 모든 SHViT 모델에 대해 부분 비율을 1 / 4.67로 설정하여 최적의 속도-정확도 균형을 달성
SHSA 레이어는 MHSA 레이어보다 더 나은 속도-정확도 균형을 제공하며, SHSA를 제거하면 속도는 빨라지지만 정확도가 크게 떨어진다. 최적의 속도-정확도 균형을 위해 부분 비율을 1/4.67로 설정. 너무 작은 값은 성능 향상에 제한적이며, 너무 큰 값은 비용 대비 성능 향상이 충분하지 않다.
4. Related Work
- 자원이 제한된 장치에서 CNN을 활용하려는 연구가 많은 관심을 받고 있다. - MobileNets에서의 합성곱 분해, ShuffleNets에서의 채널 셔플링, GhostNet에서의 저비용 선형 변환, EfficientNet에서의 복합 스케일링 법칙, 그리고 다양한 연구에서의 구조적 재파라미터화 등 - ViT(Vision Transformer)영역에서도 다양한 장치에서 추론 속도를 가속화하기 위한 효율적인 설계에 대한 노력이 진행 중 1. CNN의 로컬 우선순위를 통합하는 새로운 ViT 아키텍처를 설계 => 주로 후반 단계에서 Attention을 사용하여 큰 계산 부담 없이 글로벌 정보를 효율적으로 추출 (다른 방법들은 단일 토큰 믹서 내에서 또는 블록별로 Attention과 합성곱을 병렬로 사용하여 풍부한 특징 집합을 결합) 2. MHSA의 계산 복잡도를 줄인다. => 최소한의 공간적 중복으로 토큰을 구성하는 것을 우선시함. - 최근 연구들은 여러 헤드가 유사하게 작동하며 성능에 큰 영향을 미치지 않고 가지치기할 수 있음을 보여주었다. - EfficientViT는 풀 채널을 다른 부분으로 나눠 헤드에 제공하여 어텐션 다양성을 개선하는 방법을 제안 - 해당 논문에서는 다중 헤드 중복을 줄이는 대신 단일 헤드 구성을 사용하여 다중 헤드 중복을 본질적으로 방지하면서 계산 비용도 절감
5. Conclusion
- 이 논문에서는 기존 모델들이 일반적으로 사용하는 아키텍처 설계의 공간 및 채널 차원에서 중복성을 조사함. - 이후 3단계 계층적 표현을 갖춘 16X16 패치 임베딩과 계산 중복을 줄이기 위한 Single Head Self Attention을 제안함. - macro/micro 설계를 기반으로 다양한 장치와 비전 작업에 매우 빠른 추론 속도와 높은 성능을 달성하는 SHViT모델을 제시
이 연구는 모델 설계에서 발생하는 공간 및 채널 중복 문제를 해결하기 위해 16×16 패치 임베딩과 단일 헤드 자기어텐션을 사용한 SHViT를 제안. SHViT는 macro/micro 설계를 통해 빠른 추론 속도와 높은 성능을 달성하며, 다양한 장치와 비전 작업에서 유용하게 활용될 수 있다.
