단일 경로에서 자기 주의를 컨볼루션 레이어로 가공하기

비전 트랜스포머(Vision Transformers, ViTs)는 다양한 컴퓨터 비전 작업에서 뛰어난 성능을 달성하였습니다. 그러나 다중 헤드 자기 주의(multi-head self-attention, MSA) 층을 사용하여 전역 상관 관계를 모델링하는 것은 두 가지 널리 인식된 문제를 초래합니다: 막대한 계산 자원 소모와 로컬 시각 패턴을 모델링하기 위한 내재적 귀납 편향의 부족입니다. 이러한 두 문제를 해결하기 위해, 우리는 간단하면서도 효과적인 방법인 단일 경로 비전 트랜스포머 프루닝(Single-Path Vision Transformer pruning, SPViT)을 제안합니다. 이 방법은 사전 학습된 ViTs를 효율적으로 자동으로 압축하여 적절한 로컬성이 추가된 컴팩트 모델로 변환합니다.구체적으로, 우리는 MSA와 컨볼루션 연산 사이에 새로운 가중치 공유 방식을 제안하여 모든 후보 연산을 인코딩할 수 있는 단일 경로 공간(single-path space)을 제공합니다. 이를 통해 각 MSA 층에서 사용할 매개변수의 부분 집합을 찾는 연산 검색 문제(operation search problem)가 크게 줄어들어 계산 비용과 최적화 난이도가 감소하며, 사전 학습된 MSA 매개변수를 사용하여 컨볼루션 커널이 잘 초기화될 수 있습니다. 단일 경로 공간에 기반하여, 우리는 학습 가능한 이진 게이트(learnable binary gates)를 도입하여 MSA 층에서의 연산 선택을 인코딩합니다.비슷하게, 우리는 학습 가능한 게이트를 더 활용하여 피드포워드 신경망(Fully-Fed Forward Network, FFN) 층의 미세한 MLP 확장 비율(fine-grained MLP expansion ratios)을 인코딩합니다. 이렇게 하면 우리의 SPViT는 광범위하고 통합된 검색 공간(search space)에서 학습 가능한 게이트를 최적화하여 자동으로 탐색하고, 각 개별 밀집 모델(dense model)에 대해 MSA-FFN 프루닝 비율(pruning proportions)을 유연하게 조정할 수 있습니다.우리는 대표적인 두 개의 ViTs에 대한 광범위한 실험을 수행하였으며, 결과적으로 우리의 SPViT는 ImageNet-1k에서 프루닝 성능面에서 새로운 최고 기록(state-of-the-art, SOTA)을 달성하였습니다. 예를 들어, 우리의 SPViT는 DeiT-B 모델의 FLOPs(Floating Point Operations Per Second)를 52.0% 감소시키면서 동시에 0.6%의 Top-1 정확도 향상을 가져왔습니다. 소스 코드는 https://github.com/ziplab/SPViT 에서 확인할 수 있습니다.