[논문 리뷰] Sparse MLP for Image Recognition: Is Self-Attention Really Necessary?

Chuanxin Tang, Yucheng Zhao외 4명|arXiv (Cornell University)|2021. 09. 12.

Advanced Neural Network Applications인용 수 13

한 줄 요약

이 논문은 축 방향 전역 모델링과 가중치 공유를 사용하는 희소 MLP(sMLP) 모듈을 통해 자기주의(self-attention)를 대체하는 attention-free 비전 모델인 sMLPNet을 제안한다. 과적합을 방지하면서 파라미터 수와 FLOPs를 줄임으로써, sMLPNet은 오직 65.9M 파라미터로 ImageNet-1K에서 83.4%의 top-1 정확도를 달성한다. 이는 Swin Transformer의 성능을 25% 더 작은 모델 크기와 10% 적은 FLOPs로 재현한 것으로, 자기주의가 최첨단 이미지 인식에 필수적이지 않다는 것을 입증한다.

ABSTRACT

Transformers have sprung up in the field of computer vision. In this work, we explore whether the core self-attention module in Transformer is the key to achieving excellent performance in image recognition. To this end, we build an attention-free network called sMLPNet based on the existing MLP-based vision models. Specifically, we replace the MLP module in the token-mixing step with a novel sparse MLP (sMLP) module. For 2D image tokens, sMLP applies 1D MLP along the axial directions and the parameters are shared among rows or columns. By sparse connection and weight sharing, sMLP module significantly reduces the number of model parameters and computational complexity, avoiding the common over-fitting problem that plagues the performance of MLP-like models. When only trained on the ImageNet-1K dataset, the proposed sMLPNet achieves 81.9% top-1 accuracy with only 24M parameters, which is much better than most CNNs and vision Transformers under the same model size constraint. When scaling up to 66M parameters, sMLPNet achieves 83.4% top-1 accuracy, which is on par with the state-of-the-art Swin Transformer. The success of sMLPNet suggests that the self-attention mechanism is not necessarily a silver bullet in computer vision. The code and models are publicly available at https://github.com/microsoft/SPACH

연구 동기 및 목표

• 비전 트랜스포머에서 자기주의가 최첨단 성능을 내는 데 필수적이라는 가정에 도전한다.
• 과적합과 높은 파라미터 수라는 주요 한계를 해결함으로써, 자기주의 없이도 경쟁 가능한 성능을 내는 MLP 기반 모델이 가능한지 조사한다.
• 축 기반 계산과 가중치 공유를 통해 효율적인 전역 특징 모델링을 가능하게 하는 새로운 희소 MLP(sMLP) 모듈을 설계한다.
• Swin Transformer와 같은 최첨단 비전 트랜스포머의 성능을 따라하거나 초월하는 attention-free 네트워크(sMLPNet)를 구축한다.
• 데이터 증강이나 추가 사전학습 없이도 표준 학습 프rotocol에 따라 ImageNet-1K에서 sMLPNet의 효과성을 검증한다.

제안 방법

• sMLP 모듈은 2D 이미지 토큰의 행과 열 축을 따라 1D MLP를 적용하여 축 기반 전역 종속성 모델링을 가능하게 한다.
• 모든 행(또는 열) 간에 파라미터를 공유함으로써 모델 크기와 계산 복잡도를 크게 감소시킨다.
• 각 토큰이 sMLP 레이어에서 오직 자신의 행과 열 이웃 토큰과만 상호작용하도록 하여 희소성(spatiail sparsity)을 도입함으로써 전체 제곱형 자기주의를 피한다.
• sMLP 모듈은 ViT 유사 아키텍처에 통합되어 토큰 혼합 블록에서 자기주의 메커니즘을 대체한다.
• 지역성 편향을 유지하기 위해 토큰 혼합 블록에서 depthwise 컨벌루션을 사용하며, 채널 혼합은 표준 MLP와 동일하게 유지된다.
• 너비와 깊이 조정을 통해 모델 스케일링을 수행하였으며, 실험은 224×224 해상도의 ImageNet-1K에서 수행되었다.

실험 결과

연구 질문

• RQ1 MLP 기반 비전 모델이 자기주의를 사용하지 않고도 최첨단 성능을 달성할 수 있는가?
• RQ2 비전 트랜스포머의 자기주의 메커니즘이 다른 전역 모델링 기법보다 본질적인 이점을 제공하는가?
• RQ3 MLP에서의 희소 연결성과 가중치 공유가 과적합을 완화하고 계산 비용을 감소시키면서도 성능 유지를 가능하게 하는가?
• RQ4 지역성 편향과 축 기반 전역 모델링을 조합함으로써 더 작은 모델 크기로도 경쟁 가능한 정확도를 달성할 수 있는가?
• RQ5 Swin Transformer의 성공은 지역성에 기인한 것인가, 아니면 자기주의 덕분인가? 그리고 attention-free 모델이 그 성능을 따라할 수 있는가?

주요 결과

• sMLPNet은 오직 24.1M 파라미터로 ImageNet-1K에서 81.9%의 top-1 정확도를 달성하여, 동일한 크기 제약 조건에서 대부분의 CNN과 비전 트랜스포머를 능가한다.
• 65.9M 파라미터를 가진 sMLPNet-B는 Swin-B와 동일한 83.4%의 top-1 정확도를 달성하지만, 파라미터 수는 25% 적고 FLOPs는 10% 적다.
• 66M 파라미터에서도 과적합의 징후가 전혀 없었으며, MLP-Mixer 및 기타 MLP 기반 모델이 심각한 과적합을 겪는 것과는 대조적으로, 성능 저하 없이 안정적인 학습을 보였다.
• sMLPNet-T(19.2M 파라미터)는 5B FLOPs 이하의 모델 중에서 가장 높은 81.3%의 top-1 정확도를 기록했다.
• DeiT-S와 DeiT-B보다도 sMLPNet-T는 훨씬 적은 파라미터와 FLOPs로 더 높은 정확도를 달성했다.
• sMLP 모듈은 시퀀스 길이에 대해 계산 복잡도를 제곱형에서 선형으로 감소시켜 고해상도 입력의 효율적 처리를 가능하게 했다.