[논문 리뷰] Pruning Self-attentions into Convolutional Layers in Single Path

Haoyu He, Jianfei Cai외 5명|arXiv (Cornell University)|2021. 11. 23.

Advanced Neural Network Applications참고 문헌 64인용 수 16

한 줄 요약

SPViT는 단일 경로 프루닝 프레임워크를 제안하며, 다중 헤드 어텐션(MSA)과 피드포워드 네트워크(FFN) 프루닝을 함께 최적화함으로써, MSA와 컨볼루션 연산 간의 가중치 공유를 통해 효율적이고 자동적인 아키텍처 탐색을 가능하게 한다. 이는 낮은 계산 비용으로도 상태의 기술을 달성하며, DeiT-B에서 52.0% FLOPs를 감소시키면서 ImageNet-1k에서 상위-1 정확도를 0.6% 향상시킨다.

ABSTRACT

Vision Transformers (ViTs) have achieved impressive performance over various computer vision tasks. However, modeling global correlations with multi-head self-attention (MSA) layers leads to two widely recognized issues: the massive computational resource consumption and the lack of intrinsic inductive bias for modeling local visual patterns. To solve both issues, we devise a simple yet effective method named Single-Path Vision Transformer pruning (SPViT), to efficiently and automatically compress the pre-trained ViTs into compact models with proper locality added. Specifically, we first propose a novel weight-sharing scheme between MSA and convolutional operations, delivering a single-path space to encode all candidate operations. In this way, we cast the operation search problem as finding which subset of parameters to use in each MSA layer, which significantly reduces the computational cost and optimization difficulty, and the convolution kernels can be well initialized using pre-trained MSA parameters. Relying on the single-path space, we introduce learnable binary gates to encode the operation choices in MSA layers. Similarly, we further employ learnable gates to encode the fine-grained MLP expansion ratios of FFN layers. In this way, our SPViT optimizes the learnable gates to automatically explore from a vast and unified search space and flexibly adjust the MSA-FFN pruning proportions for each individual dense model. We conduct extensive experiments on two representative ViTs showing that our SPViT achieves a new SOTA for pruning on ImageNet-1k. For example, our SPViT can trim 52.0% FLOPs for DeiT-B and get an impressive 0.6% top-1 accuracy gain simultaneously. The source code is available at https://github.com/ziplab/SPViT.

연구 동기 및 목표

• 다중 헤드 어텐션 레이어로 인한 높은 계산 비용과 유도적 편향의 부족 문제를 해결하기 위해.
• 사전 학습된 ViT를 효율적이고 자동적으로 소형 모델로 프루닝함으로써 도입된 국소적 유도적 편향을 가능하게 하기 위해.
• NAS 기반 프루닝의 탐색 비용과 최적화 난이도를 줄이기 위해, 단일 경로 공간에서 MSA와 컨볼루션 연산을 통합함으로써.
• 목표 FLOP 제약 조건 하에 MSA 레이어와 FFN 확장 비율의 프루닝을 동시에 최적화하기 위해.
• 탐색된 아키텍처의 실험적 분석을 통해 ViT 구성 요소의 구조적 중요성을 탐색하기 위해.

제안 방법

• 컨볼루션 연산을 MSA 파라미터의 부분집합으로 표현하는 가중치 공유 기법을 제안하여, 어텐션과 컨볼루션 간의 공유 파라미터화를 가능하게 한다.
• 모든 후보 연산(MSA 및 컨볼루션)이 각 레이어의 단일 MSA 내에 인코딩된 단일 경로 탐색 공간을 구축함으로써 탐색 복잡도를 감소시킨다.
• MSA 레이어에 학습 가능한 이진 게이트를 도입하여, 어떤 MSA 파라미터 부분집합을 사용할지 선택함으로써 효과적으로 컨볼루션 연산으로 프루닝한다.
• 학습 가능한 게이트를 FFN 레이어로 확장하여 은닉 차원과 MLP 확장 비율을 동시에 프루닝하고 최적화한다.
• 학습 가능한 게이트를 종단 간 최적화하여, 목표 FLOP 제약 조건 하에 자동으로 효율적인 아키텍처를 발견한다.
• 사전 학습된 MSA 파라미터를 사용하여 컨볼루션 커널을 초기화함으로써 효과적이고 안정적인 학습을 가능하게 한다.

Figure 1: Multi-path vs. single-path search space. (a) Multi-path search space. MSA and different types of convolutional operations are added as separate trainable paths, leading to an expensive search cost. $\mathcal{P}$ and $g_{p}$ denote the ordered operation set and the learnable gate for the $p

실험 결과

연구 질문

• RQ1 통합된 단일 경로 탐색 공간은 NAS 기반 ViT 프루닝의 계산 비용과 최적화 난이도를 줄일 수 있는가?
• RQ2 MSA와 컨볼루션 연산 간의 가중치 공유가 컨볼루션 커널의 효과적이고 안정적인 초기화를 가능하게 하는가?
• RQ3 통합된 탐색 공간에서 MSA와 FFN 레이어를 동시에 프루닝하는 것이 별도의 프루닝보다 더 나은 효율-정확도 트레이드오프를 달성할 수 있는가?
• RQ4 탐색된 아키텍처의 구조적 선택(예: 浅층 블록에서 MSA를 컨볼루션으로 프루닝하는 것)은 모델의 기능적 행동과 어떻게 관련되어 있는가?
• RQ5 FFN 용량은 다양한 ViT 블록에서 어떤 역할을 하는가? 그리고 표준 및 계층적 ViT 아키텍처에서 어떻게 다를 수 있는가?

주요 결과

• SPViT는 DeiT-B에서 52.0% FLOPs를 감소시키면서 ImageNet-1k에서 상위-1 정확도를 0.6% 향상시켜, 사후 훈련 프루닝 분야에서 새로운 SOTA를 수립한다.
• 이 방법은 얕은 MSA 레이어를 컨볼루션 레이어 또는 스킵 연결로 프루닝하는 데 성공하였으며, 얕은 블록에서 국소 어텐션 패턴을 보이는 것과 일치한다.
• 모든 모델에서 깊은 MSA 레이어는 유지되며, 이는 글로벌 표현 학습에 있어 중요한 역할을 함을 확인한다.
• 표준 ViT의 깊은 레이어에서 FFN 레이어는 더 높은 부여잡힘(레이어의 임의의 부분집합이 의미 없이 존재함)을 보이며, 우선적으로 프루닝된다. 이는 후행 단계에서 낮은 용량이 요구됨을 시사한다.
• 스위니-Ti와 같은 계층적 ViT에서는 각 스테이지의 얕은 레이어에서 더 많은 은닉 차원을 유지함을 확인하였으며, 이는 토큰 병합 이후 높은 용량이 필요함을 시사한다.
• 단일 경로 탐색 공간은 다중 경로 NAS에 비해 탐색 비용과 최적화 난이도를 감소시켜, MSA와 FFN 프루닝에 대한 효율적인 종단 간 동시 탐색을 가능하게 한다.

Figure 2: Overview of our weight-sharing scheme, which expresses the bottleneck convolution with MSA parameters. First, we restrict the value and output projections in the MSA layer to only process the features within local regions. Then, we learn to ensemble the value projections with a parameteriz