안녕하세요. 지난 포스팅의 [IC2D] Resolution Adaptive Networks for Efficient Inference (CVPR2020)에서는 resolution adaptation과 depth adaptation을 동시에 달성하여 inference speed를 향상시킨 RANet에 대해서 소개하였습니다. 오늘은 Transformer에 NAS를 곁들인 NASViT를 소개하도록 하겠습니다.
Background
최근 CNN에 이어 Transformer는 self-attention을 통해 global dependency 또는 long-range dependency를 추출할 수 있다는 장점을 통해 엄청난 인기를 가지게 되었습니다. 특히, 영상 분류에서의 ViT (ICLR2021), Swin Transformer (ICCV2021), CSWin Transformer (CVPR2022) 그리고 객체 탐지에서의 DETR (ECCV2020), 의미론적 영상 분할에서의 SegFormer (NIPS2021), 마지막으로 비디오 처리 분야에서의 TimeSFormer (ICML2021) 등이 있었죠. 하지만 이러한 모델들은 너무나도 큰 대규모 모델 크기과 높은 사전학습 비용으로 인해 mobile이나 IoT 계열과 같이 컴퓨팅 리소스가 제한되는 경우에 활용하기 어렵다는 문제점이 있었죠.
이러한 문제를 해결하기 위해 CNN과 유사하게 효율적인 모델들이 나오기 시작하였습니다. 대표적으로 TinyViT 그리고 MobileViT 등이 있었죠. 하지만, 이러한 방법들이 최적의 방법이라고 볼 수 있을까요? CNN 역시 AlphaNet과 FBNetV3와 같은 Neural Archiecture Search (NAS) 기법을 적용하여 높은 효율성을 달성하였듯이 ViT 역시 가능하지 않을까요? 본 논문에서는 이러한 궁금증으로 Transformer에 NAS를 곁들여봅니다.
NAS는 실제로 정확도-효율성 사이의 trade-off를 기반으로 최적의 모델을 제공하는 AutoML 기반 방법이라고 볼 수 있습니다. 최근에는 Supernet-based NAS가 유행하고 있습니다. 대표적으로 BigNAS와 AlphaNet 등이 있죠. 이러한 방법들은 Sandwich Sampling Rule이라는 방법을 통해 모델을 샘플링하고 supernet과 subnetwork들을 동시적으로 학습 및 최적화를 진행합니다.
하지만, AlphaNet 기반의 NAS를 직접적으로 Transformer에 NAS로 적용했을 때 ViT 보다도 훨씬 낮은 성능을 가지게 되었습니다. 본 논문에서는 이러한 원인을 샘플링된 supernet과 subnetwork들 간의 gradient conflict가 발생하여 optimization이 서로 방해가 된 것으로 추측하고 이를 해결할 수 있는 gradient projection 및 Transformer-based NAS training recipe을 제안합니다.
Neural Architecture Search for Efficient Transformer
본 논문의 목표를 요약하자면 "FLOPs가 200M ~ 800M 사이의 소규모에서 중규모 사이의 효율적인 ViT 모델을 설계"하는 것입니다. 이러한 목표를 달성하기 위해 저자들은 LeViT (ICCV2021)에 주목하였습니다. 이 모델은 기본적으로 convolution과 transformer block을 결합하여 설계된 hybrid architecture입니다. convolution은 초반 feature map의 해상도가 클 때 적용되어 local computation의 장점을 살리며 transformer는 후반에 feature map의 해상도가 작을 때 적용되어 global dependency를 추출 시 self-attention의 연산량을 줄이기 위해 사용됩니다.
1) Search Space
NAS를 수행할 때는 search space를 적절하게 정의하는 것이 중요합니다. 본 논문에서는 LeViT의 성분을 바탕으로 search space를 정의하였습니다. LeViT와 마찬가지로 초기에는 convolution 연산이 주로 사용되고 후반에는 transformer 연산이 주로 사용되는 것을 볼 수 있습니다. 이 과정에서 각 성분의 세부 파라미터 (width, depth, kernel size 등)은 표 1에 제시되어 있습니다.
2) Naive Supernet-based NAS fails to find Accruate ViTs
이제 본격적으로 NAS를 Efficient Transformer를 찾기 위해 적용해보도록 하겠습니다. 기본적으로 AlphaNet의 NAS 셋팅을 그대로 차용하고 ImageNet에 360 에폭 동안 학습을 진행합니다.
다만 LeViT의 training recipe 중 Adam을 사용하는 것이 더 효고적이라는 것과 pre-trained teacher model인 EfficientNet-B5를 추가했다는 점이 다릅니다.
그림 2는 AlphaNet과 NAS를 적용한 ViT 사이의 training curve를 보여주고 있습니다. 결과적으로 CNN baseline 보다도 훨씬 낮은 성능을 보여주고 있습니다. 본 논문에서는 이러한 결과를 내는 요인을 3가지로 추정합니다.
2-1) Investigation 1: Is the search space designed poorly?
근본적으로 LeViT 기반의 Search space를 사용한 것이 잘못일수도 있을까요? 본 논문에서는 이를 검증하기 위해 FLOPs가 190M ~ 591M 사이의 sub-network 4개를 무작위로 선택한 뒤 동일한 data augmentation와 regularization을 이용하여 학습을 진행하였습니다. 결과적으로 supernet보다 더 높은 성능을 달성하는 것을 보아하니 search space는 큰 문제가 없는 것 같네요.
2-2) Investigation 2: Are the training settings suitable for ViTs?
다음으로 가능한 요인은 Training Recipe 자체가 맞지 않을 수도 있습니다. 이를 검증하기 위해 AlphaNet, DeiT, LeViT의 Training Recipe를 각각 smallest, largest subnetwork를 학습시키는 데 사용해봅니다. 실제로 최근 제안된 ViT 계열의 논문들은 더 강한 data augmentation과 regularization을 요구하기 때문에 AlphaNet 기반의 Training Recipe가 맞지 않을 수도 있을 거 같네요. 하지만, Table 3에서 보면 AlphaNet의 Training Recipe가 가장 좋은 결과를 제공하는 것을 볼 수 있죠. 따라서 Training Recipe 역시 큰 문제가 되지 않습니다.
2-3) Investigation 3: Saturated Supernet training due to gradient conflicts?
마지막 가능성은 서로 다른 supernet과 subnetwork 들의 gradient가 서로 충돌하여 optimization이 제대로 되지 않을수도 있습니다. 본 논문에서는 이를 검증하기 위해 supernet과 다른 sub-network들의 평균 gradient의 cosine similarity를 계산합니다. 여기서 negative similarity는 결국 두 gradient가 서로 다른 방향을 나타낸다는 것을 의미하기 때문에 학습에 서로 방해가 될 것 입니다. 실제로 AlphaNet의 negative 비율이 ViT 보다 훨씬 적은 것을 볼 수 있죠. 따라서 본 논문에서는 gradient conflict로 인해 발생하는 성능 하락을 해결하기 위한 방법론을 제시합니다.
3) Gradient Conflict-Aware Supernet Training
위에서 언급한 문제를 해결하기 위해 본 논문에서 제시하는 방식은 supernet의 gradient를 가장 우선시 되는 sub-network의 gradient로 projection 시켜주는 것입니다. 이 방식은 위 방정식을 기반으로 동작하며 만약 두 gradient의 cosine similarity가 양수라면 동일하게 supernet의 gradient를 변경하지 않고 cosine similarity가 음수라면 supernet의 gradient를 projection 시켜줍니다.
4) Switchable Scaling Layer
여기에 더해 Switchable Scaling Layer를 추가적으로 도입하여 모델의 학습 속도를 증가시켜줍니다. 이는 단순하게 learnable parameter $\mathbf{w}_{[c, d]}$를 입력 특징 맵 $\mathbf{x}_{[c, d]}$에 channel-wise로 곱해주어 전체적으로 값을 scaling 해주는 과정입니다. 이 과정으로 모델의 capacity가 더욱 크게 증가한다고 하네요.
5) Reduced Data Augmentation and Regularization
또한, AlphaNet에서 사용한 AutoAugment에 비해 상대적으로 약한 data augmentation인 RandAugment를 추가적으로 도입하였으며 다른 regularization 기법들은 아예 사용하지 않은 것을 볼 수 있습니다.
Experimental Results
1) ImageNet-1K Classification
2) CityScape and ADE20K
